[论文解读] Molformer: Motif-based Transformer on 3D Heterogeneous Molecular Graphs

Fang Wu, Dragomir Radev+1 more|arXiv (Cornell University)|Oct 4, 2021

Machine Learning in Materials Science被引用 7

一句话总结

Molformer 提出了一种新颖的三维异质分子图（HMG）框架，通过功能基团和强化学习实现基团层次的图结构发现，整合了原子层次和基团层次的节点。该方法采用具有多尺度建模和注意力最远点采样（AFPS）的异质自注意力 Transformer，实现了在量子化学、生理学和生物物理学基准上的最先进或具有竞争力的性能表现。

ABSTRACT

Procuring expressive molecular representations underpins AI-driven molecule design and scientific discovery. The research mainly focuses on atom-level homogeneous molecular graphs, ignoring the rich information in subgraphs or motifs. However, it has been widely accepted that substructures play a dominant role in identifying and determining molecular properties. To address such issues, we formulate heterogeneous molecular graphs (HMGs), and introduce a novel architecture to exploit both molecular motifs and 3D geometry. Precisely, we extract functional groups as motifs for small molecules and employ reinforcement learning to adaptively select quaternary amino acids as motif candidates for proteins. Then HMGs are constructed with both atom-level and motif-level nodes. To better accommodate those HMGs, we introduce a variant of Transformer named Molformer, which adopts a heterogeneous self-attention layer to distinguish the interactions between multi-level nodes. Besides, it is also coupled with a multi-scale mechanism to capture fine-grained local patterns with increasing contextual scales. An attentive farthest point sampling algorithm is also proposed to obtain the molecular representations. We validate Molformer across a broad range of domains, including quantum chemistry, physiology, and biophysics. Extensive experiments show that Molformer outperforms or achieves the comparable performance of several state-of-the-art baselines. Our work provides a promising way to utilize informative motifs from the perspective of multi-level graph construction.

研究动机与目标

• 为解决现有图神经网络（GNN）仅关注原子层次图结构、忽略信息丰富的分子亚结构（基团）的局限性。
• 开发一种统一框架，将原子层次和基团层次的节点整合到单一的三维异质分子图（HMG）中。
• 设计一种基于 Transformer 的架构，能够建模多层级图结构之间的相互作用，同时保留三维几何信息。
• 通过强化学习挖掘蛋白质中的任务相关基团，以及为小分子使用化学定义的功能基团，提升分子表征学习能力。
• 在多样化的分子性质预测任务中验证所提出的 HMG 和 Molformer 的有效性。

提出的方法

• 通过包含源自功能基团（小分子）或强化学习发现的四级氨基酸序列（蛋白质）的基团层次节点，构建异质分子图（HMG）。
• 提出一种异质自注意力（HSA）机制，能够在 Transformer 编码器中区分并计算不同节点类型（原子与基团）之间的注意力。
• 实现一种多尺度机制，通过逐步增大的感受野捕获局部模式，以建模精细的结构细节。
• 提出一种注意力最远点采样（AFPS）算法，对节点特征进行下采样和聚合，生成全局分子表征。
• 使用强化学习在蛋白质序列中挖掘与任务相关的基团候选，以优化下游预测性能。
• 在输入 Molformer 架构前，使用多层感知机（MLP）将原子特征嵌入可学习向量。

实验结果

研究问题

• RQ1 在三维异质图中显式引入分子基团作为节点，是否能相比同质原子层次图提升分子表征学习性能？
• RQ2 基于强化学习的基团挖掘策略在识别蛋白质序列中具有生物学意义的亚结构方面效果如何？
• RQ3 一种能区分原子与基团节点的异质自注意力机制，是否能提升分子性质预测的性能？
• RQ4 多尺度建模在多大程度上增强了对分子图中局部与长程结构依赖关系的捕捉能力？
• RQ5 所提出的注意力最远点采样（AFPS）算法与标准池化或基于注意力的聚合方法相比，在学习全局分子表征方面表现如何？

主要发现

• Molformer 在包括 QM7、BBBP 在内的多样化分子性质预测基准上，性能优于或与最先进基线方法相当。
• 消融实验表明，添加四类功能基团（烃类、卤代烃类、含氧类、含氮类）均带来一致的性能提升，其中烃类贡献最大。
• 基于强化学习的基团挖掘方法发现的基团在结合口袋中的出现率为 1.38%，几乎是其他区域的两倍，表明其具有生物学相关性。
• 异质自注意力机制有效建模了原子与基团之间的跨层级交互，提升了表征质量。
• AFPS 算法能有效聚合节点特征，生成鲁棒的全局分子表征，显著增强了模型的泛化能力。
• 采用显式基团节点的 HMG 建模方式相比简单融合多层级特征，性能有明显提升，验证了该设计选择的有效性。