Relational Invariant Learning for Robust Solvation Free Energy Prediction

yeyunchen

Abstract: Predicting the solvation free energy of molecules using graph neural networks holds significant potential for advancing drug discovery and the design of novel materials. While previous methods have demonstrated success on independent and identically distributed (IID) datasets, their performance in out-of-distribution (OOD) scenarios remains largely unexplored. We propose a novel Relational Invariant Learning framework (RILOOD) to enhance OOD generalization in solvation free energy prediction. RILOOD comprises three key components: (i) a mixup-based conditional modeling module that integrates diverse environments, (ii) a novel multi-granularity refinement strategy that extends beyond core substructures to enable context-aware representation learning for capturing multi-level interactions, and (iii) an invariant learning mechanism that identifies robust patterns generalizable to unseen environments. Extensive experiments demonstrate that RILOOD significantly outperforms state-of-the-art methods across various distribution shifts, highlighting its effectiveness in improving solvation free energy prediction under diverse conditions.

Lay Summary:

We want to help computers better predict how molecules behave in different environments — a crucial step in developing new drugs and materials. One important property is solvation free energy, which affects how a molecule dissolves in a liquid. Current AI models can make accurate predictions when the training and testing environments are the same. But in real-world science, molecules often appear in new environments, where existing models tend to fail.

To tackle this, we created a new method called RILOOD. It teaches the model to focus on stable, transferable patterns in molecules and their environments. We combine several ideas: mixing different environments during training, modeling both the fine-grained structure of individual molecules and their interactions with surrounding molecules, and learning invariant features that remain reliable across varied chemical settings.

Our results show that RILOOD does much better than existing approaches when handling diverse and unseen environments. This brings us a step closer to building AI models that can make accurate predictions in real-world, messy chemistry problems.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: Molecule relational learning, graph neural network, out of distribution generalization

Submission Number: 496