Relational Invariant Learning for Robust Solvation Free Energy Prediction
Authors: yeyunchen
Venue: ICML 2025 spotlightposter (2025)
Keywords: Molecule relational learning
graph neural network
out of distribution generalization
Primary area: applications->chemistry_physics_and_earth_sciences

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.

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Yeyun Chen 12
Abstract
solvation properties of molecules. More importantly, it ex- tends traditional molecular property prediction frameworks by explicitly incorporating solvent molecules as input fea- tures, thereby improving prediction accuracy and enhancing chemical interpretability.
Predicting the solvation free energy of molecules using graph neural networks holds significant po- tential for advancing drug discovery and the de- sign novel materials. While previous methods have demonstrated success on independent and identically distributed (IID) datasets, their perfor- mance in out-of-distribttion (OOD) scenarios re- mains largely unexplored. We propose novel Re- lational Invariant Learning framework (RILOOD) to enhance OOD generalization in solvation free energy prediction. RILOOD comprises three key components: (i) a mixuupbased conditional mod- eling module that integrates diverse environments, (ii) novel munti-granuulriity refinement strategy that extends beyond core substructures to enable context-aware representation learning for captur- ing multilevel interactions, and (iii) an invariant learning mechanism that identifies robust patterns generalizable to unseen environments. Exten- sive experiments demonstrate that RILOOD sig- nificantly outperforms state-of- the-art methods across various distribution shifts, highlighting its effectiveness in improving solvation free energy prediction under diverse conditions.
Despite significant advancements MRL, most existing methods operate under the assumption that training and test data independent and identically distributed (IID). However, real-world molecular systems exhibit diverse char- acteristics and uneven data distributions across solvents, making this assumption unrealistic in practical applications. Out-of-Diss D--istriibutin (OOD) scenarios arise when test data differs substantially from training data, as illustrated by the example Fig. This work focuses on exploring the OOD generalization of molecular solvation properties across dif- ferent environments within the MRL.
To address distribution shifts, several approaches have been proposed, including invariant learning (Wu et al., 2022a), feature disentanglement (Liu et al., 2021), and data aug- mentation (Sui et al., 2024; Jia al., 2024). Among these, invariant learning for OOD generalization (Krueger et al., 2021) has garnered significant attention due to its ability extract robust features that remain stable across different environments, even under distribution shifts. In molecu- lar modeling, molecular invariant learning is commonly employed to address distribution shifts by identifying core substructures that exhibit strong correlations with molecular properties. However, existing methods encounter notable limitations. For example, Lee al. (2023a) leverage privi- leged substructures as causal correlations in MRL. However, they do not account for the solvent-dependent nature of so- lute properties and the complex coupling effects that govern molecular behavior (Cramer & Truhlar, 2008). Similarly, Lee et al. (2023b) apply back-door adjustment to mitigate spurious correlations but fails account for solvent effects, thereby neglecting intricate solute-solvent interactions that are essential for accurately characterizing solute proper- ties. As result, these approaches lead to an incomplete understanding of solute behavior, increased susceptibility to spurious correlations, and poor generalization to unseen environments.
1. Introduction
Predicting the solvation free energy of molecules is crucial, as most chemical and pharmaceutical processes occur in solution, making it highly significant for downstream in- dustries (Chung et al., 2022; Varghese & Mushrif, 2019). This task, often referred to as Solute-Solvent Interaction in Molecular Relational Learning (MRL) (Lim & Jung, 2019; Subramanian et al., 2020; Panwar et al., 2021; Low al., 2022; Zhang et al., 2022: Lee al., 2023a;b), focuses on understanding and modeling the interactions between solutes and solvents, conceptualizing these interactions as
'Institute of Artificial Intelligence, Xiamen University, China "Shanghai Innovation Institute, China. Correspondence Yeyun Chen <ycyunchen2022@ @ www..muu.u......
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Although atomic interactions have been extensively mod- eled and have shown success in MRL, a precise under-

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Distribution of solvation free energy by dielectric consaant
Spurious correlation
Spurious correlation
Previous methods
25
Privileged
Privileged substructure
Scaflold
substruciure
Solvent
Solvent Gap
Solute distribution
Solute distribution
6.30
70-40
96.150
188-210
200
Dackstte ordaut
Dielectric constant
Distribution solvation unergy by solvent domain
Ofl
Our method
-000
Molecule interaction
Additional inputs
H
CH
Moleculc interaction
H
Relational
-00
invariance
CCUN
sehvet (b)
Cbo
(a)
Figure 1. Illustration of solvation free energy distrihution across different solvents. (a) illustrates how the solvation propertics of solute molecules vary with changes the solvent environment. The horizontal axis represents the dielectric constant solvent, and the dotted boxes indicate the distribution of solute properties in two solvents: benzyl alcohol (blue, left) and styrene (pink, right). Traditional methods predict solute properties by identifying core substructures, but core substructures may change different solvents. (b) presents the distribution of solvation energy across different solvent dielectric constants (top) and various solvents (bottom). highlighting the significant influence of solvent on solvation energy. Highly polar solvents, or those with intermediate dielectric constants, tend to exhibit stronger interactions with solutes, lcading broader range of solvation energies. Relational invariance refers molecular interactions such hydrogen bonding (purple) and van der Waals forces (gray). The statistical data originates from the MNSolv dataset (Marenich et al., 2012) (top) and the QMMSolv dataset (Ward et al., 2021) (bottom).
standing of the solute-solvent interaction remains elusive. Achieving both accuracy and explainability presents sig- nificant challenges, particularly in the following areas: (1) the need to accurately model solute-solvent interactions across chemically diverse solvent environments. (2) the
ery of fundamental molecular relationships in a chemically interpretable latent space. Our main contributions can be summarized as follows:
We formally formulate the out-of-distribution (OOD) gen- cralization problem in Molecular Relational Learning (MRL), establishing a rigorous foundation for studying model robustness across diverse chemical environments.
inherent complexity multilevel molecular interactions, which hinders the extraction invariant features and the construction of robust, generalizable representations. To address these issues, it is essential develop models that accurately represent multilevel molecular interactions, effectively capturing the complex one-to-many relationships between solutes and properties.
We propose RILOOD, a relational invariant learning framework for solvation free energy prediction, featuring three key components: mixup- based conditional model- ing module, a multi-grann granularity refinement strategy, and an invariant learning mechanism.
Based on the aforementioned analysis, in this work, we propose novel Relational Invariant Learning framework for Oft-oo-DDstribueee Distribution Generalization (RILOOD) in MRL. Unlike traditional methods, framework explicitly cap- tures invariant relationships in molecular pairs and achieves more generalized representation of solute-solvent interac- tions. Specifically, we first employ Graph Neural Network (GNN) to encode molecular structures, followed by cross- attention module to map atom-level interactions. then incorporate mixup-enhanced Conditional Variational Mod- eling to facilitate cross-environment invariance, leveraging multi-granularigg context-aware interaction mechanism and environment diversity inference. This enables learning of interaction invariance (Xie et al., 2024), allowing the discov-
We conduct extensive experiments across multiple distri- bution shifts, demonstrating that RILOOD consistently outperforms state-of-the-art methods, significantly ad- vancing OOD generalization in molecular property pre- diction.
2. Related Works
2.1. Molecular Relational Learning
Molecular Relational Learning (Lim & Jung, 2019; Pathak et al., 2020; Subramanian et al., 2020; Lee et al., 2023a;b), which aims study the relationship between molecules,
2

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
can be divided into molecular interaction prediction and Drug-Drug Interaction prediction. Molecular interaction prediction, i.e., solvent-based molecular property prediction, includes solvent free energy prediction, solubility predic- tion, chromophore ahsorption prediction, and so on. Unlike traditional molecular property prediction, the model need predict the properties exhibited by the same molecule ex- posed to multisolvent. Recent works (Ramani & Karmakar, 2024; Du al., 2024) leverage the merged graph to encode atomic interaction and further improve interpretability and reducing redundancy.
are difficult obtain, often relying on handcrafted rules that provide insufficient causal structure. (2) Invariant pat- terns and spurious correlations are entangled with shortcut features, complicating the identification of stable representa- tions. Addressing these challenges is crucial for improving the robustness and generalization of MRL models across diverse molecular environments.
3. Preliminaries
We define the uppercase letters (e.g.. 9) as random variables, and the blackboard typefaces (e.g., G) denote the sample spaces. Let g = (V, E G denote a graph, where V = (V1, V2,,., Vn) is the set of nodes and of edges.
2.2. Out of Distribution Generalization
V the set
Generalizing well-trained models to unseen environments with different data distributions remains a key challenge in machine learning. To address OOD generalization, three main approaches are typically employed: invariant learning (Li et al., 2022), causal inference (Dawid, 2000), and disentangled learning (Mo et al., 2023). Invariant learning aims extract stable features across distribution shifts, but ZIN (Lin et al., 2022) argues that identifying invariance in Euclidean data is impossible without envi- ronment labels, proposing auxiliary information as solu- tion. Causal inference approaches utilize Structural Causal Models (SCM) (Chen et al., 2022; Lu et al., 2021) and In- dependent Causal Mechanisms (ICM) (Peters et al., 2017: Gui et al., 2024) filter spurious correlations and enhance robust feature discovery. Disentangled learning separates features into invariant factors, which generalize across distri- butions, and spurious factors, which exhibit unstable corre- lations. While effective, it relies on strong prior assumptions and carefully curated datasets. These diverse strategies col- lectively tackle OOD generalization by distinguishing stable predictive patterns from environment-dependeer variations, yet significant challenges remain in accurately identifying and effectively leveraging invariant features.
3.1. Molecular Relational Learning.
The goal of MRL task is to predict the target label y given the associated input molecular pairs (G,, G2). can be formulated as modeling the conditional distribution p()|s1,
Notations. Given dataset D = {((G) )2)) where 91 E G is solute molecule, and G2 E G2 is solvent molecule, each molecular pair is associated with a target la- bel V. N the total number of samples. The objective is to train a model to predict y based on the input (21,92). The model should effectively learn the relationships between the input features and the target variable, leveraging the infor- mation from both G1 and G2 to accurately predict ) The model's performance will be evaluated based on the RMSE of the predicted output y in comparison to the ground truth labels y.
Molecular Representation. We implement our method based on (Pathak et al., 2020), which is message pass- ing architecture devised for the solute and solvent molecule interaction. Given pair of molecules G1 = (V2, &g) and
2.3. Invariant Learning in Molecular Relational Learning
G2 = (V2, &g). We first obtain the node representation of each molecule as follows: h = &1) h2 = GNN(V2, E2.. To capture inter-molecular interactions at the atomic level, the interaction map is constructed as fol- lowing: I = h1 h2" where is matrix multiplication, I E R N N2 Here, N1 and N₂ denote the number of atoms in molecule G and G2, respectively. We obtained a represen- tation hr E N xD of the solvent's 's interaction on the solute and a representation hig E TR. N2 of the solute's 's interaction on the solvent through a shared interaction map according to the following equations: hy = T 112, hi = 1" -hg. H2 generated by concatenating two representations hn and h1, [h,, ha]. The overall graph representation is obtained using readout layer Reatate H.)). which set the READOUT function as Set2Set (Vinyals et al., 2015).
Research on invariant learning in molecular representation learning remains sparse. One approach identifies core sub- structures using the graph information bottleneck to extract minimal task-relevant information (Lee et al., 2023a). An- other method leverages causal intervention learn causal substructures and mitigate distribution shifts (Lee et al., 2023b). In OOD settings, generalization is typically eval- uated by partitioning datasets into scenarios like "unseen solvent" or "unseen domain", where test sets exhibit spe- cific biases. However, many studies remain confined to intra-domain frameworks, failing to capture real-world com- plexities. Despite successes in graph-based invariant learn- ing (Wu et al., 2022a; Yang et al., 2022; Li et al., 2022), two key challenges persist: (1) Environmental labels for graphs
i.e., H₁ concat
3

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
3.2. OOD Generalization.
diverse environments. This approach reduces prediction errors and minimizes dependence on environmental varia- tions. A predictor performing well across multiple, varied environments is more likely to generalize robustly to un- seen distributions. Our primary objective is to develop a model that is robust to domain shifts, ensuring the mapping from molecule pairs to labels remains stable irrespective of environmental changes.
In this work, we mainly focus on OOD generalization in graph-level prediction tasks. Our aim is to train the model with limited labels to infer the domain distribution from unseen data in Die-
Problem formulation. Given a molecular pairs dataset, D =[((f) g)). ( ) colleet collect from multiple environ- ments E, which were considered as drawn independently from an identical distribution Pe. i.e., DID = {(91, 92) E D 91 E GID A G2 E GLD}. The training and test datasets are denoted as Der = {((9), ()),))))))).))))))) ) and Die = ((G)) (2), ) Our goal is to find an opti- mal predictor th (G), G2) + Y that performs well on all environments. Formally, the learning objectives can be formulated as:
Assumption 4.1. Given a molecular pair (Gi, 92), each pair is associated with R surrounding environments. We assume the existence of invariant interaction patterns that facili- tate generalizable OOD predictions across all environments. The optimal predictor (-) should satisfy the properties of Invariance and Sufficiency, as detailed in Appendix A.3.
Specifically, we further decompose ((-) into two key com- ponents as ((-) = o f G2): (a) A Conditional Varia- tional Autoencoder (CVAE) ʃ. which models the prior dis- tribution of the solute representation H₃ ~ (=(e)e across different environments. Here, z low-dimensional, con- tinuous representation of the solute in the latent space. (b) A multi-granularity context-aware learner g. which refines relational features by mapping (H1, H2) to a context-aware representation H., ie., 9 (H,, H2) + H..
min max E ((G1 [e (9) (5) ]*]]
+b eEE
(1)
Definition 3.1. (Data generation process) The OOD
distribution can be sampled according to DOOD =
(2) E D (G) E GOOD A E (GooD)((r E
GOOD A E V(G) E GIDA E GOOD)). The data generation process is as follows: Let & de- note all possible environments, supp(Ner) C supp(E), sampled train data from p(()) 92), V). Out of Distribu- tion indicate that Pe((S1,92) ()) + Pel I2), i.e., Durain = ((G) (2), ) C supp (). Drea = (((C), ) e' E
Building on Eq. 1, we reformulate the OOD generalization problem for molecular pairs as:
min max R "1eh,, (g (Gi, 9b) ,3') g EEg
(2)
where denotes the support environments, and E(:,-) repre- sents the loss function.
Details of the OOD dataset splitting are provided in the Appendix B.
4.2. Mixup-enhanced Conditional Variational Modeling
4. Methodology
Empirically, acquiring explicit environmental labels for so- lute-solved pairs is often impractical, presenting major obstacle to learning solute representations that generalize across diverse solvent conditions. This challenge is further compounded by limited data coverage, which may fail to capture distributional shifts induced by variations in solva- tion environments. Fundamentally, the difficulty in obtain- ing environmental labels arises from the fact that solvation energy is governed by complex, multi-scale interactions between solutes and solvents, including non-covalent in- termolecular forces and functional group-specific depen- dencies. In the absence of explicit labels, these underly- ing physicochemical mechanisms are difficult to accurately model or disentangle. To address this challenge, ZIN (Lin et al., 2022) introduces method for inferring latent envi- ronmental partitions using auxiliary information. Inspired by this approach, we propose conditioning on auxiliary in- formation to implicitly capture solute representations across the environment.
In this section, we introduce a Relational Invariant Learning framework designed to address oitt--f--Dtribttionn general- ization (RILOOD) in solvation free energy prediction. An overview of the proposed method is provided in Fig. 2. We detail the motivations and technical aspects of the three key components in RILOOD: Mixup-enhanced Conditional Variational Modeling (Section 4.2), Multi-granularity Context-Aware Refinement (Section 4.3), and Invariant Relational Learning Mechanism (Section 4.4).
4.1. The Overall Framework.
In molecular relation learning, substructure identification methods are widely employed. However, these approaches often fail to account for the variability in solute's 's behavior across different solvents, as solute-solvent interactions can differ significantly.
To overcome this limitation, invariant learning aims to iden- tify stable features or patterns that remain consistent across
Mixup Enhanced, To further improve generalization, we
4

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Training set
TML
Solvent
Solvent 30
Anm-sit Mise Sastroc
Solune
Pe e)
H
MMa
H
Intenction
Solventing n-1
Solvent
Solvent
Iuference
leatures
Alom-sorie
Solvent Testre H2
Fine-grained Interaction
Local
OOD Test set
Lissean
Auciliary information
Mixup-crlanned soluic fcature
Coars--graiee
O Fine-grained aware feature
Coase-grined aware
Figure 2. Our proposed RILOOD framework adopts an architecture that integrates mixup-cnhanced conditional variational modeling and guati-granularity context interaction mechanisms provide efficient way to adapt across different scenarios.
adopt mixup. which generates interpolated representations to enhance robustness across environments. We propose Mixup-enhanced CVAE (MCVAE), module designed to model molecular distributions using paired solvent infor- mation and infer the latent distribution (z(2)/1 (z)/f1, e) across diverse environments. Furthermore, we introduce uncer- tainty constraints to regularize the latent space, thereby im- proving stability and generalization across varying solvent conditions.
Here, (-)2 the uncertainty variance, and we constrain it. Consequently, the regression loss be reformulated as follows:
Lrey N (HHi) 11 ( log HH))
(4)
Variational Inference for Mixup Representations. To model the conditional log- likeliiidoood log p(H, e).
We assume that solvents belong R discrete categories, denoted as 3 = [e"]. 1, where each solvent type e" is rep- resented as R-dimensional one-hot vector e" E {0, 1) with its r-th dimension set to To enhance generaliza- tion, we employ mixup augmentation interpolate between environmental conditions, allowing the model learn con- tinuous latent representations rather than relying solely discrete labels. This encourages better adaptation to unseen domains.
introduce variational inference, reformulating the objective as a variational lower bound by approximating the posterior distribution q(> H);e):
max 13 AnDD [E (ootelieee) [logpo(0i 2, 2,ee]] 0,0
DEL (go(z) fhe |l po(= #.))
(5)
where is a threshold, and the KL divergence constraint ensures that the approximate posterior remains close to the prior, preventing latent space collapse and enhancing gener- alization across diverse solvent conditions.
Given the molecular representations H1 and H2 for molecules 9i and G2, respectively, we construct enhanced samples using the following mixup formulation:
H A- H1 +
e e= x + (3)
In practice, as the number of solvent categories grows, in- dependent modeling becomes computationally prohibitive. To enhance scalability, a learnable function could be em- ployed to assign soft environment labels to solvents, rather than modeling each solvent separately. However, given the limited solvent diversity in our dataset (e.g.. only five sol- vents in QM9Solv dataset), refrain from adopting this approach to ensure fair comparison.
where Hi denotes the mixed molecular representation and e signifies the interpolated solvent condition. The mixing coefficient is sampled from a Beta distribution, specif- ically À Beta(a, a), ensure smooth interpolation between different environments. This modeling strategy can be readily extended to other contexts, such as scaffold-base modeling. The derivations for the upper and lower bounds of this mixed representation are detailed in Appendix A.
Environment Inference via MCVAE Optimization. By optimizing MCVAE, we aim to infer the solute distribu- tion within a latent environment, offering a novel strategy for learning environment-awwre molecular representations.
For the regression task, we incorporate an uncertainty con- straint mitigate noise introduced by the mixup technique.
5

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Specifically, minimize the divergence between the ap- proximate posterior distribution (o(2 e) of the latent variable and the true posterior probability (H e, for specific environment Leveraging the rich prior knowledge embedded in conditional encoders, our training objective comprises two main components: (1) Encourage the R-group latent distribution produced by the encoder to approximate a standard normal distribution as closely as possible. (2) Sample z from the conditional distribution go(= e), and ensure that the reconstructed solute molecular features closely resemble the original features.
We begin by constructing an initial embedding E that ag- gregates relevant molecular features: F concat[z, H2], where and H2 represent specific molecular features. To model global interactions, we project E into query (Q), key (K), and value (V) matrices using learnable linear transfor- mations:
Q,K, = EWq. EWK, EWv
(7)
where WQ: WK. Wv are trainable projection matrices. A scaled dot-product self-attention mechanism is then applied to these matrices to capture long-rang -rangee dependencies across molecular entities:
LMCVAE(,, o; Bl,e) KL
(=( (=1 H1.e) elpo)] Hi)
1 12 Hall log o ( (H,))
Og = Attention (Q, K, V)
11
[QKT Vdle
(8)
(6)
= W. SSofmax
V
where 96 e,), with e (0, I). and is the identity matrix. Detailed proof is in Appendix A.2.
where We is learnable transformations that enhance atten- tion expressiveness. In parallel, local intra-molecular inter- actions are captured via a non-linear transformation:
4.3. Multi--ramularity Context-Aware Refinement.
Although the solvent type serves as a useful proxy for en- vironmental context, directly relying on can lead to pre- dictive shortcuts. To address this, we propose a joint opti- mization framework that simultaneously learns: (i) an envi- ronment modeling function from auxiliary information, (ii) an interaction-ane aware feature extractor, and (iii) mutual infor- mation constraints. The goal is to encode a context-aware molecular representation that dynamically captures solute behavior under specific conditions while accurately model- ing the context-dependent relevance of solute molecules.
Oj = PW bL)
(9)
where WL and by are learnable parameters for local fea- ture transformation. The PReLU is activation function to enhance the expressiveness of the local representations. Fi- nally, the context-aware representation is obtained by per- forming a Hadamard product between the global and local interaction features:
=z=0000
(10)
Previous studies have predominantly focused on atomic- level interactions within molecules (Pathak al., 2020). However, since solvents interact globally with solutes, multi-granularity interaction strategy is better suited to cap- ture their complex effects. Notably, solute-solvent interac- tions are non-covalent and are often neglected in conven- tional modeling approaches. To address this limitation, we employ a eeffaattention mechanism to dynamically iden- tify intermolecular interactions, allowing representations to adapt more effectively across diverse environments.
where denotes the Hadamard product, enabling the inte- gration of multi-scale contextual information into a unified and expressive feature representation.
To improve feature extraction, we maximize mutual informa- tion to retain essential features while reducing redundancy and noise. Specifically, we optimize the mutual information between the solute representation H1 and the context-aware feature H.. The solute feature H1, may be influenced by spurious correlations that hinder generalization. In contrast, the context-aware feature H. captures invariant and mean- ingful correlations, leading to more robust representations across diverse environments. We therefore denote He as Hi to emphasize its role in learning invariant features. To this end, we formulate the optimization objective as follows:
Multi-granularity Interactions with Context. To effec- tively model solute-solvent interactions, we propose a Multi- granularity Context-Aware Refinement (MCAR) strategy, designed to capture hierarchical interaction patterns across multiple levels. The MCAR mechanism is implemented in two key steps: (1) Context learning: Both coarse graiined molecular- level contexts and fine-graine fine-grained feature-level con- texts are simultaneously captured to jointly learn comprehen- sive contextual information. (2) Pattern refinement: Invari- ant interaction patterns are refined through matrix multipli- cation between the coarse- coarse-grained and fine- fine-graine grained feature representations.
max I (Bime: Fe. S.t. E
1 (Amm: Hi )) (11)
arg max
Amnv= (H)), Binc|
Finally, contrastive learning provides a practical solution for
6

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
the approximation, the learning objective is defined as:
Welling, 2016), GAT' (Veličkovié al., 2017), to obtain molecular embedding through concatenation and then enter the prediction layer, (2) there is interaction layer between molecular encoders, including ERM ( (Vapnik, 2013), Group- DRO (Sagawa et al., 2019), MixUp (Zhang et al., 7a), MolMerger (Ramani & Karmakar, 2024), CIGIN (Pathak et al., 2020), CGIB (Lee et al., 2023a), CMRL (Lee et al., 2023b).
DMI
explaimi (f Hi))
log
M
explsiml Mmo H.) 2ey cxp/simm Oo IE))
(12)
4.4. Invariant Relational Learning Mechanism
Optimization Objective. Eq. 2 clarifies the training objec- tive of OOD generalization. However, directly optimizing Eq. 2 is not impracticable. Instead, we formulate a joint optimization framework:
Metrics. We choose widely-used metries in previous works, the performance of the molecular interaction prediction task is evaluated in terms of RMSE (Pathak et al., 2020). Lower error indicate better prediction performance. AUROC (Lee et al., 2023b) for DDI prediction.
L Line + aLMvVAE + 3Lml
(13)
5.2. Main Results (RQ1)
where a and | are weight hyperparameters for LMCVAE and fyi respectively. The term Lan represents the prediction loss, measuring the discrepancy between the model's output and the ground truth. For regression tasks, Eq. 4 can be used instead of Lin--
Real-world Dataset. To assess the generalization perfor- mance our method, we conducted comprehensive exper- iments on three datasets, demonstrating the effectiveness of the proposed approach. To further explore distribution shifts across diverse environments, we evaluated model per- formance under two distinct settings: Solvent and Scaffold. The overall results, summarized in Tab. 1, lead to the fol- lowing key observations:
Proposition 4.2. Given the auxiliary enviroument e, our goal is to build model po HH e, that learns the feature H1 E I. conditioned on e. Optimizing Eq. 6 ensures that z exhibits sufficient predictive power, thereby allowing the model to satisfy the Sufficient condition in Assumption 4.1. Furthermore, minimizing Eq. 13 encourages the model to satisfy the Invariance condition in Assumption 4.1.
Our method consistently outperforms baseline models, achieving superior results across all datasets. Traditional approaches exhibit limitations, as they primarily rely on substructure- based invariance, which often introduces spu- rious correlations in MRL. The notable improvement ob- served in RILOOD stems from its ability to capture multi- granularity interactions and identify invariant patterns, en- abling the model to effectively adapt to domain shifts. Fur- ther discussions on extending this method to the L.D., set- ting are provided in Appendix C.2, with additional results in Tab. 4.
5. Experiments
In this section, we conduct extensive experiments to answer the research questions:
RQ1: How to evaluate the effectiveness of the model in OOD scenarios?
Synthetic Dataset. To evaluate model robustness under distribution shifts, we apply dataset-specific shift strategies, introducing spurious features to construct synthetic datasets. Following (Li et al., 2022; Wu et al., 2022b), spurious cor- relations are injected by controlling the variant distribution. Further details provided in Appendix C.2. Specifically, we manually introduce spurious relationships of varying degrees between environment and the label ) in the train- ing set, setting degree d = {0.25, 0.33, 0.5, 0.75}. The results, presented in Fig. (a), indicate that as d increases, performance generally improves due to a greater degree of distribution shift. However, our proposed method exhibits the highest stability, effectively mitigating the effects of spurious correlations.
RQ2: How effective is RILLOD in discovering invari- ant features and improving generalization?
5.1. Experimental Settings
Datasets. We use six datasets to evaluate our method. Specifically, the Minnesota Solvation Database (MN- Solv) (Marenich et al., 2012), QM9Solv (Ward et al., 2021), CompSoly (Moine et al., 2017), CombiSoly (Vermeire & Green, 2021), MolMerger (Ramani & Karmakar, 2024), and Abraham (Grubbs et al., 2010). The detailed statisties and descriptions are given in Appendix B. More experiments provided in Appendix C.2.
Baselines. For a comprehensive comparison, we adopt two types of baselines: (1) There is no interaction layer between molecular encoders. We use three commonly used GNN models, including GIN (Xu et al., 2018), GCN (Kipf &
Generalization on Graph Classification. To evaluate the applicability of our method to molecular pair data and classi- fication tasks, we conducted experiments on the DDI dataset.
7

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Table Performance comparison with baselines on aut-of-distribution real-world datasets from MNSolv, CompSolv. QM9Solv terms of RMSE. Different dataset splits by specific shift (solvent split and scaffold split), and details can be find in Appendix B. The best and the runner-up results are highlighted bolded and underlined respectively.
MNSolv
C
QM9Solv
Interaction
Method
Solvent
Scaffold
Solvent
Scaffold
Solvent
Scaffold
GCN
00.024
1..052+0.022
0.764440.024
0.9598
0.9115
1.031940.046
x
GIN
0.7723 1.3685+0.049
0,5927 0.113
0.600440.022
1.0928
1.0762+0002
GAT
0.7022=0.012
1.4566+0.022
0.5338+0.0m 064
0,5526440.02
1.0513
1.0415+00012
ERM
0.7503+0.00 22.026
1.3478=0.001
0.5360+0.000 002
0.633440.003
0.7471
0. 0.7261+00000
GroupDRO
0.7839+0..m 003
1.4322+0.003
0.5857=0.00 013
0.7459+0.012
0.8259 00.007
0.8503 #0.021
MixUp
0.7135=0.011
1.3843 012
0.5772+ 026
0.5604+0.017
0.7227+0.000 4000
0.7490
MolMerger
0.927640.088
1.6225520.0e
2.1115
1.2854
0.8773
1.3799 #0.041
CIGIN
0.7662+0.017
1.3649
0.5574 .002
0.6383
0.7503
0.8642##.012
CGIB
0.8312 10.017
2.2118 100002
0.3886 f 1..025
0.547610.02
1.4525 10.013
0.7894 10.006
CMRL
0.806340.00
2. 2.1524
0.3777 0077770..023 0.333
0.6672
1.4425 10.016
0.7894 2.0002
RILOOD
0.6784 10.007
1.0780 10.013
0.3660 10.022
0.5209 1..014
0.7001
10.001
0.6991
10.003
OCR
1.00
Dasp00)
MCAA
1.1
MCAI
0.5
1.0
RMSE
AUROC
0.9
0.84
0.84
0.84
0.00
0.82
0.8
0.1
0.7
0.3
0.4
0.5
0.6
0.7
RILGOD
CMRL
Solvent
Scaffold
Assay
Sizes
(a)
(b)
Figure 3. (a) Performance under different spurious correlation levels, where the strength of spurious correlation is defined as d Nunther samples with sparicus festures higher d in the training set indicates stronger spurious correlations with the underlying environments. (b) Total namber of samples
Results three DDI datasets with domain shifts, comparing approach against three SOTA methods. (c) The impact of different interaction modes on the CompSolv dataset, where w/ MCAR denotes a multi-gaauulartty interaction mode and w/o MCAR indicates a node-level interaction mode.
As illustrated in Fig. 3(b), RILOOD consistently outper- forms existing approaches under OOD conditions. This per- formance gain can be attributed to RILOOD's 's enhanced gen- cralization capability, which facilitates effective knowledge transfer from known molecular interactions to structurally similar compounds and previously unseen scaffolds. Such transferability improves the model's robustness distribu- tional shifts, thereby ensuring adaptability across diverse molecular structures.
results are summarized in Tab. 2. From Tab. 2, we observe the following: (1) Incorporating MCAR improves baseline performance, highlighting the importance of context-aware interactions in enhancing model robustness.
(2) While con- ditional modeling significantly affects performance, the in- dividual contributions of LMCVAE and LMi are smaller com- pared joint training.
(3) Removing MCAR leads performance drop; however, the model still surpasses the baseline due to the co-optimization of all losses.
Feature Visualization. To evaluate the effectiveness of MCAR, we used t-SNE to visualize molecular interactions in the best performing model and compare them to base- lines. As shown in Fig. 4, (1) The solute in the test set originated from different distributions than the training set, demonstrating the distribution shift; (2) MCAR enhances feature diversity, improving molecule interaction modeling; and (3) MCAR captures domain-invarian features, boosting generalization to unseen domains. These results confirm our
5.3. In-depth Analysis (RQ2)
To assess the contribution of each module, we conduct ablation study by removing specific components: Multi- granularity Aowte Refinement (MCAR) trained on the downstream task (M); mutual information loss Lm (Mi); conditional distribution modeling loss LMCVAE (MC); and MCAR removal but all loss is used (w/o MCAR). The com- plete model is jointly trained using Eq. 13 (Ours). The
8

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
solubility prediction in organic solvents and water. Nature communications, 1 111):5753, 2020.
Table Ablation study on CompSolv-* and MINSov-- by RMSE. We show the results of our method that performs best among baselines on all CompSolv-* and MNSolv- datasets, for comparison.
Chen, Y., Zhang, Y., Bian, Y., Yang, H., Kaili, M., Xie, B., Liu, T., Han, B., and Cheng. Learning causally invari- ant representations for autoof-distribution generalization on graphs. Advances in Neural Information Processing Systems, 15:22131-222111 2022.
CompSoly
MNSolv.
Method
Solvent
Scaffold
Solvent
Scaffold
Raseline [B]
0.5215
0.638320.011
0.7062-20.06
1.26-820.015
+ ERM loss [E]
0,486420.co2
0,591920.002
0,7263
1.3478
B - MCAR [M]
0.4914.0.om 201
0.5842.0.000
0,7115 10.000
1.2191 10.00
Chung, Y., Vermeire, H., Wu, H., Walker, J., Abraham, M. H., and Green, W. H. Group contribution and machine learning approaches to predict abraham solute parameters, solvation free energy, and solvation enthalpy. Journal of Chemical Information and Modeling, 62(3):433 446, 2022.
Lass [Mi]
0.5196
0.5444.0.aez
a.7279ao.ouz
1.2005
Mi Carevaa [MC]
0.753
0,5351
0,702260.00
1.1329 10.01
w/o MCAR
0 0.000
20.013
g0.05
1.3929 003
Ours
20.69.00.omm
0.5209:0.33
0.678420.007
1.0780
Cramer, C. J. and Truhlar, D. G. A universal approach to solvation modeling. Accounts of chemical research, 41 (6):760-768, 2008.
Figure 4. Visualization of the extracted features on training and validation set when the model achieves the best performance on the validation set. (a) The feature distribution the training set and the test set; (b) Effect of MCAR solute feature distribution; (c) Effect MCAR on global feature (solute + solvent) distribution.
Dawid, A. Causal inference without counterfactuals. Journal of the American statistical Association, 95(450): 407-424 207-424, 2000.
Delaney, J. S. Esol: estimating aqueous solubility directly from molecular structure. Journal of chemical informa- tion and computer sciences, 44(3):1000-100) 2004.
method's robustness against distribution shifts.
Du, W., Zhang, S., Di Wu, J. X., Zhao, Z., Fang, J., and Wang. Y. Mmgnn: A molecular merged graph neural net- work for explainable solvation free energy prediction. In Proceedings of the Thirty- International Joint Con- ference on Artificial Intelligence, pp. 3808-5814 2024.
6. Conclusion
In this work, we propose Relational Invariant Learning framework to solve out-of-dissribution in solvation free en- ergy prediction. Three tailored modules are jointly opti- mized to train the model and learn the representation of in- variant molecules diverse environments. Mixup enhanced molecular representations are used for variational model- ing of diverse environments, further capturing invariant in- teraction patterns through multi caaltiigaaularity context-awend refinement strategy. Extensive experiments and theoretical analysis prove the superiority of our method.
Grubbs, L. M., Saifullah, M., Nohelli, E., Ye, S., Achi, S. S., Acree Jr. W. E., and Abraham, M. H. Mathematical correlations for describing solute transfer into function- alized alkane solvents containing hydroxyl, ether, ester ketone solvents. Fluid phase equilibria, 298(1):48-8))) 2010.
Gui, S., Liu, M., Li, X., Luo, Y., and Ji, S. Joint learning of label and environment causal independence for graph out-of-distribution generalization. Advances in Neural Information Processing Systems, 36, 2024.
Impact Statement
This paper presents work whose goal advance the field of AI for Science. Our method may have some potential so- cietal consequences, but we don't 't't believe any need specific highlighting here.
Jia, T., Li, H., Yang, C., Tao, T., and Shi, C. Graph invariant learning with subgraph co-mixup for -isttriibittto generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 8562-8570, 2024.
References
Kipf, T. N. and Welling, M. Semi-supervised classifica- tion with graph convolutional networks. arXiv preprint TrXiv:160900200 2016.
Bemis, G. W. and Murcko, M. A. The properties of known drugs. 1. molecular frameworks. Journal of medicinal chemistry, 39(15))2887-2893, 1996.
Krasnov, L., Mikhaylov, S., Fedorov, M., and Sosnin, S. Bigsoldb: Solubility dataset of compounds in organic solvents and water a wide range of temperatures. 2023.
Boobier, S., Hose, D. R., Blacker, A. J., and Nguyen, B. N. Machine learning with physicochemical relationships:

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Bi- nas, J., Zhang, D., Le Priol, R., and Courville, A. Out- of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.
Moine, E., Privat, R., Sirjean, B., and Jaubert, J.-N. Estima- tion of solvation quantities from experimental thermody- namic data: Development of the comprehensive compsol databank for pure and mixed solutes. Journal of Physical and Chemical Reference Data, 46(3). 2017.
Lee, N., Hyun, D., Na, G. S., Kim, S., Lec, J., and Park, C. Conditional graph information bottleneck for molecular relational learning. arXiv preprint arXiv:2305.01520. 2023a.
Panwar, A., Shirazian, S., Singh, M., and Walker, G. M. Comprehensive modelling of pharmaceutical solvation energy in different solvents. Journal of Molecular Liq- uids, 1411111390. 2021.
Lee, N., Yoon, K., Na, G.S. S., Kim, S., and Park, C. Shift- robust molecular relational learning with causal substruc- ture. arXiv preprint arXiv:2305. 18845, 1,88,,,,,200 2023b.
Pathak, Y., Laghuvarapu, S., Mehta, S., and Priyakumar, U. D. Chemically interpretable graph interaction network for prediction of pharmacokinetic properties of drug- molecules. pp. 873-880, 2020.
Li, H., Zhang, Z., Wang, X., and Zhu, W. Learning invariant graph representations for out-of-distribttion generaliza- tion. Advances in Neural Information Processing Systems, 35:1 1828-1211111 1141,, 2022.
Peters, J., Janzing, D., and Schölkopf. B. Elements of causal inference: foundations and learning algorithms. The MIT Press, 2017.
Lim, H. and Jung, Y. Delfos: deep learning model for prediction of solvation free energies in generic organic solvents. Chemical science, 10(36):306-8315, 2019.
Ramani, V. and Karmakar, T. Graph neural networks for pre- dicting solubility in diverse solvents using molmerger in- corporating solute-solvent interactions. Journal of Chem- ical Theory and Computation, 20(15):6549 0600555599588,,, 2024.
Lin, Y., Zhu, S., Tan, L., and Cui, P. Zin: When and how learn invariance without environment partition? Advances in Neural Information Processing Systems, 35: 14529-24544 2022.
Ryu, Y., Kim, H. U., and Lee, S. Y. Deep learning improves prediction of drug-drup and drug-food interac- tions. Proceedings of the national academy of sciences, 115(18):E4304 E4311, 2018
Liu, C., Sun, X., Wang, J., Tang, Li, T., Qin, T., Chen, W., and Liu, T.-Y. Learning causal semantic representation for out-of-distributtin prediction. Advances Neural Information Processing Systems, 34:6155-6177 2021.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gener- alization. arXiv preprint arXiv: 1XXiv:9911.000 1911191911. 2019.
Low, K., Coote, M. L., and Izgorodina, E. I. Explainable solvation free energy prediction combining graph neural networks with chemical intuition. Journal of Chemical Information and Modeling, 62(22):5457-5470 2022.
Subramanian, V., Ratkova, E., Palmer, D., Engkvist, O., Fedorov, M., and Llinas, A. Multisolvent models for solvation free energy predictions using 3d-rism hydra- tion thermodynamic descriptors. Journal of Chemical Information and Modeling, 6076):2977,2988, 2020.
Lu, C., Wu, Y., Hernández-Lobato. J. M., and Schölkopf, B. Invariant causal representation learning for out-of- distribution generalization. In International Conference Learning Representations, 2021.
Sui, Y., Wu, Q., Wu, J., Cui, Q., Li, I.., Zhou, J., Wang, X., and He, X. Unleashing the power of graph data augmentation covariate distribution shift. Advances in Neural Information Processing Systems, 36, 2024.
Marenich, A. V., Kelly, C. P., Thompson, J. D., Hawkins, G. D., Chambers, C. C., Giesen, J., Winget, P., Cramer, J., and Truhlar, D. G. Minnesota solvation database- version 2012. 2012.
Vapnik, V. The nature of statistical learning theory. Springer science & business media, 2013.
Marinka Zitnik, S., Sosie, R., and Leskovec, J. Biosnap datasets: Stanford biomedical network dataset collection, 2018.
Varghese, J. and Mushrif, S. H. Origins of complex solvent effects on chemical reactivity and computational tools investigate them: review. Reaction Chemistry & Engineering, 4(2):165-200 2019.
Mo. Y., Lei, Y.. Shen, J., Shi, X., Shen, H. T., and Zhu, X. Disentangled multiplex graph representation learning. In International Conference on Machine Learning, pp. 11988-255003 PMLR, 2023.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv: 1710. 10903, 2017.
10

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Vermeire, F. H. and Green, W. H. Transfer learning for solvation free energies: From quantum chemistry periments. Chemical Engineering Journal, 418: 129307, 2021.
Vinyals, O., Bengio, S., and Kudlur, M. Order mat- ters: Sequence to sequence for sets. arXiv preprint arXiv:1511.06391, 2015.
Ward, L., Dandu, N., Blaiszik, B., Narayanan, B., Assary, R. S., Redfern, C., Foster, and Curtiss, L. A. Graph- based approaches for predicting solvation energy in mul- tiple solvents: open datasets and machine learning mod- els. The Journal of Physical Chemistry A, 125(27):5990- 5998, 2021.
Wu, Q. Zhang, H., Yan, J., and Wipf, D. Handling distribu- tion shifts graphs: An invariance perspective. arXiv preprint arXiv:2202.02400 2022a.
Wu, Y.-... Wang, X., Zhang, A., He, X.. and Chua, T.- S. Discovering invariant rationales for graph neural net- works. arXiv preprint 1xXiv:2201.12877 2022b.
Xie, Z., Tu, S., and Xu, L. Multilevel attention network with supe-supe domain adaptation for targettgeet prediction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 329 329-337, 2024.
Xu, K. Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? arXiv preprint arXiv:1810.000 00826, 2018.
Yang. N., Zeng, K., Wu, Q. Jia, X., and Yan, J. Learning substructure invariance for cut-of-distribution molecular representations. Advances Neural Information Process- ing Systems, 15:12964-12978. 2022.
Zhang, D., Xia, S., and Zhang, Y. Accurate prediction of aqueous free solvation energies using 3d atomic feature- based graph neural network with transfer learning. Jour- nal of chemical information and modeling, 62(8):1840 1848, 2022.
Zhang, H., Cisse, M., Dauphin, Y. N., and Lopez-Paz, D. mixup: Beyond empirical risk minimization. arXiv preprint wxXi:11110004122 2017a.
Zhang, W., Chen, Y., Liu, F., Luo, F., Tian, G., and Li, X. Predicting potential drug-drug interactions by integrating chemical, biological, phenotypic and network data. BMC bioinformatics. 18:1-12, 2017b.

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
A. Proofs
In this section, we provides detailed proofs in Section
A.1. Proof for Bound of Mixup Representations
Let H1 E Rd be the representation of a solute graph 91, and H2 E Rd be the representation of a solvent graph O2 under environment e, both obtained via GNN. We define the mixup representation as:
= 1H1 + (1 A)H2,
X [
(14)
We assume the following: The predictive model f IRd + y is Lipschitz continuous with constant by >> 0, i.e., for any H, H' E Id,
(I((H) J(H')l> s by || 12
(15)
The loss function & y x Vurzet - IR 2 O is upper-bounded by constant B:
E(Opreds Mime) B, Pmeeed Virue
(16)
Let D be the joint distribution over (H1, H2, (1, 62). where 61, Vurget. Define the expected prediction error under mixup as:
Emixup (A) = M(M ((f((t),))
(17)
where Xe (1 X)ez denotes environment, i.e., the soft target label. According the information-theores generalization bounds based on Fano's inequality, obtain:
Emixup (A) 2 (H(Qisage)) T(H:; Munges)
(18)
where H (Quuge)) is the entropy of the target Viurgt, and [(:; Munga) is the mutual information between the representation H and the soft label. Constant Ci > depends on the task structure and loss smoothness.
Generalization via Nuisance Information. From representation learning perspective, generalization error can be bounded by:
Emixup (A) E- + 02 V I(H: Cuuisance);
(19)
where Couisance denotes spurious nuisance factors, 02 > 0 depends on by and B, and Ex denotes the optimal lower bound of the achievable generalization error.
Therefore, obtain:
(H(Y) & Emixup (A) S --*+0 c2 V(H; Cuuisance)
(20)
These derivations indicate fundamental trade-off in mixup-based CVAE learning: while mixup reduces nuisance information through distributional smoothing, it must preserve sufficient task-relevant information avoid performance degradation. This analysis provides theoretical justification for using mixup representations in conditional generative models under structured environments.
A.2. Proof for Equation 6
For given mixed representation H, and environment e, the loss function is defined as:
I(O, do fh,e) e) E (e (n. log po (h 22, Drr(us(= fh,
ff))
(21)
Here, the learning objective is to find the optimal parameter set that maximizes the log- likellooddood log pe HH e). Since directly computing the true posterior (o(= H4, e) is intractable, we approximate it using an auxiliary variational distribution
12

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
[=(= e). By minimizing the KL divergence between (o(= H2 and po(= e), we derive tractable Evidence Lower
Bound (ELBO) formulation:
max E 0,0
(( log (Fr 2,
[2,e)]] s.t. DKL (go(2 Hl,e)llpo(z
Hi)) <
(22)
where 5 a threshold that ensures the learned latent representation z remains close to the true underlying data distribution. We define (2 H1,e) e) as the recognition model and po 2) the generative model. By leveraging approximate
posterior inference and the reparameterizzttion trick, the prior able to effectively capture environmental information from the posterior, thereby improving posterior alignment.
Using the ELBO decomposition, we express the log- ligeliiel as:
log pof, |e) DDxx H,, e) [po (1)) + EE (=)0))))))) [oog (F1 =
(23)
where (-Il- denotes the Kullback-Leeibler divergence between two distributions.
For regression-based solvation property prediction, we introduce an uncertainty constraint on RMSE, modifying the reconstruction as follows:
CM Φ; e) = DDL (go(2) li,, (po(z ii)))
(24)
Lo( (H1) 112 Hill log (H,,)
This term ensures that the variance of the latent variable a (22)² 222n is explicitly accounted for, promoting robust representation learning and reducing overfitting noise in the data.
A.3. Details of Assumption 4.1
We define two key properties, invariance and sufficiency, which ensure the generalizability of molecular relational learning across different environments. For given OOD environment e', we define invariance and sufficiency as:
(1) Invariance Property Ve, E supp(E), p() H.,e) p() H ee), where H (G 9b) denotes molecular pairs representations, H = (99, 92). H ( 9
(2) Sufficiency Property y (9 (5 + where + is a predictor, e a random noise.
the following conditions hold:
Conditional Independence:
4G G2) T (G1,92) (
(25)
which ensures that the learned representation is conditionally independent of the remaining molecular graph information.
For all environments e E @uupp()), there exists OOD environment e' E supp(E) satisfy:
(er ((),,))),))) = Per (b) ((9,, ( G2))
(26)
which implies that the joint distribution factorizes into invariant representations and residual molecular information.
Invariance of representation across environments:
pe (*(( pe (9(G).))))
(27)
This guarantees that the representation distribution remains stable across environments and e', reinforcing the assumption of invariant learning.
13

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Water
DX5D Etharnt
144
DMS0
ACN
0480
Acetone
Sel (kcal/mol)
(b)
Figure 5. (a) The distribution of property labels of solute molecules in five solvents (water, acetone, ethanol, ACN, and DMSO) in the QM9Solv dataset. (b) t-SNE visualization of the representation distributions solute molecules the QM9Solv dataset across the same five solvents.
B. Details on Datasets
We provide more details about the benchmark datasets that we use in our experiments. Statistics of the datasets are provided in Table 3.
MNSolv dataset includes 3,037 experimentally measured free energies of solvation or transfer for 790 distinct solutes and 92 different solvents. In this work, we remove the charged ions, and we consider 2,577 combinations of 528 unique solutes and 104 solvents, including mixed solvents.
QM9Solv dataset includes solvation energies for 130,258 molecules from the QM9 dataset, calculated in 5 solvents (acetone, ethanol, acetonitrile, dimethyl sulfoxide, and water) using an implicit solvent model.
CompSoly dataset is designed illustrate the impact of hydrogen-bonding association effects solvation energies. includes 3,548 pairings of 442 distinct solutes and 259 solvents, based on prior studies. The dataset also contains the original assay attribution of the data, which we have divided it into OOD settings.
CombiSoly integrates all data from MNSol, FreeSolv, CompSolv, and Abraham. After we removed the outliers, there were 1,000,000 entries, including 11,029 solutes and 284 solvents.
MolMerger compiled a dataset of 6,975 entries from three sources: BigSoIDB (Krasnov et al., 2023), BNNLabs Solubility (Boobier et al., 2020), and ESOL (Delaney, 2004). The data obtained from BigSolDB required cleaning to remove erroneous unreliable entries.
Abraham dataset, compiled by the Abraham research group at University College London, comprises 6,091 pairings of 1,038 unique solutes and 122 solvents, building on previous studies.
DI (Zhang et al., 2017b), ChChMiner (Marinka Zitnik et al.. 2018) and DeepDDI (Ryu et al., 2018) are datasets focused on drug- interaction prediction, widely used for evaluating molecular relation learning models.
B.1. Data splitting.
To evaluate the OOD generalization performance of molecule relational learning models, we employed both random scaffold splitting and random solvent splitting strategies.
Random Scaffold Splitting (Bemis & Murcko, 1996) is method where the dataset is grouped and randomly split based on molecular scaffolds. Specifically, the Bemis-Murcko was used to extract the scaffold structure of each molecule, and molecules with the same scaffold were grouped together. These scaffold groups were then randomly shuffled and split into training, validation, and test sets according to a predefined ratio (c.g., 8:1:1). This approach ensures that molecules with the same scaffold appear exclusively in one set, allowing the evaluation of the model's generalization ability to unseen scaffolds.
14

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
Table Details about the datasets used in OOD experiments.
Dataset
G1
G2
Target
#G1
#G2
MNSolv
Solute
Solvent
Solvation Energy
528
104
QM9Solv
Solute
Solvent
Solvation Energy
130625
5
CompSolv
Solute
Solvent
Solvation Energy
442
259
CombiSolv
Solute
Solvent
Solvation Energy
11029
284
MolMerger
Solute
Solvent
LogS
2699
13
Abraham
Solute
Solvent
Solvation Energy
1038
122
Molecular Weight of Solute Distribution
Heavy Atoms of Solute Distribution
4000
1000
3000
800
Frequency
Frequency
600
2000
1000
200
0
200
300
400
500
600
700
30
Molecular Weight
Mass of Heavy Atoms
Molecular Weight Solvent Distribution
Heavy Aloms Solvent Distribution
5000
1250
4000
1000
Frequency
Frequency
3000
750
2000
500
1000
250
100
200
300
400
500
600
700
5
10
15
20
25
30
35
Molecular Weight
Mass of Heavy Atoms
Figure Histograms the molecular weight and heavy atomic distribution of solutes and solvents the CombiSolv dataset. Based on the solvent solute ratios in Table 5, observe that the solutes and solvents in this datasct are mainly composed of small molecules and exhibit a long-tal tail distribution in molecular weight and heavy atom number. The average molecular weight of the solvent lower than that of the solute and may be mainly common small molecule solvents. This distribution can pose challenge the generalization ability of the model, especially when working with large molecules or complex solvents.
Random Solvent Splitting is a dataset partitioning strategy based on solvents, similar to random scaffold splitting, Solvent- domain shift considers the scenario which the solute distribution p(), (G2) is shifted across solvent splits. Here, solute molecules are grouped according to their corresponding solvents, and molecules with the same solvent are assigned the same group. These solvent groups are then randomly shuffled and divided into training, validation, and test sets in fixed ratio. Each solvent is assigned to only one set, ensuring that the solvents in the test set entirely unseen during training. This strategy aims to evaluate the model's performance in predicting molecular properties in solvent environments and its adaptability to nut-of-distribution solvents.
C. More Experiments and details
C.1. Implementation and Optimization Details.
The proposed method is implemented on a single NVIDA 3090 GPU with PyTorch. Following the CIGIN(Pathak et al., 2020), we use the same 3-layer GCN and MPNN as feature extractor for solute molecule and solvent molecule, respectively. More details about backbone can be found in S.c.3.1. During the training, the solute features were incorporated with node interaction features, which is the dot-production similarity between solute node features and solvent node features. Here, we
15

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
175
00
150
125
100
25
-10
-5
DeltaGsolv
Figure 7. Spurious correlation from dielectric constant eps.
Table 4. Performance molecular interaction prediction task (regression) in terms of RMSE.
Chromophore
Model
MNSolv
QM9Solv
FreeSoly
CompSolv
Abraham
CombiSoly
Absorption
Emission
Lifetime
GCN
25.75
31.87
0.675 0.866
20.17100011
1.192
_0.09
0.672_a0.022
GIN
24.9241.07
20.25
0.829 20.027
0.669 10.013
1.274=0..0
g.015
0,331 2.001
0.648
g0.024
0.595_0.014
GAT
26.1921.44
10.
0.859200016
0.731 007
1.305 021
1,280
0,387=0.010
0.798g0.005
0.662_0.021
CIGIN
19.32.0.35
25.0940.32
0.004
0.607:2.024
0.592 023
0.905
0.308
0.41120.000
0.451
CGIB
18.11 23.90
0.771
0.538
0.549 =: 026
0.852 032
0.276 0.003
0.390
0.422
CMRL
17.93
24.30
0.776 00?
0.551 017
0.288.20.01
0.815 046
0.255_0.011
0.374
0.421 =00 000000088
RILOOD
17.20
23.6110.21
0.706=20.05
0.48920.0c 009
0,246
0.82320.017
0,242_0.018
0.309220.13
0.292 0.00
using graph-level solute features and solvent features as input in our method. We select 168 for the dimension (dz) of latent variables. The learning rate was decreased on plateau by factor of 10-3 from 10-3 to 10-5
C.2. Synthetic data experiments.
We first consider the distribution shift caused by polarity bias w.r.t. eps. The invariant feature is A E R. where p(Y) H is constant, indicating a stable correlation between y and A Our goal is to learn model that relies solely on Hine. We use eps to control the degree of spurious correlation. Fig. 7 illustrates a spurious correlation arising from the dielectric constant, which plays a crucial role in determining the polarity of the solvent. As shown in Fig. 7, the correlation of molecular pairs and label y with eps=78 is unstable, i.e., p()] Hi) is unstable, p())) stable.
C.3. Generalization Analysis.
Table 4 presents the RMSE performance of various models on molecular interaction prediction tasks under IID settings. RILOOD consistently achieves the lowest RMSE across most datasets, demonstrating superior predictive capability, Traditional graph-based models (GCN, GIN, GAT) perform worse, particularly in QM9Solv, FreeSolv, and CombiSolv, indicating limitations capturing complex molecular interactions. More advanced models show notable improvements, yet RILOOD surpasses them in nearly all cases, with significant gains in QM9Solv, CompSolv, and CombiSoly. In Chromophore datasets, RILOOD also achieves the best results, particularly in Absorption and Emission, highlighting its strength in modeling spectral properties. The Multi Mulli-granularity Context- Aware Refinement (MCAR) and Mixup-bascd CVAE (MCVAE) mechanisms likely enhance feature extraction and generalization. Although RILOOD slightly underperforms CGIB in FreeSolv, its overall performance validates its effectiveness in molecular representation learning, even under IID conditions.
16

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
0,85
RILOOD
RILOOD
0.775
Best Baseline
Best Baseline
0.80
0.750
W
RMSE
RMSE
0.725
0.75
0.700
0.70
0.675
10
10-5
10-4
10 10
10
10-6
10-5 (b)
10-4
10
(4)
Figure 8. Sensitivity analysis of the hyperparameter (a) a and (b) 8 on CompSolv datasets. The solid line shows the average RMSE in the testing stage and the light blue area represents standard deviations. The dashed line represents the average RMSE of the best-performed haseline.
RILOOD
0.S2/CMRL)
0.54
RILDOD
Best Baseline
Best Basaline
CUSEECMBL)
0.52
0.52
RMSE
RMSE
0.50
0.50
0.48
0.48
10
10-3
10-4 (8)
10
10
10
10-6
10-3 (b)
10-4
10
Figure 9. Sensitivity analysis of the hyperparameter (a) and (b) 8 on MNSolv datasets. The solid line shows the average RMSE in the testing stage and the light blue area represents standard deviations. The dashed line represents the average RMSE of the best-performed baseline.
C.4. Hyperparameter Sensitivity Analysis
We perform sensitivity analysis the hyperparameters and B, which govern the trade-off among the loss components in Eq. 13. In practical settings, accurately approximating the true posterior is non-trivial, often resulting in reconstruction loss that dominates the supervised loss by several orders of magnitude. To mitigate this imbalance and stabilize optimization, we systematically vary within {10-7, 10-6, 10-5, 10-4, 10-3) and 3 within (10-5, 10-7, 10-6, 10-5, 10-4). Experiments are conducted on the MNSolv and CompSolv datasets, chosen for their diversity and representativeness.
The results, shown in Fig.8 and Fig.9, demonstrate that the choice of and critically affects the balance between environmental modeling and invariant representation learning. Notably, higher values of a tend yield significant performance improvements. These findings underscore the importance of tuning these hyperparameters for optimal performance. Following established practice, report the best-performing configuration along with its standard deviation.
D. Limitations and Future Directions
While we introduce a molecular pair invariant graph learning approach based on auxiliary information for environment partitioning, achieving strong OOD generalization, certain limitations remain. Since solvent polarity labels are difficult to obtain, auxiliary information can serve proxy label for environment partitioning.
This study examines spurious correlations by considering single spurious factor; however, in reality, multiple spurious associations may exist, complicating the distinction between false and invariant patterns. The influence of different spurious factors varies, impacting model performance in distinct ways.
Additionally, in the era of large-scale models, increasing data volume improves OOD generalization. However, in domain- specific applications, prioritizing broad generalization may reduce accuracy in specialized fields. This trade-off arises
17

