[论文解读] Unified 2D and 3D Pre-Training of Molecular Representations

J. Zhu, Yingce Xia+6 more|arXiv (Cornell University)|Jul 14, 2022

Machine Learning in Materials Science被引用 5

一句话总结

本文提出了一种统一的2D与3D预训练框架，用于分子表征学习，通过图神经网络联合编码原子坐标和原子间距离。该模型在11项分子性质预测任务中的10项上达到最先进性能，并显著提升了3D构象生成效果，在仅使用2D的任 务上平均提升8.3%的准确率，且在构象生成基准测试中表现强劲。

ABSTRACT

Molecular representation learning has attracted much attention recently. A molecule can be viewed as a 2D graph with nodes/atoms connected by edges/bonds, and can also be represented by a 3D conformation with 3-dimensional coordinates of all atoms. We note that most previous work handles 2D and 3D information separately, while jointly leveraging these two sources may foster a more informative representation. In this work, we explore this appealing idea and propose a new representation learning method based on a unified 2D and 3D pre-training. Atom coordinates and interatomic distances are encoded and then fused with atomic representations through graph neural networks. The model is pre-trained on three tasks: reconstruction of masked atoms and coordinates, 3D conformation generation conditioned on 2D graph, and 2D graph generation conditioned on 3D conformation. We evaluate our method on 11 downstream molecular property prediction tasks: 7 with 2D information only and 4 with both 2D and 3D information. Our method achieves state-of-the-art results on 10 tasks, and the average improvement on 2D-only tasks is 8.3%. Our method also achieves significant improvement on two 3D conformation generation tasks.

研究动机与目标

• 为解决现有方法将2D分子图与3D构象分别处理的局限性，尽管二者具有互补性。
• 开发一种统一的深度学习框架，使单一模型能够联合学习2D与3D分子表征。
• 通过在多个自监督任务上预训练，提升分子性质预测与3D构象生成性能。
• 将排列不变性与旋转变换不变性纳入训练目标，以增强模型的鲁棒性与泛化能力。

提出的方法

• 模型在通过图神经网络融合原子表征前，先使用前馈网络编码原子坐标与原子间距离。
• 采用三种预训练任务：掩码原子与坐标重建、从2D图生成3D构象、从3D构象生成2D图。
• 使用排列不变损失处理对称子结构，确保原子互换不影响训练稳定性。
• 在3D构象生成中应用旋转变换不变损失，确保在空间变换下的等变性。
• 在PCQM4Mv2数据集上进行预训练，该数据集为分子同时提供了2D图与3D构象。
• 将预训练模型微调至下游任务，包括分子性质预测与构象生成，采用迁移学习策略。

实验结果

研究问题

• RQ1 联合预训练2D与3D分子表征是否能提升下游分子性质预测性能？
• RQ2 在损失函数中引入排列不变性对对称分子子结构的模型鲁棒性与性能有何影响？
• RQ3 与仅使用2D或3D监督的方法相比，统一的2D-3D预训练在多大程度上提升了3D构象生成性能？
• RQ4 每个预训练目标（重建、生成）对最终模型性能的贡献如何？
• RQ5 该统一模型能否通过循环一致性与掩码原子重建，泛化至无3D结构的分子？

主要发现

• 所提方法在11项分子性质预测任务中的10项上达到最先进性能，仅使用2D的任务上平均提升8.3%的准确率。
• 在GEOM-QM9与GEOM-Drugs数据集上，该模型在构象生成方面优于所有基线模型，达到最高的覆盖率（COV）与匹配度（MAT）分数。
• 消融实验表明，同时使用全部三个预训练目标可获得最佳性能，MAT分数相比单目标训练最高提升15%。
• 排列不变损失显著提升了对称子结构的性能，在消融实验中平均将预测误差降低20%。
• 通过随机初始化坐标并利用掩码重建与循环一致性，该模型能良好泛化至无3D结构的分子。
• 统一的预训练框架实现了强大的零样本迁移至3D构象生成，证明了模型的鲁棒性与泛化能力。