Files
worldmodel/JEPA/LeJEPA/paper_reading.md
T
gaojie b5499c7ea0 Add detailed lecture plan and summary for LeJEPA four theorems
- Introduced a comprehensive lecture plan for the four main theorems in LeJEPA, including knowledge dependency graphs, detailed outlines for each topic, and corresponding Lean 4 files for formal verification.
- Created a summary document encapsulating the core insights and mathematical structures of the four theorems, emphasizing their interdependencies and implications in the context of LeJEPA.
2026-06-05 16:30:24 +08:00

18 KiB
Raw Blame History

论文精读:When Does LeJEPA Learn a World Model?

作者: David Klindt (CSHL), Yann LeCun (NYU), Randall Balestriero (Brown) 发表: arXiv:2605.26379v1, 2026年5月25日 本地 PDF 2605.26379v1.pdf 官网: https://klindtlab.github.io/lejepa-identifiability/ 代码: lejepa-identifiability/(已本地 clone 视频: https://youtu.be/EioGDo67ZDs


目录


一、论文要解决什么问题?

核心问题:LeJEPA 学到的表示,什么时候才算真正学到了"世界模型"?

JEPAJoint-Embedding Predictive Architecture)是 LeCun 提出的自监督学习框架,通过在表示空间做预测来避免像素级生成的容量浪费。但此前没有任何理论保证说 JEPA 学到的表示是否真正恢复了世界的潜在结构——表示可能把位置和颜色混在一起、把速度和纹理纠缠在一起,虽然在窄任务上表现好,但世界一变就崩。

这篇论文的目标:给 JEPA 的第一个可识别性(identifiability)定理

1.1 背景:什么是 JEPA 和 LeJEPA

JEPAJoint-Embedding Predictive Architecture

  • 训练编码器对同一内容的两个视图产生相似的嵌入
  • 用正则化器防止表示坍塌(collapse)

LeJEPA = JEPA + SIGRegSketched Isotropic Gaussian Regularization):

  • 对齐损失(Alignment): 拉近正样本对的嵌入
  • 高斯正则化(SIGReg): 强制嵌入分布接近各向同性高斯分布 (h(z) \sim \mathcal{N}(0, I_n))

1.2 核心缺口

此前没有任何 JEPA 的可识别性理论——不知道学到的表示是否真正恢复了世界的潜在结构。


二、世界模型的数学框架

2.1 世界的三条假设

