Fast-WAM 深度解析：世界动作模型真的需要推理时的未来想象吗？

1. 引言从 VLA 到 WAM 的范式跃迁在具身智能领域如何让机器人理解物理世界并做出合理决策一直是核心难题。过去两年视觉-语言-动作模型VLA凭借大规模预训练和端到端推理的优势成为机器人控制的主流范式。然而VLA 本质上是一种看图出动作的条件生成模型——它基于当前观测直接预测动作序列却缺乏对未来世界状态的显式建模能力。这意味着机器人在执行复杂长程任务时难以预判动作的后果也无法基于如果我这样做世界会怎么变来做出更优决策。世界动作模型World Action Model, WAM的出现正是为了弥补这一缺陷。WAM 的核心思想是将视频预测world model与动作预测action model联合建模让机器人不仅能感知当前状态还能想象未来世界的演化并据此规划动作。从 2024 年底到 2025 年初一批优秀的 WAM 工作集中涌现——包括 Motus、Cosmos Policy、LingBot-VA、DreamZero 等——它们在仿真基准和真机实验上展现出超越传统 VLA 的性能首次让学术界和工业界大范围意识到WAM 可能拥有比 VLA 更强的潜力。然而一个根本性的问题始终悬而未决WAM 到底为什么 work是因为推理时模型看到了自己想象的未来画面test-time imagination还是因为训练时视频预测任务提供的密集监督信号training-time co-training帮助模型学到了更好的表征Fast-WAM 正是为了回答这个问题而诞生的工作。2. 现有 WAM 的统一范式Imagine-then-Execute要理解 Fast-WAM 的贡献首先需要厘清当前 WAM 的主流做法。尽管各工作的具体实现细节不同但它们的核心范式高度一致训练阶段将 future video prediction 和 future action prediction 联合建模。模型同时学习世界接下来会怎么变和机器人接下来该怎么做。推理阶段模型先生成未来视频imagination然后基于这些想象的画面来输出动作。这就是所谓的 Imagine-then-Execute 范式。这种范式在直觉上很合理先想一想未来会怎样再决定怎么做。但它也带来一个隐含的假设——推理时的 future imagination 是 WAM 性能的关键来源。Fast-WAM 团队对此提出了质疑。3. 核心假设训练时的 Co-training 才是真正的功臣Fast-WAM 的作者观察到future video prediction 在训练和推理时扮演着两个截然不同的角色角色一训练时video prediction 为模型 backbone 提供了极其密集的监督信号。这些监督与未来动作和环境演化高度相关迫使 backbone 学习到对控制有价值的视觉表征——物体运动轨迹、接触关系、空间变化等。这本质上是一种 proxy task代理任务。角色二推理时video prediction 为 action expert 提供了 guidance——action expert 基于想象的未来画面来决策相当于一个 latent IDM逆动力学模型。现有 WAM 普遍采用 Imagine-then-Execute 的推理方式暗示了角色二的重要性。但作者提出了一个大胆的猜想角色一training-time co-training才是 WAM work 的主要原因而角色二test-time imagination可能是多余的。这个猜想并非凭空而来而是有先例可循。回顾 Pi-0.5 相对于 Pi-0 的进化Pi-0.5 引入了 VLM backbone 的 co-training包括 VQA 数据和自回归 discrete action token但推理时并不要求模型先预测 discrete action token 再让 action expert 基于这些 token 做 flow matching。这说明 co-training 的核心价值在于改善 backbone 表征而非为推理提供中间结果。Fast-WAM 将同样的逻辑推广到 video prediction 领域video co-training 可能只是一个更好的 proxy task而非推理时的必要环节。4. Fast-WAM 架构设计极简结构服务于严格解耦为了公平验证上述假设作者设计了一个极简的 Fast-WAM 架构。这里的设计哲学不是追求 SOTA 性能而是构建一个可控实验平台使得各个 component 的效果能被清晰解耦。4.1 整体架构Video DiT Action DiT MoTFast-WAM 的核心结构包含三个模块Video DiT视频扩散 Transformer基于 Wan2.2 5B 预训练模型负责处理视频 latent tokensAction DiT动作扩散 Transformer独立的轻量 Transformer负责处理动作 tokensMoTMixture of Tokens混合注意力层协调两个专家之间的信息交互这三个模块在训练时通过联合优化joint optimization共同学习但通过精心设计的 attention mask 控制信息流向。4.2 MoT 混合注意力的实现MoT 是 Fast-WAM 的核心连接件。它将 Video DiT 和 Action DiT 的 tokens 拼接后做联合注意力计算但通过 attention mask 严格控制哪些 tokens 之间可以交互。以下是 MoT 的核心前向传播逻辑摘自源码src/fastwam/models/wan22/mot.pyclassMoT(nn.Module):defforward(self,embeds_all,attention_mask,freqs_all,context_all,t_mod_all):tokens_all{k:vfork,vinembeds_all.items()}forlayer_idxinrange(self.num_layers):q_chunks,k_chunks,v_chunks[],[],[]cached{}# 1. 每个专家独立计算 Q/K/Vfornameinself.expert_order:# [video, action]expertself.mixtures[name]blockexpert.blocks[layer_idx]q,k,v,residual_x,gate_msa,shift_mlp,scale_mlp,gate_mlp,_\ self._build_expert_attention_io(expert,block,tokens_all[name],freqs_all[name],t_mod_all[name])q_chunks.append(q)k_chunks.append(k)v_chunks.append(v)cached[name]{...}# 缓存 post-block 所需的中间变量# 2. 拼接所有 tokens 做混合注意力q_cattorch.cat(q_chunks,dim1)k_cattorch.cat(k_chunks,dim1)v_cattorch.cat(v_chunks,dim1)mixedself._mixed_attention(q_cat,k_cat,v_cat,attention_mask)# 3. 分割输出每个专家独立做 post-blockFFN cross-attnstart0forname,seq_leninzip(self.expert_order,seq_lens):mixed_slicemixed[:,start:startseq_len,:]tokens_all[name]self._apply_expert_post_block(...)startseq_lenreturntokens_all关键设计点在于虽然 Q/K/V 被拼接在一起做全局注意力计算但attention_mask精确控制了信息的可见性。这使得同一个架构可以通过不同的 mask 配置模拟出多种 WAM 变体。5. Attention Mask解耦实验的核心武器Fast-WAM 最巧妙的设计在于通过 attention mask 实现了多种变体的统一框架。以下是核心的 mask 构建函数摘自src/fastwam/models/wan22/fastwam.pytorch.no_grad()def_build_mot_attention_mask(self,video_seq_len,action_seq_len,video_tokens_per_frame,device):total_seq_lenvideo_seq_lenaction_seq_len masktorch.zeros((total_seq_len,total_seq_len),dtypetorch.bool,devicedevice)# video tokens 之间first-frame causal maskmask[:video_seq_len,:video_seq_len]\ self.video_expert.build_video_to_video_mask(...)# action tokens 之间全可见self-attentionmask[video_seq_len:,video_seq_len:]True# action tokens 只能看到第一帧的 video tokens不看 future videofirst_frame_tokensmin(video_tokens_per_frame,video_seq_len)mask[video_seq_len:,:first_frame_tokens]Truereturnmask这段代码揭示了 Fast-WAM 的核心约束action tokens 在训练时被显式禁止看到 future video tokens。它们只能看到当前帧第一帧的视觉信息和自身的动作序列。这意味着即使训练过程中 video prediction 和 action prediction 是联合优化的action expert 也无法依赖想象的未来来做决策。基于这个统一框架作者通过修改 mask 配置实现了三种对比变体变体Mask 配置对应的已有工作核心特征Fast-WAM-Jointaction 可看 future video (joint denoise)Motus, DreamZero推理时 videoaction 联合去噪Fast-WAM-IDMaction 先看 clean future video (causal)LingBot-VA推理时先生成视频再生成动作Fast-WAMaction 不看 future video本文提出推理时跳过视频生成Fast-WAM-no-cotrain同上但去掉 video loss消融基线验证 co-training 的作用6. 训练流程联合 Flow Matching 的双流去噪Fast-WAM 的训练基于 Continuous Flow Matching对视频和动作分别施加独立的噪声调度然后联合优化。以下是训练损失的核心计算流程摘自training_loss方法deftraining_loss(self,sample,tiledFalse):inputsself.build_inputs(sample,tiledtiled)input_latentsinputs[input_latents]# VAE 编码后的视频 latentactioninputs[action]# 归一化后的动作序列# --- 视频流添加噪声并计算目标 ---noise_videotorch.randn_like(input_latents)timestep_videoself.train_video_scheduler.sample_training_t(...)latentsself.train_video_scheduler.add_noise(input_latents,noise_video,timestep_video)target_videoself.train_video_scheduler.training_target(input_latents,noise_video,timestep_video)# --- 动作流独立的噪声调度 ---noise_actiontorch.randn_like(action)timestep_actionself.train_action_scheduler.sample_training_t(...)noisy_actionself.train_action_scheduler.add_noise(action,noise_action,timestep_action)target_actionself.train_action_scheduler.training_target(action,noise_action,timestep_action)# --- 两个专家分别做 pre-processing ---video_preself.video_expert.pre_dit(xlatents,timesteptimestep_video,...)action_preself.action_expert.pre_dit(action_tokensnoisy_action,timesteptimestep_action,...)# --- MoT 联合前向通过 attention mask 控制交互 ---tokens_outself.mot(embeds_all{video:video_pre[tokens],action:action_pre[tokens]},attention_maskself._build_mot_attention_mask(...),...)# --- 分别计算两个流的预测和损失 ---pred_videoself.video_expert.post_dit(tokens_out[video],video_pre)pred_actionself.action_expert.post_dit(tokens_out[action],action_pre)loss_videoweighted_mse(pred_video,target_video)loss_actionweighted_mse(pred_action,target_action)loss_totallambda_video*loss_videolambda_action*loss_actionreturnloss_total,{loss_video:...,loss_action:...}几个值得注意的设计细节独立的噪声时间步video 和 action 使用各自独立采样的 timestep这意味着在同一个训练 step 中两者可能处于去噪过程的不同阶段。第一帧锁定latents[:, :, 0:1] first_frame_latents确保第一帧始终是干净的无噪声作为条件输入。可配置的损失权重loss_lambda_video和loss_lambda_action允许灵活调整两个任务的相对重要性。当loss_lambda_video 0时即退化为 Fast-WAM-no-cotrain 变体。7. 推理流程跳过视频生成的 Action-Only 推理Fast-WAM 最重要的工程贡献之一是infer_action方法——它完全跳过了 future video generation只利用第一帧的视觉信息做动作推理。这带来了显著的延迟降低。torch.no_grad()definfer_action(self,prompt,input_image,action_horizon,...):# 1. 编码第一帧为 latentfirst_frame_latentsself._encode_input_image_latents_tensor(input_image)# 2. Video expert 只处理第一帧timestep0无噪声timestep_videotorch.zeros(...)# t0 表示干净的状态video_preself.video_expert.pre_dit(xfirst_frame_latents,timesteptimestep_video,...)# 3. 预计算并缓存 video 的 K/V只需一次video_kv_cacheself.mot.prefill_video_cache(video_tokensvideo_pre[tokens],video_freqsvideo_pre[freqs],video_t_modvideo_pre[t_mod],video_attention_maskattention_mask[:video_seq_len,:video_seq_len],...)# 4. Action 的多步去噪循环复用 video KV cacheforstep_t,step_deltainzip(infer_timesteps,infer_deltas):pred_actionself._predict_action_noise_with_cache(latents_actionlatents_action,timestep_actionstep_t,video_kv_cachevideo_kv_cache,# 缓存复用attention_maskattention_mask,video_seq_lenvideo_seq_len,...)latents_actionself.infer_action_scheduler.step(pred_action,step_delta,latents_action)return{action:latents_action[0]}这段推理逻辑的核心洞察是Video expert 只需对第一帧做一次forward pass生成 KV cacheAction expert 在多步去噪过程中复用这个 cache无需反复运行 video branchAttention mask 保证 action tokens 只看第一帧 video tokens与训练时的约束完全一致与此对比传统的 Imagine-then-Execute 推理需要先完整运行视频去噪通常 20-50 步再运行动作去噪计算量几乎翻倍。7.1 与 InfinityStar 等快速视频生成模型的区别读到这里一个自然的疑问是如果推理时视频生成太慢是问题所在那么使用更快的视频生成方案比如字节跳动的 InfinityStar是否就能解决答案是否定的——因为 Fast-WAM 的结论指向一个更根本的层面。InfinityStar 是一种统一时空自回归视频生成框架NeurIPS 2025 Oral它用离散 token 的 next-token prediction 替代了 diffusion 的多步去噪将视频生成速度提升了约 10 倍。从工程角度看它确实大幅缓解了 diffusion 视频模型的延迟问题。但 Fast-WAM 的实验表明问题的本质不是视频生成太慢而是推理时的视频生成本身就不必要。两者解决的是不同层次的问题对比维度InfinityStarFast-WAM核心目标加速视频生成用 AR 替代 Diffusion证明推理时根本不需要生成视频方法论离散时空 token 自回归Attention Mask 解耦 KV Cache推理时是否生成视频是只是更快否完全跳过延迟改善来源减少生成步数从 50 步到约 5 步 AR彻底移除视频生成分支对 WAM 的启示可作为更快的 video backbone证明 video backbone 的价值在训练而非推理换言之即使未来有人把 Fast-WAM 的 Video DiT基于 Wan2.2 的 diffusion 模型替换为 InfinityStar 风格的 AR 视频模型来做 co-trainingFast-WAM 的核心结论依然成立推理时不需要真的走一遍视频生成过程只要训练时视频预测任务提供了足够的监督信号action expert 就能从 backbone 学到的表征中直接 decode 出动作。不过InfinityStar 的 AR 架构对 WAM 领域仍有潜在价值它的离散 token 天然适合作为 proxy task 的训练目标类似 Pi-0.5 中的 discrete action token co-training未来或许可以探索用预测离散视频 token替代预测连续 video latent作为 co-training objective在保持监督信号密度的同时进一步简化训练流水线。8. KV Cache 机制推理加速的工程细节…详情请参照古月居