Boosting Data Augmentation with Stochastic Weight Averaging

Abstract: 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) 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 equiviariance 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 computer vision and graph classification with both discrete and continuous symmetries.

• The paper shows that stochastic weight averaging (SWA) applied to augmented-data training improves both performance and equivariance, with the equivariance benefit asymptotically exceeding generic loss improvement under an Ornstein–Uhlenbeck approximation.
• The theoretical bound connects the relative equivariance gain to group size, representation structure, Hessian traces, and network width, while experiments report equivariance-to-performance ratios above 1 across vision and graph benchmarks.
• SWA produced accuracy gains of up to 14.68% and equivariance improvements in every tested configuration, including cases where accuracy declined or remained statistically unchanged.

Motivation and problem statement

Data augmentation is the standard mechanism for injecting task symmetries into an otherwise unconstrained network, but it yields only approximate, distributional equivariance. Prior theoretical work established that an infinite ensemble of networks trained on fully augmented data is exactly equivariant in expectation [gerken2024; nordenfors2024ensembles], yet ensembling requires repeating training for each member and is therefore computationally prohibitive. This paper investigates stochastic weight averaging (SWA) — averaging checkpoints along a single SGD trajectory [izmailov2019a] — as a cheap surrogate for ensembling, and asks whether SWA on augmented data produces an equivariance gain that exceeds what its generic performance improvement would predict.

The central technical obstacle is that SWA samples are drawn from one trajectory and hence are not independent. The authors resolve this by assuming a well-trained network near an equivariant minimum, where the dynamics reduce to an Ornstein–Uhlenbeck (OU) process governed by the Hessian $H$ of the loss.

Theoretical framework

The setting is a group-invariant classification task with finite group $G$, unitary representations $\rho_X,\rho_Y$ on input/output spaces, augmented data (i.e., $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$), and an invariant loss. Equivariant parameters form a linear subspace $E \subset H$ of parameter space, with orthogonal complement $E^\perp$. Two key structural facts drive the analysis:

• Equivariance of gradient statistics: for augmented data, the gradient covariance satisfies $g\Sigma(\theta)g^\top = \Sigma(g\theta)$ and the Hessian transforms covariantly (Lemma 1). Consequently, near an equivariant minimum $\theta_* \in E$, both $H$ and $\Sigma$ are block-diagonal with respect to $G$0.
• SGD as SDE: late-stage SGD is approximated by $G$1, which linearizes to an OU process at the minimum. The non-equivariant loss contribution is $G$2.

Two ratios quantify the benefit: the performance ratio $G$3 and the equivariance ratio $G$4, comparing unaveraged and SWA-averaged weights after averaging window $G$5.

Main results

For well-separated checkpoint times, the averaged weights have covariance shrinking as $G$6, giving $G$7. For dense sampling in a classification setting, the paper shows that near a well-fitted minimum the Hessian and gradient covariance coincide to first order in the residual error $G$8, so the dynamics decouple into one-dimensional OU processes along eigen-directions. This yields, for large $G$9,

$\rho_X,\rho_Y$0

where $\rho_X,\rho_Y$1 counts eigenvalues exceeding $\rho_X,\rho_Y$2. Both ratios diverge as $\rho_X,\rho_Y$3: SWA improves performance and equivariance simultaneously, even with densely sampled checkpoints. Their ratio converges to

$\rho_X,\rho_Y$4

To bound this quantity, the paper works in the infinite-width NTK limit, widening via direct sums of intermediate representations. A dimension count using Schur's lemma and Cauchy–Schwarz gives $\rho_X,\rho_Y$5, where $\rho_X,\rho_Y$6 bounds the growth-rate disparity across layers. For the trace ratio, the authors derive the NTK of a layerwise equivariant subnetwork (Lemma 5), generalizing prior work restricted to regular representations [pmlr-v267-misof25a] to arbitrary irrep decompositions. Under ReLU nonlinearities, the equivariant kernel trace is bounded by $\rho_X,\rho_Y$7 times the full kernel trace, yielding the main theorem:

$\rho_X,\rho_Y$8