假设 数学表述 直觉
独立性 (p(z_i) \perp p(z_j)),转移也独立 世界的各自由度互不干扰
平稳性 (p(z) = p(z')) 两个视图来自同一生成过程
加性噪声 (z'_i = m_i(z_i) + \eta_i) 扰动是叠加在信号上的噪声

2.2 高斯世界(Gaussian World

在以上假设下,选择最大熵分布——高斯分布 (z \sim \mathcal{N}(0, I_n))。

此时转移过程唯一确定Ornstein-Uhlenbeck (OU) 过程

z' = \rho z + \sqrt{1-\rho^2}\,\eta, \quad \eta \sim \mathcal{N}(0, I_n)

其中 (\rho \in (0,1)) 控制两个视图的相关性。

可验证:(\mathbb{E}[z'] = 0)(\text{Var}(z') = \rho^2 I_n + (1-\rho^2) I_n = I_n)(\text{Cov}(z, z') = \rho I_n)。

2.3 LeJEPA 的学习目标

\min_h \;\mathbb{E}[\|h(z') - h(z)\|^2] \quad \text{(对齐损失)} \text{s.t.} \quad h(z) \sim \mathcal{N}(0, I_n) \quad \text{(SIGReg 高斯约束)}

2.4 数据生成流程

真实潜空间 z ~ N(0, I)
    ↓ 非线性混合 g
观测数据 x = g(z)
    ↓ LeJEPA 编码器 h
学到的表示 h(x) = h(g(z))
    ↓ 目标
h(z) = Qz(正交等价恢复)

三、四大定理——论文的核心贡献

定理 1:线性可识别性(正向)

在高斯世界中,满足 LeJEPA 目标的最优表示 (h) 当且仅当 (h(z) = Qz)(Q \in O(n)) 为正交矩阵。

证明链条(6步):

高斯约束 + 最优对齐
        ↓
[步骤1] Hermite 展开:h_i(z) = Σ cₐ Heₐ(z)
        ↓
[步骤2] Mehler 公式:corr_i = Σ wₐ ρᵈ
        ↓
[步骤3] 关键不等式:corr_i ≤ ρ(等号 ⟺ w₁=1
        ↓
[步骤4] 最优性条件:L_align = 2(1-ρ)n → 每个 corr_i = ρ
        ↓
[步骤5] 线性性:每个 h_i 是线性函数
        ↓
[步骤6] 正交性:高斯约束 + 线性 → Q ∈ O(n)

步骤 1Hermite 展开

任意满足 (\mathbb{E}[h_i(z)^2] < \infty) 的函数可以展开:

h_i(z) = \sum_{\alpha} c_{i,\alpha} He_\alpha(z)

高斯约束的含义:

  • (\mathbb{E}[h_i(z)] = 0) → (c_{i,0} = 0)(零均值)
  • (\mathbb{E}[h_i(z)^2] = 1) → (\sum_{|\alpha|\geq 1} c_{i,\alpha}^2 |\alpha|! = 1)(单位方差)

定义谱权重(w_{i,d} = \sum_{|\alpha|=d} c_{i,\alpha}^2 d!),则 (w_{i,d} \geq 0)(w_{i,0} = 0)(\sum_d w_{i,d} = 1)。

步骤 2Mehler 公式计算相关性

\text{corr}_i := \mathbb{E}[h_i(z') \cdot h_i(z)] = \sum_{d=1}^{\infty} w_{i,d} \cdot \rho^d

步骤 3:关键不等式

由于 (\rho^d < \rho)(当 (d \geq 2, 0 < \rho < 1)):

\text{corr}_i = \sum_{d=1}^{\infty} w_{i,d} \cdot \rho^d \leq \sum_{d=1}^{\infty} w_{i,d} \cdot \rho = \rho

等号成立 ⟺ 对所有 (d \geq 2)(w_{i,d} = 0) ⟺ (w_{i,1} = 1) ⟺ (h_i) 是纯线性函数。

步骤 4:最优性条件

L_{\text{align}} = 2n - 2\sum_i \text{corr}_i \geq 2n - 2n\rho = 2(1-\rho)n

最优值当且仅当每个 (\text{corr}_i = \rho),即每个 (h_i) 都是线性的。

步骤 5-6:线性性 + 正交性

(h(z) = Az),高斯约束 (h(z) \sim \mathcal{N}(0, I_n)) 要求 (AA^T = I_n),即 (A \in O(n))。

核心直觉: OU 过程对高阶非线性成分衰减更快((\rho^d) 随 (d) 指数衰减),所以线性映射是唯一最优解。


定理 2:高斯分布的唯一性(逆向)

在满足世界假设的所有分布中,高斯分布是唯一使 LeJEPA 实现线性可识别性的分布。

证明工具:Sturm-Liouville 理论

核心链条:

第一特征函数 φ₁(z) = az + b(仿射)
        ↓ 代入特征方程
得分函数 (log p)' = αz + β(线性,斜率 < 0
        ↓ 积分
log p(z) = (α/2)z² + βz + C
        ↓ α < 0(向下抛物线)
p(z) ∝ exp(-(z-μ)²/(2σ²))  →  高斯分布!

惊人的对偶反转——LeJEPA 完全颠倒了 ICA 的结论:

场景 高斯分布 非高斯分布
线性 ICA 失败(旋转不可区分) 成功
LeJEPA(非线性) 成功 失败
  • ICA 失败的原因:高斯分布的旋转不变性使得无法区分不同旋转方向
  • LeJEPA 成功的原因:正是这种旋转不变性,使得 OU 过程的 Hermite 谱分解恰好给出线性最优解

直觉对比:

分布 得分函数 第一特征函数 可识别性
高斯 (\exp(-z^2/2)) (-z)(线性) (He_1(z) = z)(仿射)
拉普拉斯 (\exp(-|z|)) (-\text{sign}(z))(阶跃) 非仿射
均匀分布 (0)(常数) 非仿射

定理 3:近似可识别性

当条件只近似满足时,恢复误差优雅降级

\mathbb{E}[\|h(z) - Qz\|^2] \leq D + (\varepsilon + D)^2

其中 (D = \delta / (2\rho(1-\rho)))(\delta) 为对齐间隙,(\varepsilon) 为白化误差。

两个误差参数的含义:

参数 定义 含义
(\delta)(对齐间隙) (L_{\text{align}}(h) - 2(1-\rho)n \geq 0) 正样本对有多"不相似"
(\varepsilon)(白化误差) (|\text{Cov}(h(z)) - I_n|_F) 嵌入分布有多"不高斯"

界的推导(简化版):

  1. 从 (\delta) 到非线性权重:(\sum_{i}\sum_{d\geq 2} w_{i,d} \leq \delta / (2\rho(1-\rho)) = D)
  2. 从非线性权重到恢复误差:(\mathbb{E}[|h(z) - Az|^2] \leq D)
  3. 从线性近似到正交矩阵(Procrustes):(|A - Q|_F \leq \varepsilon + D)
  4. 三角不等式组合:(\mathbb{E}[|h(z) - Qz|^2] \leq D + (\varepsilon + D)^2)

数值感受((\rho = 0.9)):

(\delta) (\varepsilon) (D) 界 (D + (\varepsilon+D)^2)
0 0 0 0(完美)
0.018 0 0.1 0.11
0.018 0.1 0.1 0.14
0.018 0.5 0.1 0.46

关键发现:

  • 对齐质量 (\delta) 是主要瓶颈(通过 (D) 线性传播)
  • 白化误差 (\varepsilon) 影响是二阶的(在平方项中)
  • 谱间隙 (2\rho(1-\rho)) 越小,对对齐误差越敏感

定理 4:最优潜空间规划

若 (h(z) = Qz),则在任意 O(n)-不变代价函数下,潜空间规划与真实世界规划完全等价

\hat{V}^*(h(z_0)) = V^*(z_0), \quad \hat{a}^*_{1:T}(h(z_0)) = a^*_{1:T}(z_0)

O(n)-不变代价函数(\ell(Qz, a) = \ell(z, a)) 对所有 (Q \in O(n))。

覆盖的常见控制问题:

代价函数 形式 不变性
欧氏距离到目标 (|z - z_{\text{goal}}|^2)
LQR (z^T P z + a^T R a)(P = cI)
范数惩罚 (|z|^2)
目标到达 (\mathbb{1}[|z - z_{\text{goal}}| < r])

证明核心:

  1. 正交变换不改变代价:(\ell(Qz, a) = \ell(z, a))
  2. 动力学推前等价:(\hat{p}(\hat{z}'|\hat{z}, a) = p(Q^{-1}\hat{z}'|Q^{-1}\hat{z}, a))
  3. 总代价等价:(J(a_{1:T}; \hat{z}0) = J(a{1:T}; z_0))
  4. 最优动作序列等价:(\hat{a}^* = a^*)

世界模型的含义: 线性可识别性 = 可证明地学到了可用于最优规划的世界模型。


四、实验验证

实验 1:正向可识别性(验证定理 1

设置: 2D 潜变量,4 种非线性混合函数:

混合函数 公式 特点
spiral (g(z) = R(\pi|z|)z) 保测度旋转微分同胚
banana (x_0 = z_0, x_1 = z_1 + z_0^2) 抛物线弯曲
sinusoid (x_0 = z_0 + \sin(1.5 z_1)) 正弦剪切
nvp RealNVP 耦合层 可扩展到高维

结果: LeJEPA 在所有情况下恢复各向同性高斯结构(旋转等价)。

高维扩展(N = 2 → 1024):

N SIGReg (R^2) VICReg (R^2) InfoNCE (R^2)
2 0.999998 0.999996 0.950961
64 0.999966 0.999968 0.648496
256 0.999884 0.999889 0.696587
1024 0.999561 0.999582 0.720241

SIGReg 和 VICReg 在所有维度保持 (R^2 > 0.999)InfoNCE 在高维因固定核宽度退化。

实验 2:逆向验证(验证定理 2

扫描广义正态分布族 (p(z; \alpha) \propto \exp(-|z/\beta|^\alpha))

R²(h→z) 随 α 的变化:

α=0.5  ████░░░░░░░░░░░░░░░░  ~0.5(重尾,失败)
α=1.0  ██████░░░░░░░░░░░░░░  ~0.6(拉普拉斯,失败)
α=1.5  ████████░░░░░░░░░░░░  ~0.8(接近高斯,部分成功)
α=2.0  ████████████████████  ~1.0(高斯,完全成功!)
α=3.0  ████████░░░░░░░░░░░░  ~0.8(超高斯,部分失败)
α=5.0  ██████░░░░░░░░░░░░░░  ~0.6(接近均匀,失败)

(R^2) 在 (\alpha = 2)(高斯)处尖锐达到峰值,完美验证定理 2。

实验 3:近似界验证(验证定理 3

所有运行的实际误差均低于理论界 (D + (\varepsilon + D)^2),对齐损失是可识别性的最强预测指标。

实验 4:潜空间规划(验证定理 4

设置: DMC Reacher 环境(像素输入,2D 关节角度潜变量)。

数据类型 生成方式 分布 规划代价
OU 采样 (z' = \rho z + \sqrt{1-\rho^2}\eta) 各向同性高斯 ~1.0(与 oracle 无差异)
RL 轨迹 训练好的策略采样 非高斯、各向异性 ~1.5(显著偏高)
规划代价(越低越好,理想值=1):

Oracle(关节空间直线):  ████░░░░░░  ~1.0
OU 编码器:              ████░░░░░░  ~1.0(与 oracle 无统计显著差异)
轨迹编码器:             ██████░░░░  ~1.5(显著偏高)

三种方法的失效模式对比

方法 高斯约束强度 优势 失效场景
SIGReg 全分布(特征函数匹配) 对非高斯更鲁棒 高维时正交误差略增
VICReg 二阶矩(协方差白化) 与 SIGReg 性能相当 非高斯时下降更快
InfoNCE 隐式(核函数) 低维时好 高维核宽度不匹配 → 梯度消失

五、Lean 4 形式化验证

所有四大定理均在 Lean 4 定理证明器中机器验证(零 sorry),使用 Mathlib v4.28.0。

文件 内容 核心验证 状态
Hermite.lean 定理 1 Hermite 谱分解 + Mehler 公式 + 关键不等式
Uniqueness.lean 定理 2 Sturm-Liouville 特征方程 → 高斯唯一性
Approx.lean 定理 3 近似界装配 (D + (\varepsilon + D)^2)
Planning.lean 定理 4 代价等价 + 最优动作等价
Dirichlet.lean 附录 E Dirichlet 能量替代证明路径

Lean 4 验证的关键定理(示例):

-- 定理1核心:等号成立 ⟺ 纯线性
theorem equality_forces_degree_one ...
    (heq : ∑' d, sw.w d * ρ ^ d = ρ) :
     d, 2  d  sw.w d = 0

-- 定理2核心:双条件高斯唯一性
theorem gaussian_uniqueness (lc : LatentComponent) :
    (IsGaussianScore   仿射特征函数)
    
    ( 仿射特征函数  IsGaussianScore)

-- 定理3核心:近似界
theorem approximate_identifiability ... :
    total_error  δ / (2*ρ*(1-ρ)) + (ε + δ/(2*ρ*(1-ρ))) ^ 2

-- 定理4核心:规划等价
theorem planning_equivalence ... :
    totalCost cp E_hat a (Q z) = totalCost cp E a z

六、局限性与未来方向

6.1 当前局限

局限 说明
潜变量是否真的高斯? 中心极限定理支持宏观量趋向高斯,但无法从观测中验证
维度不匹配 ((m \neq n)) 编码器维度与真实潜变量维度不同时的行为未理论化
有限样本 定理 3 是总体层面结论,样本复杂度和训练动态未涉及
动作条件转移 本文只处理编码器侧,(\hat{p}(\hat{z}'

6.2 与 SFA 的对比

维度 Sprekeler et al. (2014) SFA 本文 LeJEPA
可识别性类 置换等价 正交等价
潜变量分布 任意独立 高斯(或 i.i.d.
转移结构 需要不同速率 需要各向同性
提取方式 顺序(贪心) 同时
函数空间 固定多项式核 学习(神经网络)
近似界 (D + (\varepsilon + D)^2)
实用算法 xSFA(脆弱,≤6 个潜变量) LeJEPA/SIGReg(可扩展)

七、论文的深层意义

四定理的完整逻辑闭环

定理1(正向):高斯世界 + LeJEPA → h(z) = Qz(线性可识别)
    ↕
定理2(逆向):高斯是唯一使可识别性成立的分布
    ↓
定理3(近似):条件近似满足时,误差有界且优雅降级
    ↓
定理4(应用):线性可识别 → 潜空间规划 = 真实世界规划

核心信息

LeJEPA 在高斯世界中可证明地学到了世界模型,且这个保证可以优雅降级到近似条件,并直接支持最优规划。这是 JEPA 框架从"经验上有效"到"数学上可证明"的关键一步。

对 WorldModel 项目的启示

  1. 探索策略的重要性:近似各向同性随机游走的探索策略能保持数据在理论覆盖范围内
  2. SIGReg 优于 VICReg:对非高斯潜变量更鲁棒,适合真实场景
  3. 对齐质量是关键瓶颈:训练中应优先减小对齐损失
  4. 线性可识别性 → 规划等价:为 PRISM 空间记忆架构中的潜空间规划提供理论保障

八、关键参考文献

论文 arXiv 说明
LeJEPA 原始论文 2511.08544 Balestriero & LeCun, 2025, 提出 SIGReg
本文 2605.26379 Klindt, LeCun & Balestriero, 2026, 可识别性理论
LeWorldModel 2603.19312 Maes et al., 2026, 像素到控制的端到端 JEPA
V-JEPA 2 2506.09985 Meta, 2025, 视频 JEPA
Causal-JEPA 2602.11389 Nam et al., 2026, 因果干预
VICReg 2105.04906 Bardes et al., 2021, 协方差正则化
SFA 可识别性 Sprekeler et al., JMLR 2014 慢特征分析的非线性盲源分离理论

相关资源