【AlphaFold3】网络架构篇（6）|Diffusion Module讲解

• 博主简介：努力学习的22级计算机科学与技术本科生一枚🌸
• 博主主页： @Yaoyao2024
• 往期回顾：【AlphaFold3】网络架构篇（5）|Template embedding & Pairformer stack
• 每日一言🌼: 勤学如春起之苗，不见其增，日有所长。辍学如磨刀之石，不见其损，日有所亏。——陶渊明🌺

https://deepwiki.com/google-deepmind/alphafold3/5-model-architecture

一、扩散模块-Diffusion module

翻译：

在AlphaFold 2中，最终结构通过使用不变点注意力(invariant point attention）的结构模块实现。而在AlphaFold 3中，我们用一个相对标准的、作用于所有原子的非等变点云扩散模型替代了它（算法18及主文章图2b）。在训练过程中，我们训练一个去噪器（denoiser），基于主干网络的特征，从所有重原子的位置中去除高斯噪声。该去噪器基于现代Transformer，但做了几处修改以更适配任务需求。主要改动如下：

• 我们通过多种方式整合来自主干网络的条件信息（conditioning) ：从单个嵌入初始化激活值，对单个条件使用自适应层归一化（Adaptive Layernorm）的变体，对成对条件使用对数偏置。
• 我们使用了标准的现代Transformer技巧（如SwiGLU）以及AlphaFold 2中使用的方法（门控机制）。
• 我们采用两级架构：先作用于原子，再作用于令牌，最后再回到原子。

值得注意的是，该Transformer仅使用单个线性层嵌入所有原子位置，并在最后用单个线性层投影更新量，不涉及任何几何偏置（如局部性或SE(3)不变性）。这与当前使用更强领域特定归纳偏置的趋势形成对比。

该架构的细节见算法20。

讲解：

🪧 一、核心功能：用扩散模型预测原子坐标

扩散模块（Diffusion Module）是AlphaFold 3生成最终分子结构的核心组件，替代了AlphaFold 2中基于不变点注意力的结构模块。其核心功能是：通过“加噪-去噪”的迭代过程，从随机噪声中逐步生成高精度的原子三维坐标，且整个过程受主干网络的特征（如进化信息、模板约束）引导。

• 与AlphaFold 2的关键差异：
  • AlphaFold 2依赖“结构模块+不变点注意力”直接预测结构，依赖SE(3)等变假设；
  • AlphaFold 3采用“非等变点云扩散模型”，不依赖几何偏置（如局部性、旋转不变性），通过数据驱动学习原子位置分布，更灵活适配蛋白质、核酸、配体等多分子类型。

👉🏻算法18：扩散采样流程：从噪声到结构的迭代

算法18描述了扩散模型的“采样过程”——如何从完全随机的噪声逐步去噪，最终得到原子坐标。步骤解析：

参数	含义与作用
{ f ∗ } \{\mathbf{f}^*\} {f∗}	分子的基础特征集合（如残基类型、链标识等）
{ s i i n p u t s } \{\mathbf{s}_i^{\mathrm{inputs}}\} {siinputs​}	InputEmbedder得到的单特征表示（每个令牌的独立特征）
{ s i t r u n k } \{\mathbf{s}_i^{\mathrm{trunk}}\} {sitrunk​}	主干网络(Pairformer)输出的单特征表示（整合进化、结构等信息）
{ z i j t r u n k } \{\mathbf{z}_{ij}^{\mathrm{trunk}}\} {zijtrunk​}	主干网络（Pairformer）输出的成对表示（描述令牌i与j的关系）
$\text{Noise Schedule }[c_0,c_1,\dots ,c_T]$	噪声调度表，控制噪声强度随时间的衰减（c0为初始噪声强度，cT为最终噪声强度）
${\gamma }_0=0.8$	噪声注入强度的初始系数
${\gamma }_{\min }=1.0$	噪声注入强度的阈值（低于此值时停止注入额外噪声）
$\text{noise scale }\lambda =1.003$	噪声缩放系数，控制每步添加的噪声幅度
$\text{step scale }\eta =1.5$z	步长系数，控制位置更新的幅度

1. 初始噪声采样（步骤1）：
   从初始噪声分布（$c_0\cdot \mathcal{N}(\vec{0},{\mathbf{I}}_3)$）采样得到原子的初始带噪声位置${\vec{\mathbf{x}}}_l$（三维向量），此时位置完全随机。

