Fast-WAM 模型架构#

1. 整体架构概览#

Fast-WAM 是一个世界动作模型 (World Action Model, WAM),论文标题本身就是它要回答的问题:"Do World Action Models Need Test-time Future Imagination?"(arXiv:2603.16666)。结论是不需要 —— 现有 WAM(如 DreamZero、IDM 类方法)在推理时要先把未来视频去噪生成出来,再从中读出动作,代价是每步决策都要跑一遍完整的视频扩散;Fast-WAM 证明这份"未来想象"只在训练时有价值,推理时可以彻底砍掉。

实现手段是一个双专家 MoT (Mixture-of-Transformers):视频专家来自 Wan2.2-TI2V-5B 的 DiT,动作专家 ActionDiT 是同层数、同头数的窄版 DiT,两者逐层共享同一次 mixed attention。关键设计在注意力掩码:动作 token 只被允许看首帧视频 token,而首帧在 first_frame_causal 掩码下又不允许看后续帧 —— 于是首帧的 K/V 与"视频有没有被生成"无关。推理时因此可以把视频专家跑一次(输入只有观测首帧,timestep=0)缓存下每层 K/V,后续 N 步动作去噪全部走缓存,视频分支一次都不用再跑。

三个变体共用完全相同的配置文件(configs/model/fastwam.yaml / fastwam_idm.yaml / fastwam_joint.yaml 逐字段一致,只有 _target_ 不同),差异全部落在 Python 类的两个方法上:_build_mot_attention_mask() 和 infer_action()。这让"未来想象是否必要"成为一个干净的受控对比。

2. 核心组件详解#

2.1 视频专家:Wan2.2-TI2V-5B DiT#

视频分支直接复用 Wan2.2 的文生视频 DiT,权重从 Wan-AI/Wan2.2-TI2V-5B 加载(src/fastwam/models/wan22/helpers/loader.py)。

配套 VAE 是 Wan2.2 VAE:时间下采样 4 倍、空间 16 倍。因此视频帧数必须满足 T % 4 == 1(build_inputs 中强校验,src/fastwam/models/wan22/fastwam.py:296),LIBERO 用 33 帧 → 9 个 latent 时间步。

fuse_vae_embedding_in_latents=true 时,首帧 latent 会在加噪后被原样写回(fastwam.py:468):

latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video)
...
if inputs["first_frame_latents"] is not None:
    latents[:, :, 0:1] = inputs["first_frame_latents"]   # 首帧永远是干净观测

首帧因此是条件而非生成目标,视频 loss 也相应跳过第 0 步(fastwam.py:536)。

2.2 视频自注意力掩码:first_frame_causal#

build_video_to_video_mask()(src/fastwam/models/wan22/wan_video_dit.py:473)支持三种模式,Fast-WAM 全系配置用 first_frame_causal:

if self.video_attention_mask_mode == "first_frame_causal":
    video_mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device)
    first_frame_tokens = min(video_tokens_per_frame, video_seq_len)
    video_mask[:first_frame_tokens, first_frame_tokens:] = False
    return video_mask

含义:首帧的 query 不允许看任何后续帧,其余帧之间双向可见。这条约束是整套加速的数学前提 —— 首帧 token 的表征只依赖首帧自己,所以"只喂首帧跑一次 DiT"得到的 K/V,与"喂完整视频跑一次 DiT"里首帧那段 K/V 逐位相等。缓存因此是精确的,不是近似。

三种模式对比:

2.3 动作专家 ActionDiT#

src/fastwam/models/wan22/action_dit.py 定义,结构上是一个"窄身宽头"的 DiT:

注意 hidden_dim=1024 但 num_heads × attn_head_dim = 3072。SelfAttention 里 q/k/v 投影写的是 nn.Linear(hidden_dim, num_heads * attn_head_dim)(wan_video_dit.py:179),即 Linear(1024 → 3072)。这不是笔误 —— 两个专家的 K/V 必须落在同一个 3072 维空间里,才能在 mixed attention 中直接 concat:

k_cat = torch.cat([k_video, k_action], dim=1)   # mot.py:426
v_cat = torch.cat([v_video, v_action], dim=1)

from_wan22_pretrained() 会对三项硬校验,不一致直接抛错(fastwam.py:140-145)。

ActionDiT 的初始化不是随机的,而是从 Wan2.2 视频 DiT 线性插值降维得到。scripts/preprocess_action_dit_backbone.py 把 3072 维权重沿最后一维 F.interpolate(mode="linear", align_corners=True) 压到 1024,并施加 alpha = sqrt(d_video / d_action) 的缩放以保持激活方差:

if apply_alpha_scaling and src.ndim >= 2 and src.shape[-1] != target.shape[-1]:
    alpha = (float(src.shape[-1]) / float(target.shape[-1])) ** 0.5
    value = value.to(torch.float32) * alpha

产物存为 checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt,内含 meta(8 个结构字段,加载时逐项校验)+ backbone_state_dict。action_encoder. 和 head. 两个前缀被排除在外(ACTION_BACKBONE_SKIP_PREFIXES),保持随机初始化 —— 因为动作维度与视频 latent 维度无对应关系,插值没有意义。

2.4 MoT 混合注意力#

src/fastwam/models/wan22/mot.py 的 MoT 不是常见的 sparse-MoE 路由,而是两套完整参数、逐层做一次联合 attention:

两个专家各自保留独立的 modulation、cross_attn(对文本)、ffn、o 投影 —— 只有 self-attention 那一次矩阵乘是共享的。这正是 MoT 与"共享主干 + 两个 head"的区别:参数完全不共享,只共享注意力这一次信息交换。

mot_checkpoint_mixed_attn=true 时,mixed attention 走 torch.utils.checkpoint,以重算换显存(mot.py:89)。

