PRISM: Prior Relational Information for Self-supervised Modeling to Enhance Solubility OOD Generalization

AI summary

Overview

Research area: Machine learning for molecular property prediction — specifically label-free (self-supervised) pretraining of molecular encoders for solubility prediction under distribution shift, at the intersection of cheminformatics and representation learning.

Technical level: Advanced. The paper assumes familiarity with graph neural networks, contrastive/self-supervised pretraining objectives, KL-divergence-based distribution matching, and out-of-distribution (OOD) evaluation protocols.

Scope (one sentence): The paper introduces PRISM, a pretraining framework that converts expert-defined chemical descriptors into "relational teachers" and trains prior-specific molecular encoders to reproduce the neighborhood geometry those priors induce, then evaluates the resulting representations on scaffold, solute, solvent, and solvent-pair out-of-distribution solubility benchmarks.

What This Paper Is About

Solubility prediction matters for drug discovery and chemical process engineering, but real use requires generalizing to unseen solutes, unseen solvents, and unseen solvent mixtures — situations where labeled data are scarce. Existing self-supervised molecular pretraining methods (masking, reconstruction, augmentation-defined pairs) do not encode any explicitly chemically meaningful notion of which molecules should be close together.

PRISM's goal is to close that gap by turning well-established physicochemical descriptors — which require no labels — into supervision signals that shape the geometry of learned molecular embeddings so that chemically similar molecules land near each other.

Key Contributions

1. A label-free relational pretraining framework (PRISM). Each expert-defined chemical prior acts as a "relational teacher": a prior-space distance is converted into a per-anchor probability distribution over neighbors, and an encoder is trained so its embedding-space neighbor distribution matches it via Kullback-Leibler (KL) divergence.

2. A concrete solubility model built from four complementary priors. SlogP_VSA (hydrophobic/lipophilic effects), SMR_VSA (polarizability/dispersion), PEOE_VSA (electrostatic), and Morgan fingerprints (substructure/functional-group patterns), each with its own D-MPNN encoder, combined only downstream by weighted-sum (late) fusion.

3. OOD evidence across four shift types. Scaffold splits on MoleculeNet, solute OOD via BigSolDB → Leeds/SolProp transfer, solvent OOD, and solvent-pair-disjoint mixed-solvent splits — with the frozen-encoder mixed-solvent setup showing both strong accuracy and markedly lower run-to-run variation.

4. A diagnostic for prior selection ("label energy"). A nearest-neighbor label-smoothness measure (n = 10) normalized against a Gaussian random prior, intended to identify task-compatible priors before pretraining, plus ablations and controls showing the gains come from relational structure rather than descriptor access, reconstruction, or model capacity.

Main Findings

• Distribution matching wins as the alignment objective. On ESOL with a matched SlogP_VSA prior and D-MPNN backbone, distribution matching reached RMSE 0.844 / 0.776 / 0.772 at checkpoints C1–C3; ℓ₂ distance matching was 0.956 / 0.865 / 0.987 (degrading late, suggesting exact distance matching suppresses graph-derived information); rank-based alignment stayed near 1.02; neighborhood-graph alignment was 2.011 / 1.241 / 1.083.

• Best results on two of three MoleculeNet solubility tasks. PRISM achieved ESOL RMSE 0.772 ± 0.002 and Lipo 0.538 ± 0.02, versus D-MPNN 1.050 ± 0.01 and 0.683 ± 0.02, and XGBoost 0.925 and 0.741. On FreeSolv it was a boundary case: PRISM 1.770 ± 0.09 versus D-MPNN's 1.673 ± 0.08 — the authors attribute this to hydration free energy depending on 3D geometry that 2D encoders cannot capture, and report that a 3D-aware SchNet backbone improves FreeSolv substantially (Appendix H.10).

• Best pure-solvent external transfer. Pretrained on ZINC and fine-tuned on BigSolDB: Leeds RMSE 0.91 ± 0.002 with 75.7 ± 0.9% of predictions within one log unit, and SolProp RMSE 0.77 ± 0.030 with 82.3 ± 0.8%. This beats Chemprop (0.99 / 0.83), Fastprop (0.95 / 0.83), MolFormer (1.07 / 0.98), and XGBoost given direct access to the same prior features (1.00 / 0.78).

• Small pretraining corpus suffices. PRISM was pretrained on only 250K molecules, versus 1.1B for MolFormer and 77M for ChemBERTa.