2. 多时间步迭代去噪（步骤2-12）：
   遍历噪声调度表（[c1, ..., cT]，控制噪声强度随时间衰减），每一步逐步降低噪声：

   • 步骤3：中心随机增强：调用👉🏻算法19，通过“中心化（减去均值）→随机旋转→随机平移”增强数据多样性，使模型对分子的整体位置和朝向不敏感，提升鲁棒性。
     

   • 步骤4-6：计算当前噪声参数：

     • 步骤4：计算噪声注入强度系数$\gamma$
       • 操作：当当前时间步c_τ大于阈值γmin时，γ取γ0=0.8（保留额外噪声注入）；否则γ=0（停止额外噪声注入）。
       • 作用：控制噪声注入的“开关”，后期（c_τ较小时）减少噪声，让模型聚焦于微调结构。
     • 步骤5：计算当前去噪时间步 $\hat{t}$
       • 操作：结合当前去噪步骤$c_t$和噪声系数$\gamma$（相乘），得到标量
     • 步骤6：生成当前时间步噪声$\overrightarrow{\xi l}$

   • 步骤7：添加噪声：将$\overrightarrow{\xi l}$加到当前原子位置${\vec{\mathbf{x}}}_l$，得到${\vec{\mathbf{x}}}_l^{\text{noisy}}$（模拟更嘈杂的状态）。

   • 步骤8：去噪预测：调用算法20（扩散模块核心），基于主干特征预测去噪后的原子位置 { x ⃗ l d e n o i s e d } \{\vec{\mathbf{x}}_{l}^{\mathrm{denoised}}\} {x ldenoised​}。

   • 步骤9-11：更新位置：计算修正量${\vec{\delta }}_l$（当前位置与去噪位置的差），结合时间步差$dt$和步长系数$\text{n}$，更新原子位置${\vec{\mathbf{x}}}_l$（向去噪方向调整）。

3. 输出最终结构（步骤13）：经过$T$步迭代后，${\vec{\mathbf{x}}}_l$收敛到预测的原子坐标。

🪧Q & A:

1. Question: 为什么要在随机采样后还要添加噪声，随机采样得到的不就是噪声了嘛？为什么还要去加额外噪声？
   Answer:随机采样得到的确实是噪声，但这只是初始噪声，而扩散模型的训练机制要求逐步添加噪声并去噪。在AlphaFold 3中，通过按照特定的噪声调度（Noise Schedule）逐步添加噪声，能够让模型学习到不同噪声水平下原子坐标的变化规律，从而使模型具备从高噪声状态逐步恢复到真实原子坐标状态的能力。这样训练出来的模型可以更好地处理各种噪声情况，提高对蛋白质等生物分子结构预测的准确性，以适应不同程度噪声干扰下的结构预测任务。

2. Question: DiffusionModule预测的是噪声还是去噪后的原子坐标？如果是去噪后的原子坐标的化，为什么还需要进行9~11步骤的去噪更新？
   Answer:DiffusionModule预测的是去噪后的原子坐标。而算法中9 - 11步骤仍然需要去噪更新，是因为DiffusionModule的预测结果只是一个中间阶段的去噪结果（Module认为是最终的去噪结果），并非最终的精确坐标。通过计算当前真实坐标（这里用初始${\vec{\mathbf{x}}}_l$近似表示）与预测去噪坐标的差异（步骤9），得到去噪方向，再根据噪声调度相关参数（步骤10），对当前的噪声坐标进行进一步的更新（步骤11）。这一过程可以不断迭代优化，逐步减少原子坐标中的噪声，使最终输出的原子坐标（步骤13返回的结果）更接近真实的分子结构，是整个扩散模型逐步精细化预测结果的重要环节。

---

👉🏻算法20：扩散模块核心：条件去噪的实现

算法20是去噪器的核心，基于主干特征预测去噪后的原子位置，分为6个关键步骤：

1. 条件处理（步骤1）：
   调用算法21（DiffusionConditioning），整合主干网络的特征（${\mathbf{s}}_i^{\mathrm{trunk}},{\mathbf{z}}_{ij}^{\mathrm{trunk}}$）、输入特征（${\mathbf{s}}_i^{\mathrm{inputs}}$）、时间步$\hat{t}$ 和相对位置编码，生成单特征 { s i } \{\mathbf{s}_{i}\} {si​}和成对表示 { z i j } \{\mathbf{z}_{ij}\} {zij​}。这些条件特征是去噪的“指导信号”（如进化保守性、残基对距离约束）。

2. 位置缩放（步骤2）：
   将带噪声的原子位置${\vec{\mathbf{x}}}_l^{\mathrm{noisy}}$除以$\sqrt{{\hat{t}}^2+{\sigma }_{\text{data}}^2}$，缩放为无量纲向量${\mathbf{r}}_l^{\mathrm{noisy}}$（近似单位方差），统一不同噪声强度下的位置分布，便于模型学习。

