5003.大模型-架构-Diffusion-5.潜空间扩散与结构生成实践

高维像素、全原子坐标或长序列上直接跑 1000 步 DDPM 成本过高;蛋白质小分子还具有**旋转平移(SE(3)/E(3))**对称性——在各向同性欧氏空间加噪会破坏几何先验。潜扩散模型(Latent Diffusion Model,LDM) 在压缩表示上扩散;等变扩散 在流形/群结构上定义噪声与去噪网络。本文串联 Stable Diffusion 式 LDM 思想与 RFdiffusionChromaEDM 等生物结构实践,给出可落地的 design pipeline 与评估指标。

段末注释:LDM 先在自编码器潜空间做扩散再解码;SE(3) 为三维刚体运动群;E(3) 含反射,适用于分子手性需谨慎。

前置阅读5003.大模型-架构-Diffusion-1.前向扩散与变分下界5003.大模型-架构-Diffusion-3.采样加速与引导生成Diffusion-Math-0 概率与扩散数学基础


1. 潜扩散(LDM)的两阶段范式

Stable Diffusion(Rombach et al., 2022)流程:

  1. Stage A:训练 VAE(或 VQ-VAE)将数据 (x) 编码为潜变量 (z = \mathcal{E}(x)),解码 (\mathcal{D}(z) \approx x);
  2. Stage B:在 (z) 上训练 DDPM,采样 (z_0) 后 (\hat x = \mathcal{D}(z_0))。

[
\text{像素空间 DDPM: } \dim(x) \sim 10^6 \quad \Rightarrow \quad \text{LDM: } \dim(z) \sim 10^4.
]

优势:扩散每步在低维运算;VAE 承担高频细节。
风险:VAE 重建误差上限;生物场景需验证 designability 是否被 VAE 瓶颈限制。

1.1 生物数据上的 LDM

领域 潜空间 (z) 解码器 代表
分子 2D/3D 图 VAE 潜向量 图解码为 SMILES/构象 部分 MolLDM
蛋白质 折叠模型潜表示 / 距离图 结构解码 研究中;全原子仍多坐标扩散
单细胞 scVI 潜空间 表达谱解码 scDiff

蛋白质领域:全原子或骨架坐标扩散仍占主导(RFdiffusion 直接在 (C_\alpha) / frame 上扩散),因结构 VAE 重建误差对 RMSD 敏感。LDM 更适合低维表型与图像化表示(密度图、距离矩阵)结合的任务。


2. SE(3) / E(3) 等变扩散

蛋白质坐标在全局旋转平移下物理等价。若在 (\mathbb{R}^{3N}) 直接加 (\mathcal{N}(0,I)) 噪声:

概念锚点 · 高斯噪声的前提:各向同性高斯假设各维独立同方差——对未中心化/未归一化的坐标,(\beta_t) 与数据尺度必须匹配(Math-0 §3、§10)。

概念图 高斯分布:直接对坐标加噪时的方差匹配问题

  • 噪声随结构旋转协变;
  • 网络需从数据「学会」忽略绝对朝向,样本效率低。

等变扩散做法:

