世界模型的自适应潜在容量 Adaptive Latent Capacity for World Models
让JEPA世界模型把预测信息集中到自适应选择的潜在前缀，用更少维度换来更高规划成功率。
前置知识
JEPA（联合嵌入预测架构）
JEPA 是一种自监督学习范式：不在像素空间重建观测，而是在表示空间中预测目标信号的嵌入。以 LeWM 为例，编码器 $\mathrm{enc}_\theta$ 把观测 $o_t$ 映射为潜在状态 $s_t$，预测器 $\mathrm{pred}_\phi$ 结合动作 $a_t$ 预测下一步嵌入 $\hat{s}_{t+1}$，用 $\|\hat{s}_{t+1}-s_{t+1}\|_2^2$ 训练。由于预测发生在抽象空间，模型可以忽略与动力学无关的像素细节，适合作为规划用的世界模型。
ALeWM 完全建立在 JEPA 类世界模型 LeWM 之上，理解编码器-预测器结构与嵌入空间预测损失，是读懂其方法设计和实验协议的前提。
防坍塌正则化与 SIGReg
嵌入预测损失存在平凡解：所有嵌入坍缩为常数。SIGReg（LeJEPA 提出）把嵌入投影到随机单位方向 $u$，要求投影后的一维分布匹配标准高斯，用 Epps–Pulley 特征函数统计量度量差异，从而保持各坐标方差非退化且各向同性。VICReg 等其他方案则约束逐坐标方差与非对角协方差。这些正则对所有坐标是对称的。
标准 SIGReg 对坐标对称对待、不产生任何排序，ALeWM 的 MixSIGReg 正是对它的改造；必须先理解原始机制才能看出改动点和动机。
嵌套丢弃（Nested Dropout）与 Matryoshka 表示
嵌套 dropout（Rippel et al., 2014）与 Matryoshka 表示学习（Kusupati et al., 2022）训练有序表示：训练时随机采样一个前缀长度 $k$，只用前 $k$ 个坐标计算损失，使任意前缀都独立可用，一个模型即可支持多个表示预算。早期坐标被采样得更频繁，自然承载更重要的信息。
ALeWM 的前缀掩码设计与嵌套 dropout 一脉相承，但其容量分布由网络按输入序列学习而非固定先验；对比这两者是理解本文创新点的关键。
Gumbel-Softmax 直通估计器
从类别分布 $q_\psi(k|w)$ 采样离散前缀长度时梯度无法回传。Gumbel-Softmax 用可微松弛 $\mathrm{softmax}((\log q+g)/\eta)$（$g$ 为 Gumbel 噪声）近似采样，前向使用硬 one-hot（本文 $\eta=0.5$），反向走软松弛的梯度，即直通（straight-through）估计。代价是梯度有偏。
容量选择是本文唯一的离散决策点，理解该估计器才能明白训练如何端到端进行，以及作者在结论中承认的梯度偏差局限。
MPC 与 CEM 潜在规划
世界模型规划采用模型预测控制：在潜在空间用预测器递推 $\hat{z}_{t+1}=\mathrm{pred}(\hat{z}_t,a_t)$，对候选动作序列计算终端潜在状态与目标嵌入的距离，用交叉熵方法（CEM）迭代优化动作序列，执行若干步后重规划。本文设置：300 个候选、30 轮精化、30 个精英，规划视野 5 个动作块。
ALeWM 的收益以规划成功率和规划容量 $\mathbb{E}[K_{plan}]$ 衡量，全部实验都是 CEM 目标条件控制；理解该协议才能正确解读所有结果表格。
特征函数与 Epps–Pulley 统计量
特征函数 $\varphi(\tau)=\mathbb{E}[e^{i\tau X}]$ 完全刻画一个概率分布。Epps–Pulley 检验在若干频率点 $\tau_j$ 上比较经验特征函数与参考分布特征函数的加权平方差来检验分布是否匹配，SIGReg 和 MixSIGReg 都以此作为可微损失。MixSIGReg 的混合参考特征函数形如 $\varphi_0(\tau)=\sum_k \pi_0(k)\exp(-\tau^2 u^\top M_k u/2)$。
MixSIGReg 的实现完全建立在特征函数匹配上，理解它才能看懂第 3.2 节的核心公式以及附录 C 中方差排序的严格证明。
研究动机
JEPA 世界模型（如 LeWM）的防坍塌目标把所有潜在坐标对称对待：SIGReg 要求嵌入整体逼近各向同性高斯，VICReg 约束逐坐标方差与协方差，二者都不给任何坐标赋予更高的预测地位，预测信息因此散布在整个表示中。这带来一个实际矛盾：作者在实验中观察到，增大嵌入维度常会降低规划性能——在受控振子玩具系统中，LeWM 在 $d=4$（与真实状态同维）时线性恢复 $R^2=0.980$，$d=6$、$d=8$ 时分别跌到 0.930 与 0.904；在 TwoRoom 上全宽 $d=192$ 的成功率为 88.17%，而 $d=32$ 反而有 95.00%。一种解释是：更宽的潜在让预测更难，而规划时预测器被递归调用，预测误差沿 rollout 传播并误导动作评估；直接缩窄维度又限制表示表达力，且需逐一调参、每个宽度训练一个模型。
本文的目标是本文目标是在不牺牲宽嵌入表达力的前提下，把对预测和规划最有用的信息集中到潜在向量的紧凑前缀（前若干个坐标）中，并让同一个训练好的模型支持多个潜在预算：不同任务、不同回合可自动选择不同的前缀长度，无需为每个宽度单独训练。具体要做到三点：(1) 前缀长度由一个容量网络按观测序列自动预测，训练中联合学习，部署后无需网格搜索；(2) 宽嵌入仍保留在编码器中，未被前缀使用的坐标可继续承载对其他下游任务有用的特征；(3) 在四个模拟视觉控制基准上，平均规划容量大幅下降的同时，成功率不低于甚至超过调优过维度的固定宽度 LeWM。
与已有工作不同的是，有序表示已有两条研究路线，但都没进入动作条件世界模型：嵌套 dropout 与 Matryoshka 学习用固定的、与输入无关的前缀采样分布训练有序表示，服务于重建与检索；Variational Nested Dropout 和 SparC-IB 虽学习输入相关的容量选择，但都依赖解码器在非时序监督任务（重建输入或预测标签）中训练。世界模型的独特挑战在于：前缀必须经受反复的动作条件递推预测和规划中的目标比较。同时，JEPA 的防坍塌正则（SIGReg/VICReg）与『后段坐标低方差、被掩码置零』的需求天然冲突。ALeWM 的切入点是把容量选择做成解码器无关的、序列条件的分布，并发明 MixSIGReg 混合参考分布取代对称高斯正则，从训练目标层面同时施加信息前置的排序压力。
核心方法
直觉上，ALeWM 把宽嵌入当作『分档可用』的资源：只读前 $k$ 个坐标也应能预测出完整的下一步嵌入，这样前缀就必须装下对预测最有价值的信息。技术路线上，它在 LeWM（预测损失 + SIGReg）之上加两个组件（图 1 橙色部分）：一是容量网络 $q_\psi(k|w)$，输入一段观测序列（训练时 4 帧，部署时初始帧与目标帧），输出在合法容量集合 $\mathcal{K}$ 上的分布；训练时按该分布用直通 Gumbel-Softmax 采样前缀长度 $k$，用掩码 $M_k$ 把 $s_t$ 截断为 $z_t(k)$，预测器从 $z_t(k)$ 和动作预测完整的下一嵌入 $s_{t+1}$；二是 MixSIGReg，把被掩码嵌入的正则参考从各向同性高斯换成『高斯前缀 + 零后缀』的先验加权混合分布。总损失为 $\mathcal{L}=\mathcal{L}_{pred}+\lambda\mathcal{L}_{MixSIG}$。部署时按初始与目标观测取众数容量并在整回合固定，用 CEM 在前缀空间做递推规划和目标比较。
核心创新是两股互补的排序压力。第一股是嵌套预测：预测器的输入是被采样的短前缀，预测目标却是完整的下一嵌入；由于早坐标被激活的频率不低于晚坐标，模型被逼着把预测有用的信息往前放。第二股是 MixSIGReg：其参考分布 $p_0(z_t)=\sum_{k\in\mathcal{K}}\pi_0(k;\alpha)\,\mathcal{N}(z_{1:k};0,I_k)\otimes\delta_0(z_{k+1:d})$，其中 $\pi_0(k_c;\alpha)\propto(C-c+1)^\alpha$（$\alpha>0$ 偏好小容量）；第 $j$ 个坐标的参考方差恰为存活概率 $\rho_j^\pi=\Pr(K^\pi\ge j)$，天然单调递减。理论上（命题 4），固定表示下的期望预言风险为 $\mathbb{E}[R_{K^\pi}]=R_0-\sum_c \rho_{k_c}^\pi\Delta_c$，其中 $\Delta_c=R_{k_{c-1}}-R_{k_c}$ 是各容量块的预测增益：把大增益的块放前面会被更多容量档『看到』，因此按增益非增序排列时期望预测误差最小。与已有方法的本质区别：标准 SIGReg 对坐标对称，嵌套 dropout 用固定先验，而这里是学习到的、输入条件的容量分布加混合参考分布的双重作用。
方法步骤详情
训练分五步。①编码：$s_t=\mathrm{enc}_\theta(o_t)$，4 帧窗口（相邻帧隔 5 个环境步），动作按 5 个一组扁平化。②容量预测：置换不变的 transformer 容量网络 $q_\psi$ 读入整段帧，输出容量集 $\mathcal{K}$ 上的分布（如 $d_{max}=192$ 时 $\mathcal{K}=\{8,16,\dots,192\}$ 八档），初始化为多项式先验 $\pi_0(k_c)\propto(C-c+1)^\alpha$。③前缀采样：温度 0.5 的直通 Gumbel-Softmax 采样 $k$，掩码得 $z_t(k)=M_k s_t$，后段坐标置零。④预测与正则：transformer 预测器从 $z_t(k)$ 与动作块预测完整 $s_{t+1}$，损失 $\mathcal{L}_{pred}$ 为预测嵌入与编码器下一嵌入的平方误差；MixSIGReg 把嵌入投影到 1024 个随机方向，用 Epps–Pulley 统计量匹配混合参考特征函数 $\varphi_0^{(p)}(\tau)=\sum_k\pi_0(k)\exp(-\tau^2 u^{(p)\top}M_k u^{(p)}/2)$；总损失 $\mathcal{L}=\mathcal{L}_{pred}+\lambda\mathcal{L}_{MixSIG}$，AdamW 端到端训练 10 个 epoch。⑤部署规划：$q_\psi$ 只看初始帧与目标帧，取众数容量 $k_r$ 整回合固定，递推 $\hat{z}_{t+1}=M_{k_r}\mathrm{pred}(\hat{z}_t,a_t)$，目标代价只算活跃维 $C=\frac{1}{k_r}\|\hat{z}_H-M_{k_r}s_g\|_2^2$，CEM（300 候选/30 轮/30 精英）优化动作序列。
技术新颖性
与相邻工作可从四个维度区分。表示结构上：并行的 LpWM 把固定宽度的码稀疏化，ALeWM 定义的是嵌套前缀，宽度本身可变。容量选择上：嵌套 dropout 与 Matryoshka 用固定、输入无关的采样分布；Variational Nested Dropout 与 SparC-IB 虽学习输入条件容量但依赖解码器且非时序任务；ALeWM 是解码器无关、序列条件、经动作条件预测训练的容量分布，用直通 Gumbel-Softmax 端到端学习。正则目标上：Sub-JEPA 正则化固定随机子空间，Rectified LpJEPA 与 LpWM 用稀疏或广义高斯目标，MixSIGReg 首次用『先验加权的高斯前缀 + 零后缀混合』作为参考并在特征函数层面匹配，解析上给出严格递减的方差轮廓 $\mathrm{Var}(z_j^0)=\rho_j^\pi$，且总方差预算恰为 $\mathbb{E}[K^\pi]$（命题 1）。理论上：命题 4 把『信息前置』与期望预测误差最小化定量联系起来，是 JEPA 世界模型中少见的坐标排序论证。另外，部署为『整回合固定、按初始+目标帧选档』也区别于逐帧自适应方案。
实验结果
玩具实验（受控双阻尼振子，真实状态 4 维、观测 10 维）验证信息集中：8 维 ALeWM 前 4 个坐标的线性恢复 $R^2=0.983$，与维度匹配的 LeWM（$d=4$，0.980）相当，显著超过更宽的 LeWM（$d=6$ 时 0.930、$d=8$ 时 0.904）；被掩嵌入的实测方差与先验存活概率轮廓吻合，白化 Procrustes 对齐显示前 4 个坐标经线性变换即可对上 4 个物理因子。主实验（TwoRoom、PushT、Reacher、OGBench-Cube 四个目标条件控制基准，3 个种子，各 200 个测试回合）：相对全宽 LeWM，ALeWM 成功率提升约 1.5–17 个百分点、平均规划容量降低约 67–98%（图 1：TwoRoom +11.8pp/−184 维，PushT +3.3pp/−138.5 维，Reacher +2.2pp/−143.4 维，Cube +10.3pp/−176 维）；相对逐维度调优的 best-d LeWM 仍领先 0.3–7.2pp，12 组对比中 11 组使用的维度更少，最大增益在 OGBench-Cube（79.00% vs 72.33%）。消融实验：MixSIGReg 相比 SIGReg 在 PushT 提升 2.50pp、在 Cube 提升 8.17pp，平均前缀从 138.67 降到 53.49、从 128.00 降到 16.00；学习容量（ALeWM）优于固定先验（嵌套 dropout）：PushT 96.00% vs 92.50%，Cube 79.00% vs 72.83%；部署时用活跃前缀算目标代价优于全宽目标（Cube 79.00% vs 75.17%）；对 8 档固定前缀逐档评估显示没有任何单一档位全面最优，印证了自适应选择的价值。
![LeWM and ALeWM mean test success rate (± SEM) and mean planning prefix E[Kplan].](/papers/2609.32921/images/table_1.jpg)
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 受控双阻尼振子（玩具系统，线性状态恢复） | 4 因子平均线性探针 $R^2$ | 0.983（8 维 ALeWM 的前 4 维前缀） | LeWM：$d=4$ 为 0.980，$d=6$ 为 0.930，$d=8$ 为 0.904 | 在 $d=8$ 设定下比 LeWM 高 7.9pp，用一半维度追平维度匹配模型 |
| TwoRoom 目标条件控制（ViT-Tiny, $d_{max}=192$） | 测试成功率 SR / 平均规划容量 $\mathbb{E}[K_{plan}]$ | 100.0% / 8 维 | 全宽 LeWM：88.17% / 192 维；调优 best-d LeWM：95.00% / 32 维 | 较全宽 +11.8pp 且容量降 96%；较调优 +5.0pp 且容量降 75% |
| PushT 目标条件控制（ViT-Tiny, $d_{max}=192$） | 测试成功率 SR / $\mathbb{E}[K_{plan}]$ | 96.00% / 53.49 维 | 全宽 LeWM：92.67% / 192 维；best-d LeWM：92.67% / 160 维 | +3.3pp，容量降约 67–72% |
| Reacher 目标条件控制（ViT-Tiny, $d_{max}=192$） | 测试成功率 SR / $\mathbb{E}[K_{plan}]$ | 85.83% / 48.59 维 | 全宽 LeWM：83.67% / 192 维；best-d LeWM：84.50% / 160 维 | +2.2pp，容量降约 70–75% |
| OGBench-Cube 目标条件控制（ViT-Tiny, $d_{max}=192$） | 测试成功率 SR / $\mathbb{E}[K_{plan}]$ | 79.00% / 16 维 | 全宽 LeWM：68.67% / 192 维；best-d LeWM：72.33% / 96 维 | 较全宽 +10.3pp（容量降 92%）；较调优 +6.7pp（容量降 83%），为四任务最大增益 |
| 四任务平均（三种骨架 12 组对比） | 相对全宽 LeWM 的平均成功率与容量 | 成功率 +1.5~17pp，规划容量降 67~98% | 全宽 LeWM / 调优 best-d LeWM | 对调优基线仍 +0.3~7.2pp，12 组中 11 组维度更少 |
局限与改进
作者承认的局限有三：容量网络与世界模型联合训练，行为随初始化等随机源波动（种子间所选容量可能不同，如 Reacher 上 seed 3074 全部选 $K=64$ 而另两个种子混合 32/64）；MixSIGReg 先验依赖容量秩上的多项式次数 $\alpha$，且方差只在容量边界跳变、块内恒定，更平滑或考虑容量数值间距的替代形式可能更好；直通 Gumbel-Softmax 的梯度有偏。此外玩具实验只支持前缀集中这一结论，不保证逐坐标解耦或真实因子数的恢复；理论结果是对排序压力的解释而非联合训练的收敛保证。我自己的观察：实验限于 2D/3D 仿真与小骨架（ViT-Tiny/Small），没有真实机器人或视频尺度的验证；每个数据集都要单独搜索 $\lambda$、$\alpha$、选择器学习率倍率（表 3），超参较敏感；环境回合间差异小，导致多数任务上每回合选中的容量几乎恒定（表 4），自适应逐回合选择的收益只在 Reacher 等任务部分体现；混合数据集训练时 TwoRoom 成功率从 99.50% 跌到 36.17%，多任务容量分配远未解决。
独立分析的弱点
弱点一：容量选择的收益对任务同质性敏感。四个基准都是低变化的受控环境，测试回合选中的容量几乎一致（表 4），真正需要『按回合调容量』的场景未被充分检验。改进方向：构造回合间难度差异大的数据集（目标距离、遮挡、干扰物数量变化），检验 $q_\psi$ 的逐回合分辨力。弱点二：部署时容量在回合开始即固定，无法在闭环中随新观测更新，长时程任务的信息需求可能漂移；可改为滚动重估容量或在线更新容量后验。弱点三：训练与部署存在上下文偏移——训练时容量网络看 4 帧，部署时只看初始帧和目标帧，作者自己也承认未建立这种偏移下的校准；改进可在训练时随机化可用上下文长度。弱点四：混合数据集训练性能塌方（DMC 保持 85.17% 但 TwoRoom 跌至 36.17%），说明共享容量网络的优化被不同任务的容量需求拉扯；可引入数据集条件化先验或专家容量头。弱点五：超参（$\lambda$、$\alpha$、$s_q$）需逐数据集网格搜索（表 3），削弱了『免调参选容量』的卖点；可研究从数据自动确定 $\alpha$（如让先验与经验容量分布匹配的自适应目标）。
未来方向
作者明确提出的方向：探索无偏的离散采样梯度估计替代直通 Gumbel-Softmax；设计考虑容量数值间距、避免容量边界严格方差跳变的先验形式，可能进一步推进表示可辨识性。基于本文成果可延伸的：把 MixSIGReg 的混合参考思想推广到 token 级表示——视觉 tokenizer（FlexTok、ElasticTok 一类可变长 token 化）和视觉-语言-动作模型的动作 token 排序（并行工作 Liu et al. 已在 VLA 中用固定先验的嵌套 dropout，可与学习容量结合）；利用未进入前缀的后段坐标承载其他下游任务特征，实现一个编码器多任务复用；闭环/滚动容量选择与不确定性感知的容量后验；把容量层加装到 DINO-WM、V-JEPA 2-AC 等基于预训练特征的世界模型上检验普适性；理论扩展方面，把命题 4 的预言风险排序分析推广到学习型选择器（式 35 中 $\delta_c$ 可为负的情形）与非高斯动力学；最后是在真实机器人和长视频上评估算力-精度权衡。
复现评估
复现条件较好。项目页为 arm-research.github.io/AAIR-ALeWM（正文未明确声明代码开源，需到项目页确认）。附录给出了罕见级别的实现细节：玩具系统的完整动力学参数（$\Delta t=0.2$、$\gamma=0.25$、$\kappa=1.03887957$、过程与观测噪声方差等）、四个基准的回合不重叠数据划分（各 100 验证 / 200 测试回合，PushT 共 18685 回合）、网络结构（ViT-Tiny/14 或 ViT-Small/14 编码器、6 层 transformer 预测器、4 层容量网络，容量网络参数量 1.78M–4.74M）、优化配置（10 epoch、batch 128、AdamW lr $5\times10^{-5}$、bf16、梯度裁剪 1.0）、SIGReg/MixSIGReg 的投影数与积分节点设置，以及表 3 的逐配置超参和表 4 的逐种子容量统计；所有主结果为 3 个种子并报告标准误。算力需求温和：训练小 ViT 单机即可，规划用 CEM 300 候选、批评估。主要门槛在于 JEPA 端到端训练的稳定性技巧与特征函数积分的实现细节；对有世界模型经验的团队，复现难度评估为中等。
论文图表
对比固定先验的嵌套 dropout（分别配 SIGReg 与 MixSIGReg，qψ=π0）与完整 ALeWM（学习容量）在各容量档的测试成功率。PushT 上嵌套 dropout 最佳为 92.50%（96 维，MixSIGReg），ALeWM 达 96.00% 且平均只用 53.49 维；Cube 上嵌套 dropout 各档在 65.83–72.83% 之间，ALeWM 达 79.00% 且只用 16 维。
证明『学习容量分布』而非沿用固定先验是性能来源之一，是对方法设计必要性的关键消融。
同一 ALeWM 框架下，被掩码嵌入分别用 MixSIGReg 与标准 SIGReg 正则（ViT-Tiny，dmax=192，取各任务最佳配置）：MixSIGReg 在 PushT 提升 2.50pp、在 Cube 提升 8.17pp，同时平均选中前缀从 138.67 降至 53.49、从 128.00 降至 16.00。
隔离验证 MixSIGReg 这一核心新正则的贡献——既提成功率又大幅压缩容量，是方法两大组件之一的价值证明。