all posts
AI技术 · ZH

LeWorldModel 论文深度解读:用两个损失函数从像素端到端学习世界模型

May 8, 2026·30 min read·by PandaTalk

LeWorldModel 论文深度解读:用两个损失函数从像素端到端学习世界模型

论文标题:LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels 作者:Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero 机构:Mila & Université de Montréal, New York University, Samsung SAIL, Brown University 论文链接:https://arxiv.org/abs/2603.19312

世界模型(World Model)是 AI 研究中一条极具野心的路线——让智能体在脑中构建一个环境的内部模型,通过"想象"来规划行动,而非一次次试错。Yann LeCun 多年来一直倡导的 JEPA(联合嵌入预测架构)就是这条路线的核心框架。

问题在于,现有的 JEPA 世界模型方法要么训练不稳定,要么依赖大量工程技巧(stop-gradient、指数移动平均、预训练编码器),要么需要调六七个超参数才能跑起来。这篇论文提出的 LeWorldModel(LeWM)用一种极简方案解决了这些问题:只用两个损失项,就能从原始像素端到端地稳定训练出一个潜在世界模型。

从像素到决策:世界模型为何重要

人工智能的一个核心目标是开发能跨任务、跨环境通用的智能体。理想情况下,智能体应该像人一样,直接从视觉感知出发理解环境、规划行动,而非依赖手工设计的状态表征或特定领域的标定。

世界模型就是实现这一目标的关键技术。一个世界模型本质上是对环境动力学的一个学习到的预测模型——给定当前状态和一个动作,它能预测下一个状态。当世界模型足够准确时,智能体可以在"想象空间"中模拟不同行动序列的后果,选出最优方案再执行,而无需与真实环境反复交互。这在离线学习场景中尤其有价值:智能体从已有的数据集中学习环境模型,然后基于这个模型做规划。

世界模型目前有两大流派:

生成式世界模型直接在像素空间预测未来的图像帧。代表工作包括 IRIS、DIAMOND、DreamerV4 等,它们本质上是学习到的视频生成器——给定历史画面和动作,生成未来画面。优势是直观且信息完整,劣势是计算开销巨大(要生成完整图像),且建模了大量与决策无关的视觉细节。

潜在世界模型则在一个压缩的潜在空间中预测未来。JEPA(Joint-Embedding Predictive Architecture)是这类方法中最受关注的框架——它不试图预测像素层面的未来,而是将观测编码到低维向量空间,在向量空间中预测未来状态的嵌入表征。这种方式天然地过滤掉无关细节,只保留与任务相关的特征,且计算效率远高于生成式方法。

JEPA 面临的核心挑战:表征坍缩

JEPA 的思路虽然简洁,但有一个致命的技术难题——表征坍缩(representation collapse)。

想象一下:如果编码器把所有不同的输入图像都映射到同一个向量,那预测器只要永远输出这个常量向量,预测损失就为零。这个"作弊方案"让模型完美地最小化了损失,但学到的表征毫无意义。

为了对抗坍缩,现有方法各显神通:

  • I-JEPA、V-JEPA 使用指数移动平均(EMA)加 stop-gradient 来稳定训练,但这些技巧的理论基础薄弱
  • DINO-WM 冻结预训练的 DINOv2 视觉编码器,避免编码器参数更新导致坍缩,但也放弃了端到端学习的灵活性
  • PLDM 使用 VICReg 正则化进行端到端训练,但需要七个损失项和六个超参数,训练不稳定且调参困难

下图总结了这些方法各自的特点与局限:

现有潜在世界模型方法的特性对比

LeWM 同时解决了上述所有类别方法的局限:端到端训练、任务无关、基于像素、无需重建、无需奖励信号、只有一个有效超参数,且具有可证明的防坍缩保证。

LeWM 的方法:简洁到极致

LeWM 的架构由两个核心组件构成:一个编码器和一个预测器。

LeWorldModel 训练流水线

编码器负责将原始像素观测映射到低维潜在表征。具体实现上采用 Vision Transformer(ViT),配置为 tiny 规模(约 500 万参数):patch size 14,12 层,3 个注意力头,隐藏维度 192。编码器取最后一层的 [CLS] token 嵌入,再经过一个带 Batch Normalization 的单层 MLP 投影,得到最终的潜在表征向量。

