Learning and Leveraging World Models in Visual Representation Learning

相关链接：arxiv

关键字：学习世界模型 、视觉表示学习 、自监督学习 、JEPA 、Image World Models

摘要

本文探索了在自监督视觉表示学习中学习和利用世界模型的方法。作者引入了图像世界模型（Image World Models, IWM），一种学习预测全局光度变换效果的新方法。作者发现有效的IWM模型取决于多种因素，包括条件设置、预测难度和容量。此外，通过微调，IWM可以适应不同的任务，甚至在图像分类和语义分割等任务上与或超越了先前的自监督方法。实验证明，通过IWM学习，可以控制所学表示的抽象级别，能够学习到如对比学习方法中的不变表示，或如遮蔽图像建模的等变表示。

核心方法

IWM建模框架JEPA（Joint-Embedding Predictive Architecture）的核心在于学习一个通过预测来利用世界模型的能力，包括以下方面：

条件设置：通过为预测器提供关于目标的几何信息（例如屏蔽tokens）和变换参数ax→y进行调整。
预测难度：通过使用数据增强来创建目标和源视图，增加了预测的复杂性。
容量：确定了深度的预测器模型以学习适当的变换，以此学习强大的世界模型。

作者通过多种实验测量了这些构建一个强大IWM的关键要素，从而在表示学习中实现良好的性能。

实验说明

以下是对该文中所提方法的实验性能比较的一个展示：

方法	Epochs	无预测器	固定编码器微调预测器	端到端
MAE	300	82.7	82.4	83.3
I-JEPA	300	83.0	82.0	82.0(预训练)
IWMInv (12,384)	300	83.3	82.7	83.3(预训练)
IWMEqui (18,384)	300	82.9	84.4	84.4(预训练)

实验中使用ImageNet进行评估，对比拉动MAE、I-JEPA等方法，以及对IWM的不变和等变世界模型的表现。从表中可以看出，良好的IWM在固定编码器的情况下微调预测器可以达到相当甚至超过单纯微调编码器的性能。

此外，作者还探索了多任务预测器微调，即预测器可以被微调应用于多个任务上，而且综合性能没有明显下降。这展示了世界模型的通用性和适用性。

结论

本文提出了图像世界模型IWM作为学习自监督视觉表示的新框架。考虑到世界模型条件设置、预测任务的复杂性和预测器的容量都对学习性能有显著影响，作者提出不抛弃在表示学习中学到的世界模型，而是在下游任务中进一步利用微调。实验证明该方法在多项视觉任务中表现良好，为自监督视觉表示学习提供了新的方向。