# Improving out-of-distribution generalization in graphs via hierarchical semantic environments

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 4 Mar 2024.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (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)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Hierarchical Graph Representation Learning with Differentiable Pooling](https://mlsanity.com/paper/1806.08804/) (2018)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Generalized ODIN: Detecting Out-of-distribution Image without Learning from Out-of-distribution Data](https://mlsanity.com/paper/2002.11297/) (2020)
- [Explainability in Graph Neural Networks: A Taxonomic Survey](https://mlsanity.com/paper/2012.15445/) (2020)
- [AllenNLP: A Deep Semantic Natural Language Processing Platform](https://mlsanity.com/paper/1803.07640/) (2018)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Delving into Out-of-Distribution Detection with Vision-Language Representations](https://mlsanity.com/paper/2211.13445/) (2022)
- [Understanding Attention and Generalization in Graph Neural Networks](https://mlsanity.com/paper/1905.02850/) (2019)
- [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)
- [GOOD: A Graph Out-of-Distribution Benchmark](https://mlsanity.com/paper/2206.08452/) (2022)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [Size-Invariant Graph Representations for Graph Classification Extrapolations](https://mlsanity.com/paper/2103.05045/) (2021)
- [Debiasing Graph Neural Networks via Learning Disentangled Causal Substructure](https://mlsanity.com/paper/2209.14107/) (2022)
- [Graph Representation Learning: A Survey](https://mlsanity.com/paper/1909.00958/) (2019)
- [Graph Rationalization with Environment-based Augmentations](https://mlsanity.com/paper/2206.02886/) (2022)

## Built on by (most cited first, 5 of 5)

- [A Survey of Deep Graph Learning under Distribution Shifts: from Graph Out-of-Distribution Generalization to Adaptation](https://mlsanity.com/paper/2410.19265/) (2024)
- [Raising the Bar in Graph OOD Generalization: Invariant Learning Beyond Explicit Environment Modeling](https://mlsanity.com/paper/2502.10706/) (2025)
- [SIGMMA: Hierarchical Graph-Based Multi-Scale Multi-modal Contrastive Alignment of Histopathology Image and Spatial Transcriptome](https://mlsanity.com/paper/2511.15464/) (2025)
- [Rethinking Graph Generalization through the Lens of Sharpness-Aware Minimization](https://mlsanity.com/paper/2602.08855/) (2026)
- [CATeye: Coupled Attribute-Topology Invariance Learning for Voucher Abuse Detection](https://mlsanity.com/paper/2609.01425/) (2026)

## Abstract (arXiv)

Out-of-distribution (OOD) generalization in the graph domain is challenging due to complex distribution shifts and a lack of environmental contexts. Recent methods attempt to enhance graph OOD generalization by generating flat environments. However, such flat environments come with inherent limitations to capture more complex data distributions. Considering the DrugOOD dataset, which contains diverse training environments (e.g., scaffold, size, etc.), flat contexts cannot sufficiently address its high heterogeneity. Thus, a new challenge is posed to generate more semantically enriched environments to enhance graph invariant learning for handling distribution shifts. In this paper, we propose a novel approach to generate hierarchical semantic environments for each graph. Firstly, given an input graph, we explicitly extract variant subgraphs from the input graph to generate proxy predictions on local environments. Then, stochastic attention mechanisms are employed to re-extract the subgraphs for regenerating global environments in a hierarchical manner. In addition, we introduce a new learning objective that guides our model to learn the diversity of environments within the same hierarchy while maintaining consistency across different hierarchies. This approach enables our model to consider the relationships between environments and facilitates robust graph invariant learning. Extensive experiments on real-world graph data have demonstrated the effectiveness of our framework. Particularly, in the challenging dataset DrugOOD, our method achieves up to 1.29% and 2.83% improvement over the best baselines on IC50 and EC50 prediction tasks, respectively.