Learning and Leveraging World Models in Visual Representation Learning

在视觉表示学习中学习和利用世界模型

摘要：联合嵌入预测架构（Joint-Embedding Predictive Architecture，JEPA）作为一种有前景的自监督方法崭露头角，它通过利用世界模型进行学习。此前，JEPA主要局限于预测输入缺失部分，而我们正在探索如何将其预测任务泛化到更广泛的一系列破坏性变化上。我们引入了图像世界模型（Image World Models，IWM）这一方法，该方法超越了掩模图像建模，能够学习在潜在空间中预测全局光度变换的效果。我们研究了学习高性能IWM的方法，并发现其依赖于三个关键方面：条件设定、预测难度和模型容量。此外，我们还展示了通过微调，由IWM学习得到的预测世界模型可以适应解决多种任务；经过微调后的IWM世界模型的表现可与之前自监督方法相媲美或超越它们。最后，我们展示了使用IWM学习能够让人们控制所学表征的抽象层次，既可以学习如对比学习方法那样的不变性表征，也可以学习如掩模图像建模那样的等变性表征。

（注释：简单来说，联合嵌入预测架构（JEPA）是一种很先进的自学方法，它就像一个聪明的学生，不需老师直接指导就能自己理解世界。最初，这个学生擅长根据给定的部分信息推测出缺失的信息。现在，科学家们想让它变得更厉害，不仅能处理简单的“填空题”，还能应对各种复杂的变化情况，比如图像颜色、亮度等全局性质的改变。

他们开发了一个新的工具——图像世界模型（IWM），这个工具比以前只能做局部图像预测的那些技术要强多了，它可以深入到图像内在的表达空间，去预测整个图像在各种光照条件下的变化效果。

为了更好地训练这种强大的IWM，研究人员发现需要关注三个重要因素：一是设置好学习的前提条件，二是让预测任务具有合适的挑战度，三是确保模型有足够的能力去承载复杂的知识。

而且，一旦IWM被有效训练后，就像一位多面手，经过针对性的微调，可以解决很多不同的问题，并且表现得甚至比其他同类自学方法还要出色。

最后，IWM的学习过程还有一个很棒的特点，就是能让使用者调控学习出来的图像特征的抽象级别，也就是说，既可以学会忽略图像表面差异抓住本质不变的特征（就像对比学习那样），也能学会保持某些特性不变，即使图像部分区域被遮挡或改变了也能识别出来（就像掩模图像建模那样）。）

1 引言

在强化学习（RL）领域，学习并利用世界模型是一种常见的实践，近年来尤其在Ha和Schmidhuber（2018）以及Hafner等人（2019, 2023）的工作中取得了显著成效。通常情况下，世界模型是通过训练网络来预测动作后果实现的，预测可以在输入空间中进行（如Yang等人，2023年），也可以在潜在空间中进行（如Hu等人，2023年；Hafner等人，2023年）。

鉴于对世界建模的广泛理解，我们试图探索在视觉表示学习中学习和应用世界模型是否也能带来益处。大量自监督学习方法基于编码器-预测器架构，其中编码器-预测器网络被训练来预测数据的各种变换，例如掩码图像建模（Bao等人，2021；He等人，2021）、联合嵌入架构（Grill等人，2020；Xie等人，2022；Assran等人，2023；Baevski等人，2022）或是等变预测目标（Gupta等人，2023；Garrido等人，2023b）。如果我们把数据的变换视为“动作”，那么我们可以很容易地将自监督学习方法与强化学习中的世界建模联系起来；参见图2。

 

图1 显示了通过学习得到的图像世界模型在潜在空间中进行预测的可视化结果。我们首先在潜在空间中对源图像施加某种操作，然后在包含256张图片的库中找到与预测表示最接近的邻居图像。结果显示，图像世界模型(IWM)能够模拟图像变换并对失真进行修复，显示出对底层图像变换机理的理解。该图来源于：ai.meta.com/blog/yann-lecun-advances-in-ai-research/。

举例来说，在掩码自编码器（He等人，2021）中的解码器网络可以被视为一种生成式图像世界模型，它学习推断“掩码动作”T(a)对图像y的影响；在这个例子中，变换参数a（图像中被遮蔽区域的位置）也被馈送到解码器网络中。基于联合嵌入预测架构（Joint-Embedding Predictive Architectures, JEPAs）的方法，如I-JEPA（Assran等人，2023）或data2vec（Baevski等人，2022），则以类似方式运作，但被视为学习一种潜在图像世界模型，它学习推断掩码动作对图像表示的影响。

如果预测器不根据变换参数进行条件化，则我们所能期望的最佳结果是学习对数据变换不变的表示，正如BYOL（Grill等人，2020）和SimSiam（Chen和He，2020）中所做的那样，其中图像变换对应于各种光度和几何数据增强。

