# Train Flat, Then Compress: Sharpness-Aware Minimization Learns More Compressible Models

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 67 of its 97 references that are in the map.
- First posted on arXiv 25 May 2022.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Beyond Accuracy: Behavioral Testing of NLP models with CheckList](https://mlsanity.com/paper/2005.04118/) (2020)
- [SemEval-2017 Task 1: Semantic Textual Similarity - Multilingual and Cross-lingual Focused Evaluation](https://mlsanity.com/paper/1708.00055/) (2017)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [A Survey of Quantization Methods for Efficient Neural Network Inference](https://mlsanity.com/paper/2103.13630/) (2021)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Patient Knowledge Distillation for BERT Model Compression](https://mlsanity.com/paper/1908.09355/) (2019)

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

- [PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning](https://mlsanity.com/paper/2306.10711/) (2023)
- [Unmasking the Lottery Ticket Hypothesis: What's Encoded in a Winning Ticket's Mask?](https://mlsanity.com/paper/2210.03044/) (2022)
- [Sharpness-Aware Minimization Leads to Low-Rank Features](https://mlsanity.com/paper/2305.16292/) (2023)
- [Normalization Layers Are All That Sharpness-Aware Minimization Needs](https://mlsanity.com/paper/2306.04226/) (2023)
- [Adaptive Sharpness-Aware Pruning for Robust Sparse Networks](https://mlsanity.com/paper/2306.14306/) (2023)
- [A Three-regime Model of Network Pruning](https://mlsanity.com/paper/2305.18383/) (2023)
- [Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning](https://mlsanity.com/paper/2405.20439/) (2024)
- [Flatness-Aware Prompt Selection Improves Accuracy and Sample Efficiency](https://mlsanity.com/paper/2305.10713/) (2023)
- [SAFE: Finding Sparse and Flat Minima to Improve Pruning](https://mlsanity.com/paper/2506.06866/) (2025)
- [Understanding Flatness in Generative Models: Its Role and Benefits](https://mlsanity.com/paper/2503.11078/) (2025)
- [Model Generalization: A Sharpness Aware Optimization Perspective](https://mlsanity.com/paper/2208.06915/) (2022)
- [Sparse Weight Averaging with Multiple Particles for Iterative Magnitude Pruning](https://mlsanity.com/paper/2305.14852/) (2023)
- [The Split Matters: Flat Minima Methods for Improving the Performance of GNNs](https://mlsanity.com/paper/2306.09121/) (2023)
- [Modality-Aware SAM: Sharpness-Aware-Minimization Driven Gradient Modulation for Harmonized Multimodal Learning](https://mlsanity.com/paper/2510.24919/) (2025)
- [Model Merging via Multi-Teacher Knowledge Distillation](https://mlsanity.com/paper/2512.21288/) (2025)
- [To Asymmetry and Beyond: Structured Pruning of Sequence to Sequence Models for Improved Inference Efficiency](https://mlsanity.com/paper/2304.02721/) (2023)
- [CrAFT: Compression-Aware Fine-Tuning for Efficient Visual Task Adaptation](https://mlsanity.com/paper/2305.04526/) (2023)
- [Long-Tailed Recognition on Binary Networks by Calibrating A Pre-trained Model](https://mlsanity.com/paper/2404.00285/) (2024)
- [S2AP: Score-space Sharpness Minimization for Adversarial Pruning](https://mlsanity.com/paper/2510.18381/) (2025)
- [Weight Concentration Regularization for Improving Pruning Robustness Under High Sparsity](https://mlsanity.com/paper/2511.14282/) (2025)

## Abstract (arXiv)

Model compression by way of parameter pruning, quantization, or distillation has recently gained popularity as an approach for reducing the computational requirements of modern deep neural network models for NLP. Inspired by prior works suggesting a connection between simpler, more generalizable models and those that lie within wider loss basins, we hypothesize that optimizing for flat minima should lead to simpler parameterizations and thus more compressible models. We propose to combine sharpness-aware minimization (SAM) with various task-specific model compression methods, including iterative magnitude pruning (IMP), structured pruning with a distillation objective, and post-training dynamic quantization. Empirically, we show that optimizing for flatter minima consistently leads to greater compressibility of parameters compared to vanilla Adam when fine-tuning BERT models, with little to no loss in accuracy on the GLUE text classification and SQuAD question answering benchmarks. Moreover, SAM finds superior winning tickets during IMP that 1) are amenable to vanilla Adam optimization, and 2) transfer more effectively across tasks.