Fast-WAM——训练时保留视频联合训练(虽同时生成但动作token不看未来视频token)，在推理时则移除显式的未来视频生成，直接在单次前向中，依托得到的潜在世界表征KV Cache预测动作

前言在本文中来自清华、星海图(4.2日听闻 星海图又斩获近20亿元B轮融资七月作为其合作伙伴之一特表祝贺)的研究者探讨世界模型代理WAM在测试阶段是否需要显式的未来想象还是其优势主要来源于训练阶段的视频建模作者通过提出Fast-WAM 来解耦训练时的视频建模与推理时的显式未来生成这两者的作用总之Fast-WAM 是一种 WAM 架构在训练过程中保留视频共训练co-training但在推理时跳过未来预测且作者进一步构建了若干 Fast-WAM变体以便对这两个因素进行可控对比在这些变体中作者发现 Fast-WAM 依然能够与“先想象后执行”imagine-then-execute类方法保持有竞争力的性能而去除视频共训练则会导致性能出现更大幅度的下降————即一方面如果需要 对应变体 在推理时依然可以先生成未来视觉状态二方面 模型训练时 最好是带着视频预测 共同训练实证结果表明在 WAM 中视频预测的主要价值可能在于训练阶段提升世界表征而非在推理阶段生成未来观测顺带提一下同期还有另一个类似的工作《GigaWorld-Policy——以动作为中心的世界动作模型为降低推理延迟训练用视频推理去视频(与Fast-WAM类似)》说明大家都开始注意这个技术范式了毕竟虽然世界模型相对VLA更懂物理但如果先生成视频 再生成动作 实在是太慢了哪怕是同时生成视频和动作——比如本博客中解读过的英伟达的DreamZero————详见《DreamZero——同时统一预测未来视觉状态与动作的世界动作模型解决当下VLA如果人类不示教则理论强但具体操作不强的弊病且提升任务泛化、本体泛化》即便DreamZero加速了38倍后在GB200显卡上依然只实现了7Hz的速度第一部分 Fast-WAM: Do World Action Models NeedTest-time Future Imagination?1.1 引言与相关工作1.1.1 引言如原论文所说现有的 WAM 大致可以归纳为两类具有代表性的范式未来视频与动作通过共享注意力机制联合去噪 [4,6,5]4-World action models are zero-shot policies6-Unified world models: Coupling video and action diffusion for pretraining on largerobotic datasets5-Motus: A unified latent action world model在生成未来视频之后再在其条件上预测动作 [3,7,8]3-Causal world modeling for robot control7-Vidar: Embodied video diffusion model for generalist manipulation8-Learning universal policies via text-guided video generation当然了大多数现有的 WAMWorld Action Models遵循“先想象再执行”的范式它们首先生成未来观测结果然后在想象出的未来条件下预测动作虽然这种设计直观易懂但由于需要反复进行视频去噪 [1,2,3,4,5]会在推理时引入显著的延迟。从更根本的层面看目前仍不清楚显式的未来想象是否真的是获得强动作性能所必需的WAM 的有效性可能来源于两个截然不同的因素训练阶段的视频预测目标这可能有助于模型学习更强的物理先验和与动作相关的表征推理阶段的显式未来生成这可能为动作预测提供额外的前瞻信息现有 WAM 系统通常将这两个因素纠缠在一起使得难以判断究竟是哪一项真正带来了性能提升来自1 IIIS, Tsinghua University、2 Galaxea AI的研究者重新审视这一设计选择并提出一个简单的问题WAM 在测试时是否真的需要去“想象”未来的观测还是它们主要受益于在训练过程中学习对未来观测进行建模作者的核心思路是将用于 WAM 训练的视频预测目标与推理阶段显式生成未来视频的过程解耦。如果世界建模的主要价值在于在训练期间塑造更优的潜在表征那么 WAM 应当能够在“不付出推理时进行未来视频合成”这一代价的前提下依然保留这种收益基于这一视角作者提出了 Fast-WAM这是一种在训练阶段保留视频协同训练videoco-training但在推理阶段跳过未来预测的 WAM 架构Fast-WAM 不再像以往那样在推理过程中使用预训练的视频生成模型迭代合成未来帧而是将一个预训练的视频 DiffusionTransformer(DiT)重新用作单次前向的世界编码器用于动作生成其对应的论文地址为Fast-WAM: Do World Action Models Need Test-time Future Imagination?其对应的作者为Tianyuan Yuan1,2 , Zibin Dong1,2 , Yicheng Liu1,2 , Hang Zhao1,2其对应的项目地址为yuantianyuan01.github.io/FastWAM其对应的GitHub地址为github.com/yuantianyuan01/FastWAM具体而言作者采用共享注意力的 Mixture-of-TransformerMoT架构构建 Fast-WAM该架构由一个视频DiT 和一个动作专家 DiT组成如图 1(C) 所示『三种具有代表性的 WAM 范式。(A) 联合建模式 WAM 将未来视频和动作token一并去噪。B因果式 WAM先生成未来观测然后在生成的未来表征上进行条件动作预测。CFast-WAM在训练期间保留视频联合训练但在推理时移除显式的未来生成直接在一次前向传播中从潜在世界表征预测动作』在训练过程中视频预测目标促使视频DiT 学习编码具有物理意义的运动和交互结构在推理过程中视频 DiT 通过一次前向传递处理观测上下文并为动作去噪提供潜在的世界表示从而避免显式的未来视频去噪实现高效的实时控制为了方便大家更好的理解我个人觉得有必要解释一下推理过程中的这个“单次前向传播”到底是啥意思“单次前向传播”指的是主干网络提取视觉世界特征的过程而不是直接一步“秒出”动作。动作本身依然是通过一个轻量级的扩散去噪Denoising过程生成的只是因为它抛弃了最耗时的视频去噪步骤所以整体速度飙升具体来说Fast-WAM 在推理阶段分为极其干脆的“两步走”参考论文中的图 1(C) 第一步提取潜在世界表征这是真正的“单次前向传播”砍掉未来分支只留第一帧在推理时模型直接把未来预测的视频分支整个“砍掉”仅仅保留当前第一帧观测画面的干净潜变量即单次过视频大模型把这个送进巨大的视频主干网络Video DiT大概 5B 参数中进行仅有一次的前向传播存入缓存KV Cache由于训练阶段视频联合训练的加持这一次前向传播后主干网络输出的特征已经包含了对物理世界的深刻理解这些特征会被存入键值缓存KV Cache作为后续动作生成的 “潜在世界表征底座 ”第二步动作去噪生成轻量级极速运行接下来沉重的主干网络任务结束轮到小巧的“动作专家”Action DiT1B 参数上场接管 动作专家接手动作专家不再去纠结未来的画面会长什么样而是直接利用刚才提取并存在 KV Cache 里的潜在世界表征特征快速去噪出动作由于动作序列本身的数据维度远小于视频帧动作专家只需要进行快速的 10 步去噪输出结果最终模型输出一个包含 32 步的动作序列块为什么它能做到 190ms 的极速没造未来的“假象”因为根本没有实例化instantiate未来的视频 token也就不存在对庞大视频序列进行多步去噪生成的恐怖算力开销大小脑分工明确用 5B 的大模型单次前向看一眼定下对当前世界的理解用 1B 的小模型动作专家根据理解快速多步推算该怎么动总结一下它的“流水线”传感器相机拍下当前画面原始像素VAE Encoder把繁杂的像素打包压缩成了此时还只是纯视觉/空间特征Video DiT拿过过一遍前向传播把物理直觉注入其中提取出深度的“潜在世界表征”即当它处理静态的时能从中“脑补”出环境的动态潜力和交互逻辑。它加工后的输出才被称作真正的“潜在世界表征Latent World Representations”————或者说白了Video DiT 拿到静态的动用它 5B 参数里的物理常识算出了当前环境的“动态潜力和交互逻辑”潜在世界表征这份极其宝贵的特征被转换成了Key (键)和Value (值)并被一次性存入了 KV Cache 中Action DiT动作专家直接拿着这份充满物理智慧的表征快速去噪推算出动作这就好比一位经验丰富的台球大师(5B主干)看一眼球台(当前帧)脑子里瞬间就有了球的走位直觉(提取特征)肌肉(1B动作专家)自然就知道该用多大的力、打哪个点。他不需要真的在脑子里把球碰撞的全过程像放电影一样一帧帧放完才能决定怎么出杆1.1.2 相关工作首先对于视觉-语言-行动策略具身基础模型的最新进展主要得益于视觉-语言-行动Vision-Language-ActionVLA策略这类模型利用大规模预训练的视觉-语言骨干网络[9,10,11,12,13,14,15,16,17,18,19]将视觉观测和语言指令直接映射为机器人动作通过继承来自网页规模预训练的强语义先验这些模型在物体、场景以及语言指令等多方面展现出很强的泛化能力然而正如近期 WAM 相关工作所指出的标准的 VLA 预训练主要基于静态图文数据并未显式建模物理世界在动作作用下如何演化 [3,4]作者的工作与这一研究方向是互补的Fast-WAM 在测试阶段同样保留了类似 VLA 的直接策略接口但在训练时额外引入了基于未来视觉预测的世界建模目标来学习这一接口其次对于世界动作模型与基于视频的机器人策略另一条平行的研究路线是通过未来视觉预测来进行机器人控制即将视频生成作为建模环境动力学并推断动作的一种方式[8,20,21,7,22,23,24,25,26,27,28,29,30,31]最新的方法在此基础上进一步扩展将未来视频和机器人动作联合建模在一个统一框架中 [6,2,32,33,1,3,5,4]作者遵循 [4] 将此类模型称为 World Action ModelsWAMs世界动作模型因为它们利用世界建模即预测未来视觉状态来支持下游动作预测大多数现有的 WAM要么遵循“想象后执行”imagine-then-execute范式先生成未来的视觉轨迹再据此进行动作预测要么在一个共享的生成过程内联合建模未来视频与动作————作者的工作与这一路线最为接近但关注点不同作者并非再提出一种新的“想象后执行”WAM而是研究 WAM的收益主要来源于训练阶段的视频联合训练还是来源于推理阶段显式的未来想象作者的工作也与最近的一些研究相关这些研究利用视频建模进行动作预测同时在测试时减少或绕过显式的视频合成。VPP [34] 将机器人策略条件化在由视频扩散模型提取的预测性视觉表征上而 UVA [35] 则联合建模视频和动作并在测试时跳过视频解码以实现更快速的推理与这些工作相比作者的重点在于在一个共享的框架下通过一系列可控变体来解耦训练阶段视频协同训练与测试阶段未来想象各自所起的相对作用1.1.3 问题表述考虑从视觉观测和语言指令中进行具身策略学习。令表示当前观测表示任务指令表示视野为的动作片段。一个标准的视觉运动策略建模条件分布它将当前的感知上下文直接映射为一系列动作世界动作模型World Action ModelsWAMs通过将未来的视觉观测作为中间环节引入来扩展这一形式化表述变量令表示在预测范围内的未来视觉观测。许多现有的WAM 遵循先想象再执行的分解方式在这种范式下模型首先预测未来观测然后在“想象”的未来基础上生成动作在实践中这通常通过两种方式实现要么在一个共享模型中联合去噪未来视频和动作要么先生成未来视频再将其输入到逆动力学或动作预测模块中一些先前的 WAM 进一步在这一公式外层包裹自回归式的滚动展开但为了简化并便于可控对比作者在此不采用该设计作者的出发点是这样一种观察WAM 的有效性可能源自两个不同的因素训练期间使用的视频预测目标它可以促使模型学习具有物理意义的潜在表示推理时的显式未来生成它可能为动作预测提供额外的前瞻性现有的WAM 形式通常将这两个因素耦合在一起因为同一个模型既从未来视频预测中学习又在推理时显式合成未来观测作者设计了 Fast-WAM 来解耦这两个因素在训练期间它保留世界建模作为一个联合训练信号然而在推理期间它并不显式生成未来观测相反Fast-WAM 直接根据当前观测和指令预测动作同时利用由视频协同训练塑造的潜在世界表征从这个意义上说Fast-WAM 在测试时具有类似于标准 VLA 策略的直接策略接口而其表征学习在训练期间仍然扎根于 WAM 风格的视频建模形式化地设表示由视频主干网络在给定当前上下文条件下生成的潜在世界表征。Fast-WAM 使用该表征来参数化动作分布此与” 先想象再执行” 的WAMs 的关键区别在于是通过一次前向编码过程获得的而不是在推理时显式地对未来观测进行采样或去噪1.2 Fast-WAM的完整方法论1.2.1 模型架构、训练目标Fast-WAM 的设计目标是在保留世界建模带来的训练收益的同时去除显式未来想象在推理阶段带来的开销在训练过程中它联合学习动作预测和视频建模从而鼓励视觉骨干网络捕获具有物理意义的运动和交互结构在推理过程中Fast-WAM 不会显式生成未来观测。相反它仅保留第一帧观测的干净潜在 token用视频模型在一次前向传播中对其进行处理并将得到的潜在世界表征用于直接生成动作这样Fast-WAM 在测试时具有直接策略接口同时在训练阶段仍然保留类似 WAM 的视频监督首先对于模型架构Fast-WAM 构建在 Wan2.2-5B [36] 的视频 Diffusion TransformerDiT之上将其作为世界建模骨干网络同时作者复用其预训练的文本编码器和视频 VAE任务语言由内置的T5 编码器编码并通过交叉注意力提供给所有 token而视觉观测则通过预训练的 VAE映射为潜空间视频 token在这一骨干之上作者引入一个动作专家 DiT用于生成动作分块action chunk完整模型被组织为一种 Mixture-of-TransformerMoT架构在视频分支与动作分支之间共享注意力如图 2a 所示『Fast-WAM 架构以及用于将视频协同训练与动作生成解耦的结构化注意力掩码』作者将输入的 token 组织成三类第一帧观测画面的干净潜在 token它们作为共享的视觉锚点未来视频帧的带噪潜在 token它们仅在训练阶段用于视频建模即训练阶段它依然保留视频联合训练继续让模型去学“动作会怎样改变视觉世界”这件事让共享的视频长出更强的世界模型表征以及由动作专家处理、用于动作生成的动作 token所有 token 组都通过交叉注意力(cross-attention)访问语言嵌入。一个结构化的注意力掩码用于控制这些组之间的信息流在训练过程中未来的带噪视频 token在视频分支内可以双向注意并且可以访问第一帧的干净 token动作 token在动作分支内可以双向注意同时也可以访问第一帧的干净token关键是动作 token 不能注意未来视频 token而干净的首帧 token 不会关注任何其他 token这样既保证了视频建模和动作预测建立在同一视觉上下文之上又能防止未来信息泄漏到动作分支作者在图 2b 中给出了完整的训练和推理掩码为了让你能一目了然地理解我们把模型里的处理单元Token想象成考场里的三种角色然后来看看论文是怎么给他们定“规矩”注意力掩码的考场里的三组角色Tokens在 Fast-WAM 的 Mixture-of-Transformer 架构中输入被分成了三拨人“题干” ()当前第一帧画面的干净特征 。它是所有人都要参考的基准“视频推演组” ()负责预测未来画面的噪声视频 token“动作执行组” ()负责预测机器人该怎么动的噪声动作 token(注语言指令像考场广播一样所有组都可以通过交叉注意力随时听到)训练阶段绝妙的“单向隔离”掩码在训练时这三组人都在同一个大教室里共享注意力机制同时做题 。论文通过一个结构化的注意力掩码矩阵Structured Attention Mask设定了严格的“看片”规则 规则一“题干”高冷独立当前帧只能看自己不能去看未来的视频也不能去看动作 。这样保证了它提取的是纯粹的当前真实世界状态规则二“视频推演组”看题干未来的视频 token () 可以在自己组内互相看双向注意力同时也可以看“题干” ()这是理所应当的必须要结合当前状态才能推演未来规则三核心 动作组与“视频推演组”被挡板隔开动作 token () 可以在自己组内互相看也可以看“题干” ()。但是绝对严禁动作 token 去看未来视频 tokenFast-WAM 在训练时强制切断了动作对未来画面的依赖 。这就逼着网络动作模型只能通过去“压榨”那个处理题干的大脑(视频主干网络Video DiT)让大脑把丰富的物理规律压缩到的特征里然后动作模型直接对着这个充满物理直觉的输出动作在推理阶段Fast-WAM 会完全移除未来视频分支只保留干净的首帧潜在tokenlatent tokens并将其仅通过一次视频主干网络以为动作专家生成潜在世界特征由于既不会实例化任何未来的带噪视频token也不会执行显式的未来视频去噪Fast-WAM 的推理开销相比标准的“先想象再执行”的 WAM 显著更低相当于把世界模型最值钱的“理解力”留下来把世界模型最拖后腿的“生成过程”拿掉正是因为训练时有了上述的“隔离”掩码到了实机执行推理阶段整个过程变得无比清爽直接砍掉“视频推演组”因为动作模型在训练时就从来没看过未来视频所以现在完全不需要实例化(创建)未来的视频 token极简前向传播只有“题干”干净的第一帧被保留送进视频主干网络走一遍直出动作动作专家拿到这遍前向传播提取出的世界表征特征开始快速去噪生成动作其次对于训练目标Fast-WAM 通过在动作token 和未来视频潜变量上的联合流匹配目标进行训练给定一个目标变量可以是动作片段或未来视频潜变量作者对高斯噪声和时间步进行采样并构造插值样本该模型通过标准的流匹配flow matching目标进行训练以预测对应的速度场作者为动作生成和视频协同训练两者实例化该目标对于动作预测设定并优化对于视频协同训练video co-training作者设其中表示由预训练 VAE 生成的未来视频帧的潜在 token并优化整体训练目标为其中在动作学习与视频协同训练之间起到平衡作用可能有同学疑问既然训练时保留视频联合训练可虽同时生成但动作token不看视频token那视频token 对动作token的生成还有帮助作用么表面上做题时不准“抄袭” (前向传播被阻断)在模型从左到右计算前向传播时Mask 挡板把动作 Token 的视线无情切断了动作 Token 看不到未来视频 Token 预测出了什么画面。这意味着在决策的瞬间动作和未来画面是完全物理隔离的暗地里共用同一个虽然两组 Token 互相不说话但它们都必须紧紧盯着同一个东西第一帧更重要的是处理的是那个 5B 参数的庞大“视频主干网络Video DiT”视频预测组和动作预测组其实是长在同一个“根”上的两个分支灵魂互助的真相逼着大脑“长脑子”反向传播的魔法当模型开始训练时真正的魔法发生在“对答案、算错题(反向传播)”的阶段视频组的疯狂施压视频 Token 发现自己预测的未来画面和真实的未来画面对不上产生了巨大的误差Video Loss这个庞大的误差信号会顺着网络往回传一直传到处理的视频主干网络里主干网络的被迫进化为了降低这个巨大的预测误差视频主干网络被迫去改变自己内部的神经元权重被迫去学习“杯子掉下会碎”、“布料拉扯会变形”这些底层物理规律脱胎换骨经过成千上万次这样的“逼迫”经过视频主干网络处理后的再也不是一张简单的静态图片特征了它变成了一个充满了物理直觉和动态潜力的“潜在世界表征Latent World Representation”动作组的“坐享其成”现在回到动作 Token。它虽然被蒙上了眼睛看不到未来的视频画面但是它被允许看。 此时它看到的已经是被视频预测任务“千锤百炼、疯狂开光”之后的巅峰状态了动作 Token 不需要知道未来长什么样因为它看一眼现在的就瞬间获得了老军师主干网络毕生的物理直觉总结视频 Token 到底帮了什么忙 如果把训练比作在少林寺挑水动作 Token是专门练“出拳打人”的视频 Token是专门练“挑水扎马步”的它们俩虽然练的不是同一个动作但它们长在“同一个身体(共享的视频主干网络Video DiT)上在 Fast-WAM 的设计中无论是“挑水”的视频 Token还是“出拳”的动作 Token它们都必须去读取当前第一帧经过处理后的特征 。 而负责处理这个的正是那个庞大的5B 参数 Video DiT(也就是这具“身体的躯干”)1.2.2 受控变体用于解耦的 WAM 设计为了回答这个核心问题——即 WAM 的优势主要来源于训练阶段的视频协同训练还是来源于推理阶段显式的未来想象作者在一个统一的实现框架下设计了一组可控的变体。且从近期的 WAM 中实例化具有代表性的“想象-执行”设计模式同时尽可能保持骨干网络、分词和训练方案的一致性这样的受控设置使得能够将测试时的未来生成带来的贡献与视频联合训练目标本身的贡献区分开来如图 1(A) 和 (B) 所示作者考虑两种具有代表性的“先想象再执行”imagine-then-execute变体用以刻画近期 WAMs 中占主导地位的设计模式第一种变体命名为Fast-WAM-Joint其遵循联合生成范式在一个共享模型中对未来视频 token 与动作 token 进行联合去噪从而在整个去噪过程中动作生成始终与未来视频建模保持耦合 [4,6,5]4-World action models are zero-shot policies6-Unified world models: Coupling video and action diffusion for pretraining on largerobotic datasets5-Motus: A unified latent action world model第二种变体命名为Fast-WAM-IDM其遵循先视频再动作(video-then-action)范式首先在当前观测与语言上下文条件下生成未来视频 token随后在所得到的未来表征基础上进行动作预测 [3,7,8]3-Causal world modeling for robot control7-Vidar: Embodied video diffusion model for generalist manipulation8-Learning universal policies via text-guided video generation为方便大家理解我再帮大家用一个表格总结下以方便大家一目了然模型变体核心范式本质 (打个比方)训练阶段 (Training)推理阶段 (Inference)动作是否依赖生成的未来视频真实机器推理延迟 (Latency)Fast-WAM(本文核心)“得其意而忘其形”平时逼着练“画图解题”练出直觉考试看题秒出答案。同时计算视频与动作的损失但通过掩码Mask严格切断动作对未来视频的注意力完全抛弃未来视频分支仅利用当前帧单次前向传播直接输出动作否动作生成在训练和推理时都不看未来视频极快(190 ms)Fast-WAM-Joint(联合生成基线)“边画边想”无论是平时还是考试动作都要和画面推演深度绑定在一起。视频与动作耦合训练允许视频和动作 Token 之间共享注意力视频与动作的 Token 在同一个共享模型中联合去噪同时生成是通过共享注意力相互交织计算慢(580 ms)Fast-WAM-IDM(因果/逆动力学基线)“先画后想”无论是平时还是考试都必须先严谨地把图画完再看着图解题。遵循因果顺序先生成未来的视频然后将动作预测条件化在生成的未来表示上分两步走先去噪生成未来视频再去噪生成动作是严格的先后因果依赖极慢(810 ms)在这两种情形下作者都保留了相应范式所定义的推理结构以及近期WAMs 中采用的关键训练策略同时在作者统一的框架内实现这些方法以便与 Fast-WAM 进行可控对比且作者进一步构建了一个不含视频联合训练的 Fast-WAM 变体。该变体在架构和推理过程上保持不变只在训练阶段去除了视频建模目标因此它可以作为直接的对照用于评估视频联合训练本身的作用结合这些受控变体作者得以将以往 WAM 中通常纠缠在一起的两个因素分离开来训练阶段的视频联合训练以及推理阶段的显式未来想象1.3 实验1.3.1 实现细节作者使用预训练的Wan2.2-5B [36-Wan: Open and advanced large-scale video generative models] 作为骨干网络包括其video DiT、text encoder 和videoVAEaction expert 与video 分支共享相同的架构但使用了较小的隐藏维度从而得到一个1 B 规模的action expert总模型规模为6 B 参数且作者将action horizon 设置为h 32。视频帧在时间维度上以4× 进行下采样从而每个chunk 包含9 帧视频(为何是9帧我直接给你数一下你就一目了然了)以及来自多个摄像机的图像在输入VAE 之前被拼接成一张图像对视频分支和动作分支都采用相同的flow-matching流匹配建模方式。按照文献 [36] 的方法作者在训练和推理阶段均使用以logit-normal分布为基础的噪声调度方案在推理时使用10步去噪并将classifier-free guidance(CFG无分类器引导)的尺度设置为1.0。优化器采用AdamW学习率为权重衰减为0.01对所有训练配置都使用余弦退火学习率调度训练在混合精度模式下进行并将梯度裁剪阈值设为1.0所有延迟(latency)指标均在单张NVIDIA RTX5090D V2 32GB GPU上测得且通过允许视频和动作token之间的注意力构建Fast-WAM-Joint对于Fast-WAM-IDM作者遵循[3]并以概率p 0.5 对真实视频token施加噪声增强// 待更