方向 3 调研：Adaptive-Computation Diffusion World Model with Learnable Early Exit

研究目标 在 diffusion world model / world action model 中，让不同 token、不同 latent frame、不同去噪步使用不同计算量。 简单说：不是所有 future token 都完整去噪到最后，而是让模型判断“这个 token 当前已经足够用于下游动作/规划”，然后提前冻结。

总体判断：可行，但不能直接照搬图像 diffusion early exit。
这个方向有明确文献支撑：NoiseGate 已经把 per-latent timestep 解释成信息门控；Diffusion Forcing 提供 per-token 独立噪声水平的训练范式；ACT/PonderNet 提供可学习 halting 机制；DeepCache、TokenCache、ASE、AdaDiff 等说明 diffusion 推理中确实存在跨步冗余。真正的创新点应该放在 world model / WAM 场景下的 token-level runtime halting，尤其是“冻结 token 如何继续参与 attention 与动作决策”。

1. 相关工作脉络

方向 代表工作 核心结论 对本方向的启发
Per-latent / per-token timestep Diffusion Forcing；NoiseGate 序列中每个 token/latent 可以有独立噪声水平；NoiseGate 进一步把 WAM 中每个 predicted latent frame 的 timestep 作为信息门控策略。 说明“不同未来帧/latent 不必同等去噪”是合理的。
Diffusion early exit ASE；AdaDiff 不同 denoising step 的 score estimation 难度不同，可以用更少计算完成部分步骤。 说明 diffusion 的迭代过程可以做自适应计算，而不是固定 N 步。
Feature / token caching DeepCache；TokenCache 相邻去噪步存在特征冗余；可以缓存复用 feature 或 token 级计算。 证明“后期/稳定 token 重算浪费”普遍存在。
Adaptive computation / halting ACT；PonderNet 网络可以学习何时停止计算，通过 halting probability 或 ponder cost 在精度和计算量之间折中。 给本方向的 early-exit head、halting loss、计算惩罚提供方法基础。
机器人 / WAM benchmark RoboTwin 2.0；RoboTwin manipulation tasks 提供可扩展、多任务、强 domain randomization 的双臂操作 benchmark。 适合作为动作成功率、轨迹质量、延迟收益的验证平台。

2. 和 NoiseGate 的真正区别

NoiseGate 的核心是：对每个 predicted latent frame 学习不同的 timestep / noise schedule，让不同 future frame 对 action token 的 Key/Value 贡献具有不同可靠性。 它更像是 per-latent schedule policy。

方向 3 要做的是更细粒度的 runtime early exit：

比较项 NoiseGate 方向 3：Token-level Learnable Early Exit
粒度 主要是 per latent frame / per latent timestep per token × per denoising step
决策时机 学习每帧时间步调度 推理过程中实时判断是否停止更新某个 token
目标 调节 future latent 对动作生成的信息可靠性 减少无效 denoising FLOPs，同时维持 action success / rollout quality
输出行为 不同 latent 使用不同噪声水平 部分 token 被冻结，不再走后续 denoising block，但仍可作为上下文参与 attention
论文卖点 Information gating for WAM Runtime token-wise adaptive computation for diffusion world models
注意：不要把方向 3 写成“NoiseGate 的简单细化”。真正要强调的是：NoiseGate 是 schedule learning，本方向是 runtime halting / token freezing。前者改变噪声水平，后者改变实际计算路径。

3. 方法设计：Token-wise Halting Diffusion World Model

3.1 基本设定

设 diffusion world model 在 latent token 序列上去噪： 

z_t^k = Denoiser(z_t^{k+1}, c, k),   k = K, K-1, ..., 1

其中 z 是 future latent tokens，c 是历史观察、动作、语言指令或状态上下文。 传统方法所有 token 都完整经过 K 个去噪步。本方向引入 token-level halting mask： 

h_i^k = HaltHead(z_i^k, context, k) ∈ [0, 1]
m_i^k = 1[h_i^k > τ]

当 token i 在 step k 满足退出条件后，其 latent 被冻结： 

if m_i^k = 1:
    z_i^{k-1} = stop_gradient(z_i^k)
else:
    z_i^{k-1} = DenoiserUpdate(z_i^k)

3.2 关键机制：Freeze but Attend

逐 token early exit 最大的坑是 attention。一个 token 停止更新后，不能简单从序列里删除，因为其他 token 和 action token 仍然需要它作为上下文。 所以推荐机制是：

Freeze but Attend：
已退出 token 不再更新 Q / FFN 分支，减少自身计算；但它的 K/V 表示保留在 attention memory 中，供未退出 token 和 action token 读取。这样既节省 token 更新计算，又不破坏全局上下文。

更稳妥的实现分三档：

