# Improving Generalization of Pre-trained Language Models via Stochastic Weight Averaging

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 39 of its 51 references that are in the map.
- First posted on arXiv 12 Dec 2022.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (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)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [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)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [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)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Improved Knowledge Distillation via Teacher Assistant](https://mlsanity.com/paper/1902.03393/) (2019)
- [Patient Knowledge Distillation for BERT Model Compression](https://mlsanity.com/paper/1908.09355/) (2019)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [DeepNet: Scaling Transformers to 1,000 Layers](https://mlsanity.com/paper/2203.00555/) (2022)
- [PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation](https://mlsanity.com/paper/2104.12369/) (2021)
- [An Alternative View: When Does SGD Escape Local Minima?](https://mlsanity.com/paper/1802.06175/) (2018)
- [Sharpness-Aware Minimization Improves Language Model Generalization](https://mlsanity.com/paper/2110.08529/) (2021)

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

- [Representation Surgery for Multi-Task Model Merging](https://mlsanity.com/paper/2402.02705/) (2024)
- [Knowledge is a Region in Weight Space for Fine-tuned Language Models](https://mlsanity.com/paper/2302.04863/) (2023)
- [Learning from models beyond fine-tuning](https://mlsanity.com/paper/2310.08184/) (2023)
- [A Universal Class of Sharpness-Aware Minimization Algorithms](https://mlsanity.com/paper/2406.03682/) (2024)
- [Model Fusion through Bayesian Optimization in Language Model Fine-Tuning](https://mlsanity.com/paper/2411.06710/) (2024)
- [Uncertainty-Aware Natural Language Inference with Stochastic Weight Averaging](https://mlsanity.com/paper/2304.04726/) (2023)
- [Parameter-Efficient Interventions for Enhanced Model Merging](https://mlsanity.com/paper/2412.17023/) (2024)
- [Fighting Randomness with Randomness: Mitigating Optimisation Instability of Fine-Tuning using Delayed Ensemble and Noisy Interpolation](https://mlsanity.com/paper/2406.12471/) (2024)
- [LABO: Towards Learning Optimal Label Regularization via Bi-level Optimization](https://mlsanity.com/paper/2305.04971/) (2023)

## Abstract (arXiv)

Knowledge Distillation (KD) is a commonly used technique for improving the generalization of compact Pre-trained Language Models (PLMs) on downstream tasks. However, such methods impose the additional burden of training a separate teacher model for every new dataset. Alternatively, one may directly work on the improvement of the optimization procedure of the compact model toward better generalization. Recent works observe that the flatness of the local minimum correlates well with better generalization. In this work, we adapt Stochastic Weight Averaging (SWA), a method encouraging convergence to a flatter minimum, to fine-tuning PLMs. We conduct extensive experiments on various NLP tasks (text classification, question answering, and generation) and different model architectures and demonstrate that our adaptation improves the generalization without extra computation cost. Moreover, we observe that this simple optimization technique is able to outperform the state-of-the-art KD methods for compact models.