When Does LeJEPA Learn a World Model?¶
会议: NeurIPS 2026
arXiv: 2605.26379
OpenReview: https://openreview.net/forum?id=UmYyYHEj9P
领域: 自监督/表示学习
关键词: 联合嵌入预测架构 (JEPA), 世界模型, 线性可辨识性, 厄米特多项式, 施图姆-刘维尔理论
一句话总结¶
本文首次在理论上证明了基于高斯正则化的联合嵌入预测架构（LeJEPA）能够在线性可辨识意义下从非线性观测中精确恢复真实隐变量，并严格确立了高斯分布是保证该线性恢复性质成立的唯一隐变量分布。
研究背景与动机¶
联合嵌入预测架构（JEPA）代表了自监督学习构建世界模型（World Model）的重要范式。与传统的自编码器或基于像素重构的生成模型不同，JEPA 主张在表征空间预测未来或缺失特征，从而避免为背景纹理等与任务无关的高频细节浪费模型容量。然而，当一个自监督表征被用作世界模型时，一个核心的理论问题悬而未决：模型学习到的隐状态何时能够成为物理世界真实自由度的忠实映射？如果学到的表征将物体位置、几何形状、材质与速度非线性地纠缠在一起，即便其在固定下游任务上取得了较高的微调分数，面对环境分布漂移或组合泛化需求时依然会迅速失效。
学术界与工业界通常使用线性探测（Linear Probing）来评估自监督表征的质量，但线性探测有效性的根本前提是表征对底层隐变量具备线性可辨识性（Linear Identifiability），即学到的嵌入空间与真实物理隐状态之间仅相差一个全局正交旋转变换。过去非线性独立成分分析（Nonlinear ICA）的经典理论指出，若缺乏额外的结构假设，从非线性混合中反解隐变量在数学上是病态且不可辨识的；而现存的自监督与 JEPA 方法多依赖动量教师网络、停止梯度或未定义显式分布的隐式防坍缩启发式策略，无法推导出任何严格的可辨识性理论保障。近期 LeJEPA 引入了基于草图各向同性高斯正则化（SIGReg）的显式防坍缩目标，使端到端像素级世界模型训练变得稳定，但其学到的表征是否真正恢复了世界的真实隐变量仍缺乏严格证明。
本文从数学机理上填补了这一空白。研究团队建立了一个由平稳过程与加性噪声演化驱动的高斯世界模型，证明了在该类世界中，LeJEPA 的对齐损失与高斯正则化目标足以迫使编码器在全局最优处实现精确的线性正交恢复。核心 idea：利用奥恩斯坦-乌伦贝克转移算子的厄米特多项式谱分解，证明任何阶数大于一的非线性形变都会严格削弱正样本对之间的时序相关性，从而迫使最优编码器成为纯线性正交映射，并通过施图姆-刘维尔理论证明高斯分布是促成线性可辨识性的唯一隐变量分布。
方法详解¶
整体框架¶
本文构建了世界生成过程（The World）与学习者目标（The Learner）之间的数学对齐桥梁。世界生成模型假设潜变量 \(z \in \mathbb{R}^n\) 服从各向同性标准高斯分布 \(z \sim \mathcal{N}(0, I_n)\)，并通过未知非线性混合映射 \(g: \mathbb{R}^n \to \mathcal{X}\) 生成观测数据 \(x = g(z)\)；正样本对 \((x, x')\) 由连续奥恩斯坦-乌伦贝克（OU）平稳加性噪声扩散过程产生。学习者采用参数化编码器 \(f: \mathcal{X} \to \mathbb{R}^n\)，其实际复合映射记为 \(h = f \circ g: \mathbb{R}^n \to \mathbb{R}^n\)。LeJEPA 的优化目标是在满足各向同性高斯先验（SIGReg 正则）的约束下，最小化正样本对表征的欧氏距离（对齐损失）。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["真实隐状态生成与OU转移<br/>z ~ N(0, I_n) 平稳演化为 z'"] --> B["非线性物理混合与观测<br/>x = g(z), x' = g(z')"]
B --> C["高斯正则化约束 (SIGReg)<br/>约束表征协方差与边际分布归一化"]
C --> D["正样本对齐损失最大化相关性<br/>厄米特谱衰减严格惩罚非线性分量"]
D --> E["施图姆-刘维尔第一特征函数线性约束<br/>严格排除所有非高斯隐变量分布"]
E --> F["线性正交世界模型与最优潜空间规划<br/>恢复真实自由度 h(z) = Qz, Q ∈ O(n)"]关键设计¶
1. 厄米特多项式谱分解与前向线性可辨识性证明：利用非线性衰减抑制非线性形变 传统分析难以直接限制复杂的深度神经网络映射类，本文借助高斯测度下的自然标准正交基——概率论厄米特多项式族 \(\{H_\alpha\}_{\alpha \in \mathbb{N}^n}\)。任何满足零均值与单位方差的表征分量 \(h_i(z)\) 均可按总多项式阶数 \(d = |\alpha|\) 正交展开，并将总方差分解为各阶能量之和 \(\sum_{d \ge 1} w_d^{(i)} = 1\)。利用梅勒公式（Mehler's formula），OU 转移算子在厄米特基底上呈现严格的对角化收缩：度数为 \(d\) 的厄米特模式跨越时间步的正样本相关性精确衰减为 \(\rho^d\)（其中 \(\rho \in (0, 1)\) 为转移自相关系数）。由于对任意阶数 \(d \ge 2\) 均有 \(\rho^d < \rho\)，任意非线性失真都会把方差能量推向衰减更快的更高阶模式，从而严格降低跨视角表征的相关性：
等号成立的充要条件是 \(w_1^{(i)} = 1\) 且所有高阶非线性方差 \(w_{d \ge 2}^{(i)} = 0\)。这意味着对齐损失全局最小化必然强制每个分量 \(h_i(z)\) 严格为线性映射；配合 SIGReg 带来的白化协方差约束 \(\mathrm{Cov}(h(z)) = I_n\)，最终证明了存在全局正交矩阵 \(Q \in O(n)\) 使得 \(h(z) = Qz\)。
2. 施图姆-刘维尔算子与高斯分布唯一性反向证明：非高斯分布必然破坏线性恢复 在前向证明确立高斯隐变量足以实现线性恢复后，一个自然的问题是：该保证是否依赖于高斯分布本身？本文基于扩散过程的无限小生成元 \(\mathcal{D}\) 建立了逆向充要条件。在恒定扩散系数 \(K > 0\) 下，一维平稳隐变量过程的特征方程可写为自伴随施图姆-刘维尔算子微分形式：
根据振动理论，生成元的第一非平凡特征函数 \(\varphi_1(z_i)\) 对应最慢时序特征，且始终是单调函数。若要使最优编码器保留线性恢复（即 \(\varphi_1\) 为仿射函数 \(\varphi_1(z_i) = a z_i + b\)），其导数 \(\varphi_1' = a\) 必为常数，代入微分方程后强制得分函数 \((\log p)'(z_i) = -\frac{\lambda}{K}(z_i + b/a)\) 严格为线性。积分后唯一导出的可归一化概率密度必然是高斯分布 \(\log p(z_i) \propto -(z_i - \mu)^2\)。这证明了在平稳加性噪声系统中，除高斯分布外，任何其他隐变量先验（如拉普拉斯、均匀或重尾分布）其第一特征函数均为非线性单调映射，从而在根本上破坏线性可辨识性。
3. 近似可辨识性误差界与正交潜空间最优规划保障：架起理论与工程实践的桥梁 真实训练中对齐损失与白化正则无法达到机器精度的绝对零点。为此，定理 5.3 给出了扰动分析界：若编码器的白化误差满足 \(\|\mathrm{Cov}(h(z)) - I_n\|_F \le \varepsilon\)，对齐松弛量为 \(\delta\)，则表征偏离正交恢复的均方误差上界由非线性残差能量 \(W_{\mathrm{nl}}\) 与形变项显式给出：
分母 \(2\rho(1-\rho)\) 即一阶与二阶特征值之间的谱间隙。进一步，定理 5.4 证明了正交旋转不变量在控制规划中的传递性：任何仅依赖欧氏距离、内积或旋转协变状态转移的目标函数（如目标到达 \(\ell_T(z) = \|z - z^*\|^2\) 或线性二次调节器 LQR），在学到的隐空间内规划出的动作序列与真实环境最优动作完全等价，证明了线性可辨识性正是赋予自监督表征有效世界模型规划能力的数学根基。
损失函数 / 训练策略¶
LeJEPA 模型的端到端实证训练目标由不变性（对齐）损失 \(\mathcal{L}_{\mathrm{inv}}\) 与高斯化正则项 \(\mathcal{L}_{\mathrm{SIG}}\) 构成：
其中 \(\mathcal{L}_{\mathrm{SIG}}\) 利用一维随机投影的经验特征函数与标准高斯解析特征函数计算能量距离或最大均值差异（MMD），迫使编码器输出在各个投影方向匹配各向同性标准高斯。优化过程中学习率采用余弦退火，正则化权重 \(\lambda\) 需保持在合理区间（例如 0.05 左右），过大的 \(\lambda\) 会导致表征忽略输入特征坍缩为纯随机高斯噪声，而过小的 \(\lambda\) 则会导致维度塌陷。
实验关键数据¶
主实验¶
为了验证理论在高维复杂非线性混合中的有效性，研究在由 RealNVP 耦合层构建的可逆非线性映射下，对隐变量维度从 \(N = 2\) 扩展至 \(N = 1024\) 进行了系统测试，对比了 SIGReg、VICReg 和 InfoNCE 三种不同正则化策略恢复真实隐变量的线性决定系数 \(R^2(h \to z)\)。
| 隐变量维度 \(N\) | 原始输入基线 \(R^2(x \to z)\) | LeJEPA (SIGReg) | VICReg | InfoNCE (\(\sigma=1.0\)) |
|---|---|---|---|---|
| 2 | \(0.781 \pm 2.1 \times 10^{-3}\) | \(\mathbf{0.999998 \pm 3.4 \times 10^{-7}}\) | \(0.999996 \pm 8.4 \times 10^{-7}\) | \(0.950961 \pm 1.6 \times 10^{-3}\) |
| 8 | \(0.728 \pm 1.0 \times 10^{-2}\) | \(\mathbf{0.999993 \pm 9.0 \times 10^{-7}}\) | \(0.999988 \pm 4.8 \times 10^{-7}\) | \(0.886818 \pm 4.2 \times 10^{-2}\) |
| 32 | \(0.737 \pm 2.3 \times 10^{-3}\) | \(\mathbf{0.999981 \pm 7.2 \times 10^{-7}}\) | \(0.999981 \pm 9.4 \times 10^{-7}\) | \(0.907809 \pm 2.6 \times 10^{-2}\) |
| 128 | \(0.739 \pm 6.1 \times 10^{-4}\) | \(0.999938 \pm 3.2 \times 10^{-7}\) | \(\mathbf{0.999942 \pm 7.2 \times 10^{-7}}\) | \(0.566955 \pm 6.6 \times 10^{-3}\) |
| 512 | \(0.749 \pm 3.0 \times 10^{-4}\) | \(0.999775 \pm 6.7 \times 10^{-7}\) | \(\mathbf{0.999785 \pm 6.9 \times 10^{-7}}\) | \(0.704393 \pm 2.6 \times 10^{-4}\) |
| 1024 | \(0.763 \pm 1.7 \times 10^{-4}\) | \(0.999561 \pm 1.2 \times 10^{-6}\) | \(\mathbf{0.999582 \pm 1.1 \times 10^{-6}}\) | \(0.720241 \pm 2.0 \times 10^{-4}\) |
消融实验¶
在真实物理机器人像素控制基准 DeepMind Control Suite (Reacher，2自由度关节角隐变量 \(z = (\theta_0, \theta_1)\)) 上，测试了平稳各向同性高斯 OU 采样与真实强化学习策略交互轨迹（非高斯、各向异性）对隐状态线性可辨识性 \(R^2\) 的影响。
| 数据分布类型 | 转移参数 / 帧间隔步长 \(\delta\) | 综合线性可辨识性 \(R^2(z \to h)\) | 关节 0 恢复度 \(R^2(h \to z_0)\) | 关节 1 恢复度 \(R^2(h \to z_1)\) |
|---|---|---|---|---|
| Gaussian OU | \(\rho = 0.50\) | \(0.86 \pm 0.01\) | \(0.86 \pm 0.01\) | \(0.86 \pm 0.01\) |
| Gaussian OU | \(\rho = 0.80\) | \(0.94 \pm 0.00\) | \(0.94 \pm 0.00\) | \(0.94 \pm 0.00\) |
| Gaussian OU | \(\rho = 0.99\) | \(\mathbf{0.95 \pm 0.00}\) | \(\mathbf{0.95 \pm 0.00}\) | \(\mathbf{0.95 \pm 0.00}\) |
| RL Trajectory | \(\delta = 1\) | \(-0.39 \pm 0.10\) | \(0.71 \pm 0.03\) | \(-0.03 \pm 0.04\) |
| RL Trajectory | \(\delta = 4\) | \(-0.05 \pm 0.30\) | \(0.51 \pm 0.10\) | \(0.43 \pm 0.09\) |
| RL Trajectory | \(\delta = 16\) | \(0.44 \pm 0.04\) | \(0.63 \pm 0.02\) | \(0.87 \pm 0.02\) |
| RL Trajectory | \(\delta = 64\) | \(0.44 \pm 0.04\) | \(0.55 \pm 0.03\) | \(0.77 \pm 0.02\) |
关键发现¶
- 高维扩展性：基于批次分布统计量的显式正交化与高斯化方法（SIGReg 与 VICReg）展现出近乎完美的维数可扩展性，在高达 \(N=1024\) 维的复杂非线性混合中仍能维持 \(R^2 > 0.9995\) 的线性恢复率；而 InfoNCE 由于采用固定高斯核宽超参，在高维空间中面临维度灾难与负样本对数估计方差爆炸，线性可辨识性出现显著回落。
- 高斯假设的极端敏感性：在广义正态分布族的形状参数消融实验中，线性可辨识度指标 \(R^2\) 在严格等于高斯分布的 \(\alpha=2\) 处呈现出极其陡峭的尖峰，验证了定理 5.2 关于非高斯分布破坏线性恢复的理论判断。
- 真实策略轨迹对理论假设的双重违背：在 DMC Reacher 实验中，真实强化学习策略引导的轨迹因目标导向行为坍缩至低熵流形，呈现高度各向异性（\(\rho_0 \ne \rho_1\)）及关节极限周期边界约束，导致线性可辨识性 \(R^2\) 整体上限被锁死在 0.50 左右，且不同关节间恢复精度严重分化。
亮点与洞察¶
- 经典盲源分离理论的颠覆性反转：在线性独立成分分析（Linear ICA）传统中，高斯分布恰恰是源信号分离彻底失败的唯一奇异分布；而本文在非线性自监督时序预测框架下证明，高斯分布正是使表征能够线性解耦恢复真实物理世界的唯一正向使能分布。
- 谱几何与正交不变量的优雅结合：通过将加性平稳马尔可夫链的遍历性质与厄米特多项式生成函数严格结合，首次为 JEPA 架构为何不会发生维度坍缩且能自发对齐物理坐标系提供了无需启发式先验的封闭形式解析。
- 对无监督强化学习探索策略的深刻指导：该理论揭示了为何纯基于下游任务策略收集的数据难以训练出优秀的通用世界模型——为了满足可辨识性，无监督预训练阶段的环境探索应尽可能逼近各向同性高斯随机游走（OU 过程），以最大化状态空间的各向同性覆盖。
局限与展望¶
- 维度匹配的强前提约束：当前理论严格依赖于编码器输出维度与世界潜在自由度完全相等（\(m = n\)）的前提假设。在实际应用中，真实物理维度通常未知；当编码器维度过大（\(m > n\)）或发生隐空间特征叠加（Superposition, \(m < n\)）时，模型如何分配冗余或压缩自由度仍有待进一步刻画。
- 仅覆盖状态编码侧而非受控转移：定理 5.4 严格证明了编码器在给定正交变换下的规划等价性，但将动作 \(a_t\) 作为显式干预条件的受控状态转移模型 \(\hat{p}(\hat{z}' \mid \hat{z}, a)\) 本身的可辨识性仍依赖于控制输入的持续激励条件，尚未实现完全端到端的因果联合闭环。
相关工作与启发¶
- vs VICReg / Barlow Twins: VICReg 仅对协方差矩阵施加二阶矩去相关约束（Covariance Regularization），而 LeJEPA 借助 SIGReg 匹配完整的各向同性高斯测度；本文证明了在有限样本和复杂重尾扰动下，高阶矩的正则化约束能够提供比单纯二阶去相关更稳健的收敛鲁棒性。
- vs 慢特征分析 (Slow Feature Analysis, SFA): SFA 理论通过变分最小化时间导数求解流形特征，其在一般平稳分布下只能提供保单调变换的非线性可辨识性；本文通过引入高斯边际测度与对齐损失的强偶联，成功将 SFA 的单调同胚可辨识性严格跃升为线性正交可辨识性。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次给出 JEPA 框架的线性可辨识性封闭形式证明与高斯分布唯一性充要定理。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 理论玩具、1024 维 RealNVP 缩放、先验分布族消融以及真实机器人像素级动力学规划。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨清晰，附带 Lean 4 形式化定理检验，几何物理图像直观生动。
- 价值: ⭐⭐⭐⭐⭐ 为无重构像素自监督世界模型奠定了坚实的非线性表示学习理论基石。