which exceeds one whenever $\rho_X,\rho_Y$9. Since each factor is at most $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$0, the condition holds whenever intermediate representations mix irreps or contain non-trivial irreps; for regular representations (as in group-CNNs) the per-layer factor is $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$1. Notably, the guarantee weakens as $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$2 grows, consistent with the intuition that larger symmetry groups are harder to learn statistically.

Finally, Proposition 6 establishes that the KL-based equivariance loss $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$3 — computable without Hessian knowledge — agrees with $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$4 up to $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$5 near a well-fitted equivariant minimum, licensing its use as an empirical proxy for $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$6.

Assumptions and limitations

The theorem rests on several assumptions stated explicitly by the authors: large averaging time $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$7, large width $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$8, finite $(X,Y) \stackrel{d}{=} (\rho_X(g)X, \rho_Y(g)Y)$9, a well-fitted model (small residual $E \subset H$0), ReLU nonlinearities, and convergence to a strictly equivariant minimum $E \subset H$1 (invoking identifiability results to justify treating layerwise-equivariant parameters as well-defined). The quadratic approximation of the loss is validated empirically but is itself an approximation whose accuracy depends on the SGD trajectory remaining confined to high-curvature directions. On practical architectures the theorem's hypotheses fail (restricted parameter spaces, intractable irrep decompositions), so the experiments there only test the qualitative claim $E \subset H$2 rather than the quantitative bound. The analysis also covers discrete groups only; continuous symmetries appear in experiments solely through discretized subgroups of $E \subset H$3 and $E \subset H$4.

Experimental validation

Synthetic verification. In a teacher–student MLP setup with a regular-representation teacher under $E \subset H$5, the empirical ratio $E \subset H$6 rises above the asymptotic bound $E \subset H$7 and plateaus above it as $E \subset H$8 grows, while plateaued values decrease monotonically with group order across cyclic groups $E \subset H$9 — matching both qualitative predictions of the theorem. The ratio remains above 1 throughout, though the excess shrinks with $E^\perp$0 as predicted.

Image classification. Across MNIST, Fashion-MNIST, CIFAR-10/100, and ImageNet-100 with $E^\perp$1 rotation augmentation, and seven models (MLPs, VGG-16, PreResNet-164, WideResNet-28×10, ViT-S, ViG-Ti), OSP improves strictly positively in every configuration — including ViG-Ti on ImageNet-100, where accuracy slightly decreases ($E^\perp$2) while OSP gains $E^\perp$3. Accuracy gains reach $E^\perp$4 (ViT-S on ImageNet-100) and $E^\perp$5 (WRN on CIFAR-100). The relative ratio $E^\perp$6 exceeds 1 everywhere: MLPs yield 1.14–1.44, close to the theoretical ceiling of $E^\perp$7, while deep CNNs reach 2.21–5.09, indicating the bound is conservative outside its strict regime — plausibly because residual architectures' admissible parameter space makes $E^\perp$8 a smaller fraction of $E^\perp$9 than the unrestricted dimension count assumes.

Graph classification. On DHFR, BZR, and COX2 with GINs under tetrahedral/octahedral/icosahedral augmentation, accuracy changes are essentially noise (five of nine configurations negative, error bars comparable to means), yet OSP improves in all nine configurations ($g\Sigma(\theta)g^\top = \Sigma(g\theta)$0 to $g\Sigma(\theta)g^\top = \Sigma(g\theta)$1). This dissociation — equivariance gains persisting absent accuracy gains — directly supports the paper's claim that SWA's effect on the non-equivariant subspace is not merely a side effect of generic variance reduction.

Conclusion

This paper provides the first quantitative account of how SWA interacts with data-augmentation-induced symmetrization. By modeling late-stage SGD as an OU process at an equivariant minimum and combining Hessian-trace identities with novel equivariant NTK calculations, it proves that SWA's relative improvement in equivariance asymptotically dominates its relative improvement in loss, under explicit conditions on width, averaging time, group size, and representation structure. Experiments spanning vision and molecular graph tasks, discrete rotations in 2D and 3D, and four architecture families confirm the effect robustly, including settings where accuracy does not improve. Open questions include extending the analysis beyond ReLU and finite groups, tightening the conservative bound observed for deep convolutional models, and characterizing when the predicted equivariance gain translates into downstream accuracy improvements.