JEPA4Japan · 教程

第 8 章——一次完整的前向传播:跟随一个批次从头走到尾

1,271字 4分钟阅读 #LeWorldModel#世界模型#JEPA

在最小示例中依次理解观测编码、动作条件化、目标构造、两项损失、梯度流、钩子和形状断言。

课程进度 课程大纲 已发布 48/48 课

第 0 部分——阅读指南:我们究竟要学什么?

  1. 01 第 0 章——开始之前 已发布

第 1 部分——世界模型:智能体内部的沙盒

  1. 02 第 1 章——智能体为什么需要“想象未来” 已发布
  2. 03 第 2 章——为什么不直接预测下一张图像? 已发布
  3. 04 第 3 章——JEPA 思想:预测意义,而不是图像的复制品 已发布
  4. 05 第 4 章——用一张图看懂 LeWM 已发布

第 2 部分——把图像变成状态:LeWM 架构

  1. 06 第 5 章——轨迹数据:对模型而言,世界不是图像集合 已发布
  2. 07 第 6 章——视觉编码器:为每一帧签发“状态护照” 已发布
  3. 08 第 7 章——动力学预测器:在头脑中推动时间前进 已发布
  4. 09 第 8 章——一次完整的前向传播:跟随一个批次从头走到尾 当前课程

第 3 部分——防止模型作弊:预测损失与 SIGReg

  1. 10 第 9 章——最危险的捷径:表示崩塌 已发布
  2. 11 第 10 章——预测损失:模型怎样学习下一步 已发布
  3. 12 第 11 章——SIGReg 的直觉:让表示空间“呼吸” 已发布
  4. 13 第 12 章——只保留最少却足够的数学 已发布
  5. 14 第 13 章——原始 LeWM 的端到端训练机制 已发布
  6. 15 第 14 章——训练一个不会立即崩塌的模型 已发布

第 4 部分——让模型行动起来:在潜在空间中规划

  1. 16 第 15 章——目标条件规划:从“我在哪里?”到“我要去哪里?” 已发布
  2. 17 第 16 章——潜在欧氏距离:方便,但未必可靠 已发布
  3. 18 第 17 章——CEM:用淘汰赛搜索动作 已发布
  4. 19 第 18 章——MPC:不要一次信任模型太久 已发布
  5. 20 第 19 章——长视界滚动:小误差怎样滚成大失败 已发布
  6. 21 第 20 章——从零实现一个最小 LeWM 规划器 已发布

第 5 部分——工程复现:从论文到可运行系统

  1. 22 第 21 章——官方代码仓库与实验环境 已发布
  2. 23 第 22 章——第一个实验:TwoRoom 冒烟测试 已发布
  3. 24 第 23 章——第二个实验:复现 PushT 已发布
  4. 25 第 24 章——怎样公平评估世界模型 已发布
  5. 26 第 25 章——失败诊断手册 已发布

第 6 部分——LeWM 究竟学到了什么?

  1. 27 第 26 章——线性探针:潜在状态编码了哪些物理变量? 已发布
  2. 28 第 27 章——为潜在空间做一次“健康检查” 已发布
  3. 29 第 28 章——预期违背:模型会对“不可能事件”感到惊讶吗? 已发布
  4. 30 第 29 章——怎样严谨讨论“理解世界” 已发布

第 7 部分——为什么“预测准确”仍可能“规划很差”

  1. 31 第 30 章——训练目标与规划目标之间的差距 已发布
  2. 32 第 31 章——全局不崩塌不保证保留任务相关动力学 已发布
  3. 33 第 32 章——各向同性高斯先验何时过强? 已发布
  4. 34 第 33 章——长视界规划:预测得更远,还是规划得更聪明? 已发布
  5. 35 第 34 章——从位置距离到任务进度 已发布
  6. 36 第 35 章——多任务学习、真实机器人与视觉干扰 已发布
  7. 37 第 36 章——理论边界:何时能辨识真实状态? 已发布

第 8 部分——从复现者到研究者

  1. 38 第 37 章——设计一项可信的 LeWM 改进实验 已发布
  2. 39 第 38 章——十二个可执行的研究项目 已发布
  3. 40 第 39 章——LeWM 研究中的开放问题 已发布

附录

  1. 41 附录 A——最低限度的必需数学工具箱 已发布
  2. 42 附录 B——PyTorch 实现速查 已发布
  3. 43 附录 C——完整张量形状表 已发布
  4. 44 附录 D——实验配置卡 已发布
  5. 45 附录 E——论文时间线与证据等级 已发布
  6. 46 附录 F——术语表 已发布
  7. 47 附录 G——复现检查清单 已发布
  8. 48 附录 H——专家审阅检查清单 已发布