2.5 三个变体的注意力掩码#

这是整篇论文的实验骨架。三个类的 _build_mot_attention_mask() 只在"action → video"这一个子块上不同:

FastWAM (uncond) — fastwam.py:386

mask[video_seq_len:, video_seq_len:] = True                 # action → action 全可见
first_frame_tokens = min(video_tokens_per_frame, video_seq_len)
mask[video_seq_len:, :first_frame_tokens] = True            # action → 仅首帧

FastWAMJoint — fastwam_joint.py:29

mask[video_seq_len:, video_seq_len:] = True
mask[video_seq_len:, :video_seq_len] = True                 # action → 全部视频帧

FastWAMIDM — fastwam_idm.py:20,训练时构造三段序列(noisy video / cond video / action),动作只看 cond 那一段:

mask[cond_end:, cond_end:] = True                           # action → action
mask[cond_end:, noisy_end:cond_end] = True                  # action → cond_video only

cond 分支以 video_cond_noise_prob = 0.5 的概率被加噪(fastwam_idm.py:17),否则直接用真值 latent。这是 IDM(逆动力学)的标准 teacher forcing:训练时给模型看真实的未来,让它学"从 s_t 到 s_{t+k} 需要什么动作"。

三者的推理代价:

3. 训练流水线#

training_loss()(fastwam.py:448)对视频和动作各自独立采样 timestep,两条 flow matching 支路并行:

几个实现细节:

两个独立 timestep。 timestep_video 和 timestep_action 分别采样(fastwam.py:459 / 471),不共享。这意味着模型见过"视频很干净但动作很噪"和反过来的所有组合 —— 对 uncond 变体尤其关键,因为推理时视频侧固定是 timestep_video = 0(全干净),必须在训练中覆盖到这个区域。

padding 掩码逐样本归一。 动作和图像都可能有 padding(数据集尾部对齐),loss 按有效步数归一而不是简单 .mean():

valid = (~action_is_pad).to(...)
valid_sum = valid.sum(dim=1).clamp(min=1.0)
action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum

图像侧还要把帧级 mask 折叠到 latent 步(除以 temporal_downsample_factor=4,一个 latent 步内全 pad 才算 pad,fastwam.py:431)。

proprio 作为 context token。 本体状态不进动作序列,而是过一个 Linear(proprio_dim → text_dim) 变成一个 token 拼在文本 context 末尾(_append_proprio_to_context,fastwam.py:219),两个专家的 cross-attention 都能看到。训练时只取序列首帧的 proprio(proprio[:, 0, :])。

文本编码器可以不加载。 配置里 load_text_encoder: false —— 训练前用 scripts/precompute_text_embeds.py 把 umt5 embedding 全部预计算缓存到 text_embedding_cache_dir,训练时直接读 context / context_mask,省下 umt5-xxl 的显存。

训练入口:scripts/train.py + scripts/train_zero1.sh(DeepSpeed ZeRO-1)。优化器和冻结逻辑通过 self.dit = self.mot 这个别名接到 trainer 上(fastwam.py:47)。

4. 推理流水线#

4.1 快路径:infer_action(uncond 变体)#

这是 Fast-WAM 的核心卖点,fastwam.py:906:

代码上的关键三步:

# 1. 视频专家只在首帧上跑一次,timestep 恒为 0
timestep_video = torch.zeros((first_frame_latents.shape[0],), ...)
video_pre = self.video_expert.pre_dit(x=first_frame_latents, timestep=timestep_video, ...)

# 2. 逐层缓存 K/V
video_kv_cache = self.mot.prefill_video_cache(...)          # fastwam.py:1013

# 3. N 步动作去噪全部走缓存
for step_t_action, step_delta_action in zip(...):
    pred_action = self._predict_action_noise_with_cache(..., video_kv_cache=video_kv_cache, ...)
    latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action)

infer_action 开头有一条硬断言 —— 没有 first_frame_causal 掩码,缓存就不成立:

if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal":
    raise ValueError("`infer_action` requires `video_attention_mask_mode='first_frame_causal'`.")

注意 VAE 解码器完全没被调用 —— 不生成像素,连 latent 都不生成。返回值只有 {"action": ...}。

4.2 慢路径:infer_joint#

joint 和 idm 变体走这条。两个 scheduler 步调对齐,每步都要跑完整的 _predict_joint_noise(视频 + 动作双分支),且每步结束后把首帧 latent 重新写回:

latents_video = self.infer_video_scheduler.step(pred_video_posi, step_delta_video, latents_video)
latents_action = self.infer_action_scheduler.step(pred_action_posi, step_delta_action, latents_action)
latents_video[:, :, 0:1] = first_frame_latents.clone()      # fastwam_joint.py:232

IDM 变体更进一步 —— infer_joint 里视频是独立的第一阶段先去噪完,再算动作(fastwam_idm.py:288 起,注释写明 "video is denoised in a standalone first stage")。

4.3 仿真评测#

experiments/libero/run_libero_manager.py 和 experiments/robotwin/run_robotwin_manager.py,配置见 configs/sim_libero.yaml:

5. 关键超参数表#

模型结构#

Flow Matching#

数据(LIBERO 2-cam)#

RoboTwin 配置为 3 相机 384 分辨率(configs/task/robotwin_*_3cam_384_1e-4.yaml)。

6. 关键源文件表#

7. 与 DreamZero 的对照#

两者都是 Wan 系视频扩散改造的 WAM,放在一起看差异很清楚(DreamZero 细节见 DreamZero 架构):

Fast-WAM 的 joint 与 idm 两个变体,本质上就是把 DreamZero 那类"推理时也要想象未来"的做法做成受控基线,再用 uncond 变体证明砍掉它不掉点。