PaperlayerView exact revision on arXiv

Source-linked AI summary

Boosting Data Augmentation with Stochastic Weight Averaging

Longde Huang, Axel Flinth, Jan E. Gerken

arXiv:2608.14373v1cs.LG

TL;DR

Ensemble training for symmetry-aware learning is computationally costly because it requires repeated runs. This paper analyzes stochastic weight averaging with data augmentation and finds an equivariance boost beyond ordinary SWA’s performance improvement, supported across vision and graph tasks.

  • Problem

    Training augmented deep ensembles can yield perfect symmetry, but repeated training runs make this approach computationally costly.

  • Method

    The paper models end-of-training SWA with an Ornstein–Uhlenbeck process and uses Hessian traces and infinite-width neural tangent kernels to bound equivariance improvement.

  • Results

    In the infinite-width limit, augmented-data SWA yields a non-trivial equivariance improvement beyond ordinary SWA, verified across vision and graph tasks with discrete and continuous symmetries.

  • Takeaways & Limitations

    SWA with data augmentation improves equivariance across diverse models and tasks, including cases where image-classification accuracy does not improve.

  • Takeaways & Limitations

    A practical analysis initially depends on well-separated samples, whose required waiting time is difficult to estimate and may increase training time.

Abstract

from arXiv · show

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.

1 Introduction

The introduction presents data augmentation as a flexible way to incorporate task symmetries without specialized architectures, while noting that it yields approximate rather than exact equivariance. It motivates stochastic weight averaging (SWA) as a cheaper alternative to repeated ensemble training and outlines a theoretical and experimental analysis of its equivariance boost.

  • Motivation: Task symmetries can improve learned-model performance, and they are incorporated either through exactly equivariant architectures or through flexible data augmentation.Architectural constraints improve sample efficiency but require specialized designs; augmentation works with off-the-shelf architectures when transformations are implementable.
  • Motivation: Data augmentation is generally only approximately equivariant, whereas infinite ensembles trained on augmented data are exactly equivariant in expectation over initialization.The introduction attributes the limited theory of augmentation to the nonlinear training dynamics it modifies.
  • Approach: SWA replaces independent training trajectories with an ensemble mean from one trajectory, reducing ensemble training cost to that of a single model.The approach is motivated by SGD noise near a loss minimum and is analyzed using an Ornstein–Uhlenbeck approximation of the late training trajectory.
  • Theory: Assuming a well-trained network, the analysis expresses SWA performance and equivariance boosts through Hessian-trace quantities and derives an infinite-width bound on the non-equivariant loss ratio R⊥.The introduction identifies dependent samples from a single trajectory as the central theoretical obstacle and states that Theorem 4.11 establishes a non-trivial equivariance improvement for a wide range of models.
  • Experiments: Numerical experiments verify the bound’s dependence on group size and averaging time and demonstrate SWA’s equivariance improvement across vision, graph, discrete-symmetry, and continuous-symmetry tasks.The experiments include five vision datasets with seven models and a molecular graph-classification task invariant to 3D rotations.

2 Related Work

Related work contrasts architectural equivariance with data augmentation, surveys ensembling and weight averaging, and situates the paper’s analysis in SGD dynamics, Hessian spectra, and equivariant NTKs. The paper extends prior equivariant-architecture NTK work and uses an equivariant Hessian bound relative to the full Hessian to obtain its main theorem.

  • Equivariance and data augmentation: Architectural methods impose equivariance exactly through group-equivariant, continuous, steerable, or geometric deep learning representations, whereas augmentation leaves parameters unconstrained but yields only approximate equivariance.The architectural route requires choosing and implementing a specific group, while augmentation avoids that restriction.
  • Equivariance and data augmentation: End-to-end equivariance need not generally arise from layerwise equivariance, although supporting examples and empirical conjectures for trained CNNs motivate studying alternative realizations.Under a stronger identifiability condition, prior work rigorously connected end-to-end equivariance with layerwise realizability.
  • Ensembles and checkpoint averaging: Infinite ensembles of fully augmented networks are equivariant in expectation, but ensemble methods generally require storing and evaluating multiple models at inference.Weight averaging instead collapses the ensemble into one parameter set before inference.
  • Ensembles and checkpoint averaging: Sequential checkpoint ensembles average outputs from models saved during one training run, while SWA averages iterates into a single model and is contrasted with EMA’s different averaging scheme.These methods draw on sequential training, mode connectivity, and related ensemble constructions, but multiple-model inference motivates weight averaging.
  • Weight averaging: SWA’s interpretation is contested: some work emphasizes flat-minima discovery, whereas other studies attribute its benefit mainly to variance reduction and dependence on SGD convergence.Continuous-time SDE analyses model SGD near a well-fitted minimum through a stationary distribution governed by a quadratic approximation.
  • Neural tangent kernel: The NTK describes sufficiently wide-network training as kernel gradient descent, and prior equivariant work computed NTKs for finite-group regular representations.This paper generalizes the derivation to arbitrary intermediate representations and irrep decompositions, then bounds the trace of the equivariant Hessian 𝐻E relative to the full Hessian 𝐻.

