表示学习 · 世界模型 · 理论

LeJEPA 何时学到世界模型?

When Does LeJEPA Learn a World Model?
David Klindt (Cold Spring Harbor Laboratory)  ·  Yann LeCun (New York University)  ·  Randall Balestriero (Brown University)

本文为 Joint-Embedding Predictive Architectures(JEPA)建立首个可识别性保证:在潜变量服从高斯分布的广泛世界类中,LeJEPA(对齐损失 + 高斯正则化)可从非线性观测中线性恢复真实潜变量,且高斯分布是使该保证成立的唯一潜变量分布。

arXiv 2605.26379 理论 + 实验验证 1024 维潜变量扩展性 arxiv.org/abs/2605.26379
关键词JEPA世界模型线性可识别性自监督学习Joint-Embedding Predictive Architecture表征学习高斯正则化Ornstein-Uhlenbeck过程潜变量识别潜空间规划

01 动机

自监督学习(SSL)的愿景在于无需标注数据即可学到有用的世界表示。Joint-Embedding Predictive Architectures(JEPAs)通过训练表示使同一输入的相关视图产生相似嵌入,并辅以正则化防止塌缩。然而一个更深层的问题悬而未决:何时学到的表示才是"世界模型"——即对世界潜在结构的忠实映射?

"A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization."
——论文原文

每当通过 linear probing 评估表示质量时,隐含的问题正是:模型是否学到了潜变量的线性表示?没有线性可识别性(linear identifiability),线性探针就无法精确恢复潜变量。本文正是要给出这一性质的严格数学保证。

LeJEPA world model demo
图 1:LeJEPA 学到世界模型。(左)世界拥有独立高斯潜变量;(中)未知的非线性过程将其混合为我们观测到的数据;(右)LeJEPA 将潜变量恢复至旋转等价,即学到正确的世界模型(Thm. 5.1)。
首个JEPA 可识别性理论保证
唯一高斯分布是线性可识别的唯一潜变量分布(Thm. 5.2)
R²>0.999SIGReg/VICReg 在 N=1024 维时仍保持
R²=0.95DMC Reacher OU 数据对下最优规划表现

02 方法

论文从理论角度分析 LeJEPA:在满足独立性、平稳性、加性噪声三条假设的"世界"类中,证明带高斯正则化的对齐目标可线性恢复高斯潜变量,并给出近似可识别性界和最优规划保证。

LeJEPA theory illustration
图 2(示意):LeJEPA 理论框架。(左)世界具有干净的高斯潜变量结构,正样本对高度相关;(中)未知非线性映射产生观测数据;(右)LeJEPA 以"拉近正样本对"(对齐)加"保持嵌入分布为高斯"(SIGReg)两个目标训练编码器——理论证明学到的表示必为真实潜变量的旋转。

世界模型假设

设世界潜变量 z ∈ ℝⁿ 通过未知非线性映射 g 生成观测 x = g(z)。论文对世界施加三条假设:

在高斯潜变量世界中,满足上述假设的唯一转移为 Ornstein–Uhlenbeck(OU)过程
z′ = ρz + √(1−ρ²)η,η ∼ N(0,Iₙ),η ⊥ z
其中 ρ ∈ (0,1) 控制视图间相关性。

LeJEPA 学习目标

编码器 h = f∘g: ℝⁿ → ℝⁿ 通过最小化如下目标训练:

LeJEPA 目标

minh   L(h) = 𝔼[‖h(z′) − h(z)‖²]    s.t.   h(z) ∼ N(0, Iₙ)
= 对齐损失(Alignment) + 高斯性约束(Gaussianity / SIGReg)

在白化(Cov(h(z)) = Iₙ)条件下,目标化简为最大化正样本对相关之和:L(h) = 2n − 2Σᵢ 𝔼[hᵢ(z′)hᵢ(z)]。

核心理论结果

定理 5.1(线性可识别性 — 正方向)

对于高斯世界,设 h: ℝⁿ→ℝⁿ 满足 h(z)∼N(0,Iₙ),则 L(h) ≥ 2(1−ρ)n,等号成立当且仅当 h(z) = Qz,其中 Q ∈ O(n)(正交矩阵)。证明关键:Hermite 多项式分解使每个非线性度 d≥2 受到严格惩罚,线性映射是唯一最优解。

定理 5.2(高斯唯一性 — 逆方向)

在满足三条假设的所有世界中,若对齐约束加白化的唯一最小化器是线性映射 h(z) = Qz,则 z 必须是高斯分布。证明借助 Sturm–Liouville 谱理论,排除了全部非高斯替代方案。

定理 5.3(近似可识别性)

设近似对齐间隙为 δ,近似白化误差为 ε,令 D = δ / (2ρ(1−ρ))。则存在 Q ∈ O(n) 使得:
𝔼[‖h(z) − Qz‖²] ≤ D + (ε + D)²
恢复误差随 δ、ε 连续平滑退化。

定理 5.4(最优潜空间规划)

设 h(z) = Qz(Q ∈ O(n)),对任意旋转不变代价函数的有限水平最优控制问题,有 V̂*(h(z₀)) = V*(z₀),即在学到的潜空间与真实潜空间中规划等价。

2D simulation mixings
图 3:2D 模拟实验(图示)。颜色编码真实潜变量的极坐标角度与半径。三种非线性混合(抛物线剪切、正弦剪切、RealNVP 耦合层)下,LeJEPA 均将各向同性高斯结构恢复至旋转等价,与定理 5.1 吻合。

03 实验

实验涵盖:2D 非线性混合验证、维度扩展性(N=2→1024)、分布消融(广义正态族)、DMC Reacher 机器人像素控制。基线方法:SIGReg、VICReg、InfoNCE。评估指标:线性可识别性 R²(h→z 及 z→h)。