预测器负责在潜在空间中建模环境动力学。它是一个 6 层、16 头注意力的 Transformer(约 1000 万参数),使用 10% dropout。动作信息通过自适应层归一化(Adaptive Layer Normalization, AdaLN)注入每一层——这意味着动作不是简单地拼接在输入上,而是通过调制归一化参数来影响整个网络的计算。AdaLN 的参数初始化为零,确保动作条件的影响是逐步引入的。预测器以自回归方式,根据当前潜在状态和动作,预测下一时刻的潜在状态嵌入。

整个模型只有约 1500 万参数,在单块 GPU 上几小时就能训练完成。

训练目标:两项损失,一个超参数

LeWM 的训练目标是论文最核心的贡献,它仅由两个损失项组成:

预测损失衡量预测器输出的预测嵌入与编码器产出的真实嵌入之间的均方误差(MSE)。这个损失驱动模型学习准确的环境动力学。同时,它也激励编码器学习一种"可预测的"表征——如果表征难以预测,损失就会高,所以编码器自然倾向于提取环境中平滑变化、有规律的特征。

SIGReg 正则化用来防止表征坍缩。SIGReg(Sketched-Isotropic-Gaussian Regularizer)的核心思想是:强制潜在嵌入的分布匹配一个各向同性的高斯分布。如果所有嵌入都坍缩到同一个点,它们的分布就不是高斯分布,SIGReg 损失就会很大。

直接在高维空间中评估分布是否为高斯分布是困难的,SIGReg 采用了一个巧妙的解决方案:基于 Cramér-Wold 定理,如果一个分布在所有一维方向上的投影都是高斯分布,那么这个高维分布本身就是高斯分布。因此,SIGReg 将嵌入向量投影到 M 个随机方向上,然后对每个一维投影使用 Epps-Pulley 统计检验来衡量它与标准高斯的偏离程度,最后取平均。

完整的训练目标表达为:

L_LeWM = L_pred + λ · SIGReg(Z)

其中 λ 是唯一需要调节的超参数(默认值 0.1),M(随机投影数量)在实验中被证明对性能影响可忽略不计。

以下是 LeWM 完整训练过程的伪代码,简洁程度令人印象深刻:

LeWM 训练伪代码

与 PLDM 需要调六个超参数相比,LeWM 只需要调一个 λ。而且由于 λ 可以通过简单的二分搜索来优化(O(log n) 复杂度),调参成本远低于 PLDM 的网格搜索(O(n⁶) 复杂度)。

潜在空间中的规划

训练好世界模型后,LeWM 通过在潜在空间中进行轨迹优化来实现规划。

LeWM 潜在规划流程

规划过程如下:给定当前观测和目标观测,编码器分别将它们映射为当前潜在状态和目标潜在状态。然后随机初始化一个动作序列,通过预测器在潜在空间中滚动推演(rollout),得到终端预测状态。规划的目标是找到使终端预测状态与目标状态距离最小的动作序列。

优化采用交叉熵方法(Cross-Entropy Method, CEM)——一种基于采样的优化方法,通过迭代地选择最佳候选方案并更新采样分布来逼近最优解。为了缓解长时间滚动导致的误差累积,系统采用模型预测控制(MPC)策略:只执行规划出的前 K 步动作,然后根据新的观测重新规划。

由于 LeWM 的潜在空间只有 192 维(相比 DINO-WM 使用的预训练编码器产生的约 200 倍更多的 token),规划的计算效率极高。

实验:跨环境的全面验证

论文在四个不同的控制环境中评估 LeWM,涵盖了二维和三维场景、导航和操作任务:

实验评估环境

  • Push-T:二维操作任务,智能体需要将一个 T 形方块推到目标位置
  • OGBench-Cube:三维机械臂操作任务,视觉更丰富、更复杂
  • Two-Room:二维导航任务,智能体需要在不同房间之间移动到达目标位置
  • Reacher:二关节机械臂任务,需要到达目标构型

所有环境都采用连续动作空间,LeWM 在所有环境中使用统一的超参数设置。

规划性能对比

各环境下的规划性能对比

