DeepMind | 通过去噪来进行分子性质预测的预训练

编译 | 夏忻焱 审核 | 沈祥振 今天给大家带来的是牛津大学的Sheheryar Zaidi和DeepMind公司的Michael Schaarschmidt联合发表的文章《Pre-training via Denoising for Molecular Property Prediction》。从3D结构中预测分子性质的数据有限，这对神经网络的泛化造成了挑战。作者介绍了一种预训练技术，它利用平衡状态下3D分子结构的大型数据集来为下游任务学习有意义的表示。受最近噪声正则化的启发，作者的预训练目标是基于去噪的。依赖于去噪自动编码器和分数比配之间公认的联系，作者还表明目标对应于直接从平衡结构学习分子力场——其由高斯混合近似物理状态分布产生。实验表明，使用这个预训练目标可以大大改善多个基准测试的性能，在广泛使用的QM9数据集中达到了最先进水平。最后，作者分析了不同因素对预训练的影响，并提出了实用性的见解。

1 介绍 计算机视觉（CV）和自然语言处理（NLP）中表现最好的神经网络依赖于模型在大型数据集上的预训练，预训练可帮助学习适用于下游任务的特征。对于从3D结构（中原子核的点云）中预测分子性质，如何学习3D结构的近似表示仍有待研究。比如，那些在QM9基准测试中表现最好的模型没有一个进行了预训练，这与CV和NLP产生了鲜明的对比。有效的预训练方法可以对相关领域产生重大影响，比如药物发现和材料科学。

在这项工作中，作者关注的问题是如何利用大型3D分子结构数据集来提高下游分子性质预测任务的性能，这些任务也依赖于输入的3D结构。文章讨论的问题是：如何利用像PCQM4Mv2这样包含超过300万个结构的大型数据集来提高像DES15K这样小一个数量级的数据集性能？答案是一种自监督的预训练，为下游的预测任务提供有用的表示。

受图神经网络（GNNs）噪声正则化最新进展的启发，预训练目标是基于结构空间的去噪（因此是自监督的）。与现有的大多集中于2D图的预训练方法不同的是，作者的方法目标是下游任务涉及到的定义分子结构的3D点云的设置。依赖于去噪和分数匹配之间的公认联系，作者表明去噪目标等于（近似）特定力场的学习，揭示了它如何帮助表征学习。

文章的贡献总结如下： * 研究了一种简单有效的方法，通过在3D结构空间中的去噪来预训练，目的是改善从这类3D结构中的下游分子性质的预测。去噪目标被证明与学习一种特定力场有关。 * 实验表明通过去噪来预训练可以极大的改善在大小，任务性质和分子位置上都有所不同的具有多种挑战的任务数据集上的性能表现。这证明结构去噪成功地转移到分子性质预测上，特别是在广泛使用的QM9数据集的12个目标中的10个目标都创造了最先进的表现。图1说明了QM9其中一个目标的性能。 * 作者对常见的GNN架构进行了改进，特别地展现了如何将定制的激活转换器（TAT）应用到图网络模拟器中（GNS），这是对预训练的补充并且进一步提升了性能。 * 通过深入了解训练集大小、模型大小和结构、以及上下游数据集之间的关系，作者分析了预训练的好处。

2 方法 2.1问题设置

2.2通过去噪来预训练

3 实验 实验目的是回答以下问题。首先，相比于随机初始化训练，通过去噪对神经网络进行预训练是否可以改善在下游任务中的性能？第二，上游和下游数据集的联系如何影响预训练的效果？评估涉及到四个现实的、具有挑战性的分子数据集，包括：PCQM4Mv2、QM9、OC20、 DES15K。

4 分析 作者分析了不同因素如何影响预训练，考虑了许多因素的影响，包括不同架构、不同数据集大小、不同模型大小和固定特征只对解码器进行微调。

4.1 预训练不同的架构

4.2 不同的数据集大小 如图4（左）所示，预训练提高了所有规模数据集的下游性能。随着下游训练数据集的减少，测试MAE的差异也变大了。如图4（中）所示，上游数据增加，下游性能普遍提高，GNS-TAT性能达到饱和。更多的上游数据可以产生更好质量的表征，其结果体现在下游性能中。

4.3 不同的模型尺寸 如图4（右）所示，预训练继续受益于更大的模型规模。预训练是有用的，因为比较而言，从头开始训练的模型表现不佳：从头开始训练的1.3亿参数的模型被不到一半规模（0.53亿参数）的预训练模型超越。

4.4 固定预训练参数

5 总结 作者研究了通过3D分子空间结构去噪对神经网络进行预训练。作者表明去噪等同于学习一个特定的力场，激励其学习有用的表征，揭示了去噪在其它工作中的成功应用。这项工作可以利用现有的大型3D结构数据集，提高从3D结构预测各种下游分子性质任务的性能。更广泛的说，这弥补了CV/NLP和从结构中预测分子性质的预训练之间效用的差距。 参考资料 Zaidi, S., Schaarschmidt, M., Martens, J., Kim, H., Teh, Y. W., Sanchez-Gonzalez, A., ... & Godwin, J. (2022). Pre-training via Denoising for Molecular Property Prediction. arXiv preprint arXiv:2206.00133.