# Rethinking Graph Generalization through the Lens of Sharpness-Aware Minimization

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 57 references that are in the map.
- First posted on arXiv 9 Feb 2026.

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

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (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)
- [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)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One](https://mlsanity.com/paper/1912.03263/) (2019)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Discovering Invariant Rationales for Graph Neural Networks](https://mlsanity.com/paper/2201.12872/) (2022)
- [Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism](https://mlsanity.com/paper/2201.12987/) (2022)
- [G-Mixup: Graph Data Augmentation for Graph Classification](https://mlsanity.com/paper/2202.07179/) (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)
- [Overcoming Catastrophic Forgetting in Incremental Few-Shot Learning by Finding Flat Minima](https://mlsanity.com/paper/2111.01549/) (2021)
- [Graph Adversarial Training: Dynamically Regularizing Based on Graph Structure](https://mlsanity.com/paper/1902.08226/) (2019)
- [Graph Rationalization with Environment-based Augmentations](https://mlsanity.com/paper/2206.02886/) (2022)
- [Graph Transfer Learning via Adversarial Domain Adaptation with Graph Convolution](https://mlsanity.com/paper/1909.01541/) (2019)
- [Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization](https://mlsanity.com/paper/2303.03108/) (2023)
- [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)
- [Does Invariant Graph Learning via Environment Augmentation Learn Invariance?](https://mlsanity.com/paper/2310.19035/) (2023)
- [Mind the Label Shift of Augmentation-based Graph OOD Generalization](https://mlsanity.com/paper/2303.14859/) (2023)
- [Domain Adaptation with Adversarial Training and Graph Embeddings](https://mlsanity.com/paper/1805.05151/) (2018)
- [Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization](https://mlsanity.com/paper/2306.01103/) (2023)
- [Learning Invariant Molecular Representation in Latent Discrete Space](https://mlsanity.com/paper/2310.14170/) (2023)

## Abstract (arXiv)

Graph Neural Networks (GNNs) have achieved remarkable success across various graph-based tasks but remain highly sensitive to distribution shifts. In this work, we focus on a prevalent yet under-explored phenomenon in graph generalization, Minimal Shift Flip (MSF),where test samples that slightly deviate from the training distribution are abruptly misclassified. To interpret this phenomenon, we revisit MSF through the lens of Sharpness-Aware Minimization (SAM), which characterizes the local stability and sharpness of the loss landscape while providing a theoretical foundation for modeling generalization error. To quantify loss sharpness, we introduce the concept of Local Robust Radius, measuring the smallest perturbation required to flip a prediction and establishing a theoretical link between local stability and generalization. Building on this perspective, we further observe a continual decrease in the robust radius during training, indicating weakened local stability and an increasingly sharp loss landscape that gives rise to MSF. To jointly solve the MSF phenomenon and the intractability of radius, we develop an energy-based formulation that is theoretically proven to be monotonically correlated with the robust radius, offering a tractable and principled objective for modeling flatness and stability. Building on these insights, we propose an energy-driven generative augmentation framework (E2A) that leverages energy-guided latent perturbations to generate pseudo-OOD samples and enhance model generalization. Extensive experiments across multiple benchmarks demonstrate that E2A consistently improves graph OOD generalization, outperforming state-of-the-art baselines.