Subgraph Information Bottleneck with Causal Dependency for Stable Molecular Relational Learning

1 Introduction

Molecular Relational Learning (MRL) predicts the interaction between molecular pairs by mining structures and properties [Fang et al., 2024; Du et al., 2024]. MRL has garnered significant attention in natural science research due to its widespread applications in new material design and drug discovery [Pei et al., 2024; Zhang et al., 2024a]. The molecular structural information mining-based methods have advantages in MRL [Wang et al., 2024; Job et al., 2023; Zhang et al., 2025b]. However, these methods exhibit varying levels of performance fluctuations in real-world scenarios. Consequently, enhancing model stability is essential for the effective implementation of MRL.

∗ Corresponding author.

The Graph Information Bottleneck (GIB) [Wu et al., 2020] has recently been applied to extracting and identifying core subgraphs. GIB generates representation that represent the minimal sufficient information of the input graph by information compression [Yuan et al., 2024; Fan et al., 2022; Seo et al., 2024]. The theory provides a novel solution for MRL, with the key problem being to efficiently capture the molecular representation that are most beneficial for MRL [Li et al., 2024; Seo et al., 2023]. Related research addresses the above problem by maximizing the mutual information between molecules and prediction targets [Yang et al., 2023; Li et al., 2024]. Specifically, these methods control the information flow between molecules and the targets to generate representation that contain minimal sufficient information about molecules [Lee et al., 2023a]. The generalizability of GIB-based MRL are improved by integrating the interaction information into molecular representations.

The core challenge of MRL is the instability of model predictions caused by distribution differences between training data and real-world data [Lee et al., 2023b]. An effective way to deal with this challenge is to capture the core substructures that determine molecular reactions [Fang et al., 2024]. Nevertheless, GIB-based MRL have the following limitations: (1) Missing Substructure Causal Information in Chemical Reactions: The nature of chemical reactions is the dynamic rearrangement between substructures [Hu et al., 2017]. The substructures of the same paired molecule causally related to prediction targets in different chemical reactions may be different [Iwasaki and Nozaki, 2024]. As shown in Fig. 1(a), when molecule A reacts with B, the reactive substructure in A is -COOH, while with C, it is -NH2. GIB typically focuses on information compression and fails to model causal interactions between substructures, making it challenging to identify the substructures that truly determine chemical reactions. (2) Challenges in Capturing Nuanced Differences Between Substructures: Substructures are susceptibly impacted by active atoms within the molecule, altering their properties in chemical reactions [Zhang et al., 2025a]. As shown in Fig. 1(b), molecules D and E contain a -CO, but the substructures involved in reactions with the same molecule differ. This is because the -CI in E has the stronger electron adsorption, enhancing the reactivity of -CO, which leads to the -COCI substructure involved in chemical reaction. Considering MRL solely at the molecular level makes capturing these nuanced differences between substructures difficult. These limitations restrict the ability of GIB-based relevant methods to maintain stable relational learning.

In this paper, we first establish the causal associations between substructures and molecular relationships in MRL and analyze the key factors influencing the backdoor paths. Subsequently, we propose the Causal Subgraph Information Bottleneck (CausalGIB), which incorporates the Causallydriven Substructure Representation Learning (CauSRL) and the Synergistic Optimization of Subgraph Information Bottleneck (SOSIB). CauSRL integrates causal dependencies into the optimization of substructure representation through noise injection and interactions, effectively distinguishing causal and confounding substructures. SOSIB encourages the model to learn the maximum causal mutual information between causal substructures and prediction targets while minimizing the information among confounding substructures, thereby generating stable molecular representations. The main contributions are as follows.

• We propose CausalGIB that introduces subgraph causal information into GIB to improve the stability of MRL. To the best of our knowledge, it is the first work to address the problem of causal substructure selection in different chemical reactions.
• We provide a formal validation for the design of the loss function, grounded in theoretical reasoning, to effectively capture subtle distinctions between substruc-
• tures by minimizing the disparity between confounding and causal information.
• Experiments on two tasks in nine real-world datasets and three different data distributions demonstrate that the predictive performance and stability of CausalGIB outperforms seven state-of-the-art models.

2 RelatedWork