尽管强化学习中的世界建模与从图像进行自监督学习之间存在一些表面上的相似之处，但在强化学习中习得的世界模型通常在下游任务中发挥作用，如用于规划（Hansen等人，2022）。相比之下，在自监督学习中习得的世界模型通常在预训练后被丢弃，因为焦点往往集中在学习到的编码器网络的表示质量上。这是因为计算机视觉中的大多数下游任务与世界建模任务无关。通常感兴趣的任务侧重于区分性方面，因此即使预测器学习到了有用信息，也常常会被丢弃。我们认为在表示学习中丢弃世界模型是浪费的，就像在RL中一样，我们也可以重用这个世界模型来服务于下游任务。这就促使我们深入研究学习世界模型作为表示学习的一种范式。因此，我们引入了图像世界模型（IWM，如图2右侧所示），作为一种既能学到好的表示又能学到强大且可重用的世界模型的方式。IWM基于JEPA，并将常规的潜在空间填充扩展至包括光度变换，使我们能够展示学习一个有能力的世界模型的关键方面，这些方面包括预测器条件化的选择、变换的强度以及世界模型的容量。

接下来，我们专注于如何利用已学习的世界模型执行下游任务，并发现可以通过微调来利用这个世界模型。具体来说，我们发现在冻结编码器的基础上仅微调世界模型以适应下游任务，相比直接微调编码器可以获得性能提升，同时还能以较少的成本和微调参数的数量达到这一目标。更重要的是，只有通过IWM学习的世界模型表现出这种性能提升行为；对具有同样架构的随机初始化网络进行微调并不能提供这样的性能改善。这表明世界模型应当成为推理过程中的核心部分，而非被丢弃。受到指令微调（Wei等人，2022；Zhang等人，2023）的启发，我们进一步展示世界模型可以通过微调一次性解决多个任务，从而进一步提高效率。

我们的研究表明，在使用世界模型进行表示学习时的另一个关键方面是：赋予世界模型的容量直接影响所学表示的抽象程度。直观地说，如果预测器是恒等映射（即没有预测器，如图2中间所示），网络将会捕获高层次的语义信息，因为它只会学习编码输入y及其变换x之间的共享内容。这是对比学习背后的表现质量驱动力所在，其中选取的变换只保留图像的语义信息。相反，当预测器具有更大的容量并且能有效地反转变换的效果时，编码器的输出将能保留更多关于其输入的信息。这两种思想是等变表示学习的核心；一个能够有效应用变换的预测器是等变的，而不能做到这一点的预测器则是不变的。我们发现，对变换保持不变的世界模型在线性评估中表现更好，而等变的世界模型则与更好的世界模型微调效果相关联。这就形成了适应性和原始性能之间的权衡。因此，通过学习世界模型来学习表示为我们提供了表示性质上的灵活性，使得这种方法成为一个极具吸引力的表示学习框架。

我们的贡献可以总结如下：

？ 我们展示了如何利用联合嵌入预测架构（JEPAs）来学习图像世界模型（IWM）。关键要素包括：变换的复杂性、对变换的条件化及预测器的容量。

？ 我们证明了等变世界模型可以应用于区分性任务。微调预测器相对于编码器微调能以较低的成本获得更好的性能。受指令微调启发，我们还展示了它能够一次性针对多个任务进行微调。

？ 我们展示了控制世界模型能力可以使我们得到具有不同属性的表示。一个对变换保持不变的世界模型提供了更抽象的表示，在类似于对比学习的情况下，在线性评估中表现更优。而一个等变的世界模型能保存更多的输入信息，通过预测器微调能得到更好的峰值性能。

 

图2展示了多个具有相关架构的方法家族，其中世界模型是否进行条件化是一个关键的区别特征。生成式世界模型通过利用自动编码器框架被训练来在输入空间中反转变换。这种方式可以实例化世界建模和表示学习方法。联合嵌入方法舍弃了世界模型，但在潜在空间中通过编码变换输入之间的共性来进行操作，它是SSL（自监督学习）方法的主要类别。JEPA（Joint Embedding Predictive Architecture）世界模型可以看作是一个更通用的框架，在该框架下世界模型在潜在空间中进行训练。这一家族方法在强化学习和表示学习领域都非常成功，而图像世界模型（IWM）就属于这一范畴。

（图片中描述了三种主要的方法类别：生成式世界模型、联合嵌入方法和 JEPA（联合嵌入预测架构）世界模型。下面是对这些概念的通俗易懂的解释：

1. **生成式世界模型**：

- 这些模型就像是能够“想象”和“创造”新事物的艺术家。它们通过学习输入数据（比如图片）的变化规律，来生成新的数据样本。想象一下，如果你给它们看很多猫的图片，它们就能创造出新的猫的图片，即使这些猫在现实中并不存在。

- 这些模型通常使用一种叫做“自编码器”的结构来训练，自编码器能够学习如何将数据从一个形式转换成另一个形式，然后再变回来。

2. **联合嵌入方法**：

- 这类方法有点像是寻找数据之间共同点的侦探。它们不直接处理原始数据，而是在一个叫做“潜在空间”的地方工作。在潜在空间中，它们找到不同数据样本之间的共同特征，并将这些共同特征编码下来。

