5003.大模型-架构-Diffusion-Example-1.蛋白序列离散扩散数值例题

蛋白质一级序列是有限字母表 (\mathcal{V}) 上的离散 token 串,不能对氨基酸直接加高斯噪声。离散扩散(Discrete Diffusion)——以 D3PM(Austin et al., 2021)及 MASK 吸收态变体为代表——用转移矩阵逐步把 token「腐蚀」为 [MASK],再训练网络逐步填回。本文与 Example-0 连续 DDPM 手算 平行:同样走「采样 (x_0) → 随机 (t) → 固定前向造 (x_t) → 网络预测 → 损失反传」,但符号从 (\epsilon\in\mathbb{R}^d) 换成 分类分布 + 交叉熵

段末注释:D3PM 为 Structured Denoising Diffusion Models in Discrete State-Spaces;MASK 吸收态指 token 一旦变为 [MASK] 则保持,类似多步 MLM。

前置阅读Diffusion-4 离散扩散与序列生成Example-0 前向训练数值例题
连续 DDPM 对照Diffusion-1 §5
条件 / CFGDiffusion-3 §4–§5

插图5003.大模型-架构-Diffusion-Example/example-fig04–06*.png


1. 与 Example-0 的对照

Example-0(图像/坐标) Example-1(蛋白序列)
数据 (x_0) 向量 (\mathbb{R}^2) token 串 (\mathcal{V}^L)
前向 高斯加噪 (x_t=\sqrt{\bar\alpha_t}x_0+\cdots) 转移矩阵 (Q_t),MASK 吸收
损坏形态 连续噪声 部分位置变为 [MASK]
网络输出 预测 (\epsilon) 每位置 (|\mathcal{V}|) 维 logits(常等价预测 (x_0))
损失 MSE (|\epsilon-\epsilon_\theta|^2) 交叉熵 (-\log p_\theta(x_0\mid x_t,t))
生成 从 (x_T\sim\mathcal{N}(0,I)) 去噪 从全 [MASK] 逐步填 token

2. 玩具设定:四肽 (x_0)

取长度 (L=4) 的肽段(真实蛋白常用 (L\gg4),此处仅手算):

[
x_0 = [,\mathrm{M},,\mathrm{K},,\mathrm{G},,\mathrm{D},]
]

简化词汇表(玩具,便于表格;标准蛋白为 20 种氨基酸 + 特殊 token):

索引 0 1 2 3 4 5
token A D G K M [MASK]

故 (x_0 = [4,,3,,2,,1])(M–K–G–D)。扩散步 (T=4),与 Example-0 一致。

噪声日程(每步独立 MASK 概率 (\beta_t)):

(t) (\beta_t) 该步「变为 MASK」概率
1 0.20 20%
2 0.35 35%
3 0.50 50%
4 0.65 65%

吸收规则:某位置一旦成为 [MASK],后续步保持 [MASK](不再变回原氨基酸)。


3. 前向:MASK 吸收态如何腐蚀 (x_0)

每个位置独立,单步转移:

[
q(x_t^{(\ell)} \mid x_{t-1}^{(\ell)}) =
\begin{cases}
1-\beta_t & x_{t-1}^{(\ell)} \text{ 保持} \
\beta_t & x_{t-1}^{(\ell)} \to [\mathrm{MASK}]
\end{cases}
]

从 (x_0) 到 (x_t) 闭式:位置 (\ell) 仍为原 token 当且仅当前 (t) 步都未被打 MASK

[
\mathbb{P}(x_t^{(\ell)} = x_0^{(\ell)}) = \bar\gamma_t = \prod_{s=1}^{t}(1-\beta_s).
]

手算 (\bar\gamma_t)

(t) (\bar\gamma_t) 仍保留原 token 的概率
1 0.800 80%
2 0.520 52%
3 0.260 26%
4 0.091 9.1%

图 1 MASK 吸收扩散:\(t\) 增大时序列中 [MASK] 增多


4. 训练第 1–2 步:采样 (x_0) 与 (t)

[
x_0 = [\mathrm{M},,\mathrm{K},,\mathrm{G},,\mathrm{D}] \sim q(x_0), \qquad t = 3.
]