Molecular Relational Learning. MRL is an essential task in molecular representation, encompassing various applications [Pei et al., 2024; Zhang et al., 2024b; Job et al., 2024]. We focus on Molecular Interaction (MI) and Drug-Drug Interaction (DDI) predictions. InMI prediction, CIGIN [Pathak et al., 2020] employed message-passing network and collaborative attention to encode molecular atoms and predict solvation energy. CMRL [Lee et al., 2023b] combined molecular representation with causal relationships, identifying substructures related to chemical reactions. In DDI prediction, SA-DDI [Yang et al., 2022], SSI-DDI [Nyamabo et al., 2021], and DSN-DDI [Li et al., 2023] extracted substructure information from drugs in various ways to represent drug molecules.

Graph Information Bottleneck. GIB is an extension and development of the Information Bottleneck theory within graph theory, which has been widely applied in molecular relational learning [Sun et al., 2022], multi-agent communication [Ding et al., 2024b; Ding et al., 2024a], and graph learning interpretability [Yuan et al., 2024]. Specifically, PGIB [Seo et al., 2023] introduced prototype information into the GIB framework and used it as a compression bottleneck to improve the interpretability of graph representation learning. [Ding et al., 2024b; Ding et al., 2024a] optimized robust communication in multi-agent systems by setting mutual information between message and action choices as the GIB optimization objectives. CGIB [Lee et al., 2023a] combined molecular relational learning with the GIB framework to model the nature of chemical reactions.

Despite the advantages of the related methods in MRL, they fail to consider the causal dependencies between substructure pairs. Therefore, we aim to improve the model's stability in distribution-shifted data by introducing causal information to identify core substructures of molecules.

3 Preliminaries

In this section, we define stable MRL and construct the Causal Substructure Model (CSM). Based on the analysis results of CSM, we identify the key factors influencing the backdoor paths in MRL.

3.1 Stable MRL

For any molecule G, it can be represented as G = (V,E,X,A). Here, V = {v1, v2, · · · , vN} denotes the set of nodes. E ∈ N×N represents the connections between atoms within the molecule, which is closely related to the adjacency matrix A. If (vi, vj) ∈ E, then Aij = 1; otherwise, Aij = 0. X is the feature matrix, consisting of the atom feature representations. For the given molecular pair (Gx, Yxy,Gy), the objective of Stable MRL is to generate a stable molecular embedding representation Z(Gst ) and to use it for molecule pair interaction prediction in different data distributions. Formally, Yxy = Fs(Z(Gst x ), Z(Gst y )). Therefore, CausalGIB aims to learn stable representations of molecules and improve the model's predictive performance.

3.2 CSM for Stable MRL

We define and construct the CSM for stable MRL, as shown in Fig. 2. The causal relationships illustrated in Fig. 2 are described as follows:

• Gx → Gcx ← Gy: Gcx is the causal core substructure of molecule Gx, jointly determined by Gx and Gy, and varies with changes in Gy.
• Gy → Gcy ← Gx: Gcy is the causal core substructure of molecule Gy, jointly determined by Gy and Gx, and varies with changes in Gx.
• Gx → Gnx ← Gy: Gnx is the confounding substructure of molecule Gx, representing the substructure involved of Gx other than Gcx .
• Gy → Gny ← Gx: Gny is the confounding substructure of molecule Gy, representing the substructure involved of Gx other than Gcy .
• Gcx → Zx: Zx is the representation of the molecule Gx, obtained by encoding the causal core substructure Gcx .
• Gcy → Zy: Zy is the representation of the molecule Gy, obtained by encoding the causal core substructure Gcy .
• Zx → Y ← Zy: Y represents the interaction between molecular pairs, which depends on both Gx and Gy in molecular relationship learning.

Through backdoor path analysis, we identify the core factors influencing molecular relationships Y . The backdoor paths are Gnx ← Gy→Gcy ← Gx→Gcx →Zx→Y , Gy→ Gnx ← Gx→Gcx→Zx→Y , Gx→Gny ← Gy→Gcy→Zy→Y , and Gnx ← Gx →Gcx ← Gy →Gcy →Zy →Y . By analyzing the backdoor paths in the CSM, we observe that Gcx and Gcy serve as the backdoor criteria constraining molecular relationships Y . Considering the specificity of molecular relationship learning, i.e., the substructure is either core substructure Gc or confounding substructure Gn . Therefore, our objective is to establish stable interactions between Gc and Y by eliminating the influence of Gn on Gc .

4 Methodology

Motivated by CSM, we propose the CausalGIB with joint optimization of predictive loss and confusing information for stable MRL, as shown in Fig. 3. CausalGIB models substructure dynamic binding in CSM by noise injection and substructure interactions, utilizing it to optimize substructure embeddings (Sec. 4.1). Furthermore, we propose the optimization objective for Subgraph Information Bottleneck (SIB) to perturb confounding information while preserving the association between causal substructures and the target (Sec. 4.2).

