JEPA4Japan · 教程

附录 B——PyTorch 实现速查

1,964字 5分钟阅读 #LeWorldModel#世界模型#JEPA

快速查找数据加载、Transformer 输入、因果掩码、AdaLN、SIGReg、向量化滚动、CEM、混合精度和检查点的精炼模式。

课程进度 课程大纲 已发布 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. 保持回合完整先构造有效窗口
  2. 为每条轴命名带着含义变形
  3. 测试每条路径掩码、动作、梯度
  4. 保存整段航程状态加凭据
PyTorch 负责搬运张量;我们的工作是让张量讲的故事始终完整。

先构造尊重回合边界的窗口,再由加载器组成批次。

一个小故事

把 PyTorch 想成一套速度极快的行李系统。无论我们交给它什么,它都会照样分批、变形、广播和相乘。它不会知道两个相邻帧其实属于不同回合,不会知道一个 H 表示图像高度而另一个 H 表示规划视界,也不会知道某个目标本应继续连接到共享编码器。

所以,本附录是一张驾驶舱检查卡,而不是官方实现的逐行替代品。它只问一个简单问题:每次快速移动张量之后,哪些不变量仍应成立?

技术背包

数据与形状契约

Dataset 回答“第 417 个样本是什么?”;DataLoader 则调度这些样本并组成批次。数据集或采样器必须保留回合归属、时间顺序、动作对齐和窗口长度。加载器可以打乱完整窗口,但绝不能把一个窗口内不同回合的帧混在一起,凭空制造转移。

存储格式不等于科学身份。HDF5、Lance、文件夹或视频可以承载同一套轨迹契约;两个 HDF5 文件也可能在预处理、划分、跳帧或动作对齐上不同。容器格式和数据契约都要记录。

在调用 Transformer 前,先准备一张口头登机牌:

batch = independent examples
time = ordered latent positions
feature = learned coordinates
visibility = this position may read only itself and earlier positions

尺寸相等不代表角色相等。如果候选数和时间长度恰好都是 4,交换两条轴后打印出来的形状仍可能不变。做一次可丢弃的变形测试,让易辨认的回合、时间和候选标记一路穿过它。

掩码与动作条件化

因果掩码会阻止当前位置直接看到之后的 token。修改一个较晚的上下文 token,验证较早位置的输出保持不变。通过这项测试,证明的是计算可见性规则,而不是因果物理。

AdaLN 通过平移、缩放和门控调制,为动作提供一条进入 Transformer 块的路径。在冻结实现中,它最终的调制生成器通过零初始化从中性状态开始。这并不保证训练后的模型真的使用动作。固定潜在历史,交换会造成不同后果的动作,再检查预测。

因果帘幕与 AdaLN 动作面板承担不同工作。

原始 LeWM v3 使用一个共享且可训练的视觉编码器。移位后的下一观测目标仍与计算图连接:没有停止梯度、EMA 教师或预训练冻结编码器。一个将目标分离的重新实现会拥有相同的前向形状,却是另一套学习系统。

SIGReg 参考路径

把冻结的 SIGReg 实现读成五个房间:

collect projected observation embeddings
keep time positions separate; inspect each one across the batch
sample random one-dimensional directions
compare empirical fingerprints with the analytic standard-Gaussian fingerprint
average discrepancies into a differentiable training penalty

把批次和时间展平会改变总体问题。跨设备聚合会改变经验总体;而冻结的 LeWM 类明确是单 GPU 的。分离嵌入会改变梯度。把结果称作 p 值会改变它的用途。有限的方向、小批次和积分点意味着这只是近似,不是高斯性证书。记录这项惩罚时,还要同时记录坐标离散度、协方差谱、有效秩和留出近邻。

滚动与 CEM 路径

可以把候选计划堆叠起来,在 GPU 上并排评估。每个候选内部的时间仍是自回归的:第二步使用第一步预测出的状态。环境批次、候选、模型视界与潜在宽度必须保持区分。

sample sequences from the current action proposal
roll out candidates with the frozen world model
score terminal latents against the encoded goal
select the configured elite action sequences
refit only the action proposal, then repeat

精英索引必须仍指向真正获得对应分数的动作序列。CEM 改变的是提议分布,而不是编码器或预测器权重。连续且有界的动作需要明确声明裁剪/更新语义;离散动作则需要类别型提议分布。向量化节省执行时间,却不会消除时间依赖或模型漂移。

精度、检查点与可复现性

混合精度改变算术的执行方式,不改变预测损失或 SIGReg 的科学定义。监控非有限值,并与一次短暂的全精度冒烟运行比较。这项对照不必逐比特相同,但应能暴露不合理的尺度或脆弱的累积路径。

权重文件支持推理。用于恢复训练的检查点还应标识优化器、训练位置、组合后的配置,以及任何持续演化的调度器或梯度缩放器状态。如果连续性需要随机数生成器状态,就保存它;若没有,也要说明。此外还要记录数据身份、代码修订、软硬件、精度、窗口、预算、检查点规则、规划器设置、目标、随机种子和评估回合。

same data and preprocessing
same model, optimizer, and training budget
same planner and evaluation budget
same episodes and goal rule
declare every intentional difference before reading the score

骗过我们的把戏

两个实验室使用相同的提交、配置、随机种子、检查点和 GPU,损失也完全一致。后来他们发现,两边的采样器都会把一个回合的最后一次观测,与下一个回合的第一个动作配在一起。这个错误得到了完美复现。

重复只确认一套过程,并不认证它的含义。要主动挑战回合边界、单步移位、因果可见性、动作敏感性、梯度接收者、候选身份,以及世界模型训练与规划之间的分离。反过来,两次正确的随机运行也可能在数值上不同,却支持同一个有边界的结论。

证据凭据

PyTorch 文档和 Hydra 文档定义框架抽象。LeWorldModel v3 以及冻结的官方代码 8edfeb336732b5f3ce7b8b210d0ba370a09e2cac定义本课程检查过的计算图与接口。stable-worldmodel只为它后来提供的主动学习平台构成证据,不能倒过来证明原始 v3。

允许的说法是:在已记录的配置下,这些模式保留了具名的数据、形状、条件化、滚动、优化、精度和恢复不变量。它们并不表明加载器理解回合、因果掩码发现物理规律、AdaLN 保证动作得到使用、混合精度没有伤害,或一个随机种子就证明可复现。本节并未建立独立复现。

快速检查

  1. 哪一层必须拒绝跨越回合边界的窗口?
  2. “修改较晚 token”测试能证明什么,又留下了什么未决问题?
  3. CEM 期间,哪个对象会变化,模型的哪些部分保持冻结?