【AlphaFold3】网络架构篇（3）| Sequence-local atom attention（序列局部原子注意力）

一、 序列局部原子注意力

翻译：

“序列局部原子注意力”将整个结构表示为原子的扁平列表，并允许所有原子在特定的序列邻域内直接“交流”。例如，每32个原子的子集会关注附近128个原子的子集（在序列空间中邻近）。这使网络能够学习关于局部原子群的通用规则，而不受粗粒度令牌化（每个标准残基仅用一个令牌表示）的限制。

对特定序列邻域的限制虽非最优，但对于将内存和计算成本控制在合理范围内是必要的。由此产生的注意力模式相当于计算完整的亲和矩阵，并应用沿对角线分布的矩形块邻域掩码（见补充图1）。

补充图1 序列局部原子注意力：每个原子子集（行）关注更大的原子子集（列）。蓝色区域表示理论上完整的$N_{\text{atoms}}\times N_{\text{atoms}}$注意力矩阵，黄色矩形表示实际实现的注意力。

通过令牌索引$i\in N_{\text{tokens}}$和原子名称$a\in S_{\text{atom names}}$（例如{x⃗ai}\{\vec{x}_a^i\}{xai​}）的原子索引到具有原子索引$l\in N_{\text{atoms}}$的扁平索引的转换，由扁平原子索引到令牌索引$i=\text{tok\_idx}(l)$的映射表示。

讲解：

🪧 一、核心目标：让原子在局部“直接对话”
序列局部原子注意力是AlphaFold 3处理多分子类型（尤其是配体、修饰残基等非标准分子）的关键机制。其核心目标是：打破令牌化的限制，让原子级特征直接交互，从而学习更精细的局部结构规则（如化学键、空间位阻、功能位点的原子排布）。

• 为什么需要？
  传统令牌化（如每个氨基酸用一个令牌）会丢失原子级细节（如侧链原子的精确位置），而配体、核酸等分子的功能高度依赖原子间的相互作用（如药物分子与蛋白质的氢键）。通过让原子直接“对话”，模型能捕捉这些微观细节。

🪧二、“序列局部”的含义：注意力范围的限制
为平衡精度与效率，注意力并非全局（所有原子两两交互），而是限制在“序列邻域”内：

• 具体实现：如算法7所示，将原子按序列顺序（如残基索引、令牌索引）划分为多个子集，每个子集（32个原子）仅关注附近128个原子（见补充图1的黄色矩形）。例如，序列位置1-32的原子仅与1-128的原子交互，33-64的原子仅与33-160的原子交互，以此类推。
• 原因：全局注意力的计算量随原子数呈二次增长（$N^2$），而局部限制可将复杂度降至$O(N)$，避免内存和计算爆炸（尤其对含上千原子的配体或复合物）。

👉🏻算法5（原子注意力编码器）：从原子特征到令牌嵌入

该算法是将原子级特征转换为令牌级嵌入的核心，分三步实现：

1. 原子级特征初始化（步骤1-7）

   • 为每个原子生成“化学名片”：通过c_l整合原子的坐标（⃗f_ref_pos_l）、电荷（f_ref_charge_l）、元素类型（f_ref_element_l）等元数据（128维，catom=128）。
   • 为每对原子生成“关系名片”：通过p_lm编码原子间的空间偏移（⃗d_lm）、是否同属一个空间单元（v_lm，如同一残基）、距离（逆平方距离，增强近距离原子的权重）等（16维，catompair=16）。
   • 初始原子表征q_l以c_l为基础，后续将通过注意力优化。

2. 融合外部信息（步骤8-12）
   若存在主干网络的令牌嵌入（s_trunk_i）、成对表征（z_ij）或扩散过程的噪声位置（r_l），则将其“广播”到原子级：

   • 例如，令牌i的主干嵌入会被分配给该令牌包含的所有原子，使原子特征与高层语义（如进化信息）结合。

   这里的rl是与原子索引l对应的带噪声位置参数，用于向原子的表示中引入位置噪声，可能用于增强模型对位置扰动的鲁棒性或辅助学习位置相关的特征。

3. 注意力交互与聚合（步骤13-16）

   • 增强原子对关系：步骤13-14通过MLP强化p_lm，使其包含原子自身属性（c_l、c_m）和复杂的非线性关系（如“带正电的N原子与带负电的O原子易形成氢键”）。
   • 局部注意力更新：步骤15调用AtomTransformer，通过3个Transformer块和4头注意力，让每个原子在局部序列邻域内“关注”相关原子（如成键原子、空间靠近的原子），更新q_l以包含交互信息。
   • 聚合为令牌嵌入：步骤16将同一令牌包含的所有原子的q_l取平均，生成令牌级嵌入a_i（384维，ctoken=384），实现“原子级细节→令牌级特征”的转换。

4. 步骤 17：保存中间变量:
   核心目的：这些变量将作为 “跳过连接（skip connection）” 的中间状态，供后续的原子注意力解码器（算法 6） 使用，避免在编码 - 解码过程中丢失原子级的原始细节。

👉🏻算法6（原子注意力解码器）：从令牌嵌入到原子位置更新

解码器是编码器的逆过程，用于扩散模型的去噪步骤：

• 步骤1：将令牌级嵌入a_i广播到其包含的原子，并与编码器保存的中间结果（q_skip_l）相加（跳过连接，避免梯度消失）。
• 步骤2：再次通过AtomTransformer优化原子表征，捕捉去噪过程中的原子位置调整。
• 步骤3：输出原子的位置更新量r_update_l，用于修正扩散过程中的噪声原子坐标，逐步逼近真实结构。