(x_0) 来自训练集(UniRef、Pfam 等上的真实肽段/蛋白窗口);(t=3) 表示「中等腐蚀」档。


5. 训练第 3 步:构造损坏序列 (x_t)

5.1 逐步腐蚀(理解用)

操作(本例指定随机结果) 序列
0 (x_0) M – K – G – D
1 位置 1 以 (\beta_1) 被 MASK [MASK] – K – G – D
2 位置 3 被 MASK [MASK] – K – [MASK] – D
3 位置 4 被 MASK [MASK] – K – [MASK] – [MASK]

5.2 训练用写法

[
\boxed{
x_3 = [,\mathrm{[MASK]},,\mathrm{K},,\mathrm{[MASK]},,\mathrm{[MASK]},]
}
]

等价索引 ([5,,3,,5,,5])。监督标签仍是完整 (x_0)——与 Example-0 里「已知 (\epsilon)」类比:损坏样本 (x_t) 由固定规则 + 已知随机 MASK 结果构造,真实 token (x_0) 来自数据集。

段末注释:实现中常对 (x_t \sim q(x_t\mid x_0)) 采样而非逐步枚举;本例固定一条轨迹便于阅读。


6. 训练第 4 步:网络与交叉熵损失

6.1 输入 / 输出

输入:((x_3,, t=3))——损坏 token 串 + 时间步 embedding。

输出:每位置 (\ell) 一组 logits (\mathbf{z}\ell \in \mathbb{R}^{|\mathcal{V}|}),经 softmax 得 (p\theta(x_0^{(\ell)}\mid x_3,t))。

骨干常用 Transformer(双向,非 AR);(t) 注入方式与 Example-0 相同(正弦 / MLP embedding)。

图 2 蛋白序列离散扩散单步训练流程

6.2 简化训练目标((x_0)-prediction)

常见目标:直接预测干净 token(与 Diffusion-4 §3 的 ELBO surrogates 一致):

[
\mathcal{L}{\mathrm{CE}} = -\sum{\ell=1}^{L} \log p_\theta\big(x_0^{(\ell)} \mid x_3,, t=3\big).
]

情形 A · 四位置全预测正确

位置 (\ell) 真实 (x_0^{(\ell)}) (x_3^{(\ell)}) 网络预测 损失贡献
1 M [MASK] M ✓ 0
2 K K K ✓ 0
3 G [MASK] G ✓ 0
4 D [MASK] D ✓ 0

(\mathcal{L}_{\mathrm{CE}} = 0)。

情形 B · 位置 3 错猜为 A

设 (p_\theta(\mathrm{G})=0.2),(p_\theta(\mathrm{A})=0.5)(其余忽略),则第 3 位贡献 (-\log 0.2 \approx 1.61),梯度增大 (p_\theta(\mathrm{G}))。

与 Example-0 MSE 的对照

Example-0 Example-1
标签:噪声 (\epsilon) 标签:原 token (x_0^{(\ell)})
(|\epsilon-\epsilon_\theta|^2) (-\log p_\theta(x_0^{(\ell)}\mid x_t,t))
猜对噪声 = 还原 (x_0) 猜对 token = 直接还原 (x_0)

7. 一个 minibatch 里发生什么

样本 (x_0)(数据集肽段) 随机 (t) 构造 (x_t) 损失
1 MKFL… 512 (q(x_t\mid x_0)) 采样 (\sum_\ell -\log p_\theta(x_0^{(\ell)}\mid x_t,t))
2 AGVT… 8
3 CPQY… 301

同一套 (\theta) 在所有序列、所有 (t) 上共享——学的是蛋白序列分布 (q(x_0)),不是「每条序列一个模型」。(t) 经 embedding 注入、非每档独立网络,与 Example-0 §6.3 连续 DDPM 相同。


8. 条件生成:motif 与「提示词」的蛋白版

无条件训练学 (p(x_0))。Motif scaffolding / 条件 design 学 (p(x_0\mid c)):