扩展性实验:不同正则化方法对比(Table 1 摘录)

共享 RealNVP 混合与匹配编码器,5 种随机种子,均值 ± 标准差:

潜变量维度 NSIGReg R²(h→z)VICReg R²(h→z)InfoNCE R²(h→z)
20.999998 ± 3.4e-60.999996 ± 8.4e-60.950961 ± 1.6e-3
640.999966 ± 7.4e-60.999968 ± 8.1e-60.648496 ± 3.1e-2
2560.999884 ± 7.9e-60.999889 ± 7.2e-60.696587 ± 4.9e-3
5120.999775 ± 6.7e-60.999785 ± 6.9e-60.704393 ± 2.6e-3
10240.999561 ± 1.2e-50.999582 ± 1.1e-50.720241 ± 2.0e-3

结论:SIGReg 和 VICReg 在 N=1024 时仍维持 R²>0.999;InfoNCE 在固定核宽 σ=1 下随维度增大明显退化。

RL 轨迹实验:OU 数据对 vs 策略轨迹(Table 2 摘录)

数据类型相关系数 ρR²(z→h)R²(h→z)
OU(高斯正样本对)0.300.67 ± 2e-20.67 ± 2e-2
OU(高斯正样本对)0.900.95 ± 7e-40.95 ± 7e-4
OU(高斯正样本对)0.990.95 ± 4e-40.95 ± 4e-4
RL 轨迹(非高斯)stride δ=1-0.39 ± 1e-10.71 ± 3e-2
RL 轨迹(非高斯)stride δ=640.44 ± 4e-20.55 ± 3e-2

结论:高斯 OU 数据对下,R² 随 ρ 单调上升;RL 策略轨迹因各向异性(ρ₀ ≠ ρ₁)和非高斯转移导致可识别性显著下降,与理论预测一致。

experimental results summary
图 4:实验结果总览。 (a) 定理 5.3 界的验证:SIGReg 跨网格搜索、2D 混合、扩展性和分布消融的运行结果均在对角线以下,确认近似可识别性界成立(少数接近零的离群点源于有限样本估计噪声)。 (b) 高斯最优性:线性恢复 R²(h→z) 在广义正态分布 α=2(高斯)处达到尖锐峰值,与定理 5.2 吻合;SIGReg 的高斯化对非高斯潜变量分布比白化更鲁棒。 (c) 控制代价(30 组随机起始-目标对,路径长度 ≥1,理想值为 1):高斯编码器与 oracle 在统计上无差异;Trajectory 编码器代价系统性偏高。 (d) 控制代价随线性可识别性 R² 单调下降,支持定理 5.4。
DMC Reacher environment
图:DMC Reacher 机器人控制环境。像素观测下的连续控制任务,用于验证理论在实际视觉输入上的适用性。
latent space planning interpolation
图 5:线性可识别性使能潜空间规划。 在三种编码器的潜空间中对固定起始帧和目标帧做插值,并通过最近邻检索解码。 上:Oracle(关节空间直线); 中:高斯编码器(OU,ρ=0.99)紧密追随 oracle; 下:RL 轨迹编码器(stride δ=8)出现明显偏差。 此结果直观验证了定理 5.4:线性可识别性使潜空间规划与真实潜空间规划等价。

分布消融

对广义正态分布族(generalized normal family,形状参数 α)进行扫描:线性恢复 R² 在 α=2(高斯)处达到尖锐峰值,SIGReg 和 InfoNCE 对重尾潜变量比 VICReg 拥有更宽的平坦区。这与定理 5.2 的"高斯唯一性"完全吻合。

2D 仿真

三种非线性混合(抛物线剪切、正弦剪切、RealNVP 耦合层):学到的表示将各非线性混合逆转至旋转等价,与定理 5.1 一致。网格搜索显示:过强的高斯性(λ=0.5)会压缩表示,最优恢复在 低 λ + 高 ρ 时出现。

planning scatter: control cost vs identifiability
图:规划代价 vs 线性可识别性散点图。各模型的控制代价与线性可识别性 R² 呈单调递减关系,支持"线性可识别表示可实现最优潜空间规划"(定理 5.4)的理论结论。

04 局限性

说明:以下局限性由论文作者在正文中明确陈述(stated by the authors),非推断。
真实世界的潜变量是否服从高斯分布?

高斯分布是给定均值与协方差情况下最大熵分布,是最少假设的先验。然而从观测数据中无法验证真实潜变量是否高斯。有一个尺度论据:单个微观变量可能非高斯,但任务相关的潜变量往往是大量微观变量的聚合,中心极限定理倾向于使其趋向高斯。尽管如此,这一假设本质上无法从观测中证伪。

嵌入维度不匹配时的行为(Dimension Mismatch)

定理假设编码器输出维度 m 等于真实潜变量维度 n(m=n)。当 m<n 时,高斯性约束不能确定选取哪个子空间;当 m>n 时,多余维度必须塌缩或编码冗余。理解维度不匹配对可识别性的影响是一个重要的开放问题,对 JEPA 实践设计有直接影响。

有限样本与优化效果(Finite Samples & Optimization)

理论结果是关于全局最优的总体水平(population-level)陈述。定理 5.3 证明保证随对齐间隙和协方差偏差连续退化,但未说明这些量如何随样本量或训练动态缩放。实验中观察到少量界违反情况(图 4a),与 ε 和 δ 的有限样本估计噪声一致。

动作条件下的动态学习(Action-Conditioned Dynamics)

本理论针对编码器学习,未涉及动作条件下的动态学习(action-conditioned dynamics),这是将 LeJEPA 从表示学习扩展到完整世界模型学习的重要后续方向。