- 这些方法通常用于无监督学习（SSL），这是一种不依赖于标签数据的学习方式。它们通过比较不同数据样本的相似性来学习。

3. **JEPA 世界模型**：

- JEPA 世界模型可以看作是一种更通用的框架，它在潜在空间中训练一个世界模型。这种模型不仅能够理解数据的变化，还能够预测未来的变化。

- 这类方法在强化学习和表示学习中非常成功。强化学习是一种让机器通过与环境互动来学习的方法，而表示学习则是教机器如何更好地理解和表示数据。

- 图片中提到的“图像世界模型（IWM）”就是 JEPA 世界模型的一个例子。IWM 通过学习图像的潜在表示，来理解和预测图像可能的变化。

图片三种不同的方法，它们都在尝试理解和表示数据，但是从不同的角度和使用不同的技术。生成式世界模型关注于创造新数据，联合嵌入方法关注于找到数据之间的共同点，而 JEPA 世界模型则关注于预测数据的变化。这些方法在机器学习和人工智能的不同领域中都有广泛的应用。）

2 相关工作

2.1 增强不变自监督学习

对比方法的核心在于增强不变性。一张图像的多种增强视图应在潜在空间中产生相同的表示。这类方法的关键就在于如何避免这些表示发生坍缩。样本对比方法（Chen等人，2020a; He等人，2020; Chen等人，2020b; Caron等人，2021; Chen等人，2021; Yeh等人，2021; HaoChen等人，2021; Oquab等人，2023）通过将来自其他数据点的表示推开，从而避免坍缩现象。维度对比方法（Bardes等人，2021; Zbontar等人，2021; Ermolov等人，2021; Li等人，2022; Bardes等人，2022）从整体上考虑表示，并鼓励最大化信息内容，以此避免坍缩。研究显示，维度对比和样本对比两种方法都能导致非常相似的表示（Garrido等人，2023a）。基于预测的方法（Grill等人，2020; Chen和He，2020）通过预测增强后的表示来学习，但由于未对变换进行条件化处理，它们也导致了不变的表示。

2.2 视觉表示学习中的世界建模

尽管在强化学习（Hafner等人，2019, 2023）或视频预测（Yang等人，2023; Hu等人，2023）领域，世界建模是一种成功的范式，但在表示学习方面的明显优势尚未显现。然而，从这一角度重新审视，有多类方法可以关联起来。等变自监督学习方法（Devillers和Lefort, 2022; Park等人，2022; Garrido等人，2023b; Gupta等人，2023; Dangovski等人，2021）旨在当数据变换构成群组时，预测数据的变换。掩码图像建模（He等人，2021; Bao等人，2021; El-Nouby等人，2024; Xie等人，2022）通过预测图像被掩蔽部分来学习表示。虽然这些方法是在像素空间中进行预测，但其解码器可以被视为世界模型的具体实现。同样地，JEPAs（Assran等人，2023; Baevski等人，2022）在潜在空间中预测图像的掩蔽部分。最近，生成方法也被应用于表示学习（Hudson等人，2023; Clark和Jaini，2023; Chen等人，2024），尽管这些方法看似有潜力，但其性能仍然低于对比学习或MIM（Masked Image Modeling）方法。近期的研究也显示生成质量与表示质量之间存在负相关性（Chen等人，2024）。上述工作的共同之处在于，世界模型（预测器或解码器）要么在评估阶段被丢弃，要么仅用于增强数据（如Hudson等人，2023）。我们提议超越这些实践，展示我们可以通过学习一个既能用于下游任务复用又能同时学习高质量表示的世界模型。

3.方法

现在我们将介绍图像世界模型（IWM）。它遵循联合嵌入预测架构（Joint Embedding Predictive Architecture，简称JEPA）框架（LeCun，2022年），类似于I-JEPA（Assran等人，2023年）。在这个框架中，预测器是世界模型的具体体现。我们认为一个有能力的世界模型能够在潜在空间中应用变换，从而学习到等变表示，因此我们将这样的世界模型称为等变型1，而表现较差的世界模型则称为不变型。

使用JEPA方法的一个吸引之处在于，那些采用对比学习方法来学习等变表示的方法往往需要依赖不变损失来提高表示质量，无论是明确地（如Gupta等人，2023年；Garrido等人，2023b年），还是隐含地（如Chavhan等人，2023a年）。相反，JEPA风格的方法则没有这个缺点，因为它通过潜在空间的修复学习了表示的语义方面。在潜在空间工作进一步允许网络去除不必要的信息或难以预测的信息。这使得JEPA公式更具吸引力，因为在重构方法中，重建质量并不一定与表示质量成正比（Chen等人，2024年）。

为了训练IWM，第一步是从图像I生成源视图和目标视图——分别记为图2中的x和y。

目标视图y：目标视图通过随机水平翻转、裁剪以及颜色抖动（亮度、对比度、饱和度、色调）应用于原始图像I上生成。为了避免破坏性增强，不对目标视图应用诸如灰度化之类的增强手段，以确保目标视图尽可能包含更多信息。我们在附录C中对此选择进行了进一步阐述。

