AI Research 技术调研

When Does LeJEPA Learn a World Model?

18分钟阅读

一句话定位

一篇纯理论论文，证明 LeJEPA（对齐损失 + SIGReg 高斯正则化）在”潜变量做平稳、加性噪声转移”的世界模型族中线性可辨识（linear identifiability，学到的表示是真实潜变量的一个旋转 h(z)=Qz），并且反向证明高斯是这族世界里唯一能保证这一点的潜分布——把 LeJEPA 这套经验上好用的配方，第一次上升为可 Lean 4 机器验证的数学定理。

背景与定位

LeJEPA（Balestriero & LeCun, 2025-11）用”对齐 + SIGReg 强制嵌入服从各向同性高斯”替代了此前 JEPA 家族（i-jepa、v-jepa-2 等，均可追溯到 LeCun 2022 年的 Path Towards Autonomous Machine Intelligence 蓝图）里各种隐式防塌陷技巧（stop-gradient teacher、协方差正则等），号称训练更稳定、可端到端从像素学起（如 LeWorldModel 的动作条件控制场景）。但一个悬而未决的问题是：LeJEPA 学到的表示是否真的恢复了世界的潜变量结构，还是只是”看起来还不错”的经验解？

本文的答案框架沿用非线性 ICA / 可辨识性（identifiability）文献的经典范式：可辨识性永远是”世界”（数据生成过程的假设：独立性、平稳性、加性噪声转移）和”学习器”（这里是 LeJEPA 的目标函数）的联合陈述。此前的非线性 ICA 可辨识性结果依赖非平稳性、辅助变量、干预或监督信号（Hyvärinen 等一系列工作）；本文的独特之处是只用平稳、各向同性的高斯世界 + LeJEPA 的两个损失项，就能拿到线性可辨识性——且证明这是唯一的分布选择。文章将这一结果类比为”反转经典线性 ICA 的叙事”：线性 ICA 里高斯恰恰是源分离失效的唯一分布；而在本文的非线性设定里，高斯恰恰是让分离成立的唯一分布。

模型架构

本文没有提出新的网络架构，而是给”世界”和”学习器”各自形式化：

