[论文解读] Supervised Pretraining for Molecular Force Fields and Properties Prediction

Xiang Gao, Weihao Gao+4 more|arXiv (Cornell University)|Nov 23, 2022

Computational Drug Discovery Methods被引用 5

一句话总结

本文提出在8600万个分子上对图神经网络进行监督式预训练，输入为三维几何结构和原子电荷，目标标签为分子能量。通过迁移学习，该方法在分子性质和力场预测任务中显著提升了下游性能；线性探测结果表明，模型在未显式监督结构相关任务的情况下，仍能学习到丰富的结构表征。

ABSTRACT

Machine learning approaches have become popular for molecular modeling tasks, including molecular force fields and properties prediction. Traditional supervised learning methods suffer from scarcity of labeled data for particular tasks, motivating the use of large-scale dataset for other relevant tasks. We propose to pretrain neural networks on a dataset of 86 millions of molecules with atom charges and 3D geometries as inputs and molecular energies as labels. Experiments show that, compared to training from scratch, fine-tuning the pretrained model can significantly improve the performance for seven molecular property prediction tasks and two force field tasks. We also demonstrate that the learned representations from the pretrained model contain adequate information about molecular structures, by showing that linear probing of the representations can predict many molecular information including atom types, interatomic distances, class of molecular scaffolds, and existence of molecular fragments. Our results show that supervised pretraining is a promising research direction in molecular modeling

研究动机与目标

• 通过利用大规模标注数据集进行预训练，解决分子机器学习中的数据稀缺问题。
• 通过从预训练阶段迁移学习，提升分子性质与力场预测任务的泛化能力。
• 探究在分子能量上进行监督式预训练是否能隐式学习到对多样化下游任务有用的结构信息。
• 证明在物理上合理且准确的标签（分子能量）上进行预训练，相比随机或自监督预训练，能获得更优的下游性能。

提出的方法

• 在包含8600万个分子的数据集中预训练一个消息传递图神经网络（EGNN），以3D坐标和原子电荷作为输入，分子能量作为监督目标。
• 采用标准EGNN架构，通过基于注意力的消息传递机制编码分子图，并通过最终读出层预测能量。
• 在预训练过程中应用力正则化损失，最小化能量对原子坐标的梯度的L2范数，以促进平滑的力预测。
• 在下游分子性质与力场预测任务中，使用少量标注数据对预训练模型进行微调。
• 对学习到的表征进行线性探测，以评估其预测原子类型、原子间距离、分子骨架和功能团的能力。
• 所有线性探测任务均使用负对数似然损失，以评估学习表征的质量。

实验结果

研究问题

• RQ1 在分子能量上进行监督式预训练是否能提升分子性质与力场预测任务的下游性能？
• RQ2 在物理上合理且大规模的标签（分子能量）上进行预训练，是否能带来优于自监督或任务无关预训练的泛化能力？
• RQ3 从能量预训练中学习到的表征在多大程度上编码了分子结构信息，如原子类型、原子间距离和功能团？
• RQ4 预训练模型能否在未见过的分子骨架上实现零样本预测？

主要发现

• 与从零开始训练相比，微调预训练模型在七个分子性质预测任务和两个力场预测任务中均显著提升了性能。
• 线性探测显示，预训练模型的表征在预测原子类型、原子间距离、分子骨架和85种功能团时，误差显著低于随机权重模型。
• 模型通过深层网络有效保留了原子类型信息，表明消息传递与表征学习效果良好。
• 力正则化技术提升了力场预测的泛化能力，使模型更适用于需要精确梯度的任务。
• 预训练模型在多个下游任务中达到新的最先进性能，包括在未见分子骨架上的成功零样本泛化测试。
• 结果表明，基于分子能量的监督式预训练是自监督预训练的有力替代方案，不仅能带来更优的下游性能，还能隐式捕捉结构信息。