4.1 Causally-Driven Substructure Representation Learning (CauSRL)

CauSRL utilizes the GNN encoder and an adjacency aggregation strategy to generate representations for irregular substructures. Subsequently, we intervene with causal dependencies between substructures via noise injection to identify substructures with strong causal relationships.

Irregular Substructure Representation Learning

We utilize substructure coefficient and neighbor feature weighted aggregation to extract molecular substructures. For the molecule G, the substructure Gs consists of the central atom vc and its K-hopping neighboring nodes vk c . Formally, Z(Gs ) = ∑Kk=1 ∑vs∈vk c C kGNNk (vs). GNN(∗) can be selected from GIN [Xu et al., 2019], MPNN [Gilmer et al., 2017], GAT [Veliˇckovi´c et al., 2018], or GCN [Kipf and Welling, 2017]. The calculation of C k is defined as follows:

Here, σ is the activation function. Z(∗) is the embedding representation of node. We use Z(Gs ) to obtain the molecular representation Z(G) = Z(Gs1 )|| · · · ||Z(GsJ ). J is the number of molecular substructures. Next, we compute the importance Ri of different substructures Gsi within the molecule G, which is used for subsequent noise injection.

Causal Substructure Recognition with Dependency Intervention (CSDI)

To better identify causal relationships between substructures, we introduce causal dependence to model the interaction behaviour. Causal dependence refers to a scenario where a change in one variable, due to an external intervention, directly induces a change in another variable, with this effect being unexplained by other variables. In other words, if a change in the topology of a substructure increases its interaction probability with the paired substructure, the two substructures exhibit a higher causal dependency.

Consequently, we interfere with the topological information of substructures through noise injection and quantify causal dependency by causal contribution scores to identify and select causal substructures. Specifically, given the adjacency matrix A of molecule G, our objective is to inject noise into the adjacent atoms of the substructure Gs and reconstruct the associations between them via the molecule's adjacency matrix A to obtain an intervening topologically informed substructure Gns . Formally, Gns = Gs+ △ vns, where △ vns represents the added atomic nodes. The embedding representation of the substructure after noise injection is defined as:

where N represents the set of neighboring atoms for substructure Gs , and ε is the noise function. µ and σ2 are mean and variance of Z(vns), respectively.

We compute the causal contribution difference γ between substructures before and after noise injection by incorporating stochastic interaction behaviours. γ is represented by the ratio of causal dependencies between substructures.

where Gns x represents the substructure generated from Gsi x after noise injection. Gsj y is the paired substructure of Gsi x , belonging to molecules Gy and Gx, respectively.

The numerator in Eq (4) denotes the average causal effect of the perturbed substructure Gns x on the paired substructure Gsj y , while the denominator represents the causal dependency of Gsi x on Gsj y . The value of γ indicates whether the causal dependency between substructures changes before and after noise injection. Based on γ, we identify the causal substructures Gc between paired molecules. Specifically, if γ > 0, Gns x is classified as causal substructure; otherwise, Gsi x is considered causal substructure. The remaining substructures are designated as confounding substructures Gn .

After CSDI screening, CausalGIB leverages SIB further to optimize the representations of Gc and Gn and guide the model learning. We identify the SIB-optimized causal substructures as core substructures Gco and use them to generate stable molecular representations Z(Gst ).

4.2 Synergistic Optimization of Subgraph Information Bottleneck (SOSIB)

Considering that the nature of MRL is dynamic interaction between substructures [Lee et al., 2023b], we extend GIB to the subgraph level. We derive the optimization objective for the SIB and exhibit its transformational solution process.

Subgraph Information Bottleneck

Based on the analysis of CSM, we need to generate a stable MRL that meets the following conditions: perturbing the confounding information and maintaining the causal substructure information. Therefore, we propose the optimization objective of SIB from the subgraph level.

Definition 1. (SIB) Given a pair of graphs and their interaction (Gx, Y,Gy), (Gcx ,Gcy ) and (Gnx ,Gny ) are the causal and confounding subgraph pairs of (Gx, Gy) and Y , respectively. Under the SIB principle, the optimal causal subgraphs and the optimization objective are defined as:

Optimizing Eq (5) maximizes the preservation of causal information while minimizing the perturbation of confounding information on MRL. To simplify the computation of mutual information, we reformulate and solve the upper bounds for -I(Y ; Gcx , Gcy ) and I(Gc , Gn ) to optimize Eq (5).