世界（World）的三条假设（3.1 节）：

  • 独立性：潜变量各分量 z_i 相互独立，转移 p(z_i'|z_i) 也分量独立；
  • 平稳性：两个 view 共享同一边缘分布 p(z)=p(z')；
  • 加性噪声转移：z_i' = m_i(z_i) + η_i，η_i 与 z_i 独立。

在此基础上专门化出高斯世界（3.1.1 节）：z ~ N(0, I_n)，平稳性 + 加性噪声唯一确定了转移必须是 Ornstein-Uhlenbeck（OU）过程：

z' = ρ·z + sqrt(1-ρ²)·η,   η ~ N(0, I_n) 独立于 z

其中 ρ∈(0,1) 控制两个 view 的相关性，可验证 Cov(z,z')=ρI_n。论文指出高斯是唯一在这种加性通道下能保持边缘分布不变的分布（高斯是卷积的不动点，up to rescaling）。

学习器（Learner）即 LeJEPA（3.2 节）：复合映射 h = f∘g（g 是未知非线性观测过程，f 是要学的编码器），训练目标为

min_h  E[||h(z')-h(z)||²]   s.t.  h(z) ~ N(0, I_n)
      （对齐 Alignment）        （高斯性 Gaussianity, 由 SIGReg 强制）

白化（Cov(h(z))=I_n）后损失可化简为 L(h) = 2n − 2·Σ_i E[h_i(z')h_i(z)]，即最小化距离等价于最大化 view 间相关性。

谱分析工具（第4节）：把 z→z' 的转移看作函数空间上的线性算子 T（(Th_i)(z)=E[h_i(z')|z]），其特征函数在高斯世界下就是 Hermite 多项式（Mehler 公式给出 degree-d 特征值恰为 ρ^d），对一般（非高斯）分布则由 Sturm-Liouville 理论刻画特征函数（第一个非常数特征函数总是单调的，给出”至多到单调变换”的可辨识性；而线性可辨识性要求这个特征函数是仿射的，这一约束把分布逼到唯一的高斯）。

四条定理（第5节，均给出 Lean 4 证明骨架，App. A–E 完整证明）：

  • 定理1（正向，线性可辨识）：高斯世界中，任何满足 h(z)~N(0,I_n) 的可测映射 h 满足 L(h) ≥ 2(1-ρ)n，当且仅当 h(z)=Qz（Q 为正交矩阵）时取等号；此时 h(z')|h(z) ~ N(ρh(z), (1-ρ²)I_n)，即连转移动力学也一并恢复。
  • 定理2（反向，高斯唯一性）：在假设 3.1 定义的世界族中，若每个 Cov(h(z))=I_n 的极小值点都是线性的，则 z 必为高斯。（证明思路：仿射特征函数 ⇒ score function (log p)' 是线性的 ⇒ 解出 log p(z_i) ∝ -(z_i-μ)²，即高斯。）
  • 定理3（近似可辨识性）：当对齐损失只到 δ-近似、白化只到 ε-近似时，恢复误差有界 E[||h(z)-Qz||²] ≤ D+(ε+D)²，D=δ/(2ρ(1-ρ))——保证退化是平滑的（“gracefully”），且实践中对齐误差 δ 是主导项，白化误差 ε 几乎免费达到。
  • 定理4（最优隐空间规划）：若 h(z)=Qz（Q 正交），则对任何代价函数在 O(n) 旋转下不变的有限时域最优控制问题，学到的隐空间里求解的价值函数和最优动作序列与真实隐空间完全相等——线性可辨识性直接把学到的表示变成可用于规划的”世界模型状态”。

数据

本文是理论 + 受控仿真验证性质，没有大规模真实数据训练，具体实验数据规模：

  • 2D 合成实验：z ~ N(0,I_2)，四种非线性混合函数（spiral 保测度旋转、sinusoidal shear、parabolic shear、RealNVP coupling layer），在线生成数据（每步采新样本，“infinite data”）。
  • 高维 scaling 实验：潜变量维度 N ∈ {2,4,...,1024}（2 的 1~10 次幂），同样在线生成，RealNVP 混合 + matched inverse-NVP 编码器；对比 SIGReg / VICReg / InfoNCE 三种”高斯性约束”目标（5 个随机种子）。
  • 分布消融（generalized normal 族）：形状参数 α∈{2⁻³,...,2⁵}（α→0 重尾，α=1 Laplace，α=2 高斯，α→∞ 均匀），四种 2D 混合函数各自扫描。
  • 网格搜索：λ∈{1e-6,...,5e-1}（9 档）× ρ∈{0.3,0.5,0.7,0.8,0.9,0.95,0.99}（7 档），63 组配置 × 3 个种子 = 189 次训练。
  • 像素级 RL（DMC Reacher, hard 变体）：两个关节（肩、腕），64×64 像素、MuJoCo EGL 无头渲染；每个实验条件预渲染 100,000 图像对 + 10,000 评估图像（评估图像为 i.i.d. 高斯采样，跨所有条件共享）。
    • OU 条件：ρ∈{0.3,...,0.99}（7 档），完全满足理论假设（<0.2% 的样本发生 ±π wrap）。
    • Trajectory 条件：来自 LeWorldModel 数据集里 SAC 策略跑出的 10,000 episodes（每条 201 步），按时间跨度 δ∈{1,2,4,8,16,32,64} 每 episode 采 10 对（共 100,000 对），提取关节角重新渲染；肩关节边缘分布重尾（13.6% 样本超出 ±π），腕关节近似均匀（δ=8 时 ρ_0=0.992, ρ_1=0.982，各向异性）。

数据不涉及真实世界大规模预训练语料，“数据”维度的核心是构造违反/满足理论假设的对照条件，用于验证正向/反向定理。

训练方法

  • 目标函数：L = λ·L_SIG + (1-λ)·L_inv，L_inv 为正样本对（z, z' 经 OU 转移）之间的均方距离，L_SIG 为 Balestriero & LeCun 的 SIGReg（惩罚经验切片特征函数与标准高斯目标的偏差）。
  • 2D/gennorm/grid 实验：4 层 MLP 编码器，隐藏维 256，GELU；ℝ²→Linear(2,256)→GELU→Linear(256,256)→GELU→Linear(256,256)→GELU→Linear(256,2)；优化器 AdamW，学习率 3e-3，前半程 constant + 后半程 cosine 衰减（warmup 10k + cosine 10k，共 20,000 步），batch size 256，weight decay 0，评估用固定 10,000 点；ρ=0.95（2D/gennorm 固定），λ=1e-3（2D/gennorm 固定，grid 实验中扫描）。
  • Scaling 实验：matched inverse-NVP 编码器（4 层 coupling，tanh 激活），其余超参同上，λ=1e-6。
  • Reacher 像素实验：CNN 编码器（4 个 Conv-BN-GELU block，通道 3→32→64→128→256，4×4 stride-2 卷积，AvgPool(4) + Linear(256,256)→BN→GELU→Linear(256,2)），约 1.1M 参数；BatchNorm 对稳定训练关键——去掉后约 36% 的运行会塌陷为零方差输出；梯度裁剪 ||∇||_max=1.0；AdamW（lr 3e-3, weight decay 1e-4），cosine 退火训练 100 epochs，batch size 256，SIGReg 用 256 个随机切片方向；对 λ∈{1e-3,5e-3,1e-2,5e-2} 各跑 3 个种子，报告每个条件下最优 λ 的均值±方差；不做早停/checkpoint 选择。
  • 理论建模上的”训练”实为约束优化的解析刻画：论文并非通过训练发现定理，而是先证明”在 LeJEPA 的两个约束（对齐最小 + 嵌入分布为标准高斯）下，全局最优解一定是正交旋转”，再用训练实验去验证这个解析预测。

Infra（训练 / 推理工程）

  • 本文是小规模理论验证实验（2D 到 1024 维合成数据 + 64×64 像素 Reacher），未披露具体 GPU 型号、GPU 数量或训练 GPU-hours；致谢中提到 David Klindt 使用 “the CSHL GPU cluster”（Cold Spring Harbor Laboratory GPU 集群），资助来自美国 NIH Grant S10OD028632-01，但未给出集群规模或本文实验消耗的算力数字。
  • 训练精度、并行方式未披露（推测为单卡/小规模多卡即可完成，因为最大规模实验也只是 1024 维向量 + 4 层 MLP/CNN）。
  • 推理侧不适用（本文不是可部署系统，无 FPS / 控制频率 / 边缘硬件相关披露）。
  • 形式化验证工程：Lean 4 v4.28.0 + Mathlib v4.28.0，8,032 个 build targets，零编译错误、零 sorry（未完成证明占位符）；作者说明公理化的 5 类背景引理（Hermite 多项式基础设施、Mazur-Ulam 定理、均匀权重 AM-GM、测度论/矩阵分析事实、轨迹前推期望代价关系）均为 Mathlib 尚未收录的经典结果，核心证明链条（相关性上界及等号刻画、Sturm-Liouville 特征函数方程到仿射 score 函数的代数链、正交 Jacobian 到线性等距的几何链等）完全机器验证。

评测 benchmark

论文用”线性可辨识性 R²”和”正交误差”作为核心度量（R²(h→z) 接近 1 表示 h 是 z 的线性函数；正交误差 ||Q̂ᵀQ̂-I||_F/√n 接近 0 表示拟合出的线性映射确为正交阵），全部数字均来自本文实验，非引用外部 benchmark：

  • Table 1（scaling to N=1024，5 种子）：RealNVP 混合下，SIGReg 与 VICReg 在 N 从 2 到 1024 全程保持 R²(h→z) > 0.999（N=1024 时 SIGReg 0.999561±12e-7，VICReg 0.999582±11e-7）；InfoNCE（固定核宽 σ=1）在低维尚可匹配（如 N=16 时 0.999880），但随维度增长退化明显（N=64 骤降到 0.648496，N=128 为 0.566955，N=1024 回升到 0.720241 但仍显著低于另两者）。作为对照，观测空间到潜变量的非线性度 R²(x→z) 全程只有 0.727~0.781，说明混合确实高度非线性。
  • Table 2（Reacher 像素，3 种子）：OU（高斯）条件下 R² 随 ρ 单调上升，ρ=0.99 时达到 R²(z→h)=R²(h→z)=0.95±4e-4；对照的 Trajectory（非高斯，SAC 策略轨迹）条件，即便在相近 ρ_0,ρ_1（如 δ=64 时 ρ_0=0.915, ρ_1=0.863）下，总 R² 从未超过 0.5，且逐维高度各向异性（例如 δ=8 时 R²(h→z_0)=0.80±2e-3，R²(h→z_1)=0.78±6e-3，而 δ=1 时腕关节维度 R²(h→z_1)=-0.03±4e-2，几乎不可辨识）。
  • 分布扫描（generalized normal, App. H.7）：线性恢复 R²(h→z) 在四种 2D 混合下均在 α=2（高斯）处尖峰，SIGReg 的高恢复区间比 VICReg 更宽（对重尾潜变量更鲁棒）。
  • 网格搜索（App. H.6，63 组×3 种子=189 次）：最佳性能出现在 λ∈{1e-3,5e-3} 且 ρ∈{0.9,0.95}，R²>0.97，正交误差 ≈0.15（图注给出的更宽范围是 λ∈[1e-3,1e-2] 且 ρ∈[0.8,0.95] 时 R²>0.96，正交误差 <0.2）；λ=0.5 时表示塌陷（R²≈0）；ρ=0.99 且 λ 中等时因对齐损失趋零导致 SIGReg 主导、性能反而下降。
  • 规划实验（Fig. 4c/d, K=30 起止点对）：高斯（OU, ρ=0.99）编码器的隐空间直线插值规划，控制代价与 oracle（真实关节空间直线）在统计上不可区分；Trajectory 编码器的规划代价系统性偏高；跨全部模型看，控制代价与线性可辨识性 R² 单调相关，验证定理4。
  • 论文自陈的定量局限：Fig. 4a 中有两个”接近零的离群点”被归因于有限样本噪声，而非定理3不成立。

创新点与影响

  • 第一个 JEPA 类目标函数的可辨识性证明：此前非线性 ICA 的可辨识性依赖非平稳性、辅助变量、干预或监督信号；本文只用”平稳 + 加性噪声 + 高斯”这一相对温和的假设集合，就把 LeJEPA 的经验配方升级为数学保证，填补了”linear probing 有效性为何成立”背后缺失的理论基础。
  • 反转经典 ICA 叙事：线性 ICA 中高斯是源分离失效的唯一分布；本文证明在这个非线性谱设定下，高斯恰恰是让可辨识性成立的唯一分布——用同一套 Hermite/Sturm-Liouville 谱工具解释了两者看似矛盾的结论。
  • 给”规划”一个免费的正确性保证：定理4 说明只要线性可辨识（正交）成立，隐空间里对任意 O(n)-不变代价函数做规划，价值函数和最优动作序列与真实世界完全相等——为”JEPA 表示可直接拿去做隐空间规划”提供了此前缺失的理论支撑。
  • Lean 4 形式化验证：把全部四条定理的核心证明链条机器验证（8,032 build targets、零 sorry），这是该论文区别于一般理论论文的工程亮点，公开在 GitHub 仓库中可复现。
  • 论文自述的局限（第7节）：(1) 现实世界的潜变量是否真的高斯”从观测数据本身无法验证”，作者只给出”任务相关潜变量常是许多微观变量的聚合，根据中心极限定理趋于高斯”这一启发式论证，并承认经典 ICA 的非高斯性假设同样面临这个不可验证性问题；(2) 定理假设编码器输出维度与真实潜变量维度相等（m=n），m<n（欠参数化/叠加）或 m>n（过参数化，多余维度如何塌陷或编码冗余）未被理论覆盖，作者称这是”直接影响 JEPA 设计的重要开放问题”；(3) 定理是关于全局最优的总体（population-level）陈述，未涉及有限样本或训练动力学如何影响 ε, δ 的实际取值——只有定理3给出”退化程度”的界，没有给出”样本数需要多大”的界；(4) 实践启示上，论文指出 Reacher 实验里同一个物理系统在各向同性采样（OU）下可辨识、但在目标导向策略采样下不可辨识，提示自监督预训练阶段的探索策略应尽量接近各向同性随机游走。

原始链接

一手源存档（sources/）

关系图谱

反向链接