编译 | 程志祥
今天给大家介绍的是来自北京大学机器学习基础研究实验室和微软在arXiv发表的预印本《ONE TRANSFORMER CAN UNDERSTAND BOTH 2D & 3D MOLECULAR DATA》。为了实现跨模态处理分子任务，作者开发了一种新的基于Transformer的分子模型Transformer-M。它可以将2D或3D格式的分子数据作为输入，并生成有意义的语义表示。通过对2D和3D分子数据进行训练，Transformer-M学习来自不同数据模式的知识并正确捕获表征。Transformer-M可以同时在2D和3D任务中实现较强的性能，表明其具有广泛的适用性。
简介
化学与视觉和语言等传统领域的一个关键区别是数据的多模态。在视觉和语言中，数据通常以特定的形式进行描述。例如，图像被定义为像素网格中的RGB值，而句子被定义为句子中的字词。与视觉和语言数据通常具有固定的格式不同，分子可以使用不同的化学方法进行表征。比如，可以将分子视为2D图，也可以将其定义为位于3D空间中的原子集合。然而，以往的工作大多集中于设计2D或3D结构的神经网络模型，导致一种形式的模型无法应用到另一种形式的任务中。因此作者开发了Transformer-M，一个基于Transformer的多功能分子模型，它在2D和3D分子表征学习中都表现良好。
Transformer-M
如图1所示，Transformer-M的主干网络是由标准的Transformer块组成的。作者开发了两个单独的通道来编码2D和3D结构信息。2D通道使用度编码、最短路径距离编码和从2D图结构中提取的边编码。最短路径距离编码和边编码反映了一对原子之间的空间关系和键的特征，在softmax注意中用作偏置项。度编码被添加到输入层的原子特征中。
对于3D通道，作者使用3D距离编码对3D几何结构中原子之间的空间距离进行编码。每个原子对的欧氏距离通过高斯核函数编码，并将在softmax注意中用作偏置项。对于每个原子，我们将其与所有其他原子之间的3D距离编码相加，并将其添加到输入层的原子特征中。
除了两个结构通道中的参数外，Transformer-M中的所有其他参数（例如，自注意和前馈网络）对于不同的数据模式是共享的。在训练过程中，当batch中的数据是2D图时，2D通道将被激活，3D通道将被禁用。类似地，当batch中的数据是3D几何结构时，3D通道将被激活，而2D通道将被禁用。当二维和三维信息同时给出时，两个通道都将被激活。
图1：Transformer-M模型架构。在经典主干Transformer上建立了两个通道。红色通道对2D Graph的数据激活，包含度、最短路径距离和边缘信息。紫色通道对3D几何结构的数据激活，利用欧氏距离信息。不同的编码位于适当的模块中。
实验
作者研究了三个具有代表性的任务：PCQM4Mv2 (2D)、PDBBind (2D&3D)和QM9 (3D)。使用OGBLSC（包含340万个2D和3D形式的分子）中的PCQM4Mv2数据集来预训练Transformer-M。预训练任务有两个，一个是训练模型预测2D和3D格式分子的HOMO-LUMO能级，另一个是专门设计针对3D数据的3D位置去噪任务。
PCQM4Mv2 性能（2D）
在预训练之后，使用PCQM4Mv2的验证集上评估Transformer-M。PCQM4Mv2的验证集只包含2D分子数据。因此，可以用它来评估Transformer-M在2D分子数据上的表现。该任务的目标是预测HOMU-LUMO能级，评价指标是平均绝对误差（MAE）。
由于训练目标包括HOMO-LUMO能级预测任务，所以不需要对模型参数进行微调。结果如表1所示，可以很容易地看到，Transformer-M性能大大超过了所有baseline，例如，与之前的最佳模型（GraphGPS）相比，MAE相对减少了8.2%。Transformer-M和Graphormer之间的唯一区别是，Graphormer仅对2D数据进行训练，而Transformer-M同时使用2D和3D结构信息进行训练。因此，可以得出结论，Transformer-M在2D分子数据上表现良好，共享参数的2D- 3D联合训练确实有助于模型学习更多的化学知识。
表1 PCQM4Mv2验证集的结果
Pdbbind性能（2D&3D）
为了验证Transformer-M的兼容性，进一步对PDBBind数据集进行测试。PDBBind数据集由蛋白-配体复合物组成，这些数据是在生物测定实验中获得的，与pKa（或−log Kd，−log Ki）相关。每个数据提供了3D几何结构，并通过预定义规则构造2D图形结构。在PDBBind数据集上对微调模型进行了评估。评价指标包括Pearson相关系数（R）、平均绝对误差（MAE）、均方根误差（RMSE）和标准差（SD）。
将Transformer-M与经典方法、基于CNN的方法和GNN的模型进行比较。所有实验都用不同的种子重复五次，报告平均性能。结果如表2所示，Transformer-M在所有评估指标上的表现都明显优于所有baseline，例如，Pearson相关系数（R）的绝对改善率为3.3%。值得注意的是，PDBBind数据集是蛋白-配体复合物，而Transformer-M是在简单分子上预训练的，这证明了Transformer-M的可迁移性。
表2 PDBBind数据集的结果
QM9性能（3D）
使用QM9数据集评估Transformer-M在3D分子任务上的表现。QM9是由134万个稳定的有机小分子组成的量子化学基准。每个分子都与12个目标相关联，涵盖了它的能量、电子和热力学性质。作者随机选择了10,000和10,831个分子进行验证和测试评估，剩下的分子被用来微调Transformer-M。
结果如表3所示，与这些baseline相比，Transformer-M取得了具有竞争力的性能，表明模型与3D分子数据兼容。特别地，Transformer-M在HUMO、LUMO和HUMO-LUMO预测上表现最好。这说明在预训练任务中学习到的知识能够更好地迁移到类似的任务中。该模型在其他一些任务上执行得不太好，作者认为Transformer-M可以在几个方面进行改进，包括采用精心设计的输出层或使用更多的自监督训练策略进行预训练。
表3 QM9数据集的结果
消融实验
预训练任务的影响：
Transformer-M通过两个任务在PCQM4Mv2训练集上进行预训练：
（1）预测2D和3D格式分子的HOMO-LUMO能级。
（2）3D位置去噪。
作者对PCQM4Mv2和QM9数据集进行了消融实验，以检查这两个目标是否有利于下游任务。作者进行了另外两个实验。
（1）在PCQM4Mv2和QM9上使用其2D数据和3D几何数据从头训练Transformer-M模型，以检查预训练方法的效果。
（2）对Transformer-M进行预训练，不使用3D去噪任务，以研究2D-3D联合预训练方法的有效性。
结果如表4所示。可以看出，联合预训练显著提高了在两个数据集上的性能。此外，3D位置去噪任务也是有益的，特别是在QM9数据集上。
表4 预训练任务的影响
模式分布的影响：
为了训练Transformer-M，作者为每个数据提供了三种模式：
（1）激活2D通道和禁用3D通道（2D模式）；
（2）激活3D通道，禁用2D通道（3D模式）；
（3）激活两个通道（2D&3D模式）。
在训练过程中，每个数据的模式是根据预定义的分布随机绘制的，实现方法类似于Dropout。p2D, p3D, p2D& 3D为模式的概率，实验研究了不同模式分布对模型性能的影响。作者选择三个模式分布，p2D, p3D, p2D& 3D分别为1:1:1、1:2:2和1:2:1，结果如表5所示。可以看到模型在三种模式分布上的性能都显著超过了以前的工作。此外，在3D模式上使用稍微大一点的概率可以达到最好的效果。
表5 模式分布的影响
总结
在这项工作中，作者向通用分子模型迈出了第一步。提出的Transformer-M提供了一种很有前途的方法来处理2D和3D格式的分子任务。作者使用两个单独的通道对2D和3D结构信息进行编码，并将它们集成到主干Transformer中。当输入数据为特定格式时，相应的通道将被激活，而另一个通道将被禁用。通过对2D和3D分子数据的训练，模型能自动学习不同数据格式的化学知识，并正确捕获表征。大量的实验表明，Transformer-M可以同时在2D和3D任务中实现较强的性能。
作者相信这是一个起点，未来有更多的可能性去探索。例如，在这项工作中，作者使用了一种简单的方法，将2D和3D结构的结构信息进行线性组合，作者相信应该有其他有效的方法来融合这种编码。此外，模型还可以与以往的多视角对比学习方法相结合。如何使用这些方法对模型进行预训练是值得研究的。
参考资料
Shengjie Luo, Tianlang Chen, Yixian Xu,Shuxin Zheng, Tie-Yan Liu, Liwei Wang, Di He A. (2022). One Transformer Can Understand Both 2D & 3D Molecular Data
https://doi.org/10.48550/arXiv.2210.01765
代码链接：
https://github.com/lsj2408/Transformer-M