• Strongest and most stable mixed-solvent OOD results. With frozen PRISM encoders plus a composition-weighted pooling head, solvent-pair-disjoint OOD gave RMSE 0.438 ± 0.030, R² 0.846 ± 0.031, and 96.6 ± 0.9% within one log unit — better than Triple-CIGIN (0.674), Triple-MMGNN (0.639), the D-MPNN aggregation controls (0.457–0.587), and XGBoost (0.459), with much smaller split-to-split spread than most baselines. Under the random split most methods saturate (PRISM 0.077 ± 0.006; Triple-CIGIN 0.052 ± 0.004).

• Relational pretraining, not descriptors or capacity, drives the gain. Leeds RMSE: PRISM 0.910 ± 0.002 vs. D-MPNN with concatenated priors 0.990, pointwise reconstruction 1.076, single shared encoder over all prior losses 0.930, 4× D-MPNN trained from scratch 1.001, and PRISM-Distilled 0.950. Similar ordering holds on SolProp and mixed-solvent OOD.

• Late fusion beats early fusion. Training one shared embedding against all priors (equivalent to a GraphMSL-style early fusion) underperforms, and post-hoc distillation into a single encoder recovers only part of the branch-specific signal.

• Random-prior control degrades but does not collapse. With an uninformative Gaussian prior, ESOL RMSE was 1.24 ± 0.025, FreeSolv 2.79 ± 0.058, Lipo 1.07 ± 0.014 — worse than scratch in the first two cases — showing prior quality matters, while the non-collapse shows distribution matching still permits graph-derived information.

• Priors are measurably task-aligned. Label energy ratios (ρ < 1 means smoother than random): ESOL 0.107 (VSA) / 0.150 (fingerprint); FreeSolv 0.072 / 0.090; Lipo 0.305 / 0.315.

• More priors is not better. On CYP2D6, using all candidate priors produced ROC-AUC 42.08 ± 9.35, whereas selecting the label-energy-compatible Pharm2D prior raised it to 73.31 ± 1.02.

• Complementary to physics-based methods. On the independent Cenci et al. dataset (pure, binary, and ternary solvent systems), PRISM improves on UNIFAC–PLS, and PRISM combined with UNIFAC gives the strongest overall and ternary results.

• Interpretability is prior-grounded but not causal. Adaptive fusion weights deviated from a uniform baseline in chemically plausible ways: SlogP-related contributions for glucose and cyclopentane-1,2-dicarboximide in water, and PEOE-related attribution for the nonpolar π-rich perylene. The authors explicitly frame these as prior-space contributions rather than causal decompositions. In a temperature-trend case study, PRISM preserved the correct solute ordering where Fastprop reversed it, distinguishing the two solutes only in the SlogP prior space.

Methodology in Plain English

The intuition is the old chemistry rule "like dissolves like": molecules with similar interaction profiles tend to dissolve each other. The researchers pick four standard, cheap descriptor families that each capture a different kind of interaction — hydrophobicity, polarizability/dispersion, electrostatics, and substructure patterns — and treat each one as an answer key for who should be close to whom.

For each descriptor family they run a separate graph-neural-network encoder. During pretraining, for every molecule in a mini-batch they convert descriptor distances into a probability distribution over its batch neighbors (closer molecules get higher probability), and do the same for the learned embedding distances. The loss pulls the embedding distribution toward the descriptor distribution using KL divergence. Because each anchor is normalized separately, these distributions are conditional and asymmetric, which keeps anchor-specific neighborhood geometry rather than forcing a single global picture.

No solubility labels are used in this stage — only unlabeled molecular graphs. A second design choice is to keep the four encoders separate during pretraining and merge them only afterward by a learned weighted sum, so that each branch can preserve its own distinct geometry. A gating network can produce molecule-specific weights, or weights can be shared across the task.

Downstream, task-specific heads handle three input shapes: single molecule, solute + solvent + temperature, and mixed solvent. For mixtures, the authors deliberately freeze the pretrained encoders and train only the fusion module plus a small head, pooling two solvent embeddings by mole fraction (α₁z₁ + α₂z₂ with α₁ + α₂ = 1). This low-capacity linear pooling is intentional — it tests whether the frozen embeddings already contain the relational information needed for solvent–solvent interactions.

To check whether a prior is even worth using, they measure "label energy": among a molecule's n nearest neighbors in the prior space, how much do the labels vary on average? They then normalize against a Gaussian random prior. A ratio below 1 means the prior organizes molecules more smoothly than random geometry, so it is likely a useful teacher.

