5003.大模型-架构-Diffusion-3.采样加速与引导生成

扩散模型训练稳定,但默认 1000 步逆向采样使推理成本成为部署瓶颈;生物 design 任务还要求对二级结构结合位点对称性等条件精确控制。DDIM 提供跳步确定性采样;**分类器引导(Classifier Guidance)无分类器引导(Classifier-free Guidance,CFG)**通过调节条件得分,在样本质量与条件服从度之间权衡——RFdiffusion 等蛋白质生成器 heavily 依赖 CFG 式条件化。

段末注释:CFG 在训练时随机丢弃条件,推理时用有条件与无条件得分差做外推;DDIM 为 Denoising Diffusion Implicit Models 的缩写。

前置阅读5003.大模型-架构-Diffusion-2.得分匹配与随机微分方程Diffusion-Math-0 概率与扩散数学基础


1. 采样加速的需求

设单次网络前向耗时 (c),步数 (T) 则生成一个蛋白质骨架约需 (T \cdot c)。若 (T=1000)、(c=50,\mathrm{ms}),单样本约 50 s——高通量 design 不可接受。目标:在 (T’ \ll T) 步内保持可接受的 designability / RMSD

加速路线概览:

方法 思路 典型加速比
DDIM 非马尔可夫、确定性子序列 10–50×
高阶 ODE 求解器 更大 (\Delta t)、更少步 5–20×
蒸馏(Distillation) 学生网络模仿多步教师 10–100×
潜空间扩散 LDM 低维 (z) 上扩散 2–4×(每步更轻)

2. DDIM:非马尔可夫跳步采样

Song et al.(2020)构造一族非马尔可夫前向过程,使相同的 (q(x_t\mid x_0)) 边缘分布与 DDPM 一致,但逆向可确定性、且只在子序列 ({\tau_1,\ldots,\tau_S}) 上更新。

2.1 更新公式((\eta=0) 确定性)

设 (\hat\epsilon = \epsilon_\theta(x_t,t)),预测 (\hat x_0 = \frac{1}{\sqrt{\bar\alpha_t}}(x_t - \sqrt{1-\bar\alpha_t},\hat\epsilon))。从步 (t) 跳到更早步 (t’):

[
x_{t’} = \sqrt{\bar\alpha_{t’}},\hat x_0 + \sqrt{1-\bar\alpha_{t’} - \sigma_t^2},\hat\epsilon + \sigma_t z,
]

(\eta \in [0,1]) 控制随机性:(\eta=0) 时 (\sigma_t=0),完全确定性;(\eta=1) 退化为 DDPM 类随机采样。

关键:可在 (S=50) 个均匀(或余弦)选取的 (t) 上运行,而训练仍用完整 (T=1000) 的 (\epsilon_\theta)。

2.2 与概率流 ODE 的关系

(\eta=0) 的 DDIM 等价于对概率流 ODE(第 2 篇)的一种离散化:同一 (x_T) 总得到同一 (x_0),便于插值编辑(inpainting 时在已知区域每步替换为 (q(x_t\mid x_0^{\mathrm{known}})) 的采样)。

概念锚点 · Inpainting 与条件分布:inpainting 每步将已知区域重置为 (q(x_t\mid x_0^{\mathrm{fixed}})) 的样本——即固定 (x_0) 条件下的前向加噪,与 Math-0 §4 的条件高斯一致。

概念图 条件分布 \(q(x_t \mid x_0)\):inpainting 的理论基础

图 1 DDPM 多步随机采样与 DDIM 跳步确定性采样对比


3. Classifier Guidance(分类器引导)

Dhariwal & Nichol(2021)在类条件扩散 (\epsilon_\theta(x_t,t,c)) 上,用预训练分类器 (p_\phi(c\mid x_t)) 引导采样 toward 目标类 (c):

[
\hat\epsilon = \epsilon_\theta(x_t,t,c) - w,\sqrt{1-\bar\alpha_t},\nabla_{x_t}\log p_\phi(c\mid x_t),
]

(w>0) 为引导强度(guidance scale)。直觉:在去噪方向叠加「使分类器更确信为 (c)」的梯度——本质是在 (\epsilon) 空间做条件得分的外推(与 Math-0 §11 得分方向一致)。

概念图 得分/梯度方向:CFG 与 Classifier Guidance 的几何直觉

生物场景:若存在预测「是否为 (\alpha)-螺旋丰富」的分类器,可引导生成高螺旋蛋白——但需可微、且在 noisy (x_t) 上训练稳定的分类器,工程成本高。


4. Classifier-free Guidance(CFG)

Ho & Salimans(2022):不单独训练分类器,训练时以概率 (p_{\mathrm{uncond}}) 将条件 (c) 置空(如全零嵌入、特殊 null token),学习同一网络 (\epsilon_\theta(x_t,t,c)) 同时建模条件与无条件分布。

