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

arXiv: https://arxiv.org/abs/2505.06283 · HTML: https://mlsanity.com/paper/2505.06283/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 55 references that are in the map.
- First posted on arXiv 7 May 2025.

## Builds on (most cited first, 30 of 31)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [XGNN: Towards Model-Level Explanations of Graph Neural Networks](https://mlsanity.com/paper/2006.02587/) (2020)
- [Discovering Invariant Rationales for Graph Neural Networks](https://mlsanity.com/paper/2201.12872/) (2022)
- [Graph Neural Networks for Graphs with Heterophily: A Survey](https://mlsanity.com/paper/2202.07082/) (2022)
- [Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism](https://mlsanity.com/paper/2201.12987/) (2022)
- [Graph Information Bottleneck](https://mlsanity.com/paper/2010.12811/) (2020)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs](https://mlsanity.com/paper/2202.05441/) (2022)
- [AdaRNN: Adaptive Learning and Forecasting of Time Series](https://mlsanity.com/paper/2108.04443/) (2021)
- [Graph Information Bottleneck for Subgraph Recognition](https://mlsanity.com/paper/2010.05563/) (2020)
- [Understanding and Extending Subgraph GNNs by Rethinking Their Symmetries](https://mlsanity.com/paper/2206.11140/) (2022)
- [Graph Structure Learning with Variational Information Bottleneck](https://mlsanity.com/paper/2112.08903/) (2021)
- [OOD-GNN: Out-of-Distribution Generalized Graph Neural Network](https://mlsanity.com/paper/2112.03806/) (2021)
- [SUGAR: Subgraph Neural Network with Reinforcement Pooling and Self-Supervised Mutual Information Mechanism](https://mlsanity.com/paper/2101.08170/) (2021)
- [Subgraph Neural Networks](https://mlsanity.com/paper/2006.10538/) (2020)
- [Learning from Counterfactual Links for Link Prediction](https://mlsanity.com/paper/2106.02172/) (2021)
- [DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery -- A Focus on Affinity Prediction Problems with Noise Annotations](https://mlsanity.com/paper/2201.09637/) (2022)
- [Deciphering Spatio-Temporal Graph Forecasting: A Causal Lens and Treatment](https://mlsanity.com/paper/2309.13378/) (2023)
- [Improving Subgraph Recognition with Variational Graph Information Bottleneck](https://mlsanity.com/paper/2112.09899/) (2021)
- [Does Invariant Graph Learning via Environment Augmentation Learn Invariance?](https://mlsanity.com/paper/2310.19035/) (2023)
- [Environment-Aware Dynamic Graph Learning for Out-of-Distribution Generalization](https://mlsanity.com/paper/2311.11114/) (2023)
- [Optimizing OOD Detection in Molecular Graphs: A Novel Approach with Diffusion Models](https://mlsanity.com/paper/2404.15625/) (2024)
- [Graph Invariant Learning with Subgraph Co-mixup for Out-Of-Distribution Generalization](https://mlsanity.com/paper/2312.10988/) (2023)

## Abstract (arXiv)

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.