版本 实现 优点 风险
V0：Step-level early exit 整步提前停止：若所有 token 置信度达标，则结束剩余去噪 最容易实现，适合第一版实验 创新性一般，粒度不够细
V1：Frame-level early exit 某些 future frame 的所有 token 退出，其他 frame 继续去噪 和 NoiseGate 对齐，风险较低 仍不是真正 token-level
V2：Token-level early exit 每个 spatial/latent token 独立退出；冻结 token 继续提供 K/V 创新性最强，最符合方向 3 实现复杂，容易出现 attention mismatch

3.3 Halting 信号怎么设计

不建议一开始就用复杂 RL。第一版可以做监督/自监督 proxy，降低训练难度。

信号 定义 优点 缺点
Residual stability ||z_i^k - z_i^{k+1}|| 很小则认为稳定 训练-free 可做 baseline 稳定不等于对动作有用
Predicted denoising error 预测当前 token 到最终 clean latent 的误差 和 diffusion 目标直接相关 仍偏生成质量，不一定对应控制成功
Action sensitivity 冻结该 token 后 action logits / action trajectory 改变量 最贴合 WAM / robotics 计算成本较高，可离线生成标签
Value / reward sensitivity 冻结 token 对 predicted value / success score 的影响 适合 planning / RL 需要额外 critic 或 reward model
Learned halting probability 类似 PonderNet，输出条件 halting probability 理论上最完整 训练不稳，需要 regularization

4. 推荐训练目标

4.1 基础 loss

L = L_diffusion + λ_action L_action + λ_halt L_compute + λ_cons L_consistency

4.2 计算惩罚

L_compute = (1 / NTK) Σ_i Σ_k active_i^k

即统计 token 在多少去噪步中仍然被更新。越早退出，loss 越小。

4.3 Teacher-student 一致性

用完整 K 步 denoising 作为 teacher，early-exit model 作为 student： 

L_cons = ||a_exit - a_full||_2 + β ||z_exit - z_full||_2

其中 action consistency 比 latent consistency 更重要。因为本方向的目标不是生成最好看的未来，而是用更少计算得到足够好的动作/规划。

5. 可行性验证：先做三个小实验

实验 A：Token 收敛异质性分析

目的：证明不同 token 的去噪收敛速度不同。

做法：运行 full denoising，记录每个 token 在每步的 latent change、predicted x0 change、action sensitivity。

预期：背景 token / 静态区域更早稳定；接触物体、手部、工具附近 token 更晚稳定。

图：token convergence heatmap；step-wise active token ratio。

实验 B：Oracle early exit 上限

目的：验证理论加速空间。

做法：用 full denoising 的最终结果构造 oracle：若某 token 当前状态与最终状态误差小于阈值，则允许退出。

预期：若 oracle 可省 30%–60% token-step 且动作质量不掉，说明方向值得做。

图：compute saving vs. action error / video quality Pareto curve。

实验 C：Training-free baseline

目的：快速判断不用训练能否见到趋势。

做法：用 residual stability 或 x0 stability 作为退出准则。

预期：可能质量略掉，但应有一定计算节省。若完全崩，说明 attention/freeze 机制要先修。

图：threshold sweep；success rate vs. latency。

实验 D：Learned HaltHead

目的：验证可学习 early exit 是否优于 heuristic。

做法：训练轻量 MLP/Transformer head，输入 token feature、step embedding、frame index、action context，输出 halting probability。

预期：同等成功率下比 heuristic 更省计算，尤其在动态场景中更稳。

图：learned vs. heuristic Pareto；token exit distribution。

6. 建议实验矩阵

实验组 方法 作用
Baseline-1 Full denoising 精度上限，延迟下限对照
Baseline-2 Uniform fewer steps 证明简单减少 denoising step 会伤质量
Baseline-3 Frame-level schedule / NoiseGate-style 对比 per-frame schedule 与 per-token runtime exit
Baseline-4 Residual heuristic early exit 证明学习型 halting 的必要性
Ours-V1 Frame-level runtime early exit 低风险版本，先保证可跑通
Ours-V2 Token-level freeze-but-attend early exit 主方法
Ours-V3 Token-level + action-sensitivity supervision 增强版，强调 WAM/robotics 任务相关性

7. 指标设计

类别 指标 说明
效率 Average active token-step ratio 核心指标：实际参与更新的 token-step 占 full denoising 的比例
效率 Wall-clock latency / FPS / TPS 必须测真实端到端延迟，不能只报理论 FLOPs
控制 Task success rate 机器人任务最重要指标
控制 Action MSE / trajectory deviation 和 full denoising 或 expert action 对比
预测 Latent MSE / FVD / LPIPS / SSIM 用于证明 world model 没有完全崩
稳定性 Failure recovery rate 早退后若触发不确定性，能否回退到 full denoising
可解释性 Exit map / exit step histogram 展示哪些 token 早退，哪些 token 被保留计算

8. 推荐数据集 / 平台

平台 适合程度 建议用途
RoboTwin 2.0 高 主实验：多任务、双臂操作、domain randomization，适合验证 WAM action success。
LIBERO 高 VLA / manipulation 常用 benchmark，适合和 diffusion policy / VLA 加速方法比较。
DMControl / Meta-World 中 适合先做低成本 proof-of-concept，验证 rollout 加速。
Atari / Diamond-style diffusion world model 中 适合验证纯 world model rollout，但和 WAM/action-coupled 的卖点弱一些。

