seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models

📄 arXiv: 2505.03176v3 📥 PDF

作者: Hafez Ghaemi, Eilif Muller, Shahab Bakhtiari

分类: cs.CV, cs.AI, cs.LG

发布日期: 2025-05-06 (更新: 2026-01-08)


💡 一句话要点

seq-JEPA:通过自回归预测学习不变-协变世界模型,解决表征权衡问题。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 自监督学习 世界模型 不变性 协变性 序列建模 Transformer 预测学习

📋 核心要点

  1. 现有自监督学习方法在不变性和协变性任务间存在性能权衡,限制了表征的灵活性。
  2. seq-JEPA通过引入架构归纳偏置,同时学习不变和协变表征,无需额外的损失项或预测器。
  3. 实验表明,seq-JEPA在各类下游任务中表现出色,尤其擅长序列观察聚合任务。

📝 摘要(中文)

联合嵌入自监督学习(SSL)通常依赖于数据增强和掩码等变换来学习视觉表征,通过强制图像两个视角之间变换的不变性或协变性来实现。这种SSL中占主导地位的双视角范式,通过在高层不变性任务(如图像分类)和更细粒度的协变性任务之间创建性能权衡,通常限制了学习到的表征对下游适应的灵活性。本文提出了seq-JEPA,一个世界建模框架,它将架构归纳偏置引入联合嵌入预测架构,以解决这种权衡。seq-JEPA无需依赖双重协变性预测器或损失项,即可同时学习两种架构上分离的表征,分别用于协变性和不变性任务。为此,我们的模型处理输入的不同视角(观察)的短序列。每个编码的视角与产生序列中下一个观察的相对变换(动作)的嵌入连接。这些视角-动作对通过Transformer编码器,输出聚合表征。然后,预测头以即将到来的动作为条件,预测下一个观察的表征。实验表明,seq-JEPA在协变性和不变性下游任务上都表现出强大的性能,而不会牺牲其中一个。此外,它擅长于需要聚合一系列观察的任务,例如跨动作的路径积分和跨眼球运动的预测学习。

🔬 方法详解

问题定义:现有的联合嵌入自监督学习方法,例如JEPA,在学习视觉表征时,通常需要在高层的不变性任务(如图像分类)和细粒度的协变性任务之间进行权衡。这意味着为了在图像分类等任务上获得良好的性能,模型可能会牺牲对细微变化的敏感性,反之亦然。这种权衡限制了学习到的表征在不同下游任务中的泛化能力。

核心思路:seq-JEPA的核心思路是通过引入架构上的归纳偏置,使得模型能够同时学习到适合不变性任务和协变性任务的表征。具体来说,模型通过处理一系列不同视角的观察,并结合视角之间的变换信息(动作),来学习一个能够预测未来状态的表征。这种方式允许模型在学习过程中区分哪些信息是不变的,哪些信息是随着视角变化的,从而避免了性能上的权衡。

技术框架:seq-JEPA的整体架构包括以下几个主要模块:1) 视角编码器:用于将每个视角的观察编码成一个表征向量。2) 动作嵌入:用于将视角之间的变换(动作)编码成一个嵌入向量。3) Transformer编码器:将编码后的视角表征和动作嵌入连接起来,并通过Transformer编码器进行聚合,得到一个综合的表征。4) 预测头:以即将到来的动作为条件,预测下一个观察的表征。整个流程可以看作是一个自回归的预测过程,模型通过不断地预测未来的状态来学习世界的模型。

关键创新:seq-JEPA的关键创新在于它通过架构上的设计,实现了不变性和协变性表征的解耦学习。与传统的双视角方法不同,seq-JEPA处理的是一系列的视角,并且显式地考虑了视角之间的变换信息。这种设计使得模型能够更好地理解世界的动态变化,从而学习到更具泛化能力的表征。此外,seq-JEPA不需要额外的损失函数或预测器来强制不变性或协变性,而是通过架构本身来实现。

关键设计:seq-JEPA的关键设计包括:1) 使用Transformer编码器来聚合视角表征和动作嵌入,Transformer的自注意力机制能够有效地捕捉视角之间的依赖关系。2) 使用自回归的预测方式来学习世界的模型,这种方式能够鼓励模型学习到更具预测性的表征。3) 动作嵌入的设计,动作嵌入需要能够准确地表示视角之间的变换,例如旋转、平移等。4) 损失函数的设计,通常使用预测表征和真实表征之间的相似度作为损失函数,例如余弦相似度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

seq-JEPA在多个下游任务上取得了显著的性能提升。例如,在需要序列观察聚合的任务(如路径积分和跨眼球运动的预测学习)中,seq-JEPA的表现明显优于现有的自监督学习方法。此外,seq-JEPA在图像分类等不变性任务和姿态估计等协变性任务中也取得了具有竞争力的结果,并且没有出现性能上的权衡。具体的性能数据需要在论文中查找。

🎯 应用场景

seq-JEPA具有广泛的应用前景,例如机器人导航、视频理解、增强现实等领域。在机器人导航中,机器人可以通过观察周围环境的变化,学习到自身的运动模型,从而更好地规划路径。在视频理解中,模型可以学习到视频中物体的运动轨迹和相互作用,从而更好地理解视频的内容。在增强现实中,模型可以根据用户的视角变化,实时渲染虚拟物体,从而提供更逼真的体验。该研究的未来影响在于,它为构建更智能、更具泛化能力的AI系统提供了新的思路。

📄 摘要(原文)

Joint-embedding self-supervised learning (SSL) commonly relies on transformations such as data augmentation and masking to learn visual representations, a task achieved by enforcing invariance or equivariance with respect to these transformations applied to two views of an image. This dominant two-view paradigm in SSL often limits the flexibility of learned representations for downstream adaptation by creating performance trade-offs between high-level invariance-demanding tasks such as image classification and more fine-grained equivariance-related tasks. In this work, we propose \emph{seq-JEPA}, a world modeling framework that introduces architectural inductive biases into joint-embedding predictive architectures to resolve this trade-off. Without relying on dual equivariance predictors or loss terms, seq-JEPA simultaneously learns two architecturally separate representations for equivariance- and invariance-demanding tasks. To do so, our model processes short sequences of different views (observations) of inputs. Each encoded view is concatenated with an embedding of the relative transformation (action) that produces the next observation in the sequence. These view-action pairs are passed through a transformer encoder that outputs an aggregate representation. A predictor head then conditions this aggregate representation on the upcoming action to predict the representation of the next observation. Empirically, seq-JEPA demonstrates strong performance on both equivariance- and invariance-demanding downstream tasks without sacrificing one for the other. Furthermore, it excels at tasks that inherently require aggregating a sequence of observations, such as path integration across actions and predictive learning across eye movements.