3. 原子→令牌的聚合（步骤3）： 🌟
   调用算法5（AtomAttentionEncoder），通过序列局部原子注意力，将原子级特征（ { f ∗ } , { r l n o i s y } \{\mathbf{f}^{*}\},\{\mathbf{r}_{l}^{\mathrm{noisy}}\} {f∗},{rlnoisy​}）聚合为令牌级表示 { a i } \{\mathbf{a}_i\} {ai​}，同时保存中间结果（ { q l s k i p } , { c l s k i p } , { p l m s k i p } \{\mathbf{q}_{l}^{\mathrm{skip}}\},\{\mathbf{c}_{l}^{\mathrm{skip}}\},\{\mathbf{p}_{lm}^{\mathrm{skip}}\} {qlskip​},{clskip​},{plmskip​}）用于后续解码。这一步实现“原子细节→令牌粗粒度特征”的转换。

4. 令牌级自注意力（步骤4-6）：

   • 步骤4：将单特征 { s i } \{\mathbf{s}_{i}\} {si​}融合到令牌表示 { a i } \{\mathbf{a}_i\} {ai​}中；
   • 步骤5：调用算法23（DiffusionTransformer），通过24个块和16头注意力，让令牌间进行全局交互（如残基间的长程依赖），优化 { a i } \{\mathbf{a}_i\} {ai​}；
   • 步骤6：层归一化稳定特征分布。

5. 令牌→原子的解码（步骤7）： 🌟
   调用算法6（AtomAttentionDecoder），将优化后的令牌表示 { a i } \{\mathbf{a}_i\} {ai​}广播回原子级，结合步骤3保存的中间结果，生成原子位置的更新量${\mathbf{r}}_l^{\mathrm{update}}$。这一步实现“令牌粗粒度特征→原子细节修正”的转换。

6. 位置重构（步骤8）：
   结合缩放系数${\sigma }_{\mathrm{data}}$和时间步$\hat{t}$，将更新量${\mathbf{r}}_l^{\mathrm{update}}$重新缩放为实际位置，与带噪声的输入位置${\vec{\mathbf{x}}}_l^{\mathrm{noisy}}$融合，得到去噪后的位置${\vec{\mathbf{x}}}_l^{\mathrm{out}}$。

辅助算法解析：条件注入与注意力机制

扩散模块的灵活性依赖于一系列辅助算法，核心包括：

1. 👉🏻 算法21（扩散条件处理）：
   

   • 成对条件(Pair conditioning)：拼接主干成对表示 { z i j t r u n k } \{\mathbf{z}_{ij}^{\mathrm{trunk}}\} {zijtrunk​}和相对位置编码，经过渡块增强，捕捉残基对的关系约束；
   • 单特征条件(Single conditioning)：拼接主干单特征${\mathbf{s}}_i^{\mathrm{trunk}}$和输入特征${\mathbf{s}}_i^{\mathrm{inputs}}$，融合时间步的傅里叶嵌入（算法22），让模型感知当前去噪阶段（时间信息对扩散至关重要）。

2. 👉🏻算法22（傅里叶嵌入）：

   将连续时间步 $\hat{t}$ 转换为高维余弦特征，使模型能学习时间相关的去噪模式加粗样式（如早期去噪粗调整，晚期微调细节）。
   

   • 关键性质：为什么用余弦函数？
     余弦函数cos(θ)是周期函数（周期为 2π），但其在这里的作用并非直接 “传递周期性”，而是通过多个不同频率、不同相位的余弦函数组合，将标量$\hat{t}$ 的数值信息编码到高维空间中。

   • 在实际应用的数值范围内（比如扩散模型中时间步$\hat{t}$ 的取值范围），不同的标量$\hat{t}$ 会对应不同的高维向量，从而让模型能通过这个向量 “区分” 出$\hat{t}$的具体值。

   • 在 Algorithm 21 中，傅里叶嵌入的输入是1/4 log(ˆt/σdata)，这说明ˆt 先经过对数变换再嵌入。这种处理进一步表明：

     • 扩散过程中ˆt 的影响可能是对数尺度的（如噪声随时间呈指数级增长 / 衰减），傅里叶嵌入能更好地捕捉对数空间中的连续变化；
     • 若直接线性扩展ˆt，会丢失对数变换后的非线性特征，导致模型难以建模噪声与时间步的真实关联。

综上，傅里叶嵌入并非简单的 “维度扩展工具”，而是通过频率成分编码时间步的非线性、连续性及与噪声的内在联系，这是直接线性扩展无法替代的。

