[论文解读] 3D-Transformer: Molecular Representation with Transformer in 3D Space

Fang Wu, Qiang Zhang+6 more|arXiv (Cornell University)|Oct 4, 2021

Machine Learning in Materials Science参考文献 79被引用 10

一句话总结

3D-Transformer 通过多尺度自注意力机制和自适应位置编码，将空间几何信息融入 3D 感知的 Transformer 架构中，实现分子表征学习。它通过基于注意力的下采样和对不同分子尺寸下原子间距离的上下文感知建模，在晶体性质预测和蛋白质-配体结合亲和力任务上达到最先进性能。

ABSTRACT

Spatial structures in the 3D space are important to determine molecular properties. Recent papers use geometric deep learning to represent molecules and predict properties. These papers, however, are computationally expensive in capturing long-range dependencies of input atoms; and have not considered the non-uniformity of interatomic distances, thus failing to learn context-dependent representations at different scales. To deal with such issues, we introduce 3D-Transformer, a variant of the Transformer for molecular representations that incorporates 3D spatial information. 3D-Transformer operates on a fully-connected graph with direct connections between atoms. To cope with the non-uniformity of interatomic distances, we develop a multi-scale self-attention module that exploits local fine-grained patterns with increasing contextual scales. As molecules of different sizes rely on different kinds of spatial features, we design an adaptive position encoding module that adopts different position encoding methods for small and large molecules. Finally, to attain the molecular representation from atom embeddings, we propose an attentive farthest point sampling algorithm that selects a portion of atoms with the assistance of attention scores, overcoming handicaps of the virtual node and previous distance-dominant downsampling methods. We validate 3D-Transformer across three important scientific domains: quantum chemistry, material science, and proteomics. Our experiments show significant improvements over state-of-the-art models on the crystal property prediction task and the protein-ligand binding affinity prediction task, and show better or competitive performance in quantum chemistry molecular datasets. This work provides clear evidence that biochemical tasks can gain consistent benefits from 3D molecular representations and different tasks require different position encoding methods.

研究动机与目标

• 解决现有几何深度学习模型在捕捉长程依赖关系以及处理分子中非均匀原子间距方面的局限性。
• 开发一种适应不同空间尺度和分子尺寸的 3D 感知注意力机制。
• 克服固定或以距离为主导的下采样方法在聚合分子表征时的不足。
• 设计一种能动态适应分子尺寸的位置编码策略，以提升表征质量。
• 在包括量子化学、材料科学和蛋白质组学在内的多样化科学领域中，验证 3D-Transformer 的有效性。

提出的方法

• 使用全连接图建模所有原子相互作用，保留 3D 空间中原子之间的直接连接。
• 引入多尺度自注意力模块，通过关注逐渐扩大的感受野，捕捉细粒度的局部模式和更广泛的上下文依赖关系。
• 采用自适应位置编码模块，根据分子大小选择不同的编码策略——小分子使用相对位置编码，大分子使用学习得到的绝对编码。
• 提出一种基于注意力的最远点采样方法，利用注意力分数选择代表性原子，优于虚拟节点和基于距离的池化策略。
• 将 3D 坐标作为 Transformer 编码器的输入特征，使注意力计算中引入空间归纳偏置。
• 使用标准损失函数，在分子性质预测任务上端到端训练模型。

实验结果

研究问题

• RQ1 通过显式建模空间几何和长程依赖关系，3D 感知的 Transformer 架构是否能提升分子表征学习？
• RQ2 多尺度注意力如何增强在不同原子间距分布下对上下文依赖模式的捕捉能力？
• RQ3 与固定编码方案相比，自适应位置编码是否能在不同尺寸的分子上带来更好的性能？
• RQ4 基于注意力的下采样是否能优于传统的基于距离或虚拟节点的池化方法？
• RQ5 在包括量子化学、材料科学和蛋白质组学在内的多样化科学领域中，引入 3D 空间信息在多大程度上能提升性能？

主要发现

• 3D-Transformer 在晶体性质预测基准上达到最先进性能，优于先前的 SOTA 模型。
• 该模型在蛋白质-配体结合亲和力预测任务上表现显著提升，展现出在生物和化学领域中的强大泛化能力。
• 在量子化学数据集上，3D-Transformer 的性能与现有模型相比具有竞争力或更优，验证了其在电子结构预测中的有效性。
• 自适应位置编码模块通过根据分子尺寸定制编码策略，显著提升了小分子和大分子的性能。
• 多尺度自注意力机制能有效捕捉长程依赖关系和细粒度的空间模式，即使在原子间距分布不均匀的分子中亦然。
• 基于注意力的最远点采样方法生成的分子嵌入比虚拟节点或基于距离的池化方法更具信息量和代表性。