（注释：比如你有一张原始的彩色照片（图片I）。现在，如果你想要用这张照片来训练一个计算机视觉模型，比如让模型学会识别图片中的对象，你可能希望模型能够在不同的条件下也能识别出对象，比如图片被翻转或者光线变化时。

为了做到这一点，你需要创建一些经过修改的图片，这些图片被称为“目标视图”（目标视图y）。你会对原始图片进行一系列的改变，比如：

- **水平翻转**：把图片左右翻转，就像在镜子里看一样。

- **裁剪**：从图片中剪掉一部分，创建一个新的图片，只包含原始图片的一部分。

- **颜色抖动**：改变图片的亮度、对比度、饱和度和色调，让图片看起来颜色有些不同，就像在不同的天气或时间拍摄的。

这些改变的目的是让模型在训练时看到各种各样的图片变化，这样它就能更好地适应真实世界中的不同情况。

然而，有些改变可能会“破坏”图片中的信息，比如把图片变成黑白的（灰度化）。这样的改变可能会让模型丢失一些重要的信息，比如颜色。为了避免这种情况，我们不会对目标视图使用这样的增强手段。我们希望目标视图能够尽可能地保留原始图片中的所有信息，这样模型就能学到更多有用的内容。）

源视图x：对于源视图，我们从目标视图y进一步转换开始。首先再次应用颜色抖动，以及破坏性增强：灰度化、模糊和太阳能化。这一组增强与对比自监督学习（SSL）中使用的相同。最后，我们还按照I-JEPA的方式对图像的部分区域进行遮罩。我们定义我们的遮罩M_x（一组索引）为4个矩形遮罩的并集。关于具体实现细节，请参阅附录A。

（注释：假如你有一个美丽的彩色风景照片（这就是我们的目标视图y）。你想用这张照片来教一个计算机程序如何更好地理解和处理图片。为了做到这一点，你需要创建一些经过特殊处理的图片，这些图片我们称之为“源视图”（图片x）。

首先，我们对目标视图y再做一些改变，让它更有趣、更具挑战性：

1. **颜色抖动**：就像之前对目标视图做的那样，我们再次改变图片的亮度、对比度、饱和度和色调，让颜色看起来有点不同。

2. **破坏性增强**：这次我们做一些可能会让图片“失去”一些信息的改变，比如：

- **灰度化**：把彩色图片变成黑白的，这样图片就只有灰度信息了。

- **模糊**：让图片变得模糊，就像透过脏窗户看东西一样。

- **太阳能化**：让图片看起来像是被太阳晒过，颜色变得非常明亮和饱和。

这些改变的目的是让计算机程序学会在面对各种不同质量和条件下的图片时，仍然能够识别和处理图片内容。

最后，我们对图片x做一些遮罩处理。遮罩就像是在图片上贴了一些透明的贴纸，遮盖住图片的一部分。我们用4个矩形遮罩来遮盖图片的不同区域。这样做是为了让计算机程序学会从被遮盖的区域推断出原本的内容，这是一种叫做“对比自监督学习”（SSL）的学习方法。）

动作a：我们用a_x→y表示由x变换到y的变换参数，即初始变换过程的逆变换。a_x→y包含了x和y之间颜色抖动差异的信息，以及是否应用了每一种破坏性增强的信息。

（注释：假如有两个版本的同一本彩色图书的封面照片：一个是原始的（我们称之为x），另一个是经过一系列改变的（我们称之为y）。这些改变可能包括调整颜色（比如让图片更亮或更暗）、把图片变成黑白的、或者让图片看起来模糊等。

现在，如果我们想要一个计算机程序来理解这两张照片之间的关系，我们需要告诉它这些改变是如何发生的。这就是“动作a”的概念。

“动作a”就像是一份说明书，告诉我们如何从原始照片（x）变成改变后的照片（y）。这份说明书包括了所有的改变细节，比如颜色是如何变化的，以及是否对图片进行了任何特殊的处理，比如变成黑白或者模糊处理。

用技术术语来说，我们把这个过程叫做“变换参数”，它就像是一组指令，记录了从x到y的所有操作步骤的逆过程。这样，如果计算机程序知道了这些变换参数，它就可以理解如何从改变后的照片（y）恢复到原始照片（x）的样子。

在机器学习中，我们经常用这样的方法来教计算机程序学习如何从一些输入数据（比如x）生成另一些输出数据（比如y），以及如何反向操作。这对于计算机视觉、图像处理和很多其他领域都是非常重要的。）

使用世界模型p_？进行建模：源视图和目标视图分别通过编码器f_θ及其指数移动平均版本f_θ^EMA处理，得到表示z_x = f_θ(x)和z_y = f_θ^EMA(y)。使用EMA网络至关重要，以避免解决方案坍塌。为了对作为世界模型的预测器进行条件化，我们将其输入几何信息，形式为遮罩标记和a_x→y。这些遮罩标记记为m_a，对应于M_x^C中的位置。预测器p_？随后接收嵌入后的源图像块x_c、变换参数a_x→y以及遮罩标记m_a作为输入。其目标是使p_？(z_x, a_x→y, m_a) = z？_y匹配z_y。

