[论文解读] May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations

Rui Feng, Qi Zhu+5 more|arXiv (Cornell University)|Aug 24, 2023

Machine Learning in Materials Science被引用 4

一句话总结

该论文提出ET-OREO，一种以力为中心的3D分子构象预训练框架，通过直接学习原子力，统一处理平衡态与非平衡态数据。通过将力作为统一目标——对平衡态数据使用去噪和零力正则化，对非平衡态数据使用直接力监督——该方法在分子动力学模拟中实现比未微调的等变Transformer高3倍的力预测精度，且推理速度比NequIP快2.45倍，同时在QM9性质预测任务上达到最先进性能。

ABSTRACT

Recent works have shown the promise of learning pre-trained models for 3D molecular representation. However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations. It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data. For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces. We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP. As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

研究动机与目标

• 解决现有3D分子预训练模型仅关注平衡构象、在非平衡态数据上表现不佳的局限性。
• 通过单一、物理基础明确的目标，统一基于平衡态与非平衡态分子构象的预训练。
• 通过直接从多样化构象数据中学习原子力，提升分子动力学模拟的精度与稳定性。
• 通过提升3D分子表征的泛化能力与迁移能力，增强其在下游任务（如性质预测与模拟）中的表现。
• 通过在小规模、低成本聚合物上实现有效预训练，降低数据生成成本，同时实现对大规模未见体系的外推。

提出的方法

• 提出一种以力为中心的预训练目标，将原子力作为主要信号，利用其在分子间的物理不变性与可迁移性。
• 对非平衡构象，直接监督模型预测来自从头计算方法（如DFT）获得的原子力。
• 对平衡构象，应用零力正则化与基于力的去噪方法，以近似接近平衡态的力并稳定训练过程。
• 在来自多个来源的约1500万种构象的统一数据集上训练等变Transformer模型，涵盖平衡态与非平衡态状态。
• 将预训练模型用作下游任务的强编码器，包括分子动力学模拟与QM9上的量子性质预测。
• 在平衡态数据上引入去噪目标，该目标不依赖能量，而是基于力的一致性，从而提升模型鲁棒性。

(b) Potential energy curve during simulations.

实验结果

研究问题

• RQ1 基于原子力的统一预训练目标能否有效表征平衡态与非平衡态的3D分子构象？
• RQ2 与基于能量或去噪的方法相比，以力为中心的预训练在多大程度上提升了分子动力学模拟的精度与稳定性？
• RQ3 在性质预测任务中，引入非平衡构象在多大程度上增强了3D分子表征的泛化能力？
• RQ4 在小分子和低成本DFT数据上预训练的模型，能否成功外推至大规模未见聚合物的模拟？
• RQ5 与传统的基于能量或去噪的目标相比，基于力的目标在数据效率与迁移性能方面表现如何？

主要发现

• ET-OREO实现的力预测精度约为未微调等变Transformer模型的3倍。
• 该模型在分子动力学模拟中实现稳定运行，推理速度比NequIP快2.45倍，同时与DFT参考力保持高度相关性。
• ET-OREO成功模拟了多达360个原子的大规模未见聚合物，展现出强大的泛化能力与构象采样能力。
• 在QM9 HOMO-LUMO性质预测任务上，ET-OREO达到最先进性能，与NoisyNode相当，并优于从零开始预训练的模型。
• 尽管可访问丰富的非平衡态数据，但在平衡态性质预测任务中未观察到显著精度提升，表明此类任务可能无法从非平衡态数据中获益。
• 零力正则化与基于力的去噪方法有效稳定了在平衡态数据上的训练，解决了原始等变Transformer中出现的不稳定性问题。