[论文解读] GeoTMI:Predicting quantum chemical property with easy-to-obtain geometry via positional denoising

Hyeonsu Kim, Jeheon Woo+4 more|arXiv (Cornell University)|Mar 28, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

GeoTMI 是一种新颖的训练框架，通过在正确几何结构、受损输入和目标性质之间最大化互信息，实现了使用低成本、有噪声的3D几何结构对量子化学性质进行精确预测。该方法采用位置去噪和多任务损失以提升3D GNN的性能，在分子性质预测、反应能垒和催化能垒预测任务中均实现了稳定的精度提升。

ABSTRACT

As quantum chemical properties have a dependence on their geometries, graph neural networks (GNNs) using 3D geometric information have achieved high prediction accuracy in many tasks. However, they often require 3D geometries obtained from high-level quantum mechanical calculations, which are practically infeasible, limiting their applicability to real-world problems. To tackle this, we propose a new training framework, GeoTMI, that employs denoising process to predict properties accurately using easy-to-obtain geometries (corrupted versions of correct geometries, such as those obtained from low-level calculations). Our starting point was the idea that the correct geometry is the best description of the target property. Hence, to incorporate information of the correct, GeoTMI aims to maximize mutual information between three variables: the correct and the corrupted geometries and the property. GeoTMI also explicitly updates the corrupted input to approach the correct geometry as it passes through the GNN layers, contributing to more effective denoising. We investigated the performance of the proposed method using 3D GNNs for three prediction tasks: molecular properties, a chemical reaction property, and relaxed energy in a heterogeneous catalytic system. Our results showed consistent improvements in accuracy across various tasks, demonstrating the effectiveness and robustness of GeoTMI.

研究动机与目标

• 解决在实际量子化学应用中获取高水平3D几何结构的不现实性。
• 克服在使用受损或低水平几何结构（例如MMFF或DFT优化错误）时3D GNN性能下降的问题。
• 开发一种与模型无关的框架，在推理阶段无需访问正确3D几何结构的情况下提升预测性能。
• 通过在正确几何结构、受损输入和目标性质之间最大化互信息，引导鲁棒学习。
• 在包括分子性质、反应能垒和催化体系在内的多样化量子化学任务中展示泛化性与鲁棒性。

提出的方法

• 提出一种三重互信息最大化目标，涉及正确几何结构 $X$、受损输入 $ ilde{X}$ 和目标性质 $Y$。
• 引入一个可学习的位置更新机制 $ ilde{ heta}$，在前向传播过程中促使GNN将 $ ilde{X}$ 优化为 $X$，实现去噪。
• 采用多任务损失，结合从 $ ilde{X}$ 预测性质、从 $ ilde{X}$ 重建 $X$，以及从 $X$ 辅助预测 $Y$ 的任务。
• 通过去噪目标的软正则化，隐式校正几何误差，而无需在推理阶段使用真实标签几何结构。
• 可无缝集成至现有3D GNN架构（如 Equiformer 和 SchNet），实现即插即用的部署。
• 采用渐进式去噪调度策略，通过逐步优化几何表示来稳定训练并提升收敛性。

实验结果

研究问题

• RQ1 在3D GNN中，通过在正确几何结构、受损输入和性质之间最大化互信息，是否能提升预测精度？
• RQ2 在GNN推理过程中，显式的位置优化是否能提升使用低质量几何结构时的性能？
• RQ3 在真实量子化学任务中，GeoTMI 与 Noisy Nodes 等现有去噪基线方法相比表现如何？
• RQ4 GeoTMI 在包括分子性质、反应能垒和催化能垒在内的多样化数据集和性质类型上，其泛化能力如何？
• RQ5 GeoTMI 框架中的哪些组件——去噪损失、辅助性质头或位置更新——对性能最为关键？

主要发现

• 与仅在受损几何结构上训练的基线模型相比，GeoTMI 在九种QM9分子性质上的平均MAE降低了17.6%。
• 在OC20数据集的IS2RE任务中，GeoTMI在OOD吸附能预测方面相比基线实现了53.8%的MAE降低。
• 在IS2RE基准测试中，GeoTMI优于Noisy Nodes方法，在弛豫能预测的MAE上实现了58.0%的提升。
• 消融实验表明，若移除任一组件（去噪损失、辅助正确几何结构预测头或位置更新），性能均会下降，其中移除正确几何结构头后在能垒高度预测中性能下降最为显著。
• 位置更新机制对性能贡献显著，尤其在反应能垒高度预测中，相比基线将误差降低了17.1%。
• 尽管使用了受损几何结构，GeoTMI在EwT（能量误差在0.02 eV以内）任务上的表现与全精度模型相当，甚至在OOD吸附能预测中表现更优。