先看全景

  1. 编码四个z0 z1 z2 z3
  2. 对齐三个a0 a1 a2
  3. 预测三个p1 p2 p3
  4. 拉合目标p1↔z1, p2↔z2, p3↔z3
形状正确还不够。每个时间索引和梯度路径都必须讲述正确的故事。

四个观测潜在状态构成三组移位后的预测—目标样本对。

一个小故事

工厂可以生产尺寸正确、内部标签却完全错误的箱子。LeWM 的 batch 也可能这样。三个预测和三个目标在形状上可以完全匹配,即使每个预测都与当前状态而不是下一状态进行比较,只训练了最后一组样本,或目标已被 detach。

本章作出一个简单承诺:确保来源可追踪。对于每个位置,我们都能指出对应的观测、动作块、预测和目标。

技术背包

冻结默认值让整个前向传播变得具体:batch B=128、history_size=3、num_preds=1、潜在宽度 D=192,以及 frameskip=5。

pixels                 [B,4,C,H,W]
flattened pixels       [B×4,C,H,W]
observation latents    [B,4,192]
raw action blocks      [B,4,5×A]
encoded actions        [B,4,192]
state/action context   [B,3,192]
connected targets      [B,3,192]
predictions            [B,3,192]
SIGReg input           [4,B,192]
losses                 scalar

论文报告 TwoRoom 的 history 为一,PushT/OGBench-Cube 为三;冻结的全局默认值为三。这些形状描述的是冻结默认配置,而不是所有 LeWM。

最短而又忠实的前向传播卡片如下:

z = encode_each_frame(images)              # [B,4,D]
u = encode_each_action_block(actions)       # [B,4,D]
z_pred = causal_predict(z[:,:3], u[:,:3])  # [B,3,D]
z_next = z[:,1:]                            # [B,3,D], do not detach
L_pred = mean_square_gap(z_pred, z_next)
L_sig = sigreg(time_first(z))               # [4,B,D]
L_total = L_pred + lambda * L_sig

全部三组移位后的单步样本对都参与计算。num_preds=1 表示目标偏移一个位置,并不是只有最后一个输出。在这些数据配置中,一个模型步跨越五个环境动作。

预测分支会直接训练视觉编码器、编码器投影器、动作编码器、因果预测器和预测投影器。因为 z_next 保持连接,预测梯度还会经过共享视觉编码器在目标侧的使用路径返回。SIGReg 只直接训练视觉编码器及其投影器。一个优化器更新整套连通的学习系统;不存在目标编码器优化器或 EMA 更新。

SIGReg 在 batch 中查看每一个相对时间位置。不能把 B×T 展平为一个群体,因为时间变化与样本间多样性不是一回事。冻结代码使用 1,024 个随机投影和从 0 到 3 的 17 个频率节点。论文附录采用了不同的求积表述。论文方法的权重是 0.1,冻结 YAML 则是 0.09。务必保留每个数值的来源标签。

容易骗过我们的把戏

故意构造一条错误分支:

wrong targets: [z0,z1,z2]
right targets: [z1,z2,z3]

二者都是 [B,3,D]。错误分支可以学会复制当前状态,并报告令人满意的损失。使用明显的时间校验和,并打印 prediction_slot / last_visible_context / target_source。正确的三行分别是 0/0/1、1/1/2、2/2/3。

然后测试计算图。在一次性探针中,只对正确的目标切片执行 detach,确认目标侧梯度路径消失,而上下文路径仍然存在。仅仅发现编码器有某种梯度还不够。

五项廉价检查可以捕获绝大多数管线错误:精确形状、移位来源、因果未来突变、所有预期模块中有限且非空的梯度,以及没有 NaN 或无穷值。通过这些检查只能证明机械结构正确,并不能证明动力学或规划有用。

证据凭据

LeWorldModel v3、冻结的 train.py、jepa.py 和 module.py,支持这里所述的冻结对齐方式和连通计算图。论文展示的训练伪代码中,F.mse_loss 调用格式有误,因此在精确索引问题上,应以可执行的 train.py 为准。

不要把下一状态目标称作规划目标,不要出于 JEPA 的习惯将其 detach,不要在 SIGReg 中混合 batch 与时间,也不要把形状匹配当作正确性的证明。

快速检查

  1. 四个观测会产生哪三组预测—目标样本对?
  2. 预测损失可以通过视觉编码器的哪两条使用路径传播?
  3. 为什么错误的时间目标可以通过所有普通形状断言?