[Paper Review] 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 Science3 citations

TL;DR

GeoTMI is a novel training framework that enables accurate prediction of quantum chemical properties using low-cost, corrupted 3D geometries by maximizing mutual information between the correct geometry, corrupted geometry, and target property. It employs positional denoising and a multi-objective loss to improve 3D GNN performance, achieving consistent accuracy gains across molecular property prediction, reaction barriers, and catalytic energy prediction tasks.

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.

Motivation & Objective

• Address the practical infeasibility of obtaining high-level 3D geometries in real-world quantum chemistry applications.
• Overcome the accuracy degradation in 3D GNNs when using corrupted or low-level geometries (e.g., from MMFF or DFT-optimization errors).
• Develop a model-agnostic framework that improves prediction performance without requiring access to correct 3D geometries during inference.
• Maximize mutual information between the correct geometry, corrupted input, and target property to guide robust learning.
• Demonstrate generalization and robustness across diverse quantum chemistry tasks including molecular properties, reaction barriers, and catalytic systems.

Proposed method

• Proposes a three-term mutual information maximization objective between the correct geometry $X$, corrupted input $ ilde{X}$, and target property $Y$.
• Introduces a denoising loss $ ilde{ heta}$ that encourages the GNN to refine $ ilde{X}$ toward $X$ through learnable positional updates during forward pass.
• Employs a multi-task loss combining property prediction from $ ilde{X}$, reconstruction of $X$ from $ ilde{X}$, and auxiliary prediction of $Y$ from $X$.
• Uses a soft regularization via a denoising objective to implicitly correct geometric errors without requiring ground-truth geometry at inference.
• Integrates seamlessly into existing 3D GNN architectures such as Equiformer and SchNet, enabling plug-and-play deployment.
• Applies a gradual denoising schedule to stabilize training and improve convergence by progressively refining the geometry representation.

Experimental results

Research questions

• RQ1 Can mutual information maximization between correct geometry, corrupted input, and property improve prediction accuracy in 3D GNNs?
• RQ2 Does explicit positional refinement during GNN inference enhance performance when using low-quality geometries?
• RQ3 How does GeoTMI compare to existing denoising baselines like Noisy Nodes in real-world quantum chemistry tasks?
• RQ4 To what extent does GeoTMI generalize across diverse datasets and property types, including molecular properties, reaction barriers, and catalytic energies?
• RQ5 Which components of the GeoTMI framework—denoising loss, auxiliary property head, or position update—are most critical for performance?

Key findings

• GeoTMI achieved a 17.6% average improvement in MAE across nine QM9 molecular properties compared to baseline models trained only on corrupted geometries.
• On the IS2RE task of the OC20 dataset, GeoTMI improved OOD adsorption energy prediction by 53.8% in terms of MAE reduction compared to the baseline.
• In the IS2RE benchmark, GeoTMI outperformed the Noisy Nodes method, achieving a 58.0% improvement in MAE for relaxed energy prediction.
• Ablation studies confirmed that removing any component—denoising loss, auxiliary correct-geometry prediction head, or position update—degraded performance, with the largest drop observed in barrier height prediction when the correct-geometry head was removed.
• The position update mechanism contributed significantly to performance, especially in reaction barrier height prediction, reducing error by 17.1% compared to baseline.
• Despite using corrupted geometries, GeoTMI achieved performance comparable to full-precision models on EwT (energy within 0.02 eV threshold) and even surpassed them in OOD adsorption prediction.