Solvaformer: an SE(3)-equivariant graph transformer for small molecule solubility prediction

2025-11-12

AI Summary

To address the strong reliance on quantum descriptors, poor interpretability, and high experimental costs in small-molecule solubility prediction, this work proposes a geometry-aware SE(3)-equivariant graph neural network. The model integrates SE(3)-equivariant attention with scalar attention in a synergistic architecture to enable geometrically faithful intermolecular communication without spurious relative geometric constraints. A multi-task alternating training strategy jointly leverages quantum-mechanical computations and experimental solubility data, enhancing both generalization and attribution-based interpretability. On multiple benchmarks, our method matches the performance of DFT-augmented gradient-boosting approaches, significantly outperforms ablated EquiformerV2 variants and sequence-based models, and—through attention visualization—reveals key solvation mechanisms such as hydrogen bonding. This provides a high-accuracy, interpretable predictive tool for molecular synthesis and process optimization.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web evaluation methodologies and metrics

Abstract

Accurate prediction of small molecule solubility using material-sparing approaches is critical for accelerating synthesis and process optimization, yet experimental measurement is costly and many learning approaches either depend on quantumderived descriptors or offer limited interpretability. We introduce Solvaformer, a geometry-aware graph transformer that models solutions as multiple molecules with independent SE(3) symmetries. The architecture combines intramolecular SE(3)-equivariant attention with intermolecular scalar attention, enabling cross-molecular communication without imposing spurious relative geometry. We train Solvaformer in a multi-task setting to predict both solubility (log S) and solvation free energy, using an alternating-batch regimen that trains on quantum-mechanical data (CombiSolv-QM) and on experimental measurements (BigSolDB 2.0). Solvaformer attains the strongest overall performance among the learned models and approaches a DFT-assisted gradient-boosting baseline, while outperforming an EquiformerV2 ablation and sequence-based alternatives. In addition, token-level attention produces chemically coherent attributions: case studies recover known intra- vs. inter-molecular hydrogen-bonding patterns that govern solubility differences in positional isomers. Taken together, Solvaformer provides an accurate, scalable, and interpretable approach to solution-phase property prediction by uniting geometric inductive bias with a mixed dataset training strategy on complementary computational and experimental data.

Problem

• Predicting small molecule solubility using geometry-aware graph transformer
• Overcoming limitations of quantum-derived descriptors and poor interpretability
• Modeling solution-phase properties with SE(3)-equivariant attention mechanisms

Innovation

• SE(3)-equivariant graph transformer for molecular geometry
• Multi-task training with quantum and experimental datasets
• Combines intramolecular equivariant and intermolecular scalar attention