MotionJEPA: Preventing Temporal Feature Collapse by Capturing Visual Changes in Latent Space

Embodied AI and Robotics World Models ML SSLRO REG
联合嵌入预测架构(JEPA)是一种极具前景的学习范式,可在无需视觉重建的前提下,学习任务无关的潜在世界模型。然而,标准JEPA训练过程存在强烈的归纳偏置,倾向于提取变化缓慢的特征,从而导致特征抑制乃至潜在表征的坍缩。尽管逆动力学建模可提供时间维度上的抗坍缩机制,但它依赖于动作标签,且对嵌入通用的、无标签的动力学特征缺乏有效激励。为此,我们提出“差分图像与单帧图像嵌入正则化”(DISReg)——一种新型正则化方法:它基于类逆动力学模块,仅预测时序差分图像的嵌入表示,完全不引入任何像素级重建损失,从而促进静态特征与动态特征的均衡学习。DISReg由两部分构成:其静态项用于塑造图像嵌入的分布,并鼓励提取缓慢变化的特征;其动态项则不同于对嵌入本身施加直接约束的方式——它不对图像嵌入的形状或分布施加任何限制,而仅激励模型显式编码动态特征。我们将该正则化器融入标准JEPA框架,由此构建出新型架构MotionJEPA。潜在空间探针实验表明,MotionJEPA所生成的表征比其他方法更为完备;轨迹分析进一步证实,其潜在嵌入在几何上保持简洁性,具有较低的曲率。此外,在四个不同环境中,面对静态背景干扰物的下游规划任务中,MotionJEPA均显著提升了规划成功率。
Joint Embedding Predictive Architectures (JEPAs) are a promising paradigm for learning task-agnostic latent world models without visual reconstruction. However, standard JEPA training exhibits a strong inductive bias towards slow features, causing feature suppression and the collapse of latent representation. While inverse dynamics provides temporal anti-collapse, it relies on action labels and offers little incentive to embed general, unlabeled dynamics. We introduce Difference Image and Single image embedding Regularization (DISReg), a novel regularizer that builds on an inverse-dynamics-style module that predicts temporal difference image embeddings without any pixel reconstruction loss, encouraging balanced static and dynamic feature learning. DISReg consists of a static term that shapes the distribution of the image embedding and encourages slow features, and a dynamic term, which, unlike direct regularization on the embedding, imposes no constraint on the image embedding's shape or distribution and instead only incentivizes that dynamic features be present. By integrating this regularizer into a standard JEPA, we establish our new architecture, MotionJEPA. Latent probing demonstrates that MotionJEPA produces more complete representations than other methods, and our trajectory analysis shows it maintains geometrically simple latent embeddings with low curvature. We further show that MotionJEPA improves downstream planning success under static-background distractors across four environments.
许愿