DreamZero (WAM) 模型架构

1. 整体架构概览

DreamZero 是一个世界动作模型 (World Action Model, WAM)，通过联合预测未来视频帧与机器人动作序列，实现对未见任务的零样本泛化。其核心创新在于：基于 Wan2.1 视频扩散模型构建 Causal WAN DiT，以 Flow Matching 框架在视频 latent 空间和动作空间上同步去噪，配合多机器人体 Category-specific MLP 支持 DROID、AgiBot、YAM 等多种机器人平台，在 MolmoSpaces 和 RoboArena 双榜均位列第一（截至 2026 年 2 月）。

2. 核心组件详解

2.1 感知编码器

DreamZero 使用两个独立的预训练编码器将视觉和语言信息映射到统一的 1536-dim 特征空间：

• 图像编码器：open-clip-xlm-roberta-large-vit-huge-14（CLIP ViT-H/14），对输入视频帧逐帧提取视觉特征
• 文本编码器：umt5-xxl（Google 多语言 T5），将自然语言指令编码为序列特征

2.2 视频 VAE

沿用 Wan2.1 的视频 VAE，作为压缩/解压工具（无生成能力），让 DiT 在低维 latent 空间工作以降低计算量。VAE 编码器和解码器服务于两个完全不同的目的，输入输出的 latent 含义不同。

2.3 Causal WAN DiT（核心扩散模型）

Causal WAN DiT 是 DreamZero 的核心，基于 Wan2.1 的 DiT 架构改造而来：40 层因果 Transformer，通过 Causal Masking 保证时序因果性，同时在 latent 序列中嵌入动作 token，实现视频与动作的联合生成。

因果注意力掩码设计： 视频帧 token 只能 attend 到过去帧和当前帧，动作 token attend 到所有视频 token（全局条件），保证生成的时序一致性。

2.4 多机器人体支持

DreamZero 通过每个机器人体独立的线性层（Category-specific MLP）将不同维度的动作空间映射到统一的 DiT 隐空间，支持在单一模型中处理多种机器人平台。

相对动作计算： 各机器人体在数据加载时将绝对动作转为相对值（action - reference_state），使模型学到的动作表示更加泛化。

3. 训练流水线

DreamZero 基于 Flow Matching 框架进行训练，同时对视频 latent 和动作 token 施加扩散损失。训练使用 DeepSpeed ZeRO-2 分布式优化，并通过 LoRA 高效微调 14B 参数基础模型。

LoRA 配置： 基础模型（Wan2.1，约 14B 参数）完全冻结，仅在 DiT 的注意力层插入 LoRA 适配器（rank=4, alpha=4），大幅降低显存占用。

4. 推理流水线

推理时模型从纯随机噪声出发，通过 Flow Matching 迭代去噪恢复动作序列和视频帧。DreamZero 支持通过 WebSocket 的分布式多 GPU 推理服务。

分布式推理服务

5. 关键超参数表

6. 关键源文件表

7. Attention 机制深入解析

7.1 Q / K / V 的本质分工

Attention 的计算分四步：

1. q = W_q(x)              # x 投影成 Q
2. k = W_k(context)        # context 投影成 K
3. v = W_v(context)        # context 投影成 V（同一原材料，不同矩阵）
4. 权重 = softmax(q · kᵀ / √d)
5. 输出 = 权重 · v

K 和 V 的原材料相同（都来自 context），但投影矩阵不同，作用完全不同：

K 是抽象意义上的索引——它不直接出现在输出里，只决定权重分布，告诉模型"哪些位置比较重要"。V 才是真正被读取的内容。

7.2 Self-Attention vs Cross-Attention

区别只有一个：Q、K、V 的原材料从哪来。

DreamZero 的每个 DiT Block 里两者都有：先 Self-Attention 让视频 token 内部交流，再 Cross-Attention 读取文本和图像条件。

7.3 为什么 K 不能等于 Q

K=Q 意味着 W_k = W_q，K 和 Q 是同一个向量。此时：

权重 = softmax(q · qᵀ / √d)   # x 和自身的相似度

权重退化成"找和自己像的"，模型只能检索与自身相似的 token，无法找到"内容上互补但方向不同"的信息。

类比：图书馆里每本书的索引标签直接写的是"我想找机器人抓取的书"——所有书都在大喊需求，没有一本书在说"我能提供什么"，检索系统完全失效。

K 和 Q 必须用不同的投影矩阵，让"被检索"和"去检索"解耦。

7.4 多头注意力（Multi-Head Attention）

多头就是把 dim 切成 h 份，并行跑 h 套独立的 Q/K/V：

head_i 输出 = softmax(Q_i · K_iᵀ / √d) · V_i
最终输出 = Concat(head_0, ..., head_h-1) · W_o

多头的价值来源于 Q 的多样性——每个头问不同的问题，从 context 里提取不同侧面的信息。

Q 不能共享，K/V 可以共享（MQA）：

MQA（Multi-Query Attention）中所有 Q 头共享同一组 K/V，节省 KV Cache 显存，但不影响 Q 的多样性——40 个人用不同问题检索同一个书架，仍然能找出不同的重点。

7.5 W_q、W_k、W_v 各自如何被训练

三个矩阵的梯度路径不同，学到的东西也不同。

W_q 和 W_k 通过同一条路径（点积）互相塑造，W_v 的梯度路径完全独立，只由输出质量驱动。

7.6 QK Norm：为什么归一化能稳定训练

原始点积：

q · k = |q| × |k| × cos(θ)

模长和方向混在一起，W_q 和 W_k 需要同时控制两件事，容易震荡。

QK Norm 之后：

q_norm = q / |q|,  k_norm = k / |k|
q_norm · k_norm = cos(θ)     # 值域固定在 [-1, 1]

点积退化为纯余弦相似度，只剩方向这一个自由度，模长的干扰被彻底消除。W_q 和 W_k 只需要各自学"朝哪个方向"，优化目标从三个因素（|q|、|k|、θ）变成一个（θ），训练更稳定。

DreamZero 代码中对应实现（wan2_1_submodule.py）：

q = self.norm_q(self.q(x))        # QK Norm
k = self.norm_k(self.k(context))  # QK Norm