---

Relational Invariant Learning for Robust Solvation Free Energy Prediction
because narrow distributions focus on high-probability regions, often overlooking tail samples, leading to poor generalization in rare cases. Bridging this gap between broad generalization and domain-specifie precision presents important research challenge.
18


Reviews:
[{"review_id": "dAbHCymmsp", "reviewer": "Reviewer_Jwej", "summary": "This paper presents a novel out-of-distribution learning method for addressing the challenge of predicting solvation free energy in molecular interactions. The key innovation lies in the authors' approach to modeling the distribution of molecular interactions. They validated the effectiveness of their proposed model by testing it under out-of-distribution conditions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer Jwej:\n\nWe sincerely appreciate the reviewer's thoughtful suggestions and questions. We have provided point-by-point answers to each weakness and question.\n\n## Weakness\n\n**W1. Details about datasets:** \n\nWe provided splitting details in appendix. The method of scaffold split is available in MoleculeNet, and the solvent split is the same as that of scaffold split. As shown in definition 1, we divide the dataset according to different environments, where scaffold and solvent are based on different environments.\n\n**W2. Different from CMRL:** \n\nCMRL acts as an invariant through the substructure of the interaction between molecular pairs, while at the same time weakens the influence of solvent molecules through backdoor adjustment. However, this invariance does not apply to solvent environment cases, because the interaction between solvents and solutes is multi-level, different levels of molecular interaction affect different properties. In contrast, our method recognizes that invariance is not consistent across different solvent environments due to the complex, multihierarchical nature of solute-solvent interactions, where varying interaction scales impact molecular properties differently.\n\n**W3. The choice of hyperparameter alpha:** \n\nFor alpha, we chose 0.5 in order to get a uniform distribution.\n\n**W4. Minor issues:** Thanks for pointing out our omission. Figure 3 shows the results of the COMPSOLV dataset under 4 OOD settings. We've updated the graph annotations for fig.3 (c).\n\n## Other Comments Or Suggestions\n\n**C1. Fig. 1 uses some terms, such as polarity, dielectric constants, but they are not explained in the main text, what does it have to do with the invariance in this paper?** \n\nSome of the proper nouns on the Fig. 1 is solvent related properties such as dielectric constant, polarity of solvent, which are included in solvent. By constructing multi-level interactions, the effects of these properties on solute molecules can be learned, and invariance can be further learned.\n\n**Differences from MMGNN:**  \n\nMMGNN is a model for modeling the relationship between solvents and solutes, and our method focuses on the generalization of molecules in different solvents, which are fundamentally different. In contrast, our approach recognizes that invariance is not uniform in solvent environments due to the multihierarchical nature of solute-solvent interactions, where different interaction scales affect various molecular properties. To solve this problem, we introduce a multi-granularity context-aware refinement mechanism that captures local and global interactions to ensure robust performance under different solvent conditions.\n\n## Claims And Evidence:\n\nThe statement “the decisive interaction patterns of solute vary, leading to distinct solute properties” can be supported by references to PAR[1] and Unimatch[2]. We have included these references for further clarification.\n\n## Questions\n\n**Q1: “...varying degrees between C and the label Y” in line 410?** \n\n The spurious relationships of varying degrees between C and Y means that spurious correlations are introduced by controlling the distribution of variant subgraphs.\n\nWe hope that our response will address your concerns and better clarify the contribution and value of our work.\n\n## Reference\n\n[1] Property-Aware Relation Networks for Few-Shot Molecular Property Prediction\n\n[2] UniMatch: Universal Matching from Atom to Task for Few-Shot Drug Discovery", "claims_and_evidence": "A few statements lack adequate support, such as the one in line 277: “...the decisive interaction patterns of solute vary, leading to distinct solute properties.”", "methods_and_evaluation_criteria": "The methodology and evaluation are reasonable.", "theoretical_claims": "Yes, I checked the theoretical claims.", "experimental_designs_or_analyses": "The authors do not provide a specific description of the division of the dataset, such as the number of environments in the test set and the training set, which is not discussed in the experiment details and may make the comparisons unfair.", "relation_to_broader_scientific_literature": "This paper expands the scenarios of molecular applications and helps to expand to real-world scenarios.", "essential_references_not_discussed": "I note that MMGNN is similar to this work, but the authors does not make a comparison.\n\n[1] MMGNN: A Molecular Merged Graph Neural Network for Explainable Solvation Free Energy Prediction", "other_strengths_and_weaknesses": "**Strengths**\n- This paper addresses the OOD problem of molecular solvation free property prediction and is a pioneer work in this area.\n- The proposed method relies on capturing the feature of molecular interactions that are invariant across environments, which is novelty.\n\n**Weaknesses**\n- The authors do not provide a specific description of the dataset partition settings, such as the number of environments in the test set and the training set, which are not discussed in the experimental setup and may make the comparisons unfair.\n- This paper is similar to CMRL in that it also solves OOD problems through causality, please explain how it differs from CMRL?\n- The hyperparameter alpha appears to be redundantly defined in both line 248 and Eq. 14. Additionally, the authors do not clarify how to choose alpha in line 248. \n- Fig. 3 (c) is not detailed enough to provide a description of the specific tasks, please further elaborate the description.", "other_comments_or_suggestions": "- Fig. 1 uses some terms, such as polarity, dielectric constants, but they are not explained in the main text, what does it have to do with the invariance in this paper?\n- $\\mathbb{R} ^{N_x}$ in proposition 4.2, here, $N_x$ is missing definition.", "questions_for_authors": "- Please explain “varying degrees between C and the label Y” in line 410?\n- For other questions, please refer to the weakness.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SC0ed3c4V8", "reviewer": "Reviewer_zysq", "summary": "In this paper, the authors presents the Relational Invariant Learning framework (RILOOD) to improve OOD generalization in solvation free energy prediction. RILOOD learns invariant molecular representations in varied environments and applies mixed-enhanced molecular features for modeling environmental diversity. Extensive experiments demonstrate its superiority over existing methods under various distribution shifts.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer zysq:\n\nThank you for your thoughtful suggestions and questions. We have provided point-by-point answers to each weakness and question.\n\n## Weakness\n**W1. About theoretical innovation:** \n\nWe acknowledge that our study may not introduce entirely new theoretical concepts, we believe that the theoretical support we provided is crucial for grounding our research in established frameworks, particularly in the context of invariant learning for predicting solvation properties. \n\nOur research contributes significantly to the field by demonstrating how invariant learning can effectively capture the complex relationships between molecular structures and their solvation properties. By integrating this methodology, we not only improve prediction accuracy but also provide a more interpretable framework for understanding the underlying molecular interactions.\n\n**W2. Figure 1 is intended to illustrate the invariance of interactions, but (a) the illustration is not easy to understand.** \n\nAs shown in the figure, the horizontal axis represents solvent polarity, while the dotted box illustrates the distribution of different solute properties in the two solvents. Traditional methods predict solute properties by identifying core substructures; however, these core substructures may vary across different solvents. The key to our approach is to establish invariance by modeling interactions, which allows for a deeper understanding of the diverse characteristics of solute molecules in various solvents. The green arrows in the lower part of the figure indicate the interactions between molecules.\n\n**W3. There is no clear definition of what invariant interaction patterns are.** \n\nInvariant interaction patterns are features that remain consistent across various environments, or conditions. This involves identifying patterns that reflect interactions in diverse molecular contexts. Through invariant learning, we can extract key features from complex molecular data that capture fundamental interactions between molecules, independent of specific conditions. The model can still make accurate predictions even with unseen data or conditions. This capability is crucial for studying molecular interactions, as molecules may behave differently in various biological environments or experimental settings.\n\n**W4. Since obtaining explicit environment labels for solute-solvent pairs is often infeasible. So how is this method obtained?** \n\nWe emphasize that environmental labels are not readily available, so we use auxiliary information to obtain environmental labels, i.e., other information that is easier to obtain. For example, the type of solvent, the scaffold of the molecule, etc., rather than the molecular interaction information obtained through complex modeling calculations as environmental labels.\n\n**Minor issues:** We’ve corrected errors and updated the description in the paper.\n\n## Questions\n**Q1. Could you provide real-world examples to further illustrate the OOD generalization of RILOOD?**  \n\nWe provide solvent holdout case studies for your reference. Due to time constraints, we only selected a few cases. As can be seen from the case studies, the prediction error is still low for different types of solvents, including alkanes, alcohols, and olefins.\n\nTable. The results of RILOOD on solvents holdout test. \n|  Solvent  | RMSE | \n|-------|-------|\n| Tetraglyme  |  0.1004  |\n| Pentadecane |  0.1100  |\n| Tetradecane |  0.2462  |\n| 1-Dodecanol |  0.2315  |\n| 1-bromonaphthalene | 0.7796 |\n| Tetraglycol |  0.2015  |\n\n\n**Q2. About tsne, MCAR how to enhance feature diversity?** \n\nIn fig.4(b), the feature distribution of MCAR is more dispersed than the original distribution, indicating that it contains more information and can better reflect the differences between samples after multi-interaction.\n\nWe hope that our response will address your concerns and better clarify the contribution and value of our work.", "claims_and_evidence": "Yes. Claims made in the submission supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The authors provide a wealth of benchmarks to prove the superiority of the proposed method, and the evaluation method is effective.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "The authors provide a large number of experiments to illustrate the advanced nature of the proposed method, and the experimental analysis is reasonable.", "relation_to_broader_scientific_literature": "No.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strengths**\n1. The author offers a new invariant learning method, which, as far as I know, is novel\n2. Compared with the SOTA method, the experimental results show that the proposed method performs well.\n3. The author provides a theoretical proof.\n\n**Weaknesses**\n1. Although theoretical support is provided, there is no innovation in theory\n2. Figure 1 is intended to illustrate the invariance of interactions, but (a) the illustration is not easy to understand.\n3. There is no clear definition of what invariant interaction patterns are.\n4. Since obtaining explicit environment labels for solute-solvent pairs is often infeasible. So how is this method obtained?", "other_comments_or_suggestions": "The writing is not rigorous, and it is repeated in many places。\n- What does multiple environments in line 163 $\\varepsilon$ mean? How it is and the 115 line of $\\varepsilon$ is clearly not a representation.\n- Is the $l$ in \"drawing samples $z^{(l)} (l = 1, 2, ..., l)$\" the same as the $l$ in \"$h_{l+1} = PReLU(W_lh_l + b_l)$\" in Equtation 10? $l$ is the same as $l$ in loss, it is recommended to change the symbol.", "questions_for_authors": "1. Could you provide real-world examples to further illustrate the OOD generalization of RILOOD?\n2. About tsne, MCAR how to enhance feature diversity?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K6rnP6spDS", "reviewer": "Reviewer_9m7P", "summary": "This paper investigates the challenge of out-of-distribution generalization across different environments in molecular solvation free energy prediction and introduces the RILOOD framework. RILOOD integrates mixup-based conditional modeling, a multi-granularity refinement strategy, and an invariant relational learning module to alleviate the limitations of traditional methods that overly rely on core substructures. Experimental results demonstrate that RILOOD significantly outperforms existing methods on multiple benchmark datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer 9m7P:\n\nThank you for your thoughtful suggestions and questions.We have provided point-by-point answers to each weakness and question.\n\n**W1. RILOOD integrates several modules to capture complex molecular interactions, but this results in a relatively complex model structure:** \n\nWe understand your concerns regarding the complexity of the model. We chose this complex model structure because it allows us to better capture the intricate interactions between molecules, which is crucial for improving the accuracy and reliability of our predictions.\n\nAlthough the model structure is relatively complex, we believe that this complexity enables the model to gain a more comprehensive understanding of molecular characteristics, thereby providing better performance in practical applications. Additionally, we have also considered interpretability in our model design and have made efforts to ensure that each module of the model has a clear function. We are also considering possible simplifications in future work to reduce complexity while maintaining model performance. We would greatly appreciate any further suggestions you may have regarding this.\n\n**W2. The distinction between \"surrounding environments\" and \"support environments\" is unclear:** \n\n\"Surrounding environments\" and \"support environments\" mean the same thing. We'll change it to a unified expression.\n\n**W3. Symbols are defined repeatedly:** Thank you for your corrections, and we will double-check for further proofreading.\n\n**W4. Both Sections 3.1 and 3.2 include a problem formulation, which is confusing:** \n\nThe problem formulation in 3.1 is the molecular interaction prediction task in general, mainly under the IID setting; The problem formulation in 3.2, the purpose of OOD setting is to emphasize the difference between the distribution of the training set and the test set. We'll update the language to make it easier to understand.\n\n**Minor issues:** We’ve corrected errors and updated the description in the paper.\n\nThank you for your valuable comments, and we hope that our response will address your concerns and better clarify the contribution and value of our work.", "claims_and_evidence": "Some of the paper's claims have minor issues. The difference among environments, environmental conditions, and solvent categories is not clearly expressed.", "methods_and_evaluation_criteria": "The paper evaluates the effectiveness of the model and improving generalization in OOD scenarios. Extensive experiments and theoretical analysis prove the superiority of the method.", "theoretical_claims": "The paper appears to provide sufficient theoretical analysis.", "experimental_designs_or_analyses": "1. The proposed method is evaluated using RMSE, MAE, AUROC, and Accuracy, with experiments conducted on both real-world and synthetic datasets.\n\n2. t-SNE is used to visualize molecular interactions in the best performing model.", "relation_to_broader_scientific_literature": "This work is important. Predicting the solvation free energy of molecules in out-of-distribution (OOD) scenarios has greater practical value compared to identically distributed scenarios.", "essential_references_not_discussed": "I believe that essential references have been sufficiently discussed.", "other_strengths_and_weaknesses": "Strengths \n1. Figure 1 clearly presents the motivation of the work.\n\n2. The experiments include numerous visualizations that further validate the paper’s claims.\n\n3. The performance improvements are significant.\n\nWeaknesses\n1. RILOOD integrates several modules to capture complex molecular interactions, but this results in a relatively complex model structure.\n\n2. The distinction between \"surrounding environments\" and \"support environments\" is unclear.\n\n3. Some variables are repeatedly defined (e.g., \"K\" is used for both surrounding environments and attention keys, \"c\" represents both coarse-grained conditions and environmental conditions); please check and reduce this confusion.\n\n4. Both Sections 3.1 and 3.2 include a problem formulation, which is confusing.", "other_comments_or_suggestions": "1. Lines 161-164 and 263-266 are blank.\n\n2. There are too many variables, making the paper somewhat difficult to read.\n\n3. In Eq. (4), $G$ should be $\\mathcal{G}$.\n\n4. There is a missing cite at line 332 (\"Appendix ??\").", "questions_for_authors": "Please refer to the above comments.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "None"}}, {"review_id": "BlspiZ7Zh0", "reviewer": "Reviewer_b6ge", "summary": "The paper proposes Relational Invariant Learning framework for solvation free energy prediction. RILOOD consists of three key components: a mixed conditional modeling module to integrate data from different environments, a multi-granularity refinement strategy for context-aware representation learning, and an invariant learning module to identify robust patterns that generalize to unseen environments. Experimental results demonstrate that RILOOD significantly outperforms existing SOTA methods across various datasets, showcasing its effectiveness in capturing complex molecular interactions. The paper highlights the challenges of OOD scenarios and suggests future research directions to balance broad generalization with domain-specific precision.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer b6ge:\n\nThank you for your thoughtful suggestions and questions. We have provided point-by-point answers to each weakness and question.\n\n## Weakness\n**W1. Personally, I do not find CVAE to be particularly innovative.** \n\nWhile we understand your perspective on the novelty of the CVAE framework, our focus in this paper is not solely on theoretical innovation or technological advancement. Instead, we aim to address a valuable domain-specific problem: solvent generalization in practical applications. We believe that our work is pioneering in its application of CVAE to the context of solvent generalization, which has significant implications for real-world scenarios. By concentrating on this specific challenge, we provide insights and solutions that can enhance the predictive capabilities of models in the field of solvation properties.\n\n**W2. This paper is only focused on the task of solvation free energy prediction, and the method is difficult to extend to general molecular property prediction or DDI tasks.**\n\nAs we mentioned in the limitation, the types of environments don't increase indefinitely. Considering the scalability and computational cost of the model, when there are too many types of solvents, clustering can be performed first, and then the pseudo-labels of the clustering can be conditionally modeled, so we think that our approach can be extended to more task types in the future.\n\n**W3. More case study:**\nWe've added some solvent-holdout test cases to verify real-world cases. Due to time constraints, we only selected a few cases. As can be seen from the case studies, the prediction error is still low for different types of solvents, including alkanes, alcohols, and olefins.\n\nTable. The results of RILOOD on solvents holdout test. \n|  Solvent  | RMSE | \n|-------|-------|\n| Tetraglyme  |  0.1004  |\n| Pentadecane |  0.1100  |\n| Tetradecane |  0.2462  |\n| 1-Dodecanol |  0.2315  |\n| 1-bromonaphthalene | 0.7796 |\n| Tetraglycol |  0.2015  |\n\n**Minor issue.** We’ve corrected errors and updated the manuscript.\n\n## Question:\n\n**Q2. Why is w/o MCAR better than B+MCAR[M] in Table 2?** \n\nw/o MCAR is better than B+MCAR[M] and only performs on one dataset, which may be due to the fact that w/o MCAR only removes the interactive module and the objective function does not change, which leads to the model being reverse-optimized.\n\n**Q3. What is the difference between $c^r$, $C$, $c_1$, $c_2$?**\n\n$C$ is the collection of all solvents, $c^r$ refers to a specific class of solvents, $c_1$ refers to the type of solute, and $c_2$ refers to the type of solvent\n\nWe hope that our response will address your concerns and better clarify the contribution and value of our work.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "I don’t think the baseline comparison is strong enough due to the insufficiency of OOD datasets. Compared to Table 4, which includes a broader range of datasets, Table 1 only compares three datasets. The authors should consider including more datasets to better illustrate the effectiveness of the method.", "theoretical_claims": "I carefully reviewed the proof in the appendix, and the proof of theoretical claims is generally well structured, but there are some points that require further clarification and rigor, and addressing these issues will increase the credibility of the theoretical framework.\n\n- The theoretical rationale for the uncertainty constraint of RMSE in A.1 is interesting, but the proof provided lacks sufficient detail. It would be helpful to have a more comprehensive analysis of how this constraint affects the representation of learning, and to provide formal evidence of its effectiveness in improving the robustness of the model.\n\n- For proof A.2, acknowledging any potential limitations or assumptions that may not hold in certain scenarios would provide a more balanced view.", "experimental_designs_or_analyses": "This paper conducted a variety of experiments to illustrate the effectiveness of the method, including comparative experiments, ablation experiments, and super-parameter experiments, and gave detailed analyses.", "relation_to_broader_scientific_literature": "In this paper, the authors extend the traditional molecular property prediction to solvent environment, and consider the role of solvents, which is helpful to further expand the practical application of existing AI for science.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**\n- The motivation of the paper is clearly defined.\n- This paper is well-written and effectively organized.\n- The topics discussed are very important in practical application.\n- The experiments show that RILOOD achieves robust solvation free energy prediction and outperforms baseline models under distribution shifts.\n\n**Weaknesses**\n- Personally, I do not find CVAE to be particularly innovative, as it is often used in molecular generation[1].\n- This paper is only focused on the task of solvation free energy prediction, and the method is difficult to extend to general molecular property prediction or DDI tasks.\n- Authors should add case studies illustrating generalization performance in practical applications.\n\n[1] Lim, Jaechang, et al. \"Molecular generative model based on conditional variational autoencoder for de novo molecular design.\" Journal of cheminformatics 10 (2018): 1-9.", "other_comments_or_suggestions": "- The appendix in the capital of Table 1 is incorrect.\n- Fig.1(b), the abscissa SMILES should be the corresponding name for ease of understanding.\n- There is an overlap in the data on the Fig.3 (b).", "questions_for_authors": "- What does the molecular interaction in Fig.1(a) mean by arrows of the same color? How to account for interaction invariance?\n- Why is w/o *MCAR* better than *B+MCAR[M]* in Table 2?\n- What is the difference between c_r, C, c_1, c_2, C?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposed an approach for learning solvation free energy prediction of molecules, with a focus on OOD settings. The paper is well written and the empirical evaluation supports the claims in the paper. Reviewers have particular highlighted the OOD performance to be of pioneering value with a strong performance boost over existing work. The AC agrees that this paper is of value to the ICML community.
Decision: Accept (spotlight poster)