# Efficient Sharpness-Aware Minimization for Molecular Graph Transformer Models

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 67 references that are in the map.
- First posted on arXiv 19 Jun 2024.

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

- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [GPT-GNN: Generative Pre-Training of Graph Neural Networks](https://mlsanity.com/paper/2006.15437/) (2020)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [Molecular Graph Convolutions: Moving Beyond Fingerprints](https://mlsanity.com/paper/1603.00856/) (2016)
- [Towards Understanding Sharpness-Aware Minimization](https://mlsanity.com/paper/2206.06232/) (2022)
- [Efficient Sharpness-aware Minimization for Improved Training of Neural Networks](https://mlsanity.com/paper/2110.03141/) (2021)
- [Towards Efficient and Scalable Sharpness-Aware Minimization](https://mlsanity.com/paper/2203.02714/) (2022)
- [Label Noise SGD Provably Prefers Flat Global Minimizers](https://mlsanity.com/paper/2106.06530/) (2021)
- [Towards Deeper Graph Neural Networks with Differentiable Group Normalization](https://mlsanity.com/paper/2006.06972/) (2020)
- [Sharpness-Aware Minimization Improves Language Model Generalization](https://mlsanity.com/paper/2110.08529/) (2021)
- [Dirichlet Energy Constrained Learning for Deep Graph Neural Networks](https://mlsanity.com/paper/2107.02392/) (2021)
- [Sharpness-Aware Training for Free](https://mlsanity.com/paper/2205.14083/) (2022)
- [SMILES Transformer: Pre-trained Molecular Fingerprint for Low Data Drug Discovery](https://mlsanity.com/paper/1911.04738/) (2019)
- [Auto-GNN: Neural Architecture Search of Graph Neural Networks](https://mlsanity.com/paper/1909.03184/) (2019)
- [On the Loss Landscape of Adversarial Training: Identifying Challenges and How to Overcome Them](https://mlsanity.com/paper/2006.08403/) (2020)
- [An Adaptive Policy to Employ Sharpness-Aware Minimization](https://mlsanity.com/paper/2304.14647/) (2023)

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

- [Unifying Unsupervised Graph-Level Anomaly Detection and Out-of-Distribution Detection: A Benchmark](https://mlsanity.com/paper/2406.15523/) (2024)
- [Graph-structured Small Molecule Drug Discovery Through Deep Learning: Progress, Challenges, and Opportunities](https://mlsanity.com/paper/2502.08975/) (2025)
- [Mamba-Based Graph Convolutional Networks: Tackling Over-smoothing with Selective State Space](https://mlsanity.com/paper/2501.15461/) (2025)
- [COSCO: A Sharpness-Aware Training Framework for Few-shot Multivariate Time Series Classification](https://mlsanity.com/paper/2409.09645/) (2024)
- [Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise](https://mlsanity.com/paper/2509.18001/) (2025)
- [Convergence of Sharpness-Aware Minimization Algorithms using Increasing Batch Size and Decaying Learning Rate](https://mlsanity.com/paper/2409.09984/) (2024)
- [A Generalisation Signal Need Not Be a Model-Selection Signal](https://mlsanity.com/paper/2609.39099/) (2026)

## Abstract (arXiv)

Sharpness-aware minimization (SAM) has received increasing attention in computer vision since it can effectively eliminate the sharp local minima from the training trajectory and mitigate generalization degradation. However, SAM requires two sequential gradient computations during the optimization of each step: one to obtain the perturbation gradient and the other to obtain the updating gradient. Compared with the base optimizer (e.g., Adam), SAM doubles the time overhead due to the additional perturbation gradient. By dissecting the theory of SAM and observing the training gradient of the molecular graph transformer, we propose a new algorithm named GraphSAM, which reduces the training cost of SAM and improves the generalization performance of graph transformer models. There are two key factors that contribute to this result: (i) textitgradient approximation: we use the updating gradient of the previous step to approximate the perturbation gradient at the intermediate steps smoothly (increases efficiency); (ii) textitloss landscape approximation: we theoretically prove that the loss landscape of GraphSAM is limited to a small range centered on the expected loss of SAM (guarantees generalization performance). The extensive experiments on six datasets with different tasks demonstrate the superiority of GraphSAM, especially in optimizing the model update process. The code is in:https://github.com/YL-wang/GraphSAM/tree/graphsam