3 Problem Setting

The problem setting formalizes group-invariant data augmentation, equivariant neural networks, and the induced group action on parameter space. It then specifies SGD training and stochastic weight averaging as an estimator motivated by batch-sampling noise.

  • Dataset and Data Augmentation: The dataset is a distribution on input-output spaces supported on a compact set, with input and output represented by a group G.The input and output spaces are V_X = R^n_X and V_Y = R^n_Y.
  • Dataset and Data Augmentation: Data augmentation imposes invariance of the joint distribution under simultaneous input-output group actions for every g ∈ G.The support is also invariant under the corresponding product action.
  • Network: The model is an MLP whose admissible parameters are restricted to an affine subspace, accommodating architectures such as CNNs, transformers, and residual structures.The network parameters comprise layer weights and nonlinearities, with layerwise admissible parameter spaces L_l.
  • Group action on the parameter space: Layerwise representations induce a group action on parameter space, and equivariant parameters are invariant under this associated representation.The admissible space is assumed group-invariant, and the equivariant parameter space E is non-empty.
  • Learning Algorithm: SGD uses stochastic mini-batch gradients, while weight averaging estimates an optimal solution by averaging previous weights despite batch-sampling noise.The performance ratio R(T) exceeds 1 when averaging improves the loss, with larger values indicating a greater boost.

4 Theoretical Results · 4.1 Equivariance properties of the stochastic gradients

The theoretical analysis approximates small-step SGD by a stochastic differential equation and establishes equivariance properties of its dynamics under augmented-data training. In particular, invariant initial parameter distributions yield G-equivariant ensemble averages at every positive time, while the result is limited for single models.

  • 4 Theoretical Results: The analysis replaces finite-step SGD with a stochastic differential equation, ignoring finite-step-size effects.The approximation is motivated for small η.
  • 4 Theoretical Results: The SDE uses Brownian motion with covariance I and a factor P(θ) satisfying P(θ)P(θ)^T = Σ(θ).The sample gradient has mean ∇_θL(θ(t)) and covariance Σ(θ(t)).
  • 4 Theoretical Results: For small η, the stochastic gradient increments approximately follow the SDE through a central-limit-theorem argument and an Euler–Maruyama discretization.The finite-step limit can also be justified more rigorously via [40, Corollary 10].
  • 4.1 Equivariance properties of the stochastic gradients: Under augmented-data training, the Hessian ∇^2L and covariance matrix Σ obey transformation rules that drive the equivariance analysis.These relations are stated in Lemma 4.2.
  • 4.1 Equivariance properties of the stochastic gradients: The stochastic flow Φ_t inherits an equivariance property from the transformation rules for the Hessian and covariance.The flow maps the initial parameter measure μ_0 to the time-t measure.
  • 4.1 Equivariance properties of the stochastic gradients: For an invariant initial parameter distribution μ_0, the independently trained ensemble’s parameter distribution remains governed by the equivariant flow.The condition is ḡ_*μ_0 = μ_0 for all g∈G.
  • 4.1 Equivariance properties of the stochastic gradients: For every t > 0, the ensemble-averaged network is G-equivariant.This is the conclusion of Proposition 4.3 for ensembles trained using the SDE.
  • 4.1 Equivariance properties of the stochastic gradients: The ensemble result is not informative for a single model unless its Dirac-delta initial distribution is concentrated at a point θ_0∈E.In that single-model setting, the theorem recovers results about single networks from.

