Cosmos 3 · 生成模态 ▸ 扩散采样循环内部

推理采样 —— 一步步把噪声磨成 latent

模型前向图底部那条「扩散采样 ×N」回环的放大版。生成从纯噪声起步,循环 num_steps(默认 35)次:每步跑一遍 MoT gen 分支预测速度 v、用 CFG 把「有条件/无条件」两次预测混合、再交给求解器v 换成下一步 latent; 收敛后 VAE 解码成像素。文本侧是自回归(见模型图),这里只讲扩散侧。

生成 latent(去噪中) 循环 / 采样控制 v 速度预测(rectified flow) CFG classifier-free guidance
纯噪声 latent
按每个样本的 seed 生成 x ~ 𝒩(0, I),形状同目标 latent
cond / uncond 文本
正 prompt 与空/负 prompt 各 tokenize 一份,供 CFG 两支
可选 · prompt 上采样
Reasoner tower 先把 caption 自回归改写得更详细,再进采样 _maybe_apply_prompt_upsampling
↻ 去噪循环 step = 0 … num_steps−1 · σ 从 1 → 0(噪声 → 干净)
跑 MoT gen 分支 → 预测速度 v
把当前 latent + timestep 打包成序列,过一遍 N 层 MoT trunk(即模型前向图的 gen 通路), 输出端 llm2vae 得到速度场 v_θ(x, t)🧠 这一步 = 整张模型前向图 →
▼ 同一 latent,跑两次(cond / uncond)
CFG · 有条件 / 无条件混合
vcond = MoT(x, t, 正 prompt)
vuncond = MoT(x, t, 空/负 prompt)
v = vuncond + guidance · (vcond − vuncond) guidance 默认 1.5
guidance_interval 只在指定 t 区间开 CFG(省一半算力);normalize_cfg 可把结果模长拉回 cond。 cfgp 并行时两支分到两张卡、各算一次再 P2P 交换。
▼ 拿到融合速度 v,交给求解器
求解器 · 把 v 变成下一步 latent(三选一)
FixedStep 蒸馏模型 · SDE
x₀ = x − σ·v
x ← (1−σ')·x₀ + σ'·ε
少步一次性去噪,每步回注新噪声 ε
UniPC base 默认 · ODE
多步
预测-校正
高阶 ODE 求解,少步高质量
EDM Karras · ODE
σ_max
= 80
Karras σ 调度,经典 diffusion 求解
可选 · 条件注入(i2v / inpaint)
把已知区域(首帧 / 参考)硬贴回去:x ← mask·ref + (1−mask)·x,保证条件帧永远不被噪声污染。
↺ 回到 ②,σ 降一档,重复 —— 共 num_steps
收敛 → 解码输出
干净 latent x₀
→ VAE decode →
🖼🎞 像素 · 🔊 音频 · 🦾 动作
一句话 采样 = 反复调用整张模型前向图,每次只前进一小步 σ。贵的是那 35 次(× CFG 两倍 = 70 次)完整 trunk 前向—— 这也是为什么要 cfgp / cp 并行把单次前向摊到多卡。