# Boosting Data Augmentation with Stochastic Weight Averaging

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 42 references that are in the map.
- First posted on arXiv 14 Aug 2026.

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

- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [Spherical CNNs](https://mlsanity.com/paper/1801.10130/) (2018)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [Ensemble deep learning: A review](https://mlsanity.com/paper/2104.02395/) (2021)
- [Learning Steerable Filters for Rotation Equivariant CNNs](https://mlsanity.com/paper/1711.07289/) (2017)
- [Frequency Bias in Neural Networks for Input of Non-Uniform Density](https://mlsanity.com/paper/2003.04560/) (2020)
- [Label Noise SGD Provably Prefers Flat Global Minimizers](https://mlsanity.com/paper/2106.06530/) (2021)
- [Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations](https://mlsanity.com/paper/1811.01558/) (2018)
- [Roto-Translation Covariant Convolutional Networks for Medical Image Analysis](https://mlsanity.com/paper/1804.03393/) (2018)
- [Provably Strict Generalisation Benefit for Equivariant Models](https://mlsanity.com/paper/2102.10333/) (2021)
- [Geometric Deep Learning and Equivariant Neural Networks](https://mlsanity.com/paper/2105.13926/) (2021)
- [Swallowing the Bitter Pill: Simplified Scalable Conformer Generation](https://mlsanity.com/paper/2311.17932/) (2023)
- [Does equivariance matter at scale?](https://mlsanity.com/paper/2410.23179/) (2024)
- [Improved Generalization Bounds of Group Invariant / Equivariant Deep Networks via Quotient Feature Spaces](https://mlsanity.com/paper/1910.06552/) (2019)
- [Stochastic Weight Averaging Revisited](https://mlsanity.com/paper/2201.00519/) (2022)
- [Checkpoint Ensembles: Ensemble Methods from a Single Training Process](https://mlsanity.com/paper/1710.03282/) (2017)
- [Equivariance versus Augmentation for Spherical Images](https://mlsanity.com/paper/2202.03990/) (2022)
- [The asymptotic spectrum of the Hessian of DNN throughout training](https://mlsanity.com/paper/1910.02875/) (2019)
- [Optimization Dynamics of Equivariant and Augmented Neural Networks](https://mlsanity.com/paper/2303.13458/) (2023)
- [Emergent Equivariance in Deep Ensembles](https://mlsanity.com/paper/2403.03103/) (2024)
- [A Defense of the Quadratic Model](https://mlsanity.com/paper/2607.21716/) (2026)
- [Spectral alignment of stochastic gradient descent for high-dimensional classification tasks](https://mlsanity.com/paper/2310.03010/) (2023)

## Abstract (arXiv)

The symmetries of a learning task have become an important factor in designing modern deep learning solutions. Data augmentation is a straightforward and effective way of incorporating symmetries into a generic neural network. Recent results show that infinitely large deep ensembles show perfect symmetry when trained on augmented data. However, since training ensembles requires repeating the training process many times, this method is costly. In this work, we study stochastic weight averaging (SWA) applied to classification as an alternative ensembling technique that does not require repeated training runs. We analyze SWA by approximating the stochastic training trajectory at the end of training with an Ornstein--Uhlenbeck process. We show that in the infinite-width limit, SWA on augmented data provides an equivariance boost that goes beyond what could be expected from the performance increase due to SWA alone. We verify our results with extensive numerical experiments on numerous models spanning image and graph classification with both discrete and continuous symmetries.