推理时:

[
\hat\epsilon = \epsilon_\theta(x_t,t,\varnothing) + w,\big(\epsilon_\theta(x_t,t,c) - \epsilon_\theta(x_t,t,\varnothing)\big),
]

(w=1) 为标准条件生成;(w>1) 放大条件方向,提高条件服从度,但过大时多样性下降、样本可能「过拟合条件」或物理不合理。

4.1 在蛋白质 design 中的对应

RFdiffusion 将条件(二级结构、block adjacency、motif 约束等)编码为额外通道或嵌入,训练时使用 CFG dropout;推理时 (w \approx 1!-!10) 调节 scaffold 与 motif 的保真度。

条件类型 编码方式 CFG 作用
二级结构 逐残基 SS 标签 控制 (\alpha/\beta) 比例
Motif scaffolding 固定坐标 inpainting 保留活性位点几何
对称性 对称复制 + 群平均 寡聚体 design
结合靶点 靶点结构 cross-attention 结合剂 design

5. 其他条件注入机制

除 CFG 外,条件信息进入 (\epsilon_\theta) 的常见方式:

(1)通道拼接(concatenation)
将条件图(SS、距离图、序列 profile)与 (x_t) 在通道维拼接——U-Net 常用。

(2)交叉注意力(cross-attention)
条件序列(靶点残基、文本描述)作为 Key/Value,noisy 结构作为 Query——Stable Diffusion、部分 Chroma 模块采用。

(3)AdaGN / FiLM
条件向量调制各层 GroupNorm 的 scale/shift:(\mathrm{GN}(h) \cdot (1+s(c)) + b(c))。

(4)Inpainting / 硬约束
每步将已知区域替换为 (x_t^{\mathrm{known}} \sim q(x_t\mid x_0^{\mathrm{fixed}})),未知区域继续去噪——motif scaffolding 的标准做法。


6. 步数与引导强度的调参

步数 (S)

  • 过少:结构未收敛,RMSD 高、clash 多;
  • 过多:收益递减,线性增加延迟。
  • 经验:骨架扩散 50–200 步 + DDIM/Heun 是常见起点。

引导强度 (w)

  • (w=1):标准条件;
  • (w \in [2,10]):更强条件,motif 保真度↑,多样性↓;
  • (w) 过大:模式崩溃、非物理结构。

建议在验证集上画 designability vs. (S)motif RMSD vs. (w) 曲线,而非照搬图像生成的 (w=7.5)。


7. 蒸馏与其他加速(简述)

Progressive Distillation(Salimans & Ho):训练学生网络用 2 步模仿教师 4 步,迭代减半——最终 4–8 步生成。
Consistency Models(Song et al.):学习任意 (t) 到 (0) 的一步映射 (f(x_t,t))。
生物领域:蒸馏工作相对图像少,工业 design 仍以 DDIM + 中等步数 为主。


8. 生物信息学案例:条件生成工作流

motif scaffolding 为例:

  1. 固定活性位点残基坐标 (x_0^{\mathrm{motif}});
  2. 初始化全结构 (x_T)(motif 区域加噪至与周围一致,或 inpainting 掩码);
  3. 每步 DDIM:预测 (\hat\epsilon),应用 CFG(条件 = SS + motif 掩码);
  4. Inpainting:motif 位置重置为 (q(x_t\mid x_0^{\mathrm{motif}}));
  5. 输出 (x_0),用 ProteinMPNN 设计序列,AlphaFold2 验证 self-consistency

图 2 Motif scaffolding 条件生成工作流(DDIM + CFG + Inpainting)


9. 小结

DDIM 在保持 DDPM 训练不变的前提下,用确定性跳步将推理步数降一个数量级;CFG 用条件与无条件预测之差实现免分类器引导,成为 RFdiffusion 等工具的条件控制核心。部署生物扩散 pipeline 时,(S)(w) 是与 学习率 同等重要的超参,应针对 designabilitymotif 保真 联合标定。

段末注释:inpainting 指仅生成部分区域、其余区域由已知数据约束;self-consistency 指设计序列经折叠模型预测结构与生成结构一致。


本篇数学概念索引

概念 出现章节 延伸阅读
条件分布 / inpainting §2.2 Math-0 §4
得分 / 梯度引导(CFG) §3–§4 Math-0 §11

参考与延伸阅读

  • Song et al., Denoising Diffusion Implicit ModelsDDIM).
  • Dhariwal & Nichol, Diffusion Models Beat GANs on Image SynthesisClassifier Guidance).
  • Ho & Salimans, Classifier-Free Diffusion GuidanceCFG).
  • Watson et al., De novo design of protein structure and function with RFdiffusion.
  • Karras et al., Elucidating the Design Space of Diffusion-Based Generative Models(求解器与日程).
-------------本文结束感谢您的阅读-------------