9. Ablation 设计

Ablation 要回答的问题 预期结论
w/o learned HaltHead 学习型退出是否必要？ heuristic 可以省计算，但成功率/稳定性更差。
w/o action-sensitivity loss 只看 latent 稳定是否足够？ 可能视频质量不差，但动作质量下降，说明 WAM 需要 action-aware halting。
Remove freeze-but-attend 冻结 token 是否仍需参与 attention？ 直接删除 token 会破坏上下文，性能明显下降。
Frame-level vs token-level 细粒度 token exit 是否带来额外收益？ token-level 在复杂场景更省计算，frame-level 更稳但粗糙。
Different compute penalty λ 速度-精度 trade-off 是否可控？ λ 越大越快，但过大导致成功率下降，形成 Pareto 曲线。
Static vs dynamic scenes 模型是否真的按场景难度自适应？ 静态/背景区域早退多，交互/遮挡/接触区域计算多。

10. 关键风险与解决方案

风险 1：token 早退不等于动作可靠。
token latent 稳定可能只是生成层面的稳定，但对动作决策仍然重要。解决方案：加入 action sensitivity / action consistency loss，而不是只用 latent MSE。
风险 2：逐 token 退出可能没有真实加速。
如果实现仍然把所有 token 送进 dense attention，理论 FLOPs 省了但 wall-clock 不降。解决方案：第一版可以报 active token-step，但最终必须实现 token compaction、block sparse attention 或至少 grouped token update。
风险 3：attention mismatch。
冻结 token 若完全移除，其他 token 的上下文会变。解决方案：freeze-but-attend，退出 token 保留 K/V memory，但停止自身更新。
风险 4：训练不稳定。
halting 是离散决策，直接 RL 可能抖。解决方案：从 oracle label / teacher-student consistency / Gumbel-Sigmoid straight-through 开始，不要一上来 bi-level RL。

11. 最小可投稿版本

如果目标是一篇小论文，建议不要一开始做过大系统。最小可投稿版本可以这样定义：

题目雏形：
Token-wise Adaptive Halting for Efficient Diffusion World Action Models

核心贡献三点：

  1. 问题定义：首次系统研究 WAM / diffusion world model 中的 token-wise runtime early exit。
  2. 方法：提出 freeze-but-attend 的 token halting 机制，让退出 token 不再更新但继续作为上下文。
  3. 实验：证明在机器人操作或 world model rollout 中，以较小成功率损失换取显著 token-step / wall-clock latency 降低。

12. 实验优先级

优先级 实验 为什么先做
P0 Full denoising token convergence analysis 最快证明“token 异质性”存在，是整篇论文的动机。
P0 Oracle early exit 判断上限。如果 oracle 都没收益，方向基本不值得继续。
P1 Residual heuristic early exit 作为最简单 baseline，验证工程链路。
P1 Learned HaltHead + teacher consistency 形成主方法。
P2 Action-sensitivity supervision 增强 WAM 特色，提升论文区分度。
P2 Real wall-clock optimized sparse implementation 决定论文能不能从“理论省 FLOPs”变成“真实加速”。

13. 最终建议

建议做，但路线要收窄。
最稳的切入不是“提出一个通用 early-exit diffusion 框架”，而是：

面向 World Action Model 的 action-aware token-level early exit。

这样能和普通图像 diffusion early exit 拉开差距，也能和 NoiseGate 拉开差距。论文主线应围绕一个问题展开：

哪些 future tokens 已经足够支撑动作决策，哪些还必须继续精细去噪？

参考资料

  1. NoiseGate: Learning Per-Latent Timestep Schedules as Information Gating in World Action Models, arXiv:2605.07794, 2026. https://arxiv.org/abs/2605.07794
  2. Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion, OpenReview / project page. https://openreview.net/forum?id=yDo1ynArjj
  3. DeepCache: Accelerating Diffusion Models for Free, CVPR 2024. CVPR open access
  4. Token Caching for Diffusion Transformer Acceleration, arXiv:2409.18523. https://arxiv.org/abs/2409.18523
  5. Adaptive Score Estimation / Early Exiting for Accelerated Inference in Diffusion Models. OpenReview PDF
  6. AdaDiff: Accelerating Diffusion Models through Step-Wise Adaptive Computation, ECCV 2024. ECCV PDF
  7. Adaptive Computation Time for Recurrent Neural Networks, Graves, 2016. https://arxiv.org/abs/1603.08983
  8. PonderNet: Learning to Ponder, OpenReview, 2021. https://openreview.net/forum?id=1EuxRTe0WN
  9. RoboTwin 2.0: A Scalable Data Generator and Benchmark with Strong Domain Randomization for Robust Bimanual Robotic Manipulation, arXiv:2506.18088. https://arxiv.org/abs/2506.18088