Cosmos3-Nano Policy · fwd4xl LIBERO checkpoint

网络结构:一次去噪步里发生了什么

数字来自 A6000 上实际加载的 transformer/config.json 与 Diffusers 源码 transformer_cosmos3.py。一次动作预测要把这张图重复 30 次(30 个去噪步)。

联合序列(一次前向,共 388 个 token,每个 4096 维) UND · 文字 122 当前画面 50(干净) 未来画面 4 帧 × 50 = 200(加噪,模型同时预测) 动作 16(加噪) ← 理解通道 und_seq[0:122]:只做因果自注意力,看不到画面 生成通道 gen_seq[0:266] → 每个 token 都能看到全部 388 个 token 的 K/V gen 行 250–265 文字输入 {"actions":[{"task":"pick up the milk and place it in the basket", ...}]} Qwen 分词器 embed_tokens 词表 151936 → 122×4096;目标名在第 48–49 个 token 画面输入 第三视角 腕部相机 各 256×256,旋转后拼成 512×256 Wan VAE48 通道 · 16× 下采样 2×2 patch · proj_in 320×160 内容 → 10×5 个 patch 192 维 patch → 4096:当前帧 50 个 token 注:policy 模式不走 Qwen 视觉编码器 噪声输入 未来画面 latent 噪声→ proj_in → 200 token 动作噪声 16×64 action_proj_in(按机器人域,LIBERO=5) 64 → 4096,再加 action_modality_embed 时间步嵌入 t ×0.001 cos © sin TimestepEmbedding(FP32) Linear 256→4096 SiLU Linear 4096→4096 加到所有加噪 token MoT 解码层(Qwen3-VL 文本层) × 36 两套权重共享一次注意力:文字用“理解”权重,画面和动作用“生成”权重(mixture-of-transformers) 理解通道 UND(文字 122) RMSNorm to_q / to_k / to_vQK-RMSNorm · 3D mRoPE 因果自注意力只看文字 to_out + 残差 RMSNorm MLP(SwiGLU)+ 残差4096 → 12288 → 4096 输出 → 下一层文字表示。看不到画面,所以牛奶在哪、布局如何,文字表示完全不变。 文字 K/V 交给生成通道 语言影响动作的唯一入口(k_norm_und_for_gen) 生成通道 GEN(画面 250 + 动作 16) RMSNorm_gen add_q / add_k / add_v32 个 Q 头 · 8 个 KV 头 · 128 维 全注意力:Q 来自生成通道 文字 K/V当前画面未来画面动作 K/V = [文字 122 ; 生成 266],不加因果遮挡 to_add_out+ 残差 RMSNorm_gen mlp_moe_gen(SwiGLU)+ 残差4096 → 12288 → 4096,独立于文字 MLP 我们之前做过替换的位置:第 12、30 层的动作 token;第 36 层的文字 K/V 后缀;第 36 层 mlp_moe_gen 的响应。 每层两条通道各有独立的 RMSNorm、QKV/输出投影和 MLP 权重; 只有注意力这一步把两条通道连起来,而且方向是单向的:生成通道读文字,文字不读生成通道。 输出与执行 norm_moe_gen(最终 RMSNorm) proj_out未来画面预测 action_proj_out4096→64,域 5取前 10 维 Flow UniPC 采样器 输出是“去噪速度”,不是机器人速度 30 步:更新噪声 → 回到左边再算一次 动作块 16 × 10 xyz 3rot6d 6夹爪 按 q01/q99 反归一化,逐帧相对 → 增量 转成 7D OSC 命令,裁剪到 [−1, 1] MuJoCo / LIBERO 执行 16 步 拍新画面,再预测下一块(共 8 块 = 128 步) 参数规模(config.json): · 36 层,隐藏维 4096,MLP 12288 · 32 个注意力头,8 个 KV 头(GQA),头维 128 · 3D mRoPE(时间/高/宽 = 24/20/20) · 动作头按 32 个机器人域分别存权重 · BF16;时间嵌入 FP32;guidance 1 动作和未来画面在同一序列里一起去噪, 所以动作 token 也能读到“模型自己预测的未来画面”。
文字 / 理解通道画面 / 生成通道动作物理执行红色箭头:语言进入动作的唯一路径
读图要点

这张图对诊断意味着什么

1. 文字表示与画面无关。理解通道只做因果自注意力,看不到画面。所以不管牛奶在原位还是被移远,“牛奶”这个词在 36 层里的表示都是同一份。“文字没读懂画面”不可能是原因;问题只能出在生成通道怎么使用这份文字。

2. 语言只有一个入口。每层生成通道的注意力里,动作和画面 token 用自己的 Q 去读文字的 K/V。之后画面和动作走的是独立的 mlp_moe_gen。如果“抓哪个物体”是靠画面 token 之间的注意力(动作读当前画面、读未来画面)决定的,文字那一小部分 K/V 很容易被盖过去。

3. 动作会读模型自己预测的未来画面。未来画面和动作在同一序列里一起去噪。如果预测的未来画面里,机械臂已经朝奶酪盒的位置走,动作 token 读到的也是这个画面。所以要区分两种可能:目标是由“当前画面+文字”决定的,还是由“自己想象出来的未来”锁定的。