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

原文链接: arXiv:2603.19312

PDF

Project Page

摘要

LeWorldModel (LeWM) 是首个能够从原始像素端到端稳定训练的 Joint-Embedding Predictive Architecture (JEPA)。传统 JEPA 训练面临稳定性挑战，通常需要复杂的训练技巧和多个损失项。LeWM 仅使用两个损失项——下一时刻嵌入预测损失和正则化项（强制潜在嵌入服从高斯分布）——就实现了稳定的端到端训练。这项工作为 embodied AI 和 world modeling 提供了新的技术路径。

作者: Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero

1. 问题定义

World models（世界模型）是 embodied AI 的核心组件，使智能体能够预测环境动态并规划行动。然而，训练稳定的 world models 面临重大挑战：

“Traditional world models often suffer from training instability, requiring careful tuning of multiple loss terms and hyperparameters.”

现有方法通常需要：

• 复杂的损失函数组合（重构损失、对比损失、KL 散度等）
• 精心设计的训练策略和课程学习
• 大量的工程调优才能实现稳定训练

JEPA (Joint-Embedding Predictive Architecture) 由 Yann LeCun 提出，通过预测潜在空间中的表示而非像素空间，避免了像素级重构的计算开销。但 JEPA 的稳定训练一直是一个开放问题。

2. 方法框架

LeWorldModel 的核心创新在于简化了 JEPA 的训练目标，仅使用两个关键损失项：

2.1 架构设计

LeWM 采用编码器-预测器架构：

• 编码器 (Encoder): 将输入图像映射到潜在嵌入空间
• 预测器 (Predictor): 基于当前嵌入预测下一时刻的嵌入
• 目标编码器 (Target Encoder): EMA 更新的目标网络，提供预测目标

2.2 双损失训练策略

LeWM 的关键突破在于仅使用两个损失项实现稳定训练：

损失 1: 下一时刻嵌入预测损失

L_pred = ||pred_emb - target_emb||²

预测器输出的嵌入与目标编码器产生的目标嵌入之间的均方误差。

损失 2: 高斯正则化项

L_reg = KL(emb || N(0, I))

强制潜在嵌入服从标准高斯分布，防止表示坍缩。

“The Gaussian regularization acts as a natural information bottleneck, encouraging the model to learn compact yet informative representations.”

2.3 与 V-JEPA 的关系

LeWM 与 Meta 的 V-JEPA 系列工作一脉相承：

• V-JEPA 2 (2025): 自监督视频模型，实现理解、预测和规划
• V-JEPA 2.1 (2026): 密集特征的世界建模
• LeWorldModel (2026): 首次实现端到端稳定训练的 JEPA

3. 关键创新

3.1 训练稳定性

LeWM 的核心贡献是证明了 JEPA 可以通过极简的损失设计实现稳定训练：

3.2 计算效率

由于避免了像素级重构，LeWM 具有显著的计算优势：

• 更低的内存占用: 无需存储完整的图像重构
• 更快的训练速度: 潜在空间预测比像素重构更高效
• 更好的扩展性: 可处理高分辨率视频

3.3 表示质量

高斯正则化带来了额外的优势：

• 防止表示坍缩
• 鼓励学习语义丰富的特征
• 为下游任务提供良好的初始化

4. 为什么对 AI 硬件重要

4.1 计算模式分析

LeWM 的计算特性对硬件设计有重要启示：

嵌入预测 vs 像素重构:

• 潜在空间维度通常远低于像素空间（如 256-dim vs 224×224×3）
• 矩阵乘法占比高，适合专用加速器
• 内存带宽需求显著降低

自监督训练的优势:

• 无需标注数据，降低数据预处理开销
• 连续学习特性适合边缘设备部署
• 可与其他任务共享特征提取器

4.2 硬件设计启示

内存优化:

• 潜在嵌入的紧凑表示减少片上存储需求
• 预测器的轻量设计适合移动端 NPU

计算效率:

• 编码器-预测器架构可流水线化执行
• EMA 更新适合异步硬件实现

边缘部署潜力:

• 稳定的训练过程降低了部署复杂度
• 自监督特性支持持续学习

5. 局限与未来方向

局限

• 当前主要在视觉任务上验证
• 大规模视频数据的训练效率待验证
• 与语言模型的结合尚未探索

未来方向

• 多模态扩展: 结合视觉、语言和动作的世界模型
• 机器人应用: 在真实机器人平台上的部署
• 硬件协同设计: 针对 JEPA 架构的专用加速器

6. 总结

LeWorldModel 代表了 world modeling 领域的重要进展：

1. 极简设计: 仅两个损失项实现稳定训练
2. 端到端训练: 首次实现 JEPA 的完全端到端训练
3. 计算高效: 潜在空间预测大幅降低计算开销
4. 硬件友好: 架构特性适合边缘 AI 部署

对于 AI 芯片设计，LeWM 提示了一个重要趋势：未来的加速器不仅要优化推理，还要支持高效的自监督学习和 world model 训练。

参考文献

1. Maes, L., Le Lidec, Q., Scieur, D., LeCun, Y., & Balestriero, R. (2026). LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. arXiv:2603.19312.
2. Assran, M., et al. (2025). V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning. arXiv:2506.09985.
3. Bardes, A., et al. (2024). Revisiting Feature Prediction for Learning Visual Representations. arXiv:2404.08498.