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 + MoT

Fast-WAM 的核心结构包含三个模块：

1. Video DiT（视频扩散 Transformer）：基于 Wan2.2 5B 预训练模型，负责处理视频 latent tokens
2. Action DiT（动作扩散 Transformer）：独立的轻量 Transformer，负责处理动作 tokens
3. MoT（Mixture 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.py）：

classMoT(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"]expert=self.mixtures[name]block=expert.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_cat=torch.cat(q_chunks,dim=1)k_cat=torch.cat(k_chunks,dim=1)v_cat=torch.cat(v_chunks,dim=1)mixed=self._mixed_attention(q_cat,k_cat,v_cat,attention_mask)# 3. 分割输出，每个专家独立做 post-block（FFN + cross-attn）start=0forname,seq_leninzip(self.expert_order,seq_lens):mixed_slice=mixed[:,start:start+seq_len,:]tokens_all[name]=self._apply_expert_post_block(...)start+=seq_lenreturntokens_all

关键设计点在于：虽然 Q/K/V 被拼接在一起做全局注意力计算，但attention_mask精确控制了信息的可见性。这使得同一个架构可以通过不同的 mask 配置，模拟出多种 WAM 变体。

5. Attention Mask：解耦实验的核心武器

Fast-WAM 最巧妙的设计在于通过 attention mask 实现了多种变体的统一框架。以下是核心的 mask 构建函数（摘自src/fastwam/models/wan22/fastwam.py）：

@torch.no_grad()def_build_mot_attention_mask(self,video_seq_len,action_seq_len,video_tokens_per_frame,device):total_seq_len=video_seq_len+action_seq_len mask=torch.zeros((total_seq_len,total_seq_len),dtype=torch.bool,device=device)# video tokens 之间：first-frame causal maskmask[:video_seq_len,:video_seq_len]=\ self.video_expert.build_video_to_video_mask(...)# action tokens 之间：全可见（self-attention）mask[video_seq_len:,video_seq_len:]=True# action tokens 只能看到第一帧的 video tokens（不看 future video）first_frame_tokens=min(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 配置实现了三种对比变体：

6. 训练流程：联合 Flow Matching 的双流去噪

Fast-WAM 的训练基于 Continuous Flow Matching，对视频和动作分别施加独立的噪声调度，然后联合优化。以下是训练损失的核心计算流程（摘自training_loss方法）：

deftraining_loss(self,sample,tiled=False):inputs=self.build_inputs(sample,tiled=tiled)input_latents=inputs["input_latents"]# VAE 编码后的视频 latentaction=inputs["action"]# 归一化后的动作序列# --- 视频流：添加噪声并计算目标 ---noise_video=torch.randn_like(input_latents)timestep_video=self.train_video_scheduler.sample_training_t(...)latents=self.train_video_scheduler.add_noise(input_latents,noise_video,timestep_video)target_video=self.train_video_scheduler.training_target(input_latents,noise_video,timestep_video)# --- 动作流：独立的噪声调度 ---noise_action=torch.randn_like(action)timestep_action=self.train_action_scheduler.sample_training_t(...)noisy_action=self.train_action_scheduler.add_noise(action,noise_action,timestep_action)target_action=self.train_action_scheduler.training_target(action,noise_action,timestep_action)# --- 两个专家分别做 pre-processing ---video_pre=self.video_expert.pre_dit(x=latents,timestep=timestep_video,...)action_pre=self.action_expert.pre_dit(action_tokens=noisy_action,timestep=timestep_action,...)# --- MoT 联合前向：通过 attention mask 控制交互 ---tokens_out=self.mot(embeds_all={"video":video_pre["tokens"],"action":action_pre["tokens"]},attention_mask=self._build_mot_attention_mask(...),...)# --- 分别计算两个流的预测和损失 ---pred_video=self.video_expert.post_dit(tokens_out["video"],video_pre)pred_action=self.action_expert.post_dit(tokens_out["action"],action_pre)loss_video=weighted_mse(pred_video,target_video)loss_action=weighted_mse(pred_action,target_action)loss_total=lambda_video*loss_video+lambda_action*loss_actionreturnloss_total,{"loss_video":...,"loss_action":...}

几个值得注意的设计细节：

1. 独立的噪声时间步：video 和 action 使用各自独立采样的 timestep，这意味着在同一个训练 step 中，两者可能处于去噪过程的不同阶段。
2. 第一帧锁定：latents[:, :, 0:1] = first_frame_latents确保第一帧始终是"干净的"（无噪声），作为条件输入。
3. 可配置的损失权重：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_latents=self._encode_input_image_latents_tensor(input_image)# 2. Video expert 只处理第一帧（timestep=0，无噪声）timestep_video=torch.zeros(...)# t=0 表示"干净"的状态video_pre=self.video_expert.pre_dit(x=first_frame_latents,timestep=timestep_video,...)# 3. 预计算并缓存 video 的 K/V（只需一次）video_kv_cache=self.mot.prefill_video_cache(video_tokens=video_pre["tokens"],video_freqs=video_pre["freqs"],video_t_mod=video_pre["t_mod"],video_attention_mask=attention_mask[:video_seq_len,:video_seq_len],...)# 4. Action 的多步去噪循环（复用 video KV cache）forstep_t,step_deltainzip(infer_timesteps,infer_deltas):pred_action=self._predict_action_noise_with_cache(latents_action=latents_action,timestep_action=step_t,video_kv_cache=video_kv_cache,# 缓存复用attention_mask=attention_mask,video_seq_len=video_seq_len,...)latents_action=self.infer_action_scheduler.step(pred_action,step_delta,latents_action)return{"action":latents_action[0]}

这段推理逻辑的核心洞察是：

• Video expert 只需对第一帧做一次forward pass，生成 KV cache
• Action expert 在多步去噪过程中复用这个 cache，无需反复运行 video branch
• Attention 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 的实验表明，问题的本质不是"视频生成太慢"，而是"推理时的视频生成本身就不必要"。

两者解决的是不同层次的问题：

换言之，即使未来有人把 Fast-WAM 的 Video DiT（基于 Wan2.2 的 diffusion 模型）替换为 InfinityStar 风格的 AR 视频模型来做 co-training，Fast-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 机制：推理加速的工程细节

…详情请参照古月居