👉🏻算法7（原子Transformer）：实现局部注意力的核心

Algorithm 7（Atom Transformer）是 AlphaFold 3 中实现 “序列局部原子注意力” 的核心模块，其核心功能是限制原子仅在局部序列邻域内进行注意力交互，在保证捕捉原子级细节的同时，将计算和内存成本控制在合理范围内。

🪧一、输入参数说明

Atom Transformer 的输入包含模型计算所需的核心特征和控制注意力范围的超参数，具体如下：

参数	含义与作用
$q_l$	原子级单表征集合（每个原子的特征向量，初始由原子元数据生成，后续通过注意力更新）。
$c_l$	原子级条件特征集合（整合了原子坐标、电荷、元素类型等基础属性，作为注意力计算的“上下文锚点”）。
$p_{lm}$	原子对关系特征集合（描述原子l与m的空间偏移、距离、成键可能性等，辅助注意力权重计算）。
$N_{block}=3$	Transformer块的数量（默认3个），每个块包含多头注意力和前馈网络，用于逐步优化原子表征。
$N_{head}$	多头注意力的头数（如算法5/6中默认4头），通过多组注意力并行捕捉不同类型的原子关系（如近距离键合、远距离空间作用）。
$N_{queries}=32$	每个原子子集包含的“查询原子”数量（32个），即每个局部注意力单元的“主体原子”规模。
$N_{keys}=128$	每个查询子集关注的“键原子”数量（128个），即局部注意力的“交互范围”。
$S_{subsetcentres}$	原子子集的中心位置集合（如{15.5, 47.5, 79.5, ...}），用于划分局部邻域（中心均匀分布，确保覆盖所有原子）。

🪧核心步骤解析

步骤 1：定义局部注意力掩码β_lm

• 条件解读：
  仅当原子l（查询原子）和原子m（键原子）同时落在某个中心c的邻域内时，β_lm=0（允许注意力交互）：

  • 对查询原子l：需满足|l - c| < Nqueries/2（即l在中心c周围16个原子内，因Nqueries=32，32/2=16）；
  • 对键原子m：需满足|m - c| < Nkeys/2（即m在中心c周围64个原子内，因Nkeys=128，128/2=64）。

• 掩码作用：
  若不满足上述条件，β_lm被设为-10^10（一个极小值）。在注意力权重计算中（softmax函数），极小值会使对应注意力权重趋近于0，即屏蔽非局部原子间的交互。

• 直观理解（结合补充图1）：
  整个原子集合的注意力矩阵是N_atoms×N_atoms的方阵（蓝色区域），而β_lm的作用是仅保留沿对角线分布的“矩形块”（黄色区域）——每个块对应一个中心c，横向覆盖128个键原子，纵向覆盖32个查询原子。这种模式确保每个原子只与“序列上邻近”的原子交互。

步骤2：调用DiffusionTransformer更新原子表征

• DiffusionTransformer：是一个带掩码的Transformer模块，核心功能是基于步骤1的掩码β_lm，计算原子间的注意力权重并更新q_l：
  1. 注意力计算：对每个原子l，仅关注满足β_lm=0的原子m，根据p_lm（原子对关系）和c_l（原子条件特征）计算权重（如距离近的原子权重更高）；
  2. 特征聚合：将关注的原子m的特征q_m按权重聚合，更新q_l，使q_l包含局部邻域的原子交互信息（如“该原子与附近3个氧原子可能形成氢键”）；
  3. 多层迭代：通过Nblock=3个Transformer块重复上述过程，逐步优化q_l，捕捉更复杂的局部结构规则（如多原子构成的功能位点）。

🪧 三、核心作用与设计意义
Atom Transformer的设计直接服务于“序列局部原子注意力”的目标，解决了两个关键问题：

1. 平衡细节与效率：
   分子（尤其是配体、复合物）可能包含上千个原子，若使用全局注意力（所有原子两两交互），计算量会随原子数呈平方增长（O(N²)），导致内存溢出。而局部注意力通过β_lm限制交互范围，计算量降至线性增长（O(N)），使模型能处理复杂分子。

2. 捕捉局部结构规则：
   分子的功能依赖局部原子排布（如酶的活性中心由几个关键原子构成，药物分子与靶点的结合依赖局部氢键）。Atom Transformer通过限制注意力在“序列邻近”的原子间，让模型更聚焦于这些局部规则（如“C-N键长通常在1.47Å左右”“疏水原子倾向于聚集”），提升预测精度。

总结：

序列局部原子注意力是AlphaFold 3突破蛋白质限制、实现多分子类型预测的核心创新之一：

1. 微观细节保留：通过原子级交互，捕捉配体、核酸等分子的化学键、空间排布等关键信息，而传统令牌化会丢失这些细节。
2. 效率与精度平衡：局部注意力限制避免了全局计算的高昂成本，使模型能处理含上千原子的复杂复合物。
3. 通用规则学习：通过原子群的交互模式，模型可学习跨分子类型的通用结构规则（如“氢键的原子距离通常在2.5-3.5Å”），提升对新分子的预测能力。

简言之，这一机制让模型既能“看清原子细节”，又能“高效处理复杂结构”，为多分子复合物的高精度预测奠定了基础。