Evaluation spans four shift types with RMSE and the fraction of predictions within one log unit of experimental log-solubility, plus R² for mixtures.

Why This Matters

Impact on research. The paper reframes self-supervision for molecules: instead of inventing pretext tasks from augmentations or masking, it asks where the relations in the embedding geometry should come from. It also provides a measurable, pre-pretraining criterion (label energy) for deciding which scientific priors to use and shows that stacking incompatible priors can actively hurt (the CYP2D6 result dropping to 42.08 ROC-AUC). The late-fusion finding — that fusing prior targets before learning loses information that post-hoc distillation only partly recovers — is a transferable design lesson for any multi-view pretraining setup.

Real-world applications:

• Drug discovery and formulation: predicting aqueous and co-solvent solubility for candidate compounds with limited experimental data, plus the temperature-trend behavior relevant to stability and processing.
• Crystallization and process development: selecting solvent systems for purification, where unseen solvent pairs are the norm and per-split stability matters.
• Green-solvent replacement: ranking alternative solvents for an existing solute when the substitute is not represented in training data.
• Environmental and agrochemical fate modeling: estimating how a compound distributes into water and mixed aqueous systems.

Industry relevance. Reproducing these results requires only 250K unlabeled molecules and standard descriptor software, not billion-molecule corpora or proprietary labels — a much lower barrier for labs with modest compute. The frozen-encoder mixed-solvent setup is also attractive operationally: swapping in a new solvent combination requires training only a lightweight head, and the reduced run-to-run variation reported for PRISM makes downstream decisions less sensitive to random seed. The strongest Cenci et al. results come from combining PRISM with UNIFAC, so the work is positioned as complementary to, not a replacement for, thermodynamic group-contribution methods.

Future Directions

• Backbone-dependent limits. FreeSolv remained a boundary case with 2D D-MPNN encoders; the authors report gains from a 3D-aware SchNet backbone but leave systematic integration of 3D information open, along with the framework's claims to being encoder-agnostic across more architectures.
• Beyond pairwise, linear fusion. The composition-weighted mixture pooling is explicitly a linear, low-capacity design; whether nonlinear solvent–solvent interactions can be modeled without sacrificing the stability reported here is unresolved.
• Prior selection as a general procedure. Label energy relies on some labeled molecules for diagnosis; extending it into a fully label-free or active-selection criterion — and testing whether the CYP2D6-style filtering generalizes — remains open.
• Transfer to other domains. The authors expect the strategy of turning relational priors into self-supervision signals to extend to other areas where chemically or biologically meaningful priors exist; demonstrating this beyond solubility (e.g., binding, toxicity, other thermodynamic endpoints) has not yet been shown.

Target Audience

Machine learning researchers working on self-supervised molecular representation learning and distribution-shift robustness; cheminformatics and computational chemistry practitioners building solubility or solvent-selection models; and pharmaceutical or process-engineering teams who need transferable predictions on poorly represented solutes and solvent mixtures. Readers should be comfortable with GNNs, KL divergence, and standard OOD benchmarking; the paper is not an entry-level introduction to either machine learning or thermodynamics.

Authors’ abstract

Solubility depends on intermolecular interactions between solutes and solvents, and is central to drug discovery, chemical and materials process engineering. Its prediction requires out-of-distribution (OOD) generalization to unseen solutes, solvents, and solvent mixtures, where supervised methods suffer from data scarcity, and existing self-supervised methods do not explicitly learn chemically meaningful relations. Learning chemically meaningful relational information without labels can alleviate data scarcity and improve OOD generalization, but remains challenging. We introduce Prior Relational Information for Self-supervised Modeling (PRISM) and apply it to solubility OOD generalization. PRISM is a label-free pretraining framework that converts expert-defined chemical priors into relational teachers. For solubility prediction, it trains four encoders to preserve the population-level relational information from priors capturing hydrophobic, electrostatic, polarizability, and substructural effects, using Kullback-Leibler (KL) divergence to match distance-based neighborhood distributions during pretraining. Across scaffold, solute, and solvent OOD settings, PRISM matches or surpasses supervised and self-supervised baselines. Under mixed-solvent OOD, freezing the four PRISM encoders and training a lightweight fusion-and-composition head outperforms strong baselines and substantially reduces run-to-run variation. Ablations and control studies confirm that gains stem from chemically meaningful relational structures baked into the encoders by PRISM. We expect this strategy of turning relational priors into self-supervision signals to extend beyond solubility prediction to domains where chemically or biologically meaningful relational priors are available.