Environment-aware AI model predicts molecular solubility with uncertainty
EnSol: an environment-aware graph neural network for molecular solubility prediction
Machine LearningArtificial Intelligence
Summary
Molecular solubility is important for many chemical processes but is expensive to measure experimentally. The authors developed EnSol, an AI model that represents molecules as graphs and uses a special attention mechanism to understand interactions between dissolved substances and solvents. EnSol also accounts for temperature changes and predicts a range of likely solubility values instead of a single guess, which helps capture uncertainty. It performed better than existing models on benchmark tests and reliable when tested on actual chemical mixtures.
What this means in practice
- •For chemical process engineers: Identify optimal solvents and conditions for reactions and separations by predicting solubility considering temperature and uncertainty.
- •For pharmaceutical formulation teams: Rank solvents reliably during drug formulation development to improve performance and reduce costly experiments.
Authors
Thao Nguyen, Saman Shafaei, Zhengyi Zhang, Huimin Zhao, Heng Ji
Abstract
Molecular solubility directly affects key aspects of molecular development such as reaction feasibility, formulation performance, separation efficiency, and solvent selection. However, experimental measurement across solutes, solvents, and temperatures remains costly and sparsely sampled. Existing computational models often rely on fixed-solvent assumptions, deterministic formulations, or simplified representations of solute-solvent interactions, limiting their ability to capture complex molecular interactions, continuous temperature effects, and experimental uncertainty. Here, we introduce EnSol, an environment-aware probabilistic framework for molecular solubility prediction. EnSol represents the solute and solvent as molecular graphs and learns separate representations for each before bringing them together through cross-attention to capture solute-solvent interactions. Temperature is incorporated directly into the solvent environment through feature-wise modulation, and a mixture density network predicts full solubility distributions to capture both temperature-dependent behavior and experimental uncertainty. On the independent SolProp and Leeds benchmark datasets, EnSol achieved Spearman correlations of 0.876 and 0.601, respectively, outperforming state-of-the-art solubility prediction models across both benchmarks. Beyond computational benchmarking, experimental validation across chemically diverse solute-solvent pairs showed that EnSol maintained strong predictive performance and supported reliable solvent ranking, achieving a Spearman correlation of 0.715. These results show that EnSol can support reliable solubility prediction and solvent selection across diverse chemical systems while accounting for predictive uncertainty.