[论文解读] One Transformer Can Understand Both 2D & 3D Molecular Data

Shengjie Luo, Tianlang Chen+5 more|arXiv (Cornell University)|Oct 4, 2022

Machine Learning in Materials Science被引用 4

一句话总结

本文提出 Transformer-M，一种基于 Transformer 的统一模型，通过独立但共享的注意力通道，联合编码二维分子图与三维几何结构。通过在自注意力机制中动态激活 2D 或 3D 结构编码作为偏置项，该模型在 2D 和 3D 分子性质预测任务上均表现出色，展示了在无需架构微调的情况下跨模态的泛化能力。

ABSTRACT

Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One can view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer-based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M develops two separated channels to encode 2D and 3D structural information and incorporate them with the atom features in the network modules. When the input data is in a particular format, the corresponding channel will be activated, and the other will be disabled. By training on 2D and 3D molecular data with properly designed supervised signals, Transformer-M automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M can simultaneously achieve strong performance on 2D and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https://github.com/lsj2408/Transformer-M.

研究动机与目标

• 开发一种通用深度学习模型，能够在不进行架构特化的情况下处理二维分子图与三维几何结构。
• 解决现有模型仅针对一种分子数据格式进行训练所导致的跨模态零样本迁移性能差的局限性。
• 通过设计具有模态自适应注意力机制的单一模型，实现从 2D 和 3D 数据的联合学习。
• 证明单一神经网络可通过共享参数与模态特异性结构编码，有效学习并泛化于分子数据模态之间。

提出的方法

• 模型采用标准 Transformer 主干网络，配备两条并行的、模态特异的通道，分别用于 2D 和 3D 结构编码。
• 2D 通道使用度编码、最短路径距离编码和边编码作为输入偏置，以捕捉拓扑与键合特征。
• 3D 通道使用基于高斯核的原子间 3D 欧几里得距离编码作为注意力偏置项。
• 推理过程中，根据输入模态仅激活对应通道（2D 或 3D），而共享注意力与前馈层在两者间通用。
• 采用联合训练策略，使模型可在单一批次中同时从 2D、3D 或两者数据中学习，当某一阵列不活跃时，其模态特异通道被禁用。
• 模型在 PCQM4Mv2 数据集上使用 3D 位置去噪目标进行预训练，并在下游 2D 和 3D 数据集上进行微调。

实验结果

研究问题

• RQ1 单一神经网络架构能否有效从二维分子图与三维几何结构中学习有意义的表征？
• RQ2 注意力机制中模态特异的结构编码如何影响模型在分子数据格式间的泛化能力？
• RQ3 在 2D 和 3D 数据上联合预训练在多大程度上提升了多样化分子性质预测任务的下游性能？
• RQ4 在模态间共享参数是否能带来优于模态特异性模型的泛化性能？

主要发现

• 在 PDBBind 核心数据集上，Transformer-M 的皮尔逊相关系数（R）达到 0.830，优于最佳基线模型（SIGN）和预训练的 Graphormer 模型。
• 在 QM9 数据集上，2D-3D 联合预训练策略相比从零开始训练，将 MAE 最多降低 30%，且在全部六个分子目标上均取得提升。
• 使用 DFT 生成的 3D 构象异构体带来显著性能提升（例如，R 从 0.804 提升至 0.830），而使用 RDKit 生成的构象异构体仅带来微小改进。
• 当使用 3D 位置去噪进行预训练时，模型在下游任务中表现出持续的性能增益，表明该预训练目标对 3D 分子表征学习具有有效性。
• 即使在仅使用 3D 数据进行训练的情况下，模型在 2D 任务上仍保持强劲性能，表明其具备出色的跨模态迁移能力。
• 消融实验证实，模态特异的结构编码至关重要，禁用任一通道均会导致对应模态上的性能下降。