(1)在局部 frame 上扩散(FrameDiff / RFdiffusion 思路)
每个残基用 SE(3) frame (T_i \in \mathrm{SE(3)) 表示位置与朝向;对 平移旋转(so(3) 上 IGSO(3) 噪声)分别扩散,网络用 SE(3)-Transformer / IPA 处理。

(2)E(3)-equivariant GNN(分子)
原子坐标 (x \in \mathbb{R}^{3\times M}),网络 (f) 满足 (f(Rx+b) = R f(x) + b)(或标量输出不变)。EDMGeoDiffEGNN / Tensor Field Network 上预测去噪 (\epsilon)。

(3)内坐标 / 扭转角扩散
对键长、键角、二面角扩散,天然降维并减少无效自由度;需小心 拓扑 约束与 periodicity(二面角 (2\pi))。

图 1 欧氏坐标扩散与 SE(3)/内坐标等变扩散对比


3. RFdiffusion:骨架扩散 design 解剖

RFdiffusion(Watson et al., 2023)基于 RoseTTAFold 结构模块,在 (C_\alpha) 坐标(及辅助特征)上训练扩散去噪器。

3.1 训练数据与目标

  • 数据:PDB 单体/复合物,过滤分辨率;
  • 加噪:VP 型 schedule 于坐标(常配合 centering);
  • 目标:(\epsilon)-prediction 或等价去噪;网络继承 RoseTTAFoldpair/track 更新。

概念锚点 · 坐标作为随机向量:(x_0 \in \mathbb{R}^{3N}) 经中心化后 (\mathbb{E}[x_0]\approx 0),加噪 (\epsilon \sim \mathcal{N}(0,\mathbf{I})) 才有明确 SNR 含义(Math-0 §2、§5)。

概念图 随机向量 \(x_0\):蛋白骨架坐标

3.2 推理模式

模式 条件 用途
Unconditional 新颖折叠采样
Motif scaffolding Inpainting 固定坐标 酶/结合位点 scaffold
Binder design 靶点结构 + 接口条件 结合蛋白 design
Symmetric oligomer 对称群约束 环状/二聚体

3.3 条件与 CFG

二级结构(SS)、block adjacency 等 one-hot 图作为条件;训练 CFG dropout,推理调节 guidance scale(见第 3 篇)。

3.4 标准后处理 pipeline

  1. RFdiffusion → 骨架 (C_\alpha);
  2. ProteinMPNN → 氨基酸序列;
  3. AlphaFold2 / ESMFold → 预测结构;
  4. scTM / RMSDself-consistency 筛选。

Designability:生成结构能否被「设计序列 + 折叠」还原——生物扩散的核心指标,优于纯生成似然。


4. Chroma 与可编程生成

Chroma(Ingraham et al., 2023)强调可编程条件:

  • 语义条件:自然语言式属性(通过分类器或嵌入);
  • 几何条件:对称、子结构、结合口袋形状;
  • 联合序列–结构生成模块(与纯骨架扩散 + MPNN 流水线不同)。

算法上仍属扩散/流族 + 等变网络;产品化侧重条件 API全原子输出。阅读论文时可对照:扩散变量(坐标 vs 序列 vs 联合)、条件注入层(AdaGN vs cross-attention)。


5. 分子 3D:EDM 与 GeoDiff

EDM(Karras et al.)虽出自图像社区,其 (\sigma)-空间预处理Heun 求解器GeoDiff 等分子模型借鉴:

  • 每个原子 3D 坐标;
  • EGNN 预测去噪,保证 E(3) 等变;
  • 评估:Validity(化学价)、UniquenessNSPDK 与对接 Vina

与蛋白质对比:分子 (M) 较小(十至数百原子),但化学 validity 约束强;蛋白质 (N) 大但 backbone 约束相对统一。


6. 完整 design 工作流(实践清单)

图 2 蛋白质 structure design 完整 pipeline(扩散 → 序列 → 折叠 → 筛选)

超参起点(需项目内标定):

  • 扩散步数 50–200(DDIM / Heun);
  • CFG (w \in [1, 10]);
  • Inpainting 每步重置 motif 坐标;
  • 生成 100–1000 条,MPNN 每条 8 序列,AF2 过滤 pLDDT > 70scTM > 0.5

常见失败模式

  • 步数过少 → steric clash
  • (w) 过大 → 多样性丧失、hallucinated 二级结构;
  • 忽略 centering → 扩散学绝对原点偏移,浪费容量。

7. 评估指标汇总

层级 指标 含义
结构 RMSD 与参考/motif 偏差
结构 clash score 原子重叠
设计 scTM / scRMSD 序列–结构自洽
设计 designability % 通过 AF 筛选比例
功能 对接打分 / ddG 结合 design
分子 Validity / SA / QED 可合成、类药
序列 motif 保留 离散/ inpainting 任务

8. 与系列前文的映射

概念 在结构 practice 中的体现
DDPM / (\epsilon)-pred RFdiffusion 训练目标
SDE / EDM (\sigma) 分子 GeoDiff、部分求解器
DDIM + CFG 推理默认配置
离散扩散 序列侧翼;主链多用 MPNN
LDM scRNA 等;结构域较少纯 LDM

9. 小结

潜扩散通过 VAE 降维加速,在图像与单细胞已成熟;蛋白质结构当前以等变坐标/frame 扩散为主,RFdiffusion 代表 inpainting + CFG + DDIM 的工业组合。SE(3)/E(3) 对称性应写入架构而非仅靠数据增强。落地 design 时,扩散只是 pipeline 第一段;ProteinMPNN + AlphaFold2 筛选designability 指标与算法本身同等重要。

段末注释:scTM 为 predicted TM-score 相对生成结构;IGSO(3) 为 SO(3) 上各向同性高斯噪声的推广。


本篇数学概念索引

概念 出现章节 延伸阅读
高斯噪声与坐标尺度 §2 Math-0 §3、§10
随机向量 / 中心化 §3.1 Math-0 §2、§5

参考与延伸阅读

  • Rombach et al., High-Resolution Image Synthesis with Latent Diffusion ModelsLDM).
  • Watson et al., De novo design of protein structure and function with RFdiffusion.
  • Ingraham et al., Illuminating protein space with a programmable generative modelChroma).
  • Yim et al., SE(3) diffusion model with application to protein backbone generationFrameDiff).
  • Hoogeboom et al., Equivariant Diffusion for Molecule Generation in 3DEDM).
  • Jumper et al., AlphaFold2self-consistency 评估基础).
-------------本文结束感谢您的阅读-------------