4.2 Dynamics close to equivariant minima and the Ornstein-Uhlenbeck process

This section approximates SWA dynamics near an equivariant strict local minimum by an Ornstein–Uhlenbeck process. It decomposes the dynamics into equivariant and non-equivariant components and defines a curvature-scaled measure and ratio for quantifying SWA’s equivariance gain.

  • Dynamics near equivariant minima: Near an equivariant strict local minimum, quadratic loss and constant noise approximations reduce the stochastic training dynamics to an Ornstein–Uhlenbeck process.The loss is approximated using the Hessian at the minimum, while the noise factor is evaluated there.
  • Dynamics near equivariant minima: The Hessian and noise covariance preserve the decomposition E ⊕ E⊥, so the dynamics separate into equivariant and non-equivariant blocks.Both matrices commute with the group-averaging operator and are block diagonal with respect to E and E⊥.
  • Equivariance measures: Non-equivariance is measured by the norm of the E⊥ parameter component in the Hessian-induced metric, which compares its size with the loss landscape’s curvature scale.A small L⊥ indicates that the network is close to equivariance relative to the relevant loss length scale.
  • Equivariance measures: The equivariance gain from SWA is defined as R⊥(T), the ratio of L⊥ for averaged weights to L⊥ for unaveraged weights at the end of training.The analysis targets large R⊥(T) for temporally separated samples and, in a specialized invariant-classification setting, for dense samples in the NTK limit.

4.3 Well-separated samples

When SWA samples are sufficiently separated in time, the parameter samples become approximately independent and approximately independent of initialization, approaching i.i.d. draws from a Gaussian limit distribution. This intuition is difficult to apply directly because required waiting times are hard to estimate and increase training cost.

  • Well-separated samples: Sufficiently large sampling intervals make the sequence of samples approximately i.i.d. from the limit distribution.This follows from the decay of the relevant exponential operators at large times.
  • Well-separated samples: The limiting distribution is Gaussian with mean zero and a specified covariance matrix.The provided passage identifies the limit as Gaussian and states its mean, while the covariance expression is truncated.
  • Well-separated samples: For large T, the averaged parameter also becomes Gaussian distributed, while the covariance of θ_t does not scale with T.The passage connects this scaling behavior to the SWA intuition, but the displayed covariance expression is incomplete.
  • Well-separated samples: The well-separated-sample result is not readily applicable to practical training runs because the necessary waiting time is difficult to estimate and longer waits increase training time.The next analysis removes the assumption of well-separated samples.

4.4 Dense samples in a classification setting

For group-invariant classification with dense sampling, the analysis specializes near a well-fitted minimum, where the Hessian approximately matches the gradient-noise covariance. This shared structure lets stochastic weight averaging be analyzed through decoupled eigendirections and shows that it simultaneously improves performance and encourages equivariance.

  • Dense samples in a classification setting: The classification analysis assumes a group-invariant label space, cross-entropy loss, and a loss landscape near a well-fitted minimum.A well-fitted minimum means the prediction residual is small.
  • Dense samples in a classification setting: When the residual error is small, the Hessian H is approximately equal to the covariance matrix Σ of the stochastic gradients.The difference between the matrices is linear in the residual error ε, yielding H≈Σ when ε≪1.
  • Dense samples in a classification setting: The Hessian and covariance share an eigenbasis, so the stochastic dynamics decompose into independent one-dimensional solutions.This common eigenbasis enables the simplified analysis of the averaged process and its loss components.
  • Dense samples in a classification setting: Even with dense samples, both R(T) and R⊥(T) converge to infinity as T→∞, so SWA encourages equivariance while also boosting performance.The two quantities represent the analyzed loss behaviors in the relevant and orthogonal subspaces.
  • Dense samples in a classification setting: The analysis separates generic performance improvement from equivariance gain by studying the ratio R⊥(T)/R(T), whose infinite-time behavior is then prepared for an NTK-limit bound.As T→∞, the eigenvalue counts cover all dimensions in the respective subspaces and the ε-terms vanish.

4.5 A refined analysis in the NTK limit

