[论文解读] Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning

Xiangzhe Kong, Wenbing Huang+1 more|arXiv (Cornell University)|Jun 2, 2023

Computational Drug Discovery Methods被引用 4

一句话总结

本文提出通用等变变换器（GET），一种统一的 E(3)-等变模型，将多样的 3D 分子复合物表示为集合的几何图——同时捕捉块级（如残基、核苷酸碱基）和原子级结构。GET 使用双层注意力机制和等变模块，在不同大小的块之间保留细粒度信息，实现了跨领域分子相互作用预测的最先进性能，包括零样本 RNA/DNA-配体亲和力预测。

ABSTRACT

Many processes in biology and drug discovery involve various 3D interactions between molecules, such as protein and protein, protein and small molecule, etc. Given that different molecules are usually represented in different granularity, existing methods usually encode each type of molecules independently with different models, leaving it defective to learn the various underlying interaction physics. In this paper, we first propose to universally represent an arbitrary 3D complex as a geometric graph of sets, shedding light on encoding all types of molecules with one model. We then propose a Generalist Equivariant Transformer (GET) to effectively capture both domain-specific hierarchies and domain-agnostic interaction physics. To be specific, GET consists of a bilevel attention module, a feed-forward module and a layer normalization module, where each module is E(3) equivariant and specialized for handling sets of variable sizes. Notably, in contrast to conventional pooling-based hierarchical models, our GET is able to retain fine-grained information of all levels. Extensive experiments on the interactions between proteins, small molecules and RNA/DNAs verify the effectiveness and generalization capability of our proposed method across different domains.

研究动机与目标

• 解决蛋白质、小分子和 RNA/DNA 之间多样 3D 分子相互作用缺乏统一、通用表示的问题。
• 克服领域特定模型和基于层次池化的方案在保留细粒度结构信息方面的局限性。
• 开发一种单一的等变模型，能够同时捕捉领域特定的层次结构和通用的相互作用物理规律。
• 实现在分子领域之间的零样本迁移学习，特别是针对 RNA/DNA-配体结合亲和力预测。
• 在统一的几何图表示中同时保留块级特异性和原子级可重用性。

提出的方法

• 提出一种集合的几何图表示：以分子块（如残基、核苷酸碱基）为顶层图，每个块内包含一组原子。
• 设计一种双层注意力机制，联合关注稀疏的块级表示和密集的原子级表示，以更新原子特征。
• 采用 E(3)-等变前馈网络和层归一化模块，对每个块内的排列保持不变，对 3D 旋转、平移和反射保持等变性。
• 确保块内排列不变性，并在不同大小的块之间实现大小无关的操作。
• 使用统一的变换器架构与 E(3)-等变模块，处理双层几何图，无需层次池化。
• 在混合数据集上端到端训练模型，以实现跨领域的迁移学习和零样本泛化。

实验结果

研究问题

• RQ1 一个单一的统一模型能否有效表示并从蛋白质、小分子和 RNA/DNA 的多样化 3D 分子复合物中进行学习？
• RQ2 与现有单层或层次化方法相比，双层集合几何图表示是否能更好地保留块级特异性和原子级可重用性？
• RQ3 E(3)-等变变换器是否能在不依赖基于池化的聚合方法的前提下，保留细粒度结构信息？
• RQ4 统一模型在未见领域（如 RNA/DNA-配体结合）上的泛化能力如何，是否可通过零样本迁移学习实现？
• RQ5 在预测分子相互作用亲和力方面，所提出的 GET 与领域特定、单层和层次化模型相比表现如何？

主要发现

• 在蛋白质-蛋白质亲和力预测任务中，GET 达到了最高的皮尔逊相关系数（0.697 ± 0.003）和斯皮尔曼等级相关系数（0.555 ± 0.008），优于所有基线模型。
• 在配体结合亲和力预测中，GET 实现了最低的均方根误差（1.327 ± 0.005）和最高的皮尔逊相关系数（0.620 ± 0.004）与斯皮尔曼等级相关系数（0.611 ± 0.003），并在 RNA/DNA-配体体系中展现出强大的零样本性能。
• 采用混合预训练的模型（GET-mix）在所有基准测试中均达到最先进水平，证明了其强大的泛化能力和可迁移性。
• 与依赖池化的层次化模型相比，GET 通过保留原子级和块级表示，显著提升了性能，后者通常会损失细粒度信息。
• 消融实验证实，双层注意力机制和 E(3)-等变设计对性能至关重要，尤其在跨领域迁移任务中表现突出。
• GET 在 RNA/DNA-配体结合亲和力预测中展现出强大的零样本泛化能力，表明其捕捉了超越领域特异性模式的通用相互作用物理规律。