基于内存高效等变变压器的可扩展肽设计

# 基于内存高效等变 Transformer 的可扩展肽设计
来源: https://arxiv.org/html/2606.25006

1\] 清华大学计算机科学与技术系  
2\] 清华大学人工智能产业研究院  
3\] 清华大学电子工程系  
4\] 清华大学化学系

\\contribution\[\*\] 共同第一作者  
孔祥哲\* 贾尹骏\* 张怡嘉 杨紫怡 刘洋 马剑竹†

\\[\[\[[majianzhu@tsinghua\.edu\.cn](https://arxiv.org/html/2606.25006v1/mailto:[email protected])

###### 摘要

靶标特异性肽设计需要在全原子几何约束下进行序列与结构的协同设计。隐式生成框架通过将细粒度原子结构压缩为片段级隐式表示，并在紧凑的隐空间中进行条件生成，为这一问题提供了有效途径。然而，此类系统的可扩展性很大程度上取决于其编码、解码和去噪组件中使用的几何主干。我们提出了 **MEET**（内存高效等变 Transformer），这是一个用于可扩展原子级肽建模的 E(3) 等变主干。**MEET** 维护耦合的不变标量流和等变向量流，同时围绕内存高效注意力重构几何计算。它通过全局坐标聚合初始化向量特征，通过增强的查询和键点积引入成对距离，并通过稀疏键适配器注入共价键信息。将 **MEET** 集成到用于全原子肽生成的 VAE 和隐扩散流程中，实现了随原子数量线性增长的内存缩放，并提升了生成质量，优于现有肽设计方法。在大规模 AFDB 衍生数据集上的实验进一步表明，所提出的主干支持系统的模型和数据缩放，从而带来更好的结合亲和力、物理有效性和样本多样性。

\\correspondence\\codes https://github\.com/jiaor17/MEET

## 1 引言

设计能与指定蛋白口袋结合的肽是基于结构的药物发现中的核心问题[wang2022therapeutic]。该任务是序列与结构协同设计的一种形式，其中肽序列、构象和口袋结合几何通过物理相互作用耦合。因此，精确建模需要对侧链堆积、氢键、形状互补性和局部空间兼容性进行全原子几何推理。这使得 E(3) 等变架构[thomas2018tensor, satorras2021en, deng2021vectorneuronsgeneralframework, schutt2021equivariant]成为自然选择，因为它们在直接操作分子几何的同时，尊重三维空间的对称性。

隐式生成框架[rombach2022high, kong2025unimomo]提供了一种平衡全原子保真度与生成效率的实用方法。在此公式中，VAE 将细粒度原子结构压缩为片段级隐点，从而通过重构保留局部原子细节，而非在扩散过程中显式表示。然后，在目标口袋条件下，在紧凑的隐图上进行生成，解码器将生成的隐状态映射回全原子几何。由于几何主干实例化于 VAE 和 LDM 组件中，其表达能力和内存效率成为整个生成系统可扩展性的关键。这对用于全原子肽-口袋复合物的主干提出了高要求。这类复合物通常包含数百到数千个原子。

现有的几何主干通常通过密集距离矩阵[jiao2025equivariantpretrainedtransformerunified]或显式局部分子图[schutt2018schnet, satorras2021en, schutt2021equivariant, tholke2022torchmdnet, liao2023equiformer]引入坐标依赖。这些机制虽然有效，但需要存储或重新计算成对偏差、邻接列表或边特征，这增加了内存流量，限制了可行的模型大小和批次大小。当扩展生成模型时，该问题更加突出，因为样本质量的提升通常需要更大的主干和更大的训练集。

为了解决这个主干层面的瓶颈，我们提出了 **MEET**（内存高效等变 Transformer），一个用于可扩展原子级肽建模的 E(3) 等变 Transformer 主干。**MEET** 保持耦合的标量和向量特征流，其中标量流是旋转不变的，向量流是旋转等变的。其几何操作以兼容内存高效注意力核（如 FlashAttention[dao2022flasht attention]）的形式编写。具体而言，**MEET** 通过全局坐标聚合初始化向量特征，将成对距离信息折叠进增强的查询和键点积中，并通过稀疏键适配器注入共价键信息。这些设计选择共同避免了在使用内存高效注意力时物化二次激活张量，同时保持了整个主干的 E(3) 等变性。

我们在一个受 UniMoMo[kong2025unimomo] 启发的两阶段隐式生成框架内评估该主干。在此框架内，**MEET** 作为编码器、区块类型解码器、结构解码器和 LDM 去噪器的几何主干。这一设置使我们能够测试更节省内存的等变主干是否能提升完整生成系统的性能。图1 总结了架构和主要几何模块。我们的扩展实验集中于 LDM 去噪器，其容量直接影响框架中生成肽的质量。总体而言，架构和生成评估带来了三项具体贡献。

*   **高效架构**。**MEET** 用查询和键增强替代密集距离偏置，用全局注意力聚合替代局部图向量初始化，并使用稀疏键适配器处理化学邻接。这些选择使得在固定模型大小下，与内存高效注意力核配合使用时，峰值激活内存随原子数量线性增长。
*   **大规模数据集**。从 864 万个 AFDB[varadi2022alphafold] 域开始，我们构建了约 1 亿个候选片段，并通过滑动窗口枚举、结构质量过滤、界面筛选和序列重叠聚类，采样了 10 万和 120 万个训练结构。
*   **系统扩展**。在 10 万基准测试上，**MEET** 在结合自由能和物理有效性方面超越了 PepGLAD[kong2024pepglad]、PepFlow[li2024pepflow]、UniMoMo[kong2025unimomo] 和 DiffPepBuilder[wang2024diffpepbuilder]。在 120 万数据集上，将隐式去噪器扩展到四种 DiT 风格[peebles2023scalable]模型大小，进一步提升了生成质量并恢复了样本多样性。

图 1: **MEET** 架构概览。**MEET** 处理原子坐标 **X** 以及标量特征流 **H** 和向量特征流 **V**。初始化层使用距离感知注意力产生 **H**^(0) 和 **V**^(0)。每个重复块包含一个键适配器、一个等变自注意力层和一个等变前馈层，输出最终特征 **H**^(L) 和 **V**^(L)。在自注意力中，标量和向量输入被投影为 **Q**、**K** 和 **U**，距离编码 -φ_s 和 ψ_s 被附加到 **Q** 和 **K** 上。键适配器通过稀疏消息 Δh_{i←j} 和 Δv_{i←j} 注入边属性 e_{ij}。

## 2 方法

我们的肽生成器遵循受 UniMoMo[kong2025unimomo] 启发的两阶段隐式生成框架。VAE 首先将每个全原子肽-口袋复合物编码为片段级隐变量，条件隐扩散模型从目标口袋上下文生成肽隐变量，VAE 解码器将生成的隐变量映射回肽序列和全原子几何。**MEET** 在 VAE 编码器、序列解码器、结构解码器和隐式去噪器中充当几何主干。完整的训练目标和采样程序见附录5。在接下来的小节中，我们首先在 §2.1 给出 **MEET** 架构的概述。然后描述主干共享的距离感知注意力机制 (§2.2)，接着是图1 中显示的四个主要模块：特征初始化 (§2.3)、等变自注意力 (§2.4)、等变前馈层 (§2.5) 和键适配器 (§2.6)。最后，我们在 §2.7 分析空间复杂度。

### 2.1 架构概述

**MEET** 是一个内存高效的等变 Transformer，维护两个耦合的特征流：标量和向量。我们使用与图1 相同的符号，其中 **H** 表示标量流，**V** 表示向量流，**X** 表示原子坐标。对于一个包含 N 个原子的分子，设 **X** ∈ ℝ^{N×3}，其行向量为 **x**_i ∈ ℝ^3。我们将每个原子的特征表示为一对 (**H**, **V**)，其中 **H** ∈ ℝ^{N×d} 是标量（旋转不变）特征，**V** ∈ ℝ^{N×3×d} 是向量（旋转等变）特征，遵循等变神经网络中常用的标量-向量分解[deng2021vectorneuronsgeneralframework]。在刚体变换 (**R**, **t**) 下，其中 **R** ∈ SO(3)，**t** ∈ ℝ^3，特征变换为 **H** ↦ **H** 和 **V** ↦ **R** **V**。因此，**MEET** 的所有模块在完整欧几里得群下都是 E(3)-等变的。

给定输入原子属性 **A** 和坐标 **X**，主干首先通过初始化算子 **I** 构建初始等变状态：
(**H**^(0), **V**^(0)) = **I**(**A**, **X**),    (1)
然后应用 L 个 Transformer 块产生 (**H**^(L), **V**^(L))。每个块由一个自注意力层 **A** 和一个向量-标量混合前馈层 **F** 组成，当提供了键边集 **E** 时，可选择在前面加一个键适配器 **B**。

(̄**H**^(ℓ), ̄**V**^(ℓ)) = (**H**^(ℓ), **V**^(ℓ)) + **B**^(ℓ)(Norm(**H**^(ℓ), **V**^(ℓ)), **E**),    (2)
(~**H**^(ℓ), ~**V**^(ℓ)) = (̄**H**^(ℓ), ̄**V**^(ℓ)) + **A**^(ℓ)(Norm(̄**H**^(ℓ), ̄**V**^(ℓ)), **X**),    (3)
(**H**^(ℓ+1), **V**^(ℓ+1)) = (~**H**^(ℓ), ~**V**^(ℓ)) + **F**^(ℓ)(Norm(~**H**^(ℓ), ~**V**^(ℓ))).    (4)

每一层都包裹有预归一化和残差连接，遵循标准 Transformer 实践。在下面的模块级描述中，(**H**_in, **V**_in) 和 (**H**_out, **V**_out) 指进入和离开图1 中对应面板的流。键适配器可以在每个块中应用、仅在第一块中应用，或者根据键的可用性完全禁用。

#### 归一化
对标量流，我们使用 RMSNorm[zhang2019root]。对向量流，我们使用旋转不变类似方法，通过每个原子向量特征的空间和通道轴组合均方根进行归一化：
RMSNorm_V(**V**_i) = **V**_i / sqrt( (1/(3d)) * Σ_{k,c} **V**_i,k,c^2 + ε ).    (5)
分母中的缩放因子是旋转不变的，因此输出在旋转下变换为 **R** **V**_i，并保持等变性。

### 2.2 距离感知注意力机制

初始化层 **I** 和自注意力层 **A** 共享的组件是距离感知查询-键增强。对于头 h，注意力对数包含一个不变的距离惩罚：
A_{ij}^(h) = ρ_h( ⟨**q**_i^(h), **k**_j^(h)⟩ - s_h^2 ‖**x**_i - **x**_j‖_2^2 ),    (6)
其中 **q**_i^(h) 和 **k**_j^(h) 是模块特定的查询和键特征，ρ_h 是注意力缩放因子，s_h 是可学习的距离缩放因子。根据 FlashBias[wu2026flashbias]，我们不将距离项存储为 N×N 偏置张量，而是将其吸收到查询-键点积中。对于具有分量 (x_1,i, x_2,i, x_3,i) 的坐标 **x**_i，定义：
φ_{s_h}(**x**_i) = s_h ⨁_{k=1}^3 [ x_k,i^2, 1, -2x_k,i ],
ψ_{s_h}(**x**_j) = s_h ⨁_{k=1}^3 [ 1, x_k,j^2, x_k,j ].    (7)

由于 ⟨φ_{s_h}(**x**_i), ψ_{s_h}(**x**_j)⟩ = s_h^2 ‖**x**_i - **x**_j‖_2^2，将 -φ_s 连接到查询，ψ_s 连接到键，即可重现公式 (6)：
**Q**_i^(h) = concat( **q**_i^(h), -φ_{s_h}(**x**_i) ),
**K**_j^(h) = concat( **k**_j^(h), ψ_{s_h}(**x**_j) ),
A_{ij}^(h) = ρ_h ⟨**Q**_i^(h), **K**_j^(h)⟩.    (8)