The section derives deterministic infinite-width NNGP and NTK limits for equivariant subnetworks by working in an irrep-adapted basis. It extends prior kernel calculations beyond regular representations and relates the equivariant Hessian spectrum to that of the full network near an equivariant minimizer.

  • Irrep-adapted coordinates: The analysis changes from standard coordinates to an orthogonal irrep-adapted basis, preserving inner products and the isotropic Gaussian law of unconstrained weights.After width enlargement, the basis change is applied channel-by-channel, and the transformed nonlinearity remains a collection of independent copies across channels.
  • Equivariant weights: Equivariant weights arise by projecting isotropic Gaussian weights onto intertwiner spaces, where Schur’s lemma shares weights across the d_λ copies within each irrep.This projection produces isotropic Gaussians on the equivariant subspace and induces the independence structure used in the infinite-width analysis.
  • Infinite-width equivariant kernels: As N→∞, equivariant pre-activations in irrep-adapted coordinates converge to Gaussian processes, while the equivariant subnetwork’s NTK converges to a deterministic limit.The corresponding kernels are matrix-valued and live in spaces whose ambient dimensions grow with N, so the analysis resolves convergence through fixed-size irrep blocks.
  • Kernel interpretation and novelty: The resulting NNGP kernels and NTKs can be interpreted as kernels of a layerwise equivariant architecture, extending earlier calculations restricted to regular representations of finite groups.The section presents this as a more general setting for equivariant kernel calculations.
  • Hessian comparison: Near an equivariant minimizer under cross-entropy loss, Lemma 4.10 bounds the spectral proportion of the equivariant Hessian H_E relative to the full Hessian H.The supplied passage states the bound but does not include its displayed value.

4.6 Main theoretical result

Theorem 4.11 bounds SWA’s asymptotic relative equivariance-improvement ratio under the stated assumptions and identifies when it exceeds one. The result specializes to rectangular networks and explains how representation structure and group size affect this condition.

  • Main theoretical result: The derivation assumes that the network non-linearity is the ReLU function.This is one of the assumptions listed for the theorem.
  • Main theoretical result: The main theorem gives a bound on the asymptotic relative ratio of improvement for solution equivariance.The theorem follows by combining Lemma 4.10 and equation (83) in equation (71).
  • Main theoretical result: The ratio exceeds one under the theorem’s stated condition, implying R⊥ > R as T→∞.In this regime, the expected equivariance improvement measured by R⊥ outpaces that measured by R under data augmentation.
  • Rectangular networks: For rectangular networks, Corollary 4.12 provides a simplified bound and a corresponding condition for the ratio to exceed one.Rectangular networks use layer-independent multiplicity and dimension parameters in the intermediate layers.
  • Representation dependence: The bound depends on the representation through irrep multiplicities and through the equivariant minimum around which the loss is expanded.The converged equivariant minimum selects the relevant multiplicities.
  • Representation dependence: The condition is harder to satisfy for larger groups, but it holds when sufficiently many intermediate representations mix irreps or contain non-trivial-dimensional irreps.Regular intermediate representations, as in group-CNNs, provide a notable special case.

4.7 A proxy for the orthogonal loss

The section establishes that the KL-based equivariance loss is an accurate proxy for the theoretical orthogonal loss near a well-fitted equivariant minimum. It also explains why the proxy is practical and remains valid when averaging predictive distributions instead of logits.

  • Proposition 4.14: Proposition 4.14 shows that the equivariance metric accurately proxies the orthogonal loss under the stated small-error and near-minimizer assumptions.The result applies when ε(x) = p_θ(x) − e_y(x) is small and θ lies near the equivariant minimizer θ* = 0.
  • Proposition 4.14: Near the equivariant minimizer, deviations from group-averaged logits are driven entirely by the orthogonal parameter component θE⊥, while θE does not contribute at this order.This follows from linearizing the network output and using the equivariance transformation law of its Jacobian.
  • Corollary 4.15: The resulting corollary relates equivariance losses evaluated on unaveraged and averaged weights in the regime where Proposition 4.14 applies.The passage states the ratio result but does not include its explicit formula.
  • Variant: Averaging predictive distributions rather than logits changes the averaged distribution by O(∥θ∥2) and the equivariance loss by O(∥θ∥3).Thus, the same proxy argument extends to this alternative orbit-averaging procedure.
  • Practical significance: The orthogonal loss is generally intractable to compute, whereas the equivariance loss is directly computable from predictions and therefore provides a practical proxy near a well-fitted minimum.Computing the orthogonal loss requires the Hessian and an orthogonal parameter-space projection.

