Causal-JEPA:通过对象级潜在遮蔽学习世界模型 Causal-JEPA: Learning World Models through Object-Level Latent Interventions
将JEPA从图像块扩展到对象级表示,通过遮蔽诱导因果归纳偏置,实现高效交互感知的世界建模
前置知识
联合嵌入预测架构(JEPA)
JEPA是Yann LeCun提出的一种自监督学习范式,核心思想是在表示空间而非像素空间进行预测。与重建型方法(如MAE)不同,JEPA不试图还原原始像素,而是学习预测输入数据的潜在表示。这种方法的优势在于模型专注于语义结构和预测关系,避免在像素级细节上浪费容量。I-JEPA将这一思想应用于图像,V-JEPA扩展到视频,通过预测时空tube的表示来学习视觉动态。
C-JEPA是JEPA范式在对象级表示上的首次应用,理解JEPA的核心思想(表示空间预测、避免重建)是理解本文技术创新的基础。
Slot Attention(槽注意力)
Slot Attention是一种学习对象中心表示的机制,通过竞争性注意力将特征图迭代分组到固定数量的槽(slot)中。每个槽对应场景中的一个实体,整个过程无需监督信号。对于视频,扩展版本(如SAVi、VideoSAUR)通过在帧间传递槽状态来维持时间一致性。最终得到一个集合表示 $S_t = \{s_t^1, ..., s_t^N\}$,其中每个 $s_t^i \in \mathbb{R}^d$ 对应一个对象。
C-JEPA建立在对象中心表示之上,Slot Attention是其底层编码器的核心组件,理解槽表示的特性(置换等变性、对象级抽象)对理解遮蔽策略的设计至关重要。
因果归纳偏置
因果归纳偏置是指通过模型设计或训练目标引入的先验知识,引导模型学习具有因果结构的表示,而非仅仅利用统计相关性。在世界模型中,这意味着模型应该学习对象之间真实的交互依赖关系,而不是依赖虚假的像素级相关性或对象自动态的捷径。传统的因果发现方法假设固定图结构,而本文通过遮蔽作为一种潜在干预来诱导因果偏置。
本文的核心创新就是通过对象级遮蔽作为一种潜在干预来诱导因果归纳偏置,这是论文方法论的理论基础。
模型预测控制(MPC)
MPC是一种在线规划方法,在每个时间步通过优化未来动作序列来最小化目标成本。具体来说,给定当前状态,MPC在规划范围内采样多个动作序列,使用世界模型预测每个序列的结果状态,选择成本最低的动作执行,然后在新观测到来时重新规划。常用交叉熵方法(CEM)进行动作优化:采样300个候选序列,保留30个最优作为精英,更新高斯参数,重复30次迭代。
本文的一个关键贡献是C-JEPA在MPC规划中的高效性——仅使用1.02%的潜在token就能达到与基于patch的方法可比的性能,规划速度提升8倍以上。
研究动机
当前世界模型面临一个核心矛盾:对象中心表示提供了有用的抽象,但仅靠对象表示不足以捕获交互依赖的动态。具体来说,现有研究表明,如果没有显式机制引导交互学习,模型很容易退化到依赖对象自动态或利用偶然相关性。例如,SlotFormer通过自回归滚动学习动态,但没有显式约束交互结构;C-SWM依赖固定关系图,缺乏适应性;SPARTAN通过稀疏注意力鼓励交互选择性,但不是通过学习目标本身来诱导交互结构。在CLEVRER等多对象交互场景中,这些方法在反事实推理任务上表现不佳,因为反事实问题需要理解如果某个对象不存在会发生什么,这恰恰需要建模对象间的因果依赖关系。此外,基于patch的方法(如DINO-WM)虽然性能好,但注意力机制随patch token数量二次方增长,在规划时计算开销巨大。
本文的目标是本文的目标是提出一种简单而有效的对象中心世界模型,通过对象级遮蔽作为一种潜在干预来诱导因果归纳偏置,使得交互推理成为最小化预测目标的功能性必要条件。具体可量化的目标包括:在CLEVRER视觉问答基准上提升反事实推理能力;在控制任务上实现与基于patch的方法可比的性能,同时将潜在输入特征减少到仅1%;使规划速度提升一个数量级。
与已有工作不同的是,本文的独特切入角度在于:不通过架构约束(如分解自动态和交互)或后处理(如稀疏正则化)来强制学习交互,而是通过学习目标本身来使交互推理成为必要。具体来说,遮蔽对象的历史状态创造了一种反事实式查询:当模型无法直接观察某个对象的历史时,它必须从其他对象的演化状态中推断该对象的动态,这本质上要求模型理解对象间的因果依赖。这种设计的关键洞察是:现有的patch级遮蔽方法优化的是局部patch相关性,不强制对象级交互推理;而对象级遮蔽通过控制预测器的可观察性,在学习过程中系统地引入干预,使得干预稳定的影响力邻域得以涌现。
核心方法
C-JEPA的方法可以用一个直觉来理解:想象你在观看一场台球比赛,但某个球被突然遮住了。为了预测被遮住球的运动轨迹,你不能只看它最后出现的位置(因为那已经被遮蔽了),你必须观察其他球的运动——它们是否与被遮住的球发生了碰撞?碰撞后各自如何运动?这就是C-JEPA的核心思想:通过遮蔽对象的历史状态,强制模型从其他对象的交互中推断被遮蔽对象的动态。技术路线上,C-JEPA首先使用冻结的对象中心编码器(如VideoSAUR)将视频帧编码为对象槽表示,然后在训练时随机遮蔽部分对象的历史状态(保留最早时间步作为身份锚点),使用ViT风格的双向Transformer预测器同时预测被遮蔽的历史状态和未来状态,优化遮蔽潜在预测目标。推理时不应用遮蔽,直接从完整历史预测未来。
C-JEPA的核心创新在于将对象级遮蔽作为一种潜在干预,而非简单的数据增强。与已有方法的本质区别体现在三个方面:第一,遮蔽粒度——从patch级/时空tube级提升到对象级,这使得遮蔽对应于移除对整个实体的访问,而非局部视觉区域;第二,遮蔽目标——不是为了学习重建(如MAE),而是为了创造反事实式预测问题,使得交互推理成为最小化目标的功能性必要条件;第三,因果解释——遮蔽改变的是预测器的可观察性,而非数据生成过程本身,这是一种对可观察性的潜在干预。形式上,当对象 $i$ 在时间 $t$ 被遮蔽时,遮蔽token定义为 $\tilde{z}_i^\tau = \phi(z_i^{t_0}) + e_\tau$,其中 $\phi$ 是线性投影,$z_i^{t_0}$ 是身份锚点,$e_\tau$ 是可学习的时间位置编码。这种设计确保预测器知道哪个对象被遮蔽(通过身份锚点),但不知道其当前状态(因为被遮蔽),从而必须依赖上下文来推断。
方法步骤详情
C-JEPA的完整方法包含以下步骤:(1)对象中心编码:使用冻结的编码器 $g$ 将每个观测帧 $X_t$ 映射为对象槽集合 $S_t = g(X_t) = \{s_t^1, ..., s_t^N\}$,每个槽 $s_t^i \in \mathbb{R}^d$ 对应一个实体。(2)构建实体token:将对象状态 $S_t$ 与辅助变量 $U_t$(如动作 $a_t$ 和本体感受信号 $p_t$)组合为实体token集合 $Z_t = \{S_t, U_t\}$。(3)对象级遮蔽:对历史窗口 $T$ 内的每个时间步,随机选择遮蔽索引集 $M \subset \{1,...,N\}$,将被遮蔽对象的状态替换为遮蔽token $\tilde{z}_i^\tau = \phi(z_i^{t_0}) + e_\tau$,仅保留最早时间步 $t_0$ 作为身份锚点。(4)联合预测:使用ViT风格的双向Transformer预测器 $f$ 处理遮蔽后的输入序列 $\bar{Z}_T$,同时预测历史窗口内被遮蔽的token和未来时间步的token。(5)优化目标:最小化遮蔽潜在预测损失 $\mathcal{L}_{mask} = \mathbb{E}[\sum_{\tau \in T} \sum_{i=1}^N \mathbf{1}[\bar{z}_i^\tau \neq z_i^\tau] \|\hat{z}_i^\tau - z_i^\tau\|_2^2]$,其中历史项抑制对简单自动态的依赖,未来项确保与前向世界建模对齐。(6)推理:不应用遮蔽,从完全可观察的历史 $S_T$ 前向预测未来状态。
技术新颖性
C-JEPA的技术新颖性体现在多个层面。首先,这是首次将JEPA范式与对象中心世界建模相结合,此前的JEPA工作(I-JEPA、V-JEPA、V-JEPA2)都操作在patch级表示上,而C-JEPA将其扩展到对象级,这带来了显著的效率提升——对象级表示的token数量仅为patch级的1.02%。其次,遮蔽作为潜在干预的概念是新颖的:不同于MAE等方法将遮蔽视为重建目标,也不同于因果发现中将遮蔽用于结构识别,C-JEPA将遮蔽视为对预测器可观察性的干预,创造反事实式预测问题。第三,身份锚点的设计巧妙解决了槽表示置换等变性带来的问题——由于对象维度没有位置编码,被遮蔽的对象需要最小身份信息来区分,这通过保留最早时间步的线性投影实现。最后,理论分析首次形式化了对象级遮蔽如何诱导因果归纳偏置,证明了影响力邻域的概念及其预测充分性。
实验结果
C-JEPA在视觉推理和预测控制两个互补视角上都取得了显著成果。在CLEVRER视觉问答基准上,对象级遮蔽带来了一致的性能提升:当使用VideoSAUR编码器遮蔽4个对象时,平均准确率从82.79%提升到89.40%(+6.61%),反事实推理从47.68%大幅提升到68.81%(+21.13%)。这一提升在反事实问题上尤为显著,表明遮蔽不仅增强了预测准确性,还直接强化了反事实推理能力。使用SAVi编码器时,遮蔽2个对象达到最优性能,平均准确率从77.28%提升到83.88%(+6.60%),反事实推理从41.10%提升到60.19%(+19.09%)。值得注意的是,过度遮蔽(如SAVi编码器遮蔽4个对象)会导致性能下降,表明存在最优遮蔽机制,取决于底层对象表示的鲁棒性。与基线方法的对比显示,移除重建损失会导致SlotFormer严重退化(从79.44%降至44.94%),而C-JEPA在无重建设置下达到了83.88%的最优性能。在Push-T控制任务上,C-JEPA取得了与基于patch的DINO-WM可比的成功率(88.67% vs 91.33%),但仅使用6个token×128维度,而DINO-WM使用196个token×384维度——这是1.02%的token比例。这种效率提升直接转化为规划速度:在相同设置下,C-JEPA评估50条轨迹平均需要673秒,而DINO-WM需要5,763秒,实现了8倍以上的加速。理论分析表明,对象级遮蔽诱导的影响力邻域使交互推理成为最小化预测目标的功能性必要条件,形式化证明了任何忽略影响力邻域中变量的预测器都无法达到最小可实现的期望重建误差。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| CLEVRER反事实推理(VideoSAUR编码器) | 反事实VQA准确率(每问题) | 68.81% | 47.68%(OC-JEPA无遮蔽) | +21.13%绝对提升 |
| CLEVRER平均准确率(VideoSAUR编码器) | 平均VQA准确率(每问题) | 89.40% | 82.79%(OC-JEPA无遮蔽) | +6.61%绝对提升 |
| CLEVRER反事实推理(SAVi编码器) | 反事实VQA准确率(每问题) | 60.19% | 41.10%(OC-JEPA无遮蔽) | +19.09%绝对提升 |
| Push-T控制任务 | 规划成功率 | 88.67% | 60.67%(OC-DINO-WM) | +28.00%绝对提升 |
| 规划效率 | 评估50条轨迹时间 | 673秒 | 5,763秒(DINO-WM) | 8倍加速 |
局限与改进
论文作者和我们的分析共同指出了若干局限性。首先,性能依赖于对象中心编码器的质量,这构成了性能天花板。当底层编码器的对象对齐不够好时,遮蔽的效果会减弱——这在SAVi编码器遮蔽4个对象时表现明显,性能反而下降(从77.28%降至73.28%)。其次,虽然论文形式化了影响力邻域的概念,但没有在具有显式时间因果图的数据集上直接验证,因果图恢复的实验留作未来工作。第三,论文假设了有限历史充分性(Assumption 4),但当真实系统表现出比历史窗口更长的记忆时,预测器只能学习近似,准确性可能饱和。第四,辅助变量(如动作和本体感受)的处理方式相对简单,只是作为独立实体条件化预测器,没有探索更复杂的交互建模。此外,我们的观察是:论文的理论分析假设了对象对齐的潜在表示(Assumption 3),但这在实际中很难完美满足,特别是在复杂真实场景中。最后,C-JEPA的最优遮蔽比例需要针对不同任务和编码器进行调整,缺乏自动确定机制。
独立分析的弱点
基于对论文的深入分析,我们识别出几个值得改进的弱点。第一,编码器依赖问题:C-JEPA的性能上限由冻结的对象中心编码器决定,当编码器质量不高时,遮蔽可能引入噪声而非有益的归纳偏置。改进方向是联合优化编码器和预测器,使用强预训练骨干(如DINOv2)而不导致表示崩溃。第二,遮蔽策略的固定性:当前方法在训练时随机遮蔽固定数量的对象,但最优遮蔽比例因任务和场景而异。可以探索自适应遮蔽策略,根据预测误差或对象交互强度动态调整遮蔽比例。第三,理论假设的限制:影响力邻域的定义假设了对象对齐和有限历史充分性,这些在复杂真实场景中可能不成立。需要发展更鲁棒的理论框架,处理部分对象对齐和长期依赖的情况。第四,缺乏不确定性建模:当多个对象被遮蔽时,预测可能存在多个合理结果,但当前方法使用MSE损失,无法捕获这种多模态性。引入概率预测或扩散模型可能改善这一点。第五,评估局限于合成环境:CLEVRER和Push-T都是相对受控的场景,需要在更复杂的物理环境(如PHYRE的扩展版本或真实机器人任务)中验证方法的有效性。
未来方向
论文作者提出了几个有前景的未来方向。首先,联合优化对象中心编码器:使用强预训练骨干而不导致表示崩溃,这可以通过目标编码器自举等技术实现。其次,在更复杂环境中评估C-JEPA:具有更丰富交互的真实世界场景,如多机器人协作或复杂物理模拟。第三,直接验证因果图恢复:在具有显式时间因果图的数据集上评估影响力邻域与真实因果结构的对应关系。基于论文成果,我们还可以延伸出几个方向:将C-JEPA扩展到多智能体设置,其中每个智能体维护自己的对象表示并进行通信;探索遮蔽在迁移学习中的作用,研究在源域学到的影响力邻域是否能迁移到新域;发展自监督的遮蔽比例选择方法,根据学习进度自动调整;将C-JEPA与强化学习结合,利用其高效规划能力加速策略学习;探索在自然语言推理任务中的应用,将对象级遮蔽的思想扩展到符号推理。
复现评估
论文的复现条件相对良好。代码已在GitHub开源(https://github.com/galilai-group/cjepa),基于stable-pretraining和stable-worldmodel框架构建,这些框架本身也有良好文档。数据方面,CLEVRER和Push-T都是公开可用的标准基准,无需额外数据收集。算力需求方面,论文报告所有实验在单GPU上进行,使用预提取的对象嵌入,这大大降低了计算门槛。具体来说,预测器训练30个epoch,批量大小256,在L40s GPU上评估50条轨迹需要673秒。复现难度中等:虽然论文提供了详细的超参数和训练细节(见附录D-H),但成功复现需要理解对象中心表示、JEPA框架和遮蔽策略的交互作用。特别是,预训练的对象中心编码器(VideoSAUR、SAVi)的质量对最终性能有显著影响,这可能需要额外调优。总体而言,对于有相关经验的研究者,复现应该是可行的,但需要仔细遵循论文的训练协议。
论文图表