文生图 蛋白序列离散扩散
条件 (c) = 文本 prompt(CLIP 嵌入) (c) = 固定 motif、二级结构标签、pH/温度等
训练 (\epsilon_\theta(x_t,t,c)) 训练 (p_\theta(x_{t-1}\mid x_t,t,c)) 或 (p_\theta(x_0\mid x_t,t,c))
CFG 强化 prompt CFG 强化 motif / 结构条件(Diffusion-3 §4.1

8.1 Motif Inpainting 手算示意

设计 8 肽固定活性位点 (\mathrm{H}\mathrm{D}) 于位置 3–4,侧翼 6 个位置由扩散生成:

[
x = [\text{?},,\text{?},,\mathrm{H},,\mathrm{D},,\text{?},,\text{?},,\text{?},,\text{?}]
]

采样((t=T\to1))每步:

  1. 网络对全序列输出 logits;
  2. 非固定位置按 (p_\theta) 采样 token;
  3. 固定位置强制为 (\mathrm{H},\mathrm{D})(或从 (q(x_{t-1}\mid x_0^{\mathrm{fixed}})) 重采样,与 inpainting 一致)。

图 3 Motif Inpainting:固定 HD,扩散生成侧翼序列

这与 Example-0 §11 文生图条件 同构:损坏/去噪骨架不变,额外条件 (c) 告诉网络「哪些 token 必须保留、其余填什么分布」。


9. 生成阶段:从全 MASK 到肽段

训练完成后(不再使用 (x_0))。与连续图像从 (x_T\sim\mathcal{N}(0,I))、(t=T) 起去噪(Example-0 §10)类比:

  1. 初始化 (x_T = [\mathrm{MASK},\ldots,\mathrm{MASK}])(或均匀随机 token)——对应离散扩散的最大腐蚀档
  2. (t = T,\ldots,1):同一 (p_\theta) 输出每位置分类分布,采样 (x_{t-1});
  3. 得到 (x_0) 作为新设计序列
  4. 可选:用 ESMFoldOmegaFold 等做结构一致性筛选(Diffusion-4 §6)。

10. 与 AR / MLM 的一句话选型

范式 本玩具例是否适用 备注
离散扩散(本文) motif inpainting、并行填 MASK
AR(ProGen2 等) 左→右生成,非本例 长序列生态成熟
MLM(ESM) 单步遮盖,非多步 schedule 常作表征而非生成

多数 3D 蛋白 design 仍用 RFdiffusion(连续结构扩散)+ ProteinMPNN(序列)纯序列离散扩散适合 motif 侧翼全局编辑 等场景(见 Diffusion-4 §5–§6)。


11. 小结

环节 Example-1 中的角色
(x_0 \sim q(x_0)) 真实肽段/蛋白窗口
固定前向 (q(x_t\mid x_{t-1})) MASK 吸收转移,不训练
随机 (t) 噪声档位
构造 (x_t) 部分 [MASK] + 部分可见 token
(p_\theta(x_0\mid x_t,t)) Transformer 分类去噪
(\mathcal{L}_{\mathrm{CE}}) 各位置 token 交叉熵之和
条件 (c) motif / 结构 / CFG(§8)

记忆:连续 DDPM 猜 (\epsilon);离散序列扩散猜 原 token(或 (x_{t-1}) 分布)——监督都来自已知 (x_0),前向只是「把 (x_0) 弄坏」的机制不同。


本篇概念索引

概念 出现章节 延伸阅读
MASK 吸收态 / (\beta_t,,\bar\gamma_t) §2–§3 Diffusion-4 §2.1
构造 (x_t) §5 Diffusion-4 §2
(\mathcal{L}_{\mathrm{CE}}) / (x_0)-prediction §6 Diffusion-4 §3
Motif inpainting / 条件 (c) §8 Diffusion-3 §4–§5Example-0 §11
与 AR/MLM 选型 §10 Diffusion-4 §6
连续 DDPM 对照 §1 Example-0

参考与延伸阅读

  • Austin et al., Structured Denoising Diffusion Models in Discrete State-SpacesD3PM).
  • Ho & Salimans, Classifier-Free Diffusion Guidance(CFG,条件序列/结构通用).
  • 本仓库 Diffusion-4Example-0
-------------本文结束感谢您的阅读-------------