论文

Sub-JEPA:稳定端到端世界模型的子空间高斯正则化

Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

模型训练预训练

摘要

联合嵌入预测架构 (JEPA) 为通过预测未来潜在表示来学习世界模型提供了一个简单的框架。但是,JEPA 训练受到偏差-方差权衡的影响。如果没有足够的结构约束,过多的表示方差会导致模型崩溃为微不足道的解决方案。最近的 LeWorldModel (LeWM) 表明,可以通过使用各向同性高斯简单地约束潜在嵌入来缓解此问题然而,潜在表示本质上位于高维环境空间内的低维流形上,并且直接在该环境空间中强制执行各向同性高斯先验会引入过强的偏差。在这项工作中,我们提出了ame,它通过在多个随机子空间而不是原始嵌入空间中应用高斯约束来寻求偏差方差边界上的有利操作点。这种设计在保留其抗崩溃效果的同时放宽了全局约束,导致在训练稳定性和表示灵活性之间取得更好的平衡。在四个连续控制环境中进行的广泛实验表明,该方法始终优于 LeWM,且具有非常明显的优势。我们的方法简单而有效,可以作为未来基于 JEPA 的世界模型研究的强大基线。fdefineeeemode 代码可在 https://github.com/intcomp/Sub-JEPA 获取。