（注释：假如你有两个版本的同一个地方的照片：一张是原始的（我们称之为源视图x），另一张是经过一些改变的（我们称之为目标视图y）。现在，你想创建一个“世界模型”（我们用p_？表示），这个模型能够理解这两张照片之间的关系，并能够预测如果对原始照片进行同样的改变，它会是什么样子。

为了做到这一点，你首先使用一个特殊的相机（我们称之为编码器f_θ）来拍摄这两张照片。这个相机不仅拍下了照片，还记录了一些关于照片的信息。然后，你还使用了这个相机的一个特别版本（我们称之为指数移动平均版本f_θ^EMA），它能够平滑地记录照片随时间的变化。这样，你就得到了两组信息：一组是关于原始照片的（z_x），另一组是关于改变后照片的（z_y）。

但是，为了让这个“世界模型”能够有效地工作，你需要给它一些额外的信息，这就是“遮罩标记”（m_a）和“变换参数”（a_x→y）。遮罩标记就像是你给照片上的某些部分贴上的标签，告诉你哪些部分被改变了。变换参数则是一组指令，告诉你如何将原始照片变成目标视图。

现在，你的“世界模型”p_？需要做两件事：

1. 它要查看原始照片的信息（z_x）和遮罩标记（m_a）。

2. 它要使用变换参数（a_x→y）来理解如何从原始照片（x）变成目标视图（y）。

最终目标是让“世界模型”p_？能够准确预测出改变后的照片（我们称之为z？_y）看起来会是什么样子，这样它就能够和实际的目标视图（z_y）相匹配。

这个过程就像是你在教一个机器人如何通过看两张照片来理解发生了什么变化，并能够预测如果对其他照片做同样的改变会是什么样子。这个“世界模型”p_？就是你的机器人，它通过学习和理解照片之间的变化来做出预测。）

损失函数：所使用的损失函数是预测值z？_y与其目标z_y之间L2距离的平方：

 

（注释：这个公式用来评估我们的“世界模型”（预测器）做出的预测和实际结果之间的差距。

- p_？：这是我们的预测器（世界模型）根据源视图的信息、变换参数和遮罩标记做出的预测。

- f_θ^EMA：这是我们的目标视图经过编码器的指数移动平均（EMA）版本处理后得到的表示。

- ||...||：这是绝对值的两倍，用来计算预测值和实际值之间的距离。在机器学习中，我们通常使用 L2 距离（也称为欧几里得距离）来衡量两点之间的直线距离。

L(x, y)就是预测值和实际值之间差的平方和。我们的目标是最小化这个损失函数，这样我们的预测就会尽可能接近实际值。在训练过程中，通过调整模型的参数，我们不断减少这个损失，使得模型的预测越来越准确。）

3.1 架构与命名约定

我们的编码器是一个视觉Transformer（Vision Transformer，ViT）（Dosovitskiy等人，2021年），具体来说，我们使用的是ViT-B/16架构。预测器基于相同的架构，但深度和嵌入维度不同。我们将IWM的不同实例表示为IWM^Z_X,Y，其中X是预测器的深度，Y是其嵌入维度，Z取决于世界模型的能力，要么是Inv表示不变型，要么是Equi表示等变型。例如，IWM^Equi_18,384意味着预测器有18层深，具有384维嵌入，并展现出等变行为，即已经学会了灵活多样的世界模型。

4 学习用于表征学习的图像世界模型

4.1 评估世界模型的质量

正如之前讨论的那样，学习等变表示和学习世界模型是紧密相关的问题。因此，我们可以借鉴等变性文献中的指标来评估训练好的世界模型的质量。我们依赖于均值倒数排名（Mean Reciprocal Rank, MRR）（Kipf等人，2019年）作为主要评估指标。计算MRR的过程如下：首先，我们生成一组增强的目标图像库（实践中为256张）。然后，我们把干净图像的表示通过预测器，目标是预测目标图像。接下来，我们计算预测结果与增强表示库之间的距离，由此获得目标在近邻图（NN-graph）中的排名。将多个图像和变换下的倒数排名平均后，我们就得到了MRR，它反映了世界模型的质量。若MRR接近1，则表明世界模型能够成功应用相应变换；相反，若MRR接近0，则意味着世界模型无法完成变换任务。

4.2 学习强大的图像世界模型

为了构建高性能的IWM，我们关注三个关键方面：对预测器进行变换（或动作）条件化、控制变换的复杂性以及控制预测器的容量。我们展示出如果忽视这三个方面的任何一点，都会导致得到不变表示。

**世界模型条件化**。我们研究了两种方法来对预测器进行变换信息的条件化。

**序列条件化**。一种方法是在预测器输入中直接添加代表变换的令牌。尽管这种方法看似直观，但在实施过程中需要打破变压器预测器的置换等变性。为此，每个令牌都要经过一个独特的线性层，使得网络能够将信息转换为预测器能够区分的形式。

