[论文解读] Pre-training via Denoising for Molecular Property Prediction

Sheheryar Zaidi, Michael Schaarschmidt+7 more|arXiv (Cornell University)|May 31, 2022

Machine Learning in Materials Science被引用 30

一句话总结

作者提出一种自监督预训练方法，通过去噪3D分子结构来学习表示，在QM9及其他基准上达到最先进结果。

ABSTRACT

Many important problems involving molecular property prediction from 3D structures have limited data, posing a generalization challenge for neural networks. In this paper, we describe a pre-training technique based on denoising that achieves a new state-of-the-art in molecular property prediction by utilizing large datasets of 3D molecular structures at equilibrium to learn meaningful representations for downstream tasks. Relying on the well-known link between denoising autoencoders and score-matching, we show that the denoising objective corresponds to learning a molecular force field -- arising from approximating the Boltzmann distribution with a mixture of Gaussians -- directly from equilibrium structures. Our experiments demonstrate that using this pre-training objective significantly improves performance on multiple benchmarks, achieving a new state-of-the-art on the majority of targets in the widely used QM9 dataset. Our analysis then provides practical insights into the effects of different factors -- dataset sizes, model size and architecture, and the choice of upstream and downstream datasets -- on pre-training.

研究动机与目标

• 在标注数据稀缺的情况下，动机并开发一种基于3D结构的分子性质预测的预训练方法。
• 利用去噪/分数匹配从平衡结构中学习一个有意义的力场样表示。
• 证明基于去噪的预训练在不同数据集和架构上提升下游任务的性能。
• 提供关于上游数据集规模、模型规模、架构和数据集相似性如何影响迁移的实际见解。

提出的方法

• 通过对原子坐标加入高斯噪声并训练一个GNN去预测噪声，来构建一个3D分子结构的去噪预训练目标。
• 将去噪目标理解为学习一个力场，该力场对应于围绕平衡结构的Boltzmann分布的高斯混合近似。
• 将这项预训练应用于Graph Net Simulator (GNS)及其增强版本 Tailored Activation Transform (TAT) 的变体（GNS-TAT）。
• 通过也将去噪预训练应用于TorchMD-NET（一个具有不变/协变组件的基于Transformer的模型），以展示架构无关性。
• 使用 Noisy Nodes 作为辅助损失来补充预训练并分析其影响。
• 分析数据集规模、模型规模以及上游/下游数据集关系对迁移性能的影响。

实验结果

研究问题

• RQ1 与随机初始化相比，基于去噪的预训练是否能改善下游分子性质预测？
• RQ2 去噪预训练的收益如何取决于上游（预训练）与下游数据集之间的关系？
• RQ3 去噪预训练是否能在不同架构（GNNs 和Transformer）和任务（QM9、OC20、DES15K）上有效？

主要发现

• 通过去噪的预训练显著提升下游性能，覆盖多个数据集和目标，在QM9的多数目标上达到新的最佳状态。
• 在对GNS-TAT应用去噪预训练并使用Noisy Nodes后，在12个QM9目标中有10个达到最先进的结果；某些目标可能对噪声尺度敏感。
• 在PCQM4Mv2（3.4M个3D结构）上进行的预训练通常能很好地迁移到QM9和DES15K，迁移效果随元素重叠和任务相似性而异。
• 在OC20上进行预训练时，如果是在PCQM4Mv2上进行预训练，则不会提高OC20 IS2RE性能，凸显数据集不相似带来的迁移限制；但在OC20自身上进行预训练时，收敛速度有所提升。
• 增加上游数据通常能提升下游表征；更大的模型受益于预训练，且可以超过从头训练的模型。
• 预训练也提升力场预测（例如MD17阿司匹林）并在仅微调解码器组件时加速训练。