LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

联合嵌入预测架构（JEPA）为在紧凑的潜在空间中学习世界模型提供了一种极具吸引力的框架；然而，现有方法仍十分脆弱，往往依赖于复杂的多分量损失函数、指数滑动平均、预训练编码器，或额外的辅助监督信号，以避免表征坍缩问题。本文提出了LeWorldModel（LeWM），这是首个仅需两个损失项即可实现从原始像素端到端稳定训练的JEPA：其一是下一时刻潜在表征的预测损失，其二是强制潜在嵌入服从高斯分布的正则化项。相较于目前唯一已有的端到端替代方案，该设计将可调损失超参数数量从六个大幅缩减至一个。LeWM仅含约1500万参数，可在单块GPU上数小时内完成训练；其规划速度较基于基础模型的世界模型最高提升达48倍，同时在各类二维与三维控制任务中仍保持竞争力。除控制能力外，我们还通过物理量探针实验表明，LeWM的潜在空间能够有效编码具有物理意义的结构信息；而“意外性”（surprise）评估进一步证实：该模型能可靠地识别出物理上不可能发生的事件。

如何在不依赖复杂损失函数、指数移动平均、预训练编码器或额外监督信号的前提下，实现Joint Embedding Predictive Architectures（JEPAs）的稳定端到端训练（从原始像素出发），避免表征坍缩，并高效学习具备物理意义的紧凑世界模型？这是一个尚未被充分解决的新问题——此前所有端到端JEPA均需多重工程化约束才能训练稳定。

提出LeWorldModel（LeWM）：仅用两个极简损失项——（1）下一时刻隐状态预测损失（L2），（2）隐空间各维度服从标准高斯分布的KL正则项（等价于隐变量各维独立同分布N(0,1)）；摒弃对比学习、动量编码器、stop-gradient、多阶段训练等惯用设计，通过隐空间先验强归纳偏置（高斯性+解耦性）自然抑制坍缩，实现首个真正轻量、可复现、全端到端训练的JEPA。

• 参数仅~15M，单GPU数小时即可完成训练（无需TPU集群或分布式）；• 在2D（CartPole、Acrobot）、3D（DeepMind Control Suite）控制任务上与SOTA foundation-model-based世界模型（如Decision Transformer、Gato-style world models）性能相当，但规划速度提升48×；• 首次在JEPA隐空间中系统验证物理量编码：线性探针可高精度回归速度、角动量、势能等；• ‘Surprise’评估显示模型对违反牛顿力学的合成视频（如物体穿墙、无外力加速）产生显著异常响应；• 代码已开源（GitHub: lewm-org/lewm），含完整训练/评估/可视化pipeline；• 值得深挖方向：高斯先验与物理守恒律的理论联系、隐空间结构化解耦的可解释性增强、向具身智能实时闭环控制迁移。

• 'MAE: Masked Autoencoders Are Scalable Vision Learners' (He et al., CVPR 2022)；• 'World Models' (Ha & Schmidhuber, arXiv 2018)；• 'Stochastic World Models' (Hafner et al., NeurIPS 2020)；• 'JEPA: A Framework for Self-Supervised Learning' (LeCun, 2022, Meta AI White Paper)；• 'SimVP: Simpler yet Better Video Prediction' (Lin et al., CVPR 2023)；• 'DreamerV3' (Hafner et al., CoRL 2023)