Cosmos 3 · OmniMoT ▸ 完整前向传播

模型前向 —— 一次完整前向,从原始数据到输出数据

五种形状迥异的原始数据,各走自己的编码器压成 latent、投影进同一个 hidden、打包成一条联合序列(输入端 I); 整条序列一起过 N 层 MoT trunk 做单次前向;输出端按 token 类型分流、各经反向投影回自己的模态再解码(输出端 O)。 I 与 O 是镜像对称的一对(vae2llm↔llm2vaepatchify↔unpatchifypack↔unpack),trunk 夹在正中。 文本走自回归、生成模态走扩散——两条推理回环在底部标出。

理解 / 文本通路(und) 生成通路(gen · 带噪 latent) VAE 外部预训练编码器(冻结) 推理回环(采样 ×N / AR 逐 token)
输入端 I · 异构数据 → 一条统一序列
原始数据 · 形状各不相同
🅣文本 prompt / 指令
🖼🎞图像 / 视频 像素 RGB
🔊音频 波形 / 谱
🦾动作 关节 / 位姿
"a cat on a sofa…"变长 UTF-8 字符串
[3, T, H, W]RGB 视频张量(帧 × 高 × 宽)
[wav / mel]音频波形或梅尔谱
[T, action_dim]逐帧动作向量 · 按 embodiment
各自的编码器 → latent · 外部预训练、推理时冻结
Qwen3-VL tokenizer BPE 分词 + 特殊 token text → token_ids [L] set_up_tokenizers
Wan VAE · 4×8×8 时间 ↓4 · 空间 ↓8,编成连续 latent → latent [C=16, T', H', W'] wan2pt1_vae_4x8x8.py
audio tokenizer 音频 VAE/编码器,时间压缩 → sound latent [T'', sound_dim] tokenizers/audio/
(无 VAE) 动作已是低维向量,直接进投影 [T, action_dim=32] action_dim / num_embodiment_domains
patch 化 + 投影进同一个隐空间 · 统一成 [·, hidden]
embed_tokens 词嵌入查表 token_ids → [L, hidden]
patchify → vae2llm p=2 分块,cthpwq→thwpqc [N_p, p²·C=64][N_p, hidden] + time_embedder(t) 仅带噪 token patchify_and_pack_latents vae2llm
sound2llm 线性投影 sound_dim → hidden + sound_modality_embed + time(t) sound2llm
action2llm DomainAwareLinear · 每 embodiment 一套 bias action_dim → hidden + action_modality_embed + time(t) action2llm
④ scatter → pack 成一条联合序列 [N_total, hidden]
每个模态的 token 按 sequence_indexes 散射进同一条序列 + mRoPE 位置编码 + 特殊 token 分隔 · 文本=und、latent=gen
<bos>texttext <img>img patchimg patch <vid>vid patch <snd>sound <act>action
MoT trunk · 单次前向(整条序列一起过 N 层)

Unified MoT ×N 层

language_model · MoTDecoderLayer MoTDecoderLayer
每层 = 共享 Self-Attention(gen 看 und+gen 的联合注意力,跨模态在此单向交互)+ 按 token 路由的 und / gen 双专家分支。 序列长度与 token 顺序全程不变——进出都是 [N_total, hidden]
输出端 O · 一条序列 → 各模态输出 ↕ 与输入端镜像
⑥ 按 token 类型分流(und 文本位置 | 各生成模态位置)
反向投影回各模态 · hidden → 模态 镜像 ③ 的进投影
norm → lm_head hidden → vocab logits → 采样 token
llm2vae → unpatchify hidden → patch_latent → 速度 v_vision llm2vae unpatchify_and_unpack
llm2sound → unpack hidden → sound_dim → 速度 v_sound llm2sound
llm2action → unpack hidden → action_dim → 速度 v_action llm2action
解码 → 输出数据 镜像 ②①:VAE / tokenizer 反向
🅣 文本 / 推理
detokenize
🖼🎞 图像 / 视频
clean latent → VAE 解码
🔊 音频
sound tok 解码
🦾 动作
ActionProcessor 还原
训练时 · 单次前向 生成模态在随机 t 加噪一次,走 ③→⑤→⑦v_pred,与目标 v = x₀−x₁ 算 flow-matching 损失;文本走 teacher-forcing 交叉熵。不循环
推理时 · 循环 同一张图重复跑:生成模态 ×N 步扩散去噪,文本逐 token 自回归——见上方回环带。 ⏱️ 训练/推理时序 →
连接的骨架:每个模态有一对进/出隐空间的投影—— embed_tokens/lm_headvae2llm/llm2vaesound2llm/llm2soundaction2llm/llm2action。 左半张图(① → ④,进)与右半张图(⑥ → ⑧,出)互为镜像;中间那段 trunk 所有模态共用同一条序列、同一套权重。 模态差异只在两端的编解码,以及输出端的用法(AR 采样 vs 扩散积分)。