**特征条件化**。另一种选择是在变换和遮罩令牌之间混合信息，即将条件化信息作为额外维度加入，然后通过1x1卷积神经网络将遮罩令牌馈送，以便混合遮罩令牌中的信息并映射回正确的维度。

如表1所示，未经条件化的世界模型无法应用变换，而采用序列或特征轴条件化的两种方法都能产生良好的世界模型。在实践中，由于特征条件化能带来更高的下游性能，我们选择使用这种条件化方式。

表1 预测器条件化对世界模型质量的影响。序列和特征条件化都能导致良好的世界模型。灰色是我们的默认设置。

 

**变换复杂性**。我们依赖于对比学习中常用的增强策略，包括颜色抖动（亮度、色相、对比度、饱和度）、灰度化、模糊和太阳能化，后三者因会移除信息被称为破坏性增强。除了要考虑模型化的变换集合之外，变换的强度也必须适当，以便学习到有用的世界模型。如果预测任务过于简单，预测器将无法学到有用的东西。如表2所示，增强越强，就越容易学习到强大的世界模型。在附录C中，我们提供了更多关于增强策略的详细消融实验，在更广泛的增强场景下可以看到这一趋势持续存在。

表2 预测器架构和变换对平均排名倒数（MRR）的影响。学习一个有效的世界模型需要复杂的变换和足够的预测器容量。灰色代表我们的默认设置。红色和绿色分别表示不变性和等变性行为。

 

**世界模型容量**。当变换复杂时，预测器需要更大的容量才能成功应用变换，这就说明容量是学习图像世界模型的关键因素。如表2所示，更深的预测器允许我们在更广泛的增强范围内学习到强大的世界模型，且这是IWM成功的关键所在。在附录C中，我们更深入地研究了深度对实现良好世界模型的影响。对于12层的预测器，颜色抖动等变性只在5次试验中有1次达到，而对于18层预测器，有4次试验中成功达到。因此，预测器的容量是构建强大世界模型的重要组成部分。

4.3 可视化预测结果。

与计算MRR类似，我们可以将预测出的表示与一组变换后的图像库进行比较，并查看与预测结果最近邻的图像。如图1所示，通过IWM学习到的世界模型能够在潜在空间正确地应用变换。然而，当我们反转灰度化操作时，可以观察到一些不准确之处，因为灰度化并不是完全可逆的。这些可视化结果进一步证实了IWM能够为图像变换学习到强大的世界模型。更多可视化内容请参阅附录I。

5 利用世界模型进行下游任务

在图像上学习的世界模型有一个局限性，即它们所解决的任务与大多数下游计算机视觉应用所需的任务并不对齐。我们已展示IWM能够应用颜色抖动或着色，但这并非此类模型驱动应用的核心任务。相比之下，LLMs（大型语言模型）的下一个词预测正是此类模型的主要应用之一。因此，我们研究如何在视觉领域利用世界模型解决超越简单变换的应用问题，重点关注如图像分类和图像分割等判别任务。

5.1 预测器微调

对于任何任务，评估头都需要理解所学习的潜在空间，并利用它来解决手头的问题。我们的学习到的预测器能做到这一点，这表明它已经学习到了不一定存在于编码器中的有用信息。然而，由于预测器被训练去预测另一个有效的表示，如果直接使用其输出，并不一定能带来更好的下游任务性能。因此，需要对预测器进行微调以解决判别任务。我们重点研究了与He等人（2021）提出的微调协议的比较。所有研究方法都在ImageNet数据集上预先训练和评估，并使用ViT-B/16作为编码器。

预测任务设置。在微调预测器时，仍然需要为其设定一个预测任务。在表3中，我们研究了定义预测任务的各种方式及其对性能的影响。首先注意到的一点是，使用教师网络相比于学生网络可以提高性能。是否使用随机变换并不是一个重要因素，最重要的是预测另一幅完整的图像。这样做的好处是评价更为灵活，因为我们不必重用预训练目标作为评估标准。使用CLS标记聚合信息而非预测完整图像也是一种有效策略，尽管这样做会使性能降低约半个百分点。这种技术的优势在于成本更低（N+1个标记对比2N个标记），根据应用场景的不同，它可以成为一个很好的替代方案。总体而言，最简单的策略是最好的：预测一幅未变形的完整图像。这使得微调协议易于复用，因为它不依赖于预训练任务。我们在附录D中提供了更多详细的消融实验结果。

表3 如何进行预测器微调的预测。使用教师模型可以提高性能，而具体的预测任务并不关键。空潜在表示更加灵活且表现更好。为了提高效率，不需要完整的预测，但这会导致性能的小幅下降。灰色是我们的默认设置。

 