Minimizing I(Gc , Gn ) in Eq (5)

The goal of minimizing I(Gc ,Gn ) is to extract subgraphs that contain less confounding information. Inspired by [Lee et al., 2023a], we minimize the upper bound of Isi as the conditionally independent distribution between Gc , Gn , and R in Eq (2). Since Gc , Gn are subgraphs of G, according to principles of information theory, we have Isi = I(Gc ,Gn ) ≤ min(I(Gc ,G), I(Gn ,G)). The minimization of the upper bound of Isi can be expressed as:

where p is a posterior distribution function.

Eq (6) indicates that minimizing (LcKL,LnKL) provides an upper bound for I(Gc ,Gn ). In the context of MRL, for a paired molecule (Gx,Gy), the I(Gc ,Gn ) is refined as I (Gcx , Gnx ) + I (Gcy , Gny ).

The goal of minimizing -I(Y ;Gcx ,Gcy )) is to capture the stable representation of causal substructures that determine molecular relationship learning. We maximize Ica = I(Y ;Gcx ,Gcy ) through γ in Eq (4), the causal substructures Gcx , Gcy , and the prediction target Y :

where q(Gcx ,Gcy |Y ) is the variational approximation distribution used to approximate the posterior distribution p(Gcx , Gcy |Y ).

The Eq (7) indicates that minimizing the prediction loss Lpred achieves the minimization of -I(Y ;Gcx ,Gcy ). Maximizing the mutual information between γ and Y helps the model identify the causal relationship between Gc and Y in the CSM backdoor paths.

4.3 Model Loss and Analysis

We define the model training objective as the sum of Lpred and LKL as follows:

where α, β are hyperparameters. LKL = min(LcKL,LnKL) controls the impact of confounding substructures on the causal substructure representation. Lpred is computed by Mean Absolute Error Loss for regression task and CrossEntropy Loss for classification task.

5 Experiments

In this section, we analyze the experiments to answer the following questions: RQ1: Can CausalGIB improve the stability of MRL in distribution-shifted data? RQ2: How does the performance of CausalGIB in MRL, and whether CausalGIB is susceptible to backbone? RQ3: Can the results of ReAlignFit be visually supported? RQ4: What is the key to CausalGIB's performance improvement?

Following the related research [Lee et al., 2023a; Lee et al., 2023b; Boulougouri et al., 2024], we conduct extensive MI prediction and DDI prediction experiments on nine datasets. To better simulate the real world data distribution, we set up three different data distributions (Original Partitioning (Original), Rule-based Partitioning (P1) and Graph-based Partitioning (P2)) in DDI prediction and use it for CausalGIB learning.

5.1 Stability Analysis of CausalGIB (RQ1)

We compared the stability of CausalGIB in distributionshifted data with four backbone models (GCN [Kipf and Welling, 2017], GAT [Veliˇckovi´c et al., 2018], MPNN [Gilmer et al., 2017], and GIN [Xu et al., 2019]) and seven domain models (MIRACLE [Wang et al., 2021], CIGIN [Pathak et al., 2020], SSI-DDI [Nyamabo et al., 2021], DSNDDI [Li et al., 2023], CMRL [Lee et al., 2023b], CGIB [Lee et al., 2023a] and MMGNN [Du et al., 2024]). We quantify model stability by calculating the performance retention (PR) in distribution shifts, inter-distribution performance variation (PV) for AUROC, ACC, F1, Pre, and AUPR.

where Eva represents the evaluation metrics. EvaPM and EvaOri M denote the prediction performance in the original and drift distributions, respectively. M ∈ {AUROC,ACC, F1, Pre,AUPR} and Pi ∈ {P1, P2}. It is important to clarify that we conduct model training and testing separately in three distinct data distribution scenarios (Original, P1, and P2), rather than directly applying the model trained on distribution Original for inference in distributions P1 and P2.

The experimental results are shown in Table 1 and Table 2. From Tables 1 and 2, the stability of the MRL models is influenced by data distribution.

Overall Analysis of Model Stability: As data distribution differences increase, MRL models exhibit varying degrees of performance degradation. However, it is noteworthy that CausalGIB consistently maintains the highest in different distributions. Compared to the second-best model, CausalGIB achieves approximately 3% to 6% improvements in RPD, PR, and PV. These experimental results indicate that CausalGIB learns more stable molecular representations under distribution shifts, which we consider a key factor to its success.