在 Push-T 和 Reacher 上,LeWM 持续优于 PLDM 和 DINO-WM。值得注意的是,在 Push-T 上 LeWM 比 PLDM 的成功率高出 18%,并且即使 DINO-WM 额外使用了本体感知信息(proprioceptive information),LeWM 仅凭像素输入就超过了它的表现——这说明 LeWM 的编码器确实学到了环境中与任务相关的关键物理量。

在 OGBench-Cube 上,DINO-WM 略微领先,这可能是因为三维环境的视觉复杂度更高,预训练的 DINOv2 编码器在此场景下有优势。在最简单的 Two-Room 环境中,PLDM 和 DINO-WM 优于 LeWM,这揭示了 SIGReg 正则化在低内禀维度环境中的一个潜在局限——当环境的真实信息量很低时,强制嵌入匹配高维高斯分布可能反而不利。

计算效率:48 倍加速

规划时间和固定计算量下的性能对比

计算效率是 LeWM 最亮眼的指标之一。由于编码器生成的潜在表征只有约 200 倍更少的 token(相比 DINO-WM),LeWM 的完整规划过程不到 1 秒即可完成,而 DINO-WM 需要约 47 秒——48 倍的速度差距。在固定计算量(FLOPs)的公平对比下,LeWM 在 Push-T 上的成功率达到 90%,而 DINO-WM 只有 13%;在 OGBench-Cube 上是 74% 对 48%。

这个结果传达了一个重要信息:端到端学习的紧凑表征,在规划效率上有压倒性优势。

训练稳定性

论文还展示了 LeWM 训练过程中的损失曲线,两项损失都呈现平滑、单调的收敛行为。相比之下,PLDM 的七项损失中有多项出现噪声和非单调行为,训练过程更难控制。

潜在空间真的理解物理吗

一个紧凑的 192 维潜在空间能真正捕获环境的物理结构吗?论文通过三种方式验证了这一点。

物理量探测

作者训练线性和非线性探针,从潜在嵌入中预测物理量(智能体位置、方块位置、方块角度等)。

Push-T 环境上的物理量探测结果

结果显示,LeWM 在所有指标上持续优于 PLDM,并且与 DINO-WM 保持竞争力。考虑到 DINO-WM 的编码器是在约 1.24 亿张图像上预训练的 DINOv2,而 LeWM 是从头端到端训练的 tiny ViT,这个结果相当出色。

潜在空间解码可视化

为了更直观地验证潜在表征的信息量,作者训练了一个解码器(训练过程中从未使用),将 192 维嵌入重建为像素图像:

训练过程中的解码器可视化

随着训练推进,解码出的图像越来越接近真实场景。训练初期,解码器只能捕获缓慢变化的特征(slow features),随后逐渐恢复出完整的场景结构。这证明了紧凑的潜在空间确实保留了足够的信息来重建视觉场景,尽管模型的目标从来不是重建。

论文还展示了预测器在潜在空间中的开环滚动推演,经解码器还原为图像:

预测器在 Push-T 和 OGBench-Cube 上的开环推演

给定三帧上下文和一个动作序列,预测器在潜在空间中自回归地生成未来状态,解码后与真实观测高度一致。对 Push-T 这类场景,"想象"出的轨迹几乎与真实轨迹完全吻合。OGBench-Cube 中某些精细细节(如末端执行器角度)的重建不够完美,但整体场景结构和动力学趋势都被准确捕获。

潜在空间的拓扑结构

通过 t-SNE 可视化,LeWM 的潜在空间呈现出与物理状态空间高度一致的拓扑结构:

潜在空间的 t-SNE 可视化

左图是 Push-T 环境中智能体和方块在物理空间中的均匀网格采样,右图是对应潜在嵌入的 t-SNE 投影。两者呈现出相似的空间关系——物理空间中相邻的状态,在潜在空间中也彼此临近。这说明编码器学到的映射保持了环境的空间拓扑。

物理直觉的涌现:违反预期测试

论文的最后一部分将 LeWM 与发展心理学中的"违反预期"(Violation-of-Expectation)范式联系起来——如果一个世界模型真正理解了物理规律,它应该在观察到"不可能事件"时表现出"惊讶"。

具体做法是:让世界模型观察正常轨迹和两种扰动轨迹(物体颜色突变、物体瞬间传送),然后测量模型的预测误差(surprise)。