5 Experiments

Experiments validate the quadratic-loss approximation and the theoretical predictions for SWA’s equivariance boost. Across image and graph classification, SWA improves equivariance more consistently than general performance, with effects depending on architecture, task complexity, and symmetry-group size.

  • Quadratic approximation: The Hessian’s effective rank increases with sample size before plateauing at 45, while its near-full-rank structure is explained by ReLU scale invariances.Despite slight rank deficiency, the quadratic loss approximation remains accurate along training trajectories and becomes more accurate near the minimizer.
  • Synthetic validation: For G=C4, R⊥(T)/R(T) rises from below its asymptotic bound, then decreases and plateaus above the theoretical bound as averaging continues.The experiments therefore support convergence toward the asymptotic value with increasing averaging time.
  • Synthetic validation: Larger symmetry groups reduce the dimension ratio and weaken SWA’s guaranteed excess equivariance, but the empirical ratio remains above 1 for all tested group sizes.This agrees with the predicted decrease of |G|r^2/(|G|r^2 − 1) toward 1 as |G| grows.
  • Image classification: Across image-classification architectures and datasets, SWA improves equivariance more than general performance, with R⊥/R consistently larger than 1.OSP improvement is strictly positive across all tested configurations, including ViG-Tiny on ImageNet-100 when standard accuracy decreases.
  • Image classification: MLPs achieve R⊥/R ratios of 1.14 on MNIST and 1.44 on FashionMNIST, near the theoretical C4 value of 4/3, whereas deep convolutional models reach 2.21–5.09.The 4/3 bound is conservative for the convolutional settings and can be close to tight for MLPs.
  • Graph classification: Across nine graph-classification configurations, ΔAcc ranges from −2.81% to +2.48%, while ΔOSP improves in every setting by 0.68%–2.66%.Five of nine configurations have negative ΔAcc, and equivariance effects are clearly resolved in most settings.

A Proof of Proposition 4.3

The proof establishes Proposition 4.3 by showing that the Fokker–Planck operator and time evolution commute with the push-forward symmetry operation, preserving invariant distributions and the resulting parameter symmetry.

  • Definitions: The push-forward operation S_g distinguishes transforming a loss gradient from taking the gradient of a transformed loss.Specifically, ∇L(gθ) is S_g∇L, whereas ∇[L(ḡθ)] is ∇(S_gL).
  • Operator commutation: Lemma A.2 proves that the Fokker–Planck operator A commutes with S_g for every g∈G.The proof derives the required divergence and Laplace commutation relations using the chain rule and tensor transformations.
  • Operator commutation: Using loss invariance S_gL=L and Lemma 4.2, the transformed drift and diffusion terms reduce to expressions compatible with the symmetry operation.The resulting identity explicitly rewrites both divergence and second-order diffusion contributions under the group action.
  • Proposition 4.3: Because A and ∂_t commute with S_g, the flow maps an initial distribution p_0 to a transformed solution S_gp_t initialized at S_gp_0.An invariant initialization therefore remains invariant throughout the evolution.

B Quadratic approximation with Gaussian diffusion sampling

Replacing gradient-ascent sampling with Gaussian diffusion yields comparable quadratic-approximation accuracy, while producing systematically lower losses and faster stabilization of approximation error at matched perturbation scales.

  • B Quadratic approximation with Gaussian diffusion sampling: Gaussian diffusion achieves comparable approximation accuracy to gradient-ascent sampling, demonstrating robustness of the quadratic approximation to perturbation choice.The experiment repeats Section 5.1.1 with Gaussian diffusion; results are shown in Figure 4.
  • B Quadratic approximation with Gaussian diffusion sampling: At matched perturbation scales, Gaussian diffusion produces systematically lower loss values than gradient ascent.
  • B Quadratic approximation with Gaussian diffusion sampling: The approximation error stabilizes considerably faster under Gaussian diffusion than under gradient ascent.
Loading 2608.14373v1…