Soft causal learning for generalized molecule property prediction: An environment perspective

TL;DR

The paper addresses the challenge of out-of-distribution generalization in molecular graph learning by shifting focus from invariant subgraphs to expanding molecular environments. It introduces CauEMO, a triad framework consisting of a knowledge-guided environment growth generator, an Environment-Graph Information Bottleneck for distilling environment representations, and an environment-invariance soft causal interaction module that enables dynamic coupling between environment signals and invariant substructures. Empirical results across seven datasets—including real-world drugs and synthetic motifs—demonstrate improved OOD generalization, with ablations confirming the contributions of environment modeling, disentanglement, and soft causal fusion. This work potentially enhances the reliability and interpretability of AI-driven molecular property prediction in settings with evolving chemical spaces and distribution shifts.

Abstract

Learning on molecule graphs has become an increasingly important topic in AI for science, which takes full advantage of AI to facilitate scientific discovery. Existing solutions on modeling molecules utilize Graph Neural Networks (GNNs) to achieve representations but they mostly fail to adapt models to out-of-distribution (OOD) samples. Although recent advances on OOD-oriented graph learning have discovered the invariant rationale on graphs, they still ignore three important issues, i.e., 1) the expanding atom patterns regarding environments on graphs lead to failures of invariant rationale based models, 2) the associations between discovered molecular subgraphs and corresponding properties are complex where causal substructures cannot fully interpret the labels. 3) the interactions between environments and invariances can influence with each other thus are challenging to be modeled. To this end, we propose a soft causal learning framework, to tackle the unresolved OOD challenge in molecular science, from the perspective of fully modeling the molecule environments and bypassing the invariant subgraphs. Specifically, we first incorporate chemistry theories into our graph growth generator to imitate expaned environments, and then devise an GIB-based objective to disentangle environment from whole graphs and finally introduce a cross-attention based soft causal interaction, which allows dynamic interactions between environments and invariances. We perform experiments on seven datasets by imitating different kinds of OOD generalization scenarios. Extensive comparison, ablation experiments as well as visualized case studies demonstrate well generalization ability of our proposal.

Figures (7)

• Figure 1: Motivation of CauEMO. (a) Increasingly growing environments can finally dominate the property of whole graph. (b) Two substances, ethyl alcohol and phenol are with the same functional group of hydroxyl, but are with different environmental substructures connection, resulting in different solubility properties.
• Figure 2: Framework overview of CauEMO
• Figure 3: Environment-invariance soft causal interaction
• Figure 4: The ability of CauEMO to identify 'house' in Spurious-Motif dataset.
• Figure 5: Ablation studies on CauEMO-Random and CauEMO-Subgraph