ICLR2023 | MOLE-BERT 重新思考分子图模型预训练方法

今天为大家介绍的是来自西湖大学Stan Z. Li的一篇发表在2023ICLR的分子预训练任务的论文。文章提出了一种新的基于图神经网络的分子预训练框架MOLE-BERT。作者发现目前基于AttrMask预训练任务存在的局限性，并提出了一种名为Masked Atoms Modeling (MAM)的原子级别预训练新方法。此外，作者还提出一种基于分子级别的预训练任务，triplet masked contrastive learning (TMCL)。在下游任务中Mole-BERT表现了非常卓越的性能。

分子是由原子和化学键组成的复杂结构，因此对于机器学习算法来说是一个具有挑战性的领域。为了解决分子预训练问题，研究人员已经提出了许多预训练方法，如AttrMask。具体来说，AttrMask的训练策略是先随机掩码掉分子中的部分原子，再让目标模型对掩码位置的原子类别进行预测。此方法目前被预训练领域广泛采用。然而，AttrMask在处理分子数据时存在一些限制，例如无法处理极端不平衡的原子类型分布，这导致模型在迁移到下游任务上时，表现欠佳，很多时候甚至不如未经预训练的模型。

因此，作者提出了一种新的原子级别预训练方法，称为MAM，旨在克服现有方法的局限性并提高模型性能。该方法使用VQ-VAE变体的编码器来扩大原子类型词汇表，并引入一种新的节点级别预训练任务，以更好地捕捉分子结构中的信息。

尽管MAM是一种十分有效的原子级别预训练任务，作者发现用它提取的特征无法衡量分子之间的相似度，如下图a部分所示。为了让模型提取的分子特征符合分子的相似度特征，作者提出了一种名为TMCL的基于对比学习的预训练模型，有效的改善了分子特征相似度的问题，结果如下图b部分所示。

原子级别预训练方法MASKED ATOMS MODELING (MAM)

MAM是Mole-BERT预训练方法的一个重要组成部分，流程如上图所示。它是一种节点级别的预训练任务，旨在通过随机屏蔽原子类型来预测它们。具体来说，对于每个分子，MAM会随机选择一些原子并将它们的类型替换为特殊的掩码标记。然后，GNN模型被用来预测这些被屏蔽的原子类型。尽管整体流程与AttrMask相似，MAM的预测目标并不是直接预测原子类型。它首先使用了VQ-VAE架构中的编码器编码原子信息，然后将原子特征映射到预设的词汇表中，词汇表中具体的符号名称就成为了MAM的预测目标。

在原始的原子词汇表中，一共能得到118个不同的词汇，并且不同词汇的数据量极度不平衡，如下图b所示，这使得基于AttrMask的预训练模型很快就会收敛的同时难以学习有效的原子信息表示。作者用VQ-VAE将特征空间词汇表化，此举扩展了原子的表示形式的同时，还缓解了样本不均衡的问题，使得编码器能更好的编码原子级别特征。作者用t-sne展示了跟C原子相关的不同官能团被编码后的特征，如下图c所示，可以看到不同官能团被有效的区分，扩展了更多的分子表示可能性。

分子级别预训练方法

TMCL是Mole-BERT预训练方法中的一种图级别预训练任务，旨在模拟分子之间的异质相似性。具体来说，TMCL使用图网络对分子进行编码，并将分子相似性作为预训练任务的目标，使得相似的分子在特征空间上相近，而不同的分子在特征空间上远离。与传统的对比学习不同的是，本文的对比学习增加了一个层次，即分子自身基于随机掩码比例不同的对比学习，让掩码比例小的分子特征更贴近原始分子，而掩码比例大的分子远离。通过不同层次的对比学习，Mole-BERT具备了一定的分子查询能力，使得它可以学习到如何捕捉分子之间的相似性，并在下游任务中提高模型性能。

实验结果

Mole-Bert首先在两百万的分子数据集上进行的预训练，随后主要在分子性质预测的下游任务上对模型效果进行检验。

作者在八个数据集上对模型进行了检测(这八个数据集的任务目标均为二分类的分子性质预测,由MoleculeNet提供),并且在其中的七个数据集上达到了目前最好的结果。

对于实验结果，作者还进行了深入的分析：

1. AttrMask预训练在HIV和BBBP两个数据集上表现糟糕，结果甚至低于未经预训练的模型。而本文提出的MAM方法明显优于AttrMask，大幅提高了基于原子级别预训练模型的效果。

2. 相比于目前效果最好的模型GraphMVP，Mole-BERT效果更佳。

3. 
   对于有监督形式的直接训练，MAM效果仍然要好于AttrMask。具体结果如下表所示。

4. VQ-VAE中词汇表的设置方式十分重要。将数据量过大的一些原子单独分配词汇表可以有效的防止不同类型的原子被映射到同一词汇上。除此之外，本文提出的TMCL中，基于掩码比例的对比学习十分重要，如果只采用传统的对比学习，效果会有很显著的倒退。

除了分类任务外，作者还在回归任务上对模型的能力进行了评估。评估内容集中在分子性质预测和分子亲和力预测。具体结果如下表所示。

结论

本文首先讨论了传统的预训练方法AttrMask的劣势，并且针对性的提出了全新的原子级别预训练方法MAM。MAM借助VQ-VAE的词汇表，有效的提高了预训练模型的编码能力。除此之外，本文还提出了一种全新的，基于对比学习的分子级别预训练方法TMCL。MAM和TMCL共同构建出名为Mole-BERT的全新预训练框架，此框架在下游数据上取得到了非常好的效果。对于下一步工作，作者想要将本文提出的模型框架迁移到蛋白质模型上去，探讨其在蛋白质任务上的能力。

参考资料

原文链接：

https://openreview.net/pdf?id=jevY-DtiZTR

代码

https://github.com/junxia97/Mole-BERT