三种环境中的违反预期评估

结果清晰地显示,LeWM 对物理性扰动(瞬间传送)表现出显著更高的惊讶度(预测误差飙升),而对纯视觉扰动(颜色变化)的反应则相对温和。这说明 LeWM 学到的表征更侧重于物理结构而非表面外观——它"理解"了物体不应该凭空传送这一物理常识。

跨三个环境的统计检验表明,传送扰动引起的惊讶度在所有环境中都显著高于基线(配对 t 检验,p < 0.01),而颜色扰动引起的惊讶度则不显著。

消融实验的关键发现

论文的消融实验揭示了几个值得关注的设计决策:

  • SIGReg 对内部参数不敏感:随机投影数量 M 和积分节点数对性能的影响可忽略不计,唯一有效的超参数是正则化权重 λ
  • 嵌入维度有门槛效应:维度需要足够大才能让方法有效工作,但超过某个门槛后性能迅速饱和,说明方法对编码器容量的精确选择不敏感
  • 编码器架构可替换:将默认的 ViT 替换为 ResNet-18 后性能依然保持竞争力,说明 LeWM 的有效性不依赖于特定的视觉编码器架构
  • 训练曲线的对比:LeWM 的两项损失在训练过程中呈现平滑单调的收敛,而 PLDM 的七项损失中多项出现噪声和非单调行为

局限性与未来方向

论文坦诚地讨论了当前方法的局限:

短视规划——当前潜在世界模型的规划仍限于较短的时间视野。自回归推演会累积预测误差,视野越长误差越大。层次化世界模型是应对长视野推理和规划的一个有前景的方向。

数据多样性依赖——方法依赖具有足够交互多样性的离线数据集。当数据多样性不足时(如 Two-Room 环境),SIGReg 正则化在高维空间中强制匹配高斯分布可能适得其反。在大规模自然视频数据集上预训练可能是缓解此问题的途径。

动作标注需求——当前方法需要与观测配对的动作标签,收集成本较高。通过逆动力学建模来学习动作表征,有望减少对显式动作标注的需求。

这篇论文的意义

对世界模型研究的意义

LeWM 最直接的影响在于它重新定义了这个领域的基线标准。在此之前,端到端训练 JEPA 世界模型被默认为一件需要大量工程技巧的困难工作——你需要 stop-gradient、指数移动平均、多项损失函数的精细平衡、反复的超参数搜索。LeWM 证明这些复杂性并非必需。一个预测损失加一个正则化项,就足以让训练稳定收敛。

这意味着后续的研究者不再需要以"如何稳定训练"为起点,可以直接在一个已经稳定的基础上研究更根本的问题:如何学习更好的表征、如何做更长视野的规划、如何适应更复杂的环境。LeWM 提供了一个干净的实验平台,让研究聚焦于真正重要的问题。

单 GPU 可训练这一点也值得关注。当前 AI 研究越来越向大型实验室集中,算力门槛不断提高。LeWM 将世界模型的训练拉回到了个人研究者和小型实验室可以负担的范围内。这对一个研究方向的长期健康发展至关重要——真正的科学突破往往来自广泛的探索,而非少数玩家的算力竞赛。

对 LeCun 通往 AGI 路线图的意义

理解 LeWM 的更深层意义,需要将它放在 Yann LeCun 关于通用人工智能(AGI)的整体路线图中来看。

LeCun 在 2022 年发表的论文"A Path Towards Autonomous Machine Intelligence"中,提出了一个与当前主流 LLM 路线截然不同的 AGI 架构。他认为,真正的智能系统需要一个世界模型作为核心组件——智能体通过这个内部模型来理解物理世界的运作规律、预测行动的后果、在想象中进行规划。JEPA 正是他为世界模型设计的学习框架。

在 LeCun 的构想中,这条路径有几个关键阶段:

第一步,是让 JEPA 能从原始感知输入中稳定地学习有意义的表征。I-JEPA 和 V-JEPA 分别在图像和视频上验证了这一点,但它们依赖 EMA 和 stop-gradient 等启发式技巧,理论基础不够坚实。LeWM 用 SIGReg 这一有可证明数学保证的正则化方法替代了那些技巧,让 JEPA 的理论根基更加牢固。这不仅仅是工程上的简化,更是理论上的推进——从"经验上有效"走向"数学上可解释"。