总体结果。在表4中，我们比较了预测器微调、编码器微调以及编码器和预测器端到端微调的性能，所有方法都采用了ViT-B/16作为编码器。结果显示，IWM在保持或改进了I-JEPA性能的同时，不变性行为在编码器微调中表现更好。有趣的是，等变IWM的预测器微调能够匹敌不变模型编码器微调的性能。这表明该协议具有竞争力，因为它在推理时交换了参数量，以换取更友好的计算适应性。尽管这种评估增加了推理时使用的参数数量，但它仍然节省了通过主干网络的前向传播，这是完全微调所做不到的。因此，一旦考虑到多个任务，使用微调过的预测器将提供比常规微调更高的吞吐率。

表4 在ImageNet-1k上的微调评估。我们通过微调它们的编码器、保持编码器冻结并微调它们的预测性世界模型，或者同时微调两者来评估基于预测的方法。当IWM展示出等变性行为时，微调世界模型非常有效。这种行为在其他方法中不存在或不太明显，表明了强大的世界模型的重要性。

 

当比较使用随机初始化的预测器（即，大型评估头）与预训练预测器时，我们发现在MAE上几乎没有增益。这表明MAE学习的世界模型在分类任务上并不比随机初始化的网络更有优势。对于具有不变世界模型的I-JEPA和IWM，性能提升不到1个百分点，表明世界模型不够强大，不足以充分利用。然而，当我们看到带有等变世界模型的IWM时，相比于随机预测器，性能提高了1.8个百分点。这表明预测器已经学习到了有用的属性和信息，这些信息为编码器所学内容带来了额外的收益。

通过端到端微调编码器和预测器，性能可以进一步提升，IWM能够超过所有其他微调协议。这使得我们能够从单次预训练中获取更多的性能，因为世界模型始终处于训练状态。我们推测，大部分方法在端到端微调上缺乏性能的原因可能源于训练网络一部分（编码器）的同时从零开始训练另一部分（预测器）所带来的优化复杂性。在表5中，当汇总所有协议的性能时，利用我们的IWM在冻结编码器的情况下可以获得最佳性能，即充分利用预训练的所有部分。详情见附录A。

表5 单次预训练实现的最高性能。我们比较了在冻结编码器的情况下，或者允许任何评估头部使用任何协议、是否微调，以及在编码器顶部有无预测器的情况下，ImageNet顶部1准确率。

 

表6 在ADE20k上的分割微调。与图像分类类似，我们观察到预测器微调提高了性能，并且超过了编码器微调。

 

图像分割任务。在表6中，我们研究了I-JEPA和IWM在ADE20k数据集上进行图像分割任务的性能。我们观察到与图像分类相似的趋势，其中不变模型产生了最佳的编码器。然而，使用等变模型微调预测器相比编码器微调取得了显著提升，大幅度超越后者。再次验证了我们的IWM具备应用于广泛任务的潜力。附加详细性能分析可在附录A.2中找到。

效率分析。在图3中，我们研究了预测器微调与编码器微调的效率。我们发现，当参数量可比并在多种预测器大小下，IWM的预测器微调相比于MAE提升了约1个百分点的性能，相较于IWM本身提升了1.5个百分点。这意味着预测器微调不仅在性能上具有竞争力，而且在适应效率上也表现出色。我们还在章节E中进一步研究了使用ViT-L/16的IWM的行为。当比较ViT-B的端到端微调与ViT-L的编码器微调时，我们观察到性能有所提升（84.4% vs 84.3%），而参数量仅为后者的几分之一（121M vs 307M）。这进一步证明了利用IWM学习的世界模型是多么高效，重复利用预训练的所有部分可以像扩大编码器规模一样有效。

 

图3 微调效率。当考虑到微调参数的数量时，预测器的微调比编码器的微调要高效得多。

5.2 多任务预测器调优

先前我们讨论了与编码器微调相比的效率提升，但实际上还可以更进一步提高效率。表示学习的主要目标之一是获取可用于多种任务的表示。正如预测器被训练去解决多种任务（着色、修复、改变颜色）一样，我们表明它可以针对多个任务进行微调，灵感来源于LLMs（大规模语言模型）中的前缀调优（Prefix Tuning，Li和Liang，2021年）和指令调优（Instruction Tuning，Wei等人，2022年；Zhang等人，2023年）。

基本思想是，我们通过补充图S2形象地展示，向预测器提供新的学习令牌，以指示它正在尝试解决哪个任务。这一做法让人联想到DyTox（Douillard等人，2022年）中用于连续学习的任务令牌。对于每个任务，我们都设计了一个任务令牌，以及任务特定的头部组件和/或损失函数。然后将所有任务损失组合在一起，同时更新预测器和任务特定的头部组件。我们研究了一个简单的场景，其中批次数据均匀分配给各个任务，需要注意的是，其他采样策略可能会进一步提高性能。

我们在表7中评估了预先在ImageNet上训练好的IWMEqui 18,384模型在ImageNet、iNaturalist18（Horn等人，2018年）、SUN397（Xiao等人，2010年）和Places205（Zhou等人，2014年）上的性能。对于每个任务，我们都训练了一个单任务基线，其中总迭代次数与多任务训练相同。因此，训练所有四个单任务基线的成本与多任务训练完全相同，尽管这会导致四个不同的模型而不是一个模型。多任务预测器能够实现与单任务预测器相似的性能，在大多数任务上略有下降，但在SUN397任务上实现了显著性能提升。平均来看，它达到了与单任务预测器相当的性能水平。这进一步证明了利用优秀世界模型带来的效率提升，现在参数可以在所有任务间共享，使得预测器在任何任务上的推理时间都能实现轻量级的微调。