Analysis of the Importance of Causal Dependency: As shown in Table 2, methods that do not adequately account for causal relationships (e.g., CIGIN, CGIB) exhibit at least a 1.5% lower PVPre in DDI prediction compared to methods that fully consider causal dependencies (e.g., CausalGIB). Similar trends are observed in the experimental results shown in Table 1. This underscores the critical importance of causal dependencies between prediction targets and substructures for the stability of MRL.

Adaptability to Distribution-Shifted Data: In Table 1, backbone models demonstrate poor adaptability to data shifts, showing significant performance drops and fluctuations. In contrast, CausalGIB achieves the best results in multiple evaluation metrics on different datasets compared to other MRL domain models. In P1 scenario, the performance fluctuations of CausalGIB range from 15% to 20%, which is significantly smaller than those of MMGNN and CMRL. These findings suggest that CausalGIB exhibits acceptable adaptability to data distribution shifts. We attribute this result to two main reasons: the molecular representations generated by CausalGIB contain robust causal information between substructures and prediction targets. On the other hand, CausalGIB mitigates the influence of confounding substructures on molecular representations by weakening the associations between confounding and causal substructures.

5.2 Predictive Performance Analysis (RQ2)

We further compared the predictive performance of CausalGIB and baseline methods on regular data to evaluate whether CausalGIB remains competitive in independent and identically distributed data. The experimental results are reported in Table 3. CausalGIB exhibits pronounced advantages even in independent and identically distributed data.

Comprehensive Performance Analysis: CausalGIB demonstrates the best overall performance in MI and DDI predictions. In MI prediction, CausalGIB improves RMSE by 4% to 6% in multiple datasets. CausalGIB's comprehensive performance outperforms other comparative methods in DDI prediction. This results highlights the importance of mitigating confounding information and strengthening causal substructure-target relationships to enhance predictive performance.

Substructure Importance Analysis: Methods integrating substructure information demonstrate superior predictive performance compared to those that disregard it, highlighting the pivotal role of substructure interactions in MRL. Encouraging models to capture substructure-level features significantly enhances the prediction of molecular interactions, consistent with the principle that substructures are fundamental determinants of chemical reactivity.

5.3 Visualization Analysis (RQ3)

To validate CausalGIB's ability to capture causal substructures, we visualize the node features and the interactions between substructures, as shown in Fig. 4.

CausalGIB effectively captures substructure information in drug molecules, aligning closely with actual molecular substructures and highlighting intra-molecular atomic interactions. As shown in Fig. 4, it demonstrates strong recognition of interactions between functional groups, such as the symmetrical -COOH groups (green and orange areas) in DB00548 and Group1 (benzene ring marked by red box) in DB01117. This aligns with domain knowledge that chemical reactions primarily occur between core substructures. CausalGIB identifies critical substructures and underscores their pivotal role in driving chemical reactions.

5.4 Ablation Experiment (RQ4)

To analyze the impact of confounding information and causal information on the overall performance of the model, we compared the following model variants: CausalGIBnone (no optimizations applied in the model architecture), CausalGIBcon (only minimizing I(Gc ,Gn ) ), and CausalGIBcau (only maximizing I(Y ;Gcx ,Gcy )). The experimental results are shown in Fig. 5. As shown in Fig.

5, minimizing I(Gc ,Gn ) and maximizing I(Y ;Gcx ,Gcy ) improve AUROC by 4%, 8% and 8%, 16%, respectively. Notably, the performance gain from maximizing I(Y ;Gcx ,Gcy ) significantly surpasses that of minimizing I(Gc ,Gn ). While minimizing I(Gc ,Gn ) has limited impact on stability, maximizing I(Y ; Gcx , Gcy ) enhances predictive stability, though its contribution alone remains insufficient. The best performance and stability are achieved when both objectives are optimized concurrently, highlighting SIB's role in refining substructure differentiation and enhancing predictive stability.

6 Conclusion

In this paper, we propose the Causal Subgraph Information Bottleneck to enhance the stability of MRL. CausalGIB leverages causal dependencies to guide substructure representation and integrates SIB to optimize the core substructure representation, generating stable representations. Experimental results confirm that CausalGIB significantly enhances stability in distribution shifts. Future work will explore its validation in diverse datasets and tasks to establish broader applicability.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (No.62472332, No.62276196, No.62076045), the China Scholarship Council, the Open Bidding for Selecting the Best Candidates Project of Wuhan East Lake High-Tech Development Zone (No.2024KJB322), the Hubei Provincial International Science and Technology Cooperation Project (No.2024EHA031), the 111 Project (No.D23006) and the National Foreign Expert Project of China (No.D20240244).