第二步,是让世界模型能够在潜在空间中支撑有效的规划与决策。LeWM 用实验证明了这一点:一个从头训练的 1500 万参数模型,仅凭像素输入,就能在潜在空间中做规划,并且规划速度比使用预训练大模型的方法快 48 倍。这说明 JEPA 框架从"学习表征"到"支撑行动"的通路是可行的。

第三步,也是目前尚未完成的,是将世界模型扩展到真实世界的复杂度——从受控的机器人环境走向开放世界的视觉场景,从短视野规划走向长时间跨度的推理。LeWM 在这一步上还有明显的距离(它目前只在相对简单的控制环境中验证),但它为这一步提供了一个正确的起点。

对通用人工智能路线之争的意义

当前通往 AGI 存在两条根本不同的技术路线,而 LeWM 所代表的世界模型路线与主流的大语言模型路线形成了鲜明对比。

LLM 路线认为,通过在海量文本上训练足够大的语言模型,智能可以从语言中涌现——语言已经编码了人类对世界的理解,因此足够强大的语言模型可以"理解"世界。这条路线的核心操作对象是离散的 token 序列。

世界模型路线认为,真正的智能必须建立在对物理世界的感知和建模之上。语言是人类经验的压缩和抽象,但大量的物理直觉、空间推理、因果关系无法通过语言充分传达。一个从未"看过"也从未"行动过"的系统,仅通过阅读文字来理解重力、碰撞、物体永恒性,存在根本性的瓶颈。

LeWM 的违反预期实验为世界模型路线提供了一个有力的论据:一个仅从像素和动作序列中训练的 1500 万参数模型,自发涌现了区分物理性扰动和视觉性扰动的能力——它"知道"物体不应该凭空传送,但颜色变化是相对次要的。这种物理直觉是从感知经验中直接习得的,无需任何语言中介。

这并不意味着两条路线互相排斥。更合理的前景可能是融合:语言模型提供高层次的推理和知识,世界模型提供底层的物理理解和感知规划。但 LeWM 证明了世界模型这条路线是可实现的、可扩展的,并且产出的表征确实编码了有意义的物理结构。

对"简洁性"作为研究原则的意义

在更深的方法论层面,LeWM 传递了一个在当前 AI 研究中逐渐被遗忘的信念:好的科学理论应该是简洁的。

近年来的 AI 研究有一种明显的趋势——通过不断叠加技巧、增加损失项、引入辅助目标来提升性能。每一项新增的技巧可能带来边际收益,但也引入了额外的超参数、更复杂的交互效应、更难复现的结果。PLDM 的七项损失和六个超参数就是这种趋势的典型产物。

LeWM 的做法恰好相反。它回到问题的本质——"什么是防止坍缩所需的最小充分条件?"——然后给出了一个数学上可证明的回答:强制嵌入分布匹配各向同性高斯分布。这个单一的正则化项就够了,不需要其他任何技巧。结果是一个更稳定、更易复现、更容易分析和理解的系统。

这提醒我们,当一个方法需要大量技巧才能工作时,也许问题出在方法本身的设计上,而非技巧还不够多。真正的进步往往来自找到正确的抽象和原理,让复杂性自然消解。

总结

LeWorldModel 最核心的贡献,在于证明了一件看似矛盾的事情:通过做减法——减少损失项、减少超参数、减少工程技巧——反而可以获得更稳定的训练、更高效的规划、以及更具物理结构意义的潜在表征。

在一个追求"堆更多损失项、加更多技巧"的研究趋势下,这篇论文提供了一个令人信服的反例:一个预测损失加一个正则化项,1500 万参数,单块 GPU,几小时训练——就能得到一个在多种环境中具有竞争力的世界模型,并且其潜在空间还自发涌现出了物理理解和轨迹平滑化等有意义的性质。

它不是 AGI 的答案,但它可能是通往答案的道路上一块重要的铺路石——证明了从像素到物理理解这条路径是可行的,并且比我们此前以为的要简洁得多。

━━━ fin ━━━

If you read this far — thank you.
Come tell me what you thought on X.