3. 👉🏻算法23（Diffusion Transformer）：令牌级特征优化的核心 Transformer 模块
   

   • 核心功能：通过多轮 Transformer 块迭代，基于单特征（{s_i}）和成对表示（{z_ij}）优化令牌级特征（{a_i}），捕捉令牌间的全局交互关系，为后续原子级解码提供高质量的全局约束。
   • 核心逻辑：每个 Transformer 块由 “注意力层（算法 24）” 和 “过渡层（算法 25）” 组成，通过Nblock=24次迭代逐步优化a_i：
     • 注意力层（步骤 2）：让令牌关注全局范围内的重要伙伴（受z_ij引导，如距离近的令牌权重更高）；
     • 过渡层（步骤 3）：引入非线性变换，增强特征表达能力（如捕捉 “令牌 A 与 B 交互 + 令牌 B 与 C 交互” 的组合模式）。

4. 👉🏻算法24（AttentionPairBias）：带成对偏置的扩散注意力机制
   

• AdaLN（步骤 2）：通过单特征 { s i } \{\mathbf{s}_{i}\} {si​}动态调整归一化的缩放和偏移（见算法 26），让令牌特征更适配当前条件（如时间步、残基类型）。
• 成对偏置（步骤 8）：${\mathbf{z}}_{ij}$（如残基距离、相互作用强度）经转换后作为注意力权重的 “加分项”——${\mathbf{z}}_{ij}$越大，$A_{ij}^h$（i对j的注意力）越高，确保模型优先关注重要对。
• 门控机制（步骤 9、11）：${\mathbf{g}}_i^{\mathbf{h}}$（0~1）像 “阀门”，只允许有用的信息（如j的重要特征）通过，过滤噪声；步骤 13 进一步用${\mathbf{s}}_i$控制输出强度，避免无关信息干扰。

5. 👉🏻算法25（ConditionedTransitionBlock）：带自适应归一化的过渡块
   

   • 核心功能：通过 SwiGLU 激活和自适应归一化，对注意力输出的特征进行非线性变换，增强模型对复杂模式的捕捉能力（如多令牌交互的组合特征）。
   • SwiGLU 激活（步骤 2）：比传统 ReLU 更灵活的激活函数，通过swish(x) ⊙ x捕捉特征间的乘法交互（如 “令牌 A 的疏水性 × 令牌 B 的电荷”），增强特征表达。
   • 单特征控制（步骤 3）：sigmoid(Linear(s))生成 0-1 的系数，默认因biasinit=-2.0而接近 0，仅允许重要特征通过（避免过强激活导致的噪声）。

6. 👉🏻算法26（ConditionedTransitionBlock）：带自适应归一化的过渡块
   

• 核心功能：不同于普通层归一化（固定缩放和偏移），AdaLN 用单特征{s_i}动态调整归一化参数，让特征处理更适配当前条件（如时间步、残基类型）。
• 普通 LayerNorm vs AdaLN：
  • 普通 LayerNorm：a = (a - mean)/std * γ + β（γ、β是固定参数）；
  • AdaLN：a = (a - mean)/std * sigmoid(Linear(s)) + LinearNoBias(s)（γ和β由{s_i}动态生成）。
• 作用：让单特征{s_i}（如 “该残基是保守残基”）直接控制特征的缩放（sigmoid(…)）和偏移（LinearNoBias(s)），例如：保守残基的特征缩放更强，确保其在去噪中更稳定。

总结

🪧设计亮点与优势

1. 非等变设计：不依赖SE(3)等变假设或几何偏置，完全通过数据驱动学习原子位置分布，适配配体、核酸等非蛋白质分子的不规则结构。
2. 两级架构（原子-令牌-原子）：先聚合原子为令牌（捕捉全局结构），再解码回原子（细化局部细节），平衡全局约束与局部精度。
3. 灵活的条件注入：通过自适应层归一化、注意力偏置等方式，将主干特征（进化、模板）深度融合到去噪过程，使预测受生物学约束引导。
4. 现代Transformer技巧：SwiGLU激活、门控机制提升特征表达能力，随机增强提升模型泛化性。

扩散模块是AlphaFold 3实现多分子类型结构预测的核心创新，通过“加噪-去噪”的迭代过程，从随机噪声中逐步生成原子坐标。其设计摒弃了传统几何偏置，采用两级架构和灵活的条件注入，结合现代Transformer技巧，在蛋白质、核酸、配体等复杂分子的预测中实现了高精度。相比AlphaFold 2的结构模块，扩散模块更通用、更灵活，为解析复杂生物分子相互作用提供了强大工具。