表7 多任务微调。同时在多个任务上对预测器进行微调的表现与分别在每个任务上进行微调相似。这样可以使用单一的预测头部来处理多个任务，分摊其成本。

综上所述，当学习到一个好的世界模型时，可以通过微调它来重用于下游任务。这能在较少成本下达到与编码器微调相媲美的性能。通过进行多任务微调，可以使效率更加高效，突出了这种方法的多功能性。

6 图像世界模型支持灵活的表示形式

为了完善我们对IWM在表示学习中的分析，我们研究了它在自我监督学习中常用的一些轻量级评估协议上的表现。我们专注于线性评估（Chen等人，2021年）和注意力探查（Chen等人，2023年）。

如表8所示，当IWM学习到一个不变的世界模型时，其性能表现与对比学习方法（如MoCov3）类似，在线性评估中相比于MIM或其他基于JEPA的方法有显著性能提升。同样，当IWM学习到一个等变的世界模型时，其行为类似于MAE等MIM方法，在线性评估中性能稍低，但在注意力探查中表现更具竞争力。

表8 在ImageNet-1k上的线性和注意力探测性能。IWMInv（不变性世界模型）的表现与对比学习方法相似，而IWMEqui（等变性世界模型）则与掩蔽建模方法相似。

 

这表明不同方法之间的显著区别不一定是表示质量，而是它们的抽象层次，即从它们提取信息的难易程度。线性探查是最简单的评估之一，注意力探查略为复杂，而微调则是更为复杂的协议。

在图4中，我们可以清楚地看到最适合的评估协议与世界模型的等变性之间的关联。更具不变性的世界模型在线性评估中表现出色，而更具等变性的世界模型在更大规模的评估头如预测器微调中表现出色。此外，我们注意到源自等变世界模型的丰富表示在OOD（Out-of-Distribution）数据集上表现出更好的性能（见附录F）。

 

图4 虽然等变性的水平影响线性和预测器微调设置中的性能，但它与注意力探测几乎没有相关性。这表明在表示的抽象层次方面存在权衡，不同的评估协议评估的是不同的属性。

这样，我们就可以将一系列方法放置在一个表示抽象性的谱系图中（见图5）。对比学习方法位于谱系的高抽象一端，其中信息很容易通过简单的协议提取出来，但如果忽略适应成本，其峰值性能较低，如表5所示。相反一端则是Masked Image Modeling，虽然在诸如微调等复杂评估中表现出更强的性能，但由于信息不易访问，在线性探查中则表现欠佳。通过改变世界模型的等变性，IWM能够在对比学习方法和MIM之间的谱系中占据位置，如图4和表8所示，其中IWMInv 12,384 和 IWMEqui 18,384 分别代表了IWM谱系的两个极端。

 

图5 图像世界模型允许表示的模块化。不同家族的方法提供了具有不同属性的表示，但图像世界模型（IWM）允许探索整个范围。

这一谱系可以用SSL（Self-Supervised Learning）理念“学习可预测的内容”来概括。通过弱世界模型学习意味着模型不能很好地模拟世界，编码器会移除那些不可预测的信息。另一方面，如果世界模型非常强大，那么表示就不必过于抽象或语义化，因为它可以在任何情况下找到预测表示的方法。这意味着学习世界模型提供了一种可衡量的方式来控制表示的抽象层级。

7 结论与未来展望

我们引入了IWM，这是一种借助世界模型学习自我监督视觉表示的方法。通过深入研究，我们为学习优秀的图像世界模型提供了指导原则和关键组件。对世界模型施加图像变换条件是避免退化为经典SSL行为的关键。使用强大的变换也是确保世界模型学会模拟更复杂行为并变得有用的关键。最后，足够的容量是模拟复杂行为所必需的。我们展示了只有能力强大的世界模型才能用于判别任务。这引导我们提出了预测器微调协议，它以较小的成本与编码器微调相媲美，表明世界模型是多用途的评估头。我们进一步将其改编为一次解决多个任务而不丧失性能。最后，我们研究了学习世界模型如何影响表示质量。一个能力强的世界模型能学习丰富的表示，从而改善下游任务（如图像分类和语义分割）的性能。另外，学习不变世界模型有助于在进行线性评估时获得更好的表示。虽然MIM和对比学习方法在表示抽象性上是一条谱系的两端，但图像世界模型让我们得以在这两者之间插值。因此，我们相信学习图像世界模型是视觉表示学习领域极具前景的框架。

8 更广泛影响声明

本论文提出的工作旨在推进机器学习领域的进步。我们的工作有许多潜在的社会影响，但我们认为在此无需特别强调任何一项。