Boosting Data Augmentation with Stochastic Weight Averaging

Boosting Data Augmentation with Stochastic Weight Averaging Longde Huang 𝑎

Axel Flinth∗ 𝑏

Jan E. Gerken∗ 𝑎

arXiv:2608.14373v1 [cs.LG] 14 Aug 2026

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.

1

Introduction

Many machine-learning tasks, ranging from protein folding [1] to medical image processing [2] feature symmetries that constrain the learning problem and, if taken into account, can be used to improve the performance of the trained models. In deep learning more specifically, there are two major approaches for incorporating symmetries into the learning algorithm. Firstly, the symmetry can be realized as a constraint on the neural network architecture. In particular, the layers of the network can be chosen such that they are equivariant functions with respect to the symmetry of the task at hand [3, 4]. This approach requires specialized architectures (discussed in a large body of literature) but realizes the symmetry exactly and improves sample efficiency during training. Secondly, the neural network can be trained on transformed versions of the original training data, a method known as data augmentation. Since this technique does not constrain the network architecture, it can be applied to any highly optimized off-the-shelf architecture as long as the symmetry transformation can be implemented. It is therefore straightforward and flexible. However, it does not lead to an exactly equivariant model, but rather one that is approximately equivariant (or equivariant in distribution). Theoretical studies of equivariant neural networks are plentiful due to deep connections between equivariant layers and representation theory. In particular, practically relevant special cases such as the symmetry ∗ Equal Contribution 𝑎 Department of Mathematical Sciences, Chalmers University of Technology and the University of Gothenburg, SE-412 96

Gothenburg, Sweden. Emails: longde@chalmers.se, gerken@chalmers.se 𝑏 Department of Mathematics and Mathematical Statistics, Umeå University, Linnéus väg 49, 901 87 Umeå, Sweden. Email: axel.flinth@umu.se

1

groups 𝑆 𝑛 of permutations, 𝑆𝑂 (𝑛) of rotations or 𝑆𝐸 (𝑛) of roto-translations have been studied extensively in the literature, for various different model architectures. In contrast, most studies of data augmentation are purely empirical, see e.g. [5]. This is because data augmentation modifies the training procedure and therefore one inevitably needs to study the complicated, non-linear training dynamics of neural networks. Despite this obstacle, recent advances in this area show that in expectation over the initialization distribution, neural networks are exactly equivariant when they are trained on augmented data [6, 7]. Approximately, this can be realized by training a deep ensemble. However, ensemble training is notoriously compute intensive, since the training needs to be repeated for each ensemble member. For this reason, and motivated by the recent results mentioned above, we study the effects of data augmentation on stochastic weight averaging (SWA). Instead of an ensemble over independent training trajectories, SWA takes an ensemble mean over several models from the same training trajectory [8], thereby reducing the training cost to that of a single model. Intuitively, the training trajectory of SGD is dominated by noise from batch selection when it is close to a minimum of the loss. Therefore, SWA can be expected to have a similar effect to averaging over random initializations. In this work, we investigate whether this intuition holds up. In particular, we approximate the stochastic training trajectory at the end of training with an Ornstein–Uhlenbeck process and use the ratio of the expected losses with and without SWA to quantify the performance boost due to SWA. A central obstacle to the theoretical analysis of SWA in this setting is the non-independence resulting from the samples being taken from a single training trajectory. To overcome this problem, we assume a well-trained network and express the performance boost in terms of the trace of the Hessian of the network. In this setting, we then consider the expected equivariance boost which focuses on the non-equivariant contributions to the loss. In the infinite-width limit, we compute a bound on this expected equivariance boost using the theory of neural tangent kernels (NTKs). Our calculations show that combining data augmentation with SWA leads to an equivariance boost which goes beyond the mere performance increase expected from ordinary SWA. We validate our theoretical analysis with extensive numerical experiments. In particular, we explicitly verify that the bound predicted by our theory holds and shows the right dependency on group size and averaging time on a synthetic problem over which we have full control. We furthermore demonstrate the claimed improvement on computer vision and graph classification tasks, as well as both discrete and continuous symmetry groups, for numerous different model classes. Our main contributions are as follows: • We formalize SWA at the end of training with an Ornstein–Uhlenbeck process and analyse the equivariance properties of the resulting model in the case of augmented training data. To quantify the performance boost due to SWA, we introduce the loss ratio 𝑅 in (20). Similarly, to quantify the equivariance boost due SWA, we define the loss ratio 𝑅⊥ in (41). We overcome the central obstacle of dependent ensemble members from one trajectory by assuming a well-trained network and expressing 𝑅 and 𝑅⊥ in terms of the trace of the Hessian of the network, see (66) and (69). • We then take the infinite-width limit and leverage the theory of NTKs to derive a bound on 𝑅⊥ (Theorem 4.11) which shows that SWA with data augmentation leads to a non-trivial improvement in equivariance for a wide range of models. In the process, we derive novel results about the NTK limit for equivariant models. • We validate our theoretical results numerically by considering a synthetic learning task. We verify for this task that our bound still holds for a finitely-sized network and has the predicted dependence on the group size and averaging time (Section 5.1). We furthermore test the equivariance gain from SWA in more realistic settings (Section 5.2). In particular, we demonstrate the predicted improvement on five vision datasets with seven different models (MLPs, CNNs, vision and graph 2

transformers) for discrete rotations. We also show that the same effect can be observed for a molecular graph classification task invariant with respect to 3D rotations.

2

Related Work

Equivariance and data augmentation A large body of work builds equivariance directly into the network architecture, so that the symmetry constraint is imposed exactly throughout training rather than only approximately. Group-equivariant CNNs [9] replaced standard convolutions with group convolutions over discrete groups, and were later generalized to continuous and steerable representations [10, 11], ultimately unified under the framework of geometric deep learning [3]. Taking the reverse approach, Kondor et al. [12] studied equivariant maps from first principles, proving that convolutions are the unique equivariant linear maps for compact groups. These architectures guarantee equivariance without relying on the data distribution, at the cost of restricting the parameter space to a limited subspace, which Nordenfors et al. [13] study in a general setting. This architectural restriction has been argued to improve generalization [14, 15] From this parameter restriction arises the structural question: is layerwise equivariance the only way to realize equivariance? Bökman et al. [16] first tackled this question in theory, arguing that this is not the general scenario despite the existence of supporting examples; they nevertheless conjectured that it holds empirically for trained CNNs. Shahverdi et al. [17], under a stronger condition of identifiability, showed rigorously that end-to-end equivariance always implies the existence of a parameter choice realizing the same function with layerwise equivariance. This thread of works is the premise on which we treat the equivariant subspace E as well-defined, and study the training dynamics restricted to it. Layerwise realizability does not make the architectural route costless in practice, since it requires committing to a specific group at design time, implementing specialized layers for it, and training within a restricted parameter space. As an alternative, data augmentation offers a route to equivariance that leaves the parameter space unconstrained, at the cost that the resulting network is only approximately equivariant. Whether this matters is contested. Gerken et al. [5] find that for invariant classification tasks for spherical images, sufficiently augmented models with large scale match equivariant architectures. Similarly, [18] reports that dropping the rotational inductive bias simplifies model scaling considerably. Brehmer et al. [19], in contrast, argues equivariant models outperform augmented ones at every compute budget tested. Augmentation thus appears as a viable substitute for architectural equivariance in some scenarios, including the classification tasks with finite groups considered in our work. Chen et al. [20] give a theoretical framework for this alternative, modelling data augmentation as an averaging operation over group orbits that renders the data distribution approximately invariant, which is exactly the setting of our work. Ensembles and checkpoint averaging To achieve exact equivariance from data augmentation, ensembling has emerged as a provable mechanism [6, 7]. The result of Nordenfors et al. [7], that an infinite ensemble of networks trained with fully augmented data is equivariant in expectation, appears in our work as a special case of a more general framework. This connection to ensembling motivates a closer look at the broader literature on combining multiple models. Ensembles of independently trained networks are a classical and extensively studied technique for improving generalization, and we refer to Ganaie et al. [21] for a general review. Among the broad spectrum of methods, ensembling with multiple sequentially trained models has emerged as an efficient strategy. Huang et al. [22] and Chen et al. [23] constructed ensembles by averaging the outputs of models saved at different points during a single training round. Both works are based on the hypothesis that the training trajectory explores different local minima, and averaging over them reduces overfitting. Garipov et al.

3

[24], on the other hand, explored the loss landscape further and discovered that local minima are connected by simple low-loss curves, challenging the isolated minima view. Inspired by this insight of mode connectivity, they traverse nearby low-loss regions to collect a diverse set of ensemble members. A common issue among these methods is that they still require storing and evaluating multiple models at inference time, limiting their practicality. Weight averaging avoids this cost entirely by collapsing the ensemble into a single set of weights before inference, and is exactly the backbone of the mechanism we study in our work. Weight averaging Averaging the iterates of a stochastic optimization algorithm has classical roots in Polyak–Ruppert averaging [25]. The most common form in deep learning is the exponential moving average (EMA) [26], a related but distinct technique that stochastic weight averaging (SWA) [8] is often contrasted against. The two differ in their averaging scheme, which leads to different accounts of their effect on the loss landscape: by averaging iterates sampled uniformly rather than with exponential decay, SWA is argued to land in the interior of a wide, flat optimum, which is connected to better generalization. Several works adopt this wide-minima point of view. Xie et al. [27] ground this flat-minima intuition by modeling SGD as a diffusion process whose stationary distribution concentrates exponentially more on flat minima; Damian et al. [28] similarly show that SGD with label noise biases toward flatter minimizers by analyzing how the noise affects local quadratic structure of the loss. This flat-minima account is not, however, unchallenged. Guo et al. [29] revisit SWA empirically and find its benefit highly dependent on how well the underlying SGD trajectory has converged before averaging, arguing the operation contributes primarily through variance reduction rather than flat-minima discovery. This coincides with Mandt et al. [30], who model SGD as a continuous-time SDE and show that near a well-fitted minimum, its stationary distribution is governed by a quadratic approximation of the loss with covariance set by the Hessian, which is the same picture we adopt. Unlike Mandt et al. [30], however, who treat this as a general posterior characterization of SGD, we consider the setting where the dynamics decompose into equivariant and non-equivariant components, studying the impact of SWA on each in turn. The quadratic, Hessian-governed picture is itself an active subject of empirical analysis. Meterez et al. [31] stress-test it at LLM pretraining scale and find it remains predictive over a non-trivial fraction of training, while Arous et al. [32] show that in high-dimensional classification tasks the SGD trajectory aligns with a low-dimensional outlier eigenspace of the Hessian, layer by layer. We likewise focus on classification tasks and study the spectrum of the Hessian, but apply the theory of the neural tangent kernel to derive a quantitative analysis. Neural tangent kernel The NTK [33] characterizes the training dynamics of sufficiently wide neural networks as kernel gradient descent under a fixed, architecture-determined kernel, and is based on the network’s correspondence to a Gaussian process at initialization known as the neural network Gaussian process (NNGP) [34]. Subsequent work has used the spectrum of the NTK to characterize structure in the Hessian and gradients of trained networks. Jacot et al. [35] show that the NTK itself governs the asymptotic spectrum of the Hessian throughout training. Basri et al. [36] use the NTK’s spectral decomposition to characterize which frequencies a network learns fastest, extending the classical frequency-bias result to non-uniformly sampled data. Seleznova et al. [37] build an OOD detector for classification models, based on the discovery of a related alignment phenomenon: the NTK’s approximate block-diagonal structure across classes. Closest to our setting in equivariant tasks, Misof et al. [38] compute the NTK of equivariant architectures whose intermediate representations are regular representations of a finite group. Our derivation in Section 4.5.2 generalizes this to arbitrary intermediate representations and irrep decompositions, and we

4

use the resulting equivariant NTK to bound the trace of the equivariant Hessian 𝐻 E relative to the full Hessian 𝐻, which leads to our main theorem.

3

Problem Setting

In this section, we will set the notation and formalize stochastic weight averaging (SWA). Dataset and Data Augmentation For a generic training task, we denote the true dataset as a joint distribution 𝜇 𝑋,𝑌 over 𝑉𝑋 × 𝑉𝑌 , from which the training samples are drawn. Here, 𝑉𝑋 = R𝑛𝑋 and 𝑉𝑌 = R𝑛𝑌 are the input and output spaces. We assume the distribution 𝜇 𝑋,𝑌 to be supported on a compact set Ω. Furthermore, we assume a pair of group representations (𝜌 𝑋 , 𝜌𝑌 ) of a group 𝐺 on the input and output spaces, respectively, 𝜌 𝑋 : 𝐺 → GL(𝑉𝑋 ) ,

𝜌𝑌 : 𝐺 → GL(𝑉𝑌 ) .

(1)

In this context, we consider data augmentation of the training data as a condition on the distribution. Specifically, an augmented dataset corresponds to a distribution 𝜇 𝑋,𝑌 that is invariant under the group action, namely the joint random variable (𝑋, 𝑌 ) is identically distributed with (𝜌 𝑋 (𝑔) 𝑋, 𝜌𝑌 (𝑔)𝑌 ), ∀𝑔 ∈ 𝐺. This also means (𝜌 𝑋 × 𝜌𝑌 ) (𝑔)Ω = Ω, ∀𝑔 ∈ 𝐺. For simplicity, we denote 𝜌 𝑋 (𝑔), 𝜌𝑌 (𝑔) with 𝑔 𝑋 , 𝑔𝑌 . Network As the neural network we train, we consider a multilayer perceptron (MLP) N𝜃 with parameters 𝜃 = (𝑊 0 , 𝑊 1 , . . . , 𝑊 𝐿 ) and nonlinearities 𝜎 = (𝜎0 , . . . , 𝜎𝐿 ). In terms of the preactivations (ℎ0 , . . . , ℎ 𝐿+1 ), the network function is given by 1 ℎ0 (𝑥) = 𝑊 0 𝑥 , ℎ𝑙 (𝑥) = √ 𝑊 𝑙 𝜎𝑙−1 (ℎ𝑙−1 (𝑥)) for 𝑙 = 1, . . . , 𝐿 + 1 , N𝜃 (𝑥) = ℎ 𝐿+1 (𝑥) , (2) 𝑛𝑙 where ℎ𝑙 (𝑥) ∈ R𝑛𝑙 and 𝑊 𝑙 ∈ R𝑛𝑙+1 ×𝑛𝑙 . In particular, we have 𝑛 𝑋 = 𝑛0 and 𝑛𝑌 = 𝑛 𝐿+1 . We denote the space of parameters 𝜃 of this network by H=

𝐿 Ê

R𝑛𝑙+1 ×𝑛𝑙 .

(3)

𝑙=0

To accommodate practical and modern architectures of interest, such as convolutional networks (CNN), transformers, and residual structures, we follow a standard parameter-space restriction and further restrict the admissible parameters 𝜃 to an affine subspace L ⊂ H . At layer 𝑙, we will denote space of admissible parameters by L 𝑙 . We call N𝜃 an equivariant network, if for (𝑥, 𝑦)-𝜇 𝑋,𝑌 a.e. we have N𝜃 (𝑔 𝑋 𝑥) = 𝑔𝑌 N𝜃 (𝑥) .

(4)

Loss function The loss function ℓ(N𝜃 (𝑥), 𝑦) is defined over 𝑉𝑌 × 𝑉𝑌 . We assume it is non-negative, twice differentiable, locally Lipschitz and attains zero if and only if N𝜃 (𝑥) = 𝑦. We further assume that it is group invariant, ℓ(𝑔𝑌 N𝜃 (𝑥), 𝑔𝑌 𝑦) = ℓ(N𝜃 (𝑥), 𝑦) . (5) This implies that, if the network N𝜃 itself is equivariant, then ℓ 𝜃 (𝑥, 𝑦) := ℓ(N𝜃 (𝑥), 𝑦) is invariant ℓ 𝜃 (𝑥, 𝑦) = ℓ 𝜃 (𝑔 𝑋 𝑥, 𝑔𝑌 𝑦) .

(6)

The cumulative loss, 𝐿 (𝜃), is defined as the expectation of ℓ 𝜃 (𝑥, 𝑦) with respect to the data distribution, ∫ 𝐿 (𝜃) = E 𝜇 (𝑙 (𝑥, 𝑦)) = 𝑙 𝜃 (𝑥, 𝑦)d𝜇 𝑋,𝑌 . (7) Ω

5

Group action on the parameter space For a network with a feed forward structure, one natural way to achieve equivariance is through layerwise equivariance. Therefore, we define a series of representations 𝜌𝑖 on the intermediate spaces 𝑉𝑖 = R𝑛𝑖 . We assume all representations, including the one on the input and the output space, are unitary with respect to the canonical tensor basis. The group then naturally acts on the parameter space as 𝜌¯ : 𝐺 → GL(H ) ,

𝜌(𝑔)𝜃 ¯ = 𝜌(𝑔) ¯ (𝑊 0 , . . . , 𝑊 𝐿 ) = (𝜌 𝑋 (𝑔)𝑊 0 𝜌 𝑋 (𝑔) −1 , . . . , 𝜌𝑌 (𝑔)𝑊 𝐿 𝜌 𝐿 (𝑔) −1 ). (8)

If we require the non-linearities 𝜎𝑖 to be equivariant with respect to the corresponding intermediate representations 𝜌𝑖 , 𝜌𝑖 (𝑔) ◦ 𝜎𝑖 = 𝜎𝑖 ◦ 𝜌𝑖 (𝑔), (9) Then, to say that a set of parameters 𝜃 is equivariant w.r.t. these representations, 𝑊 𝑙 𝜌𝑙 (𝑔) = 𝜌𝑙+1 (𝑔)𝑊 𝑙 ,

(10)

is to say that it is invariant under the associated representation 𝜌¯ (denote 𝜌(𝑔) ¯ as 𝑔), ¯ 𝜃 = 𝑔𝜃 ¯ .

(11)

The parameters satisfying (11) form a linear subspace H𝐺 in H . We assume the space of admissible parameters L is invariant under the group action (𝑔L ¯ ⊂ L, ∀𝑔 ∈ 𝐺), and the space of admissible equivariant parameters E = H𝐺 ∩ L is non-empty. By picking 𝜃 𝐺 ∈ E, we can associate L with 𝑇𝜃𝐺 L + 𝜃 𝐺 . Assuming (8) and (9), the group action on the data is equivalent to an action on the parameter space, N𝜃 (𝑔 𝑋 𝑥) = 𝑔𝑌 N𝑔¯ −1 𝜃 (𝑥).

(12)

This coincides with the definition of E, since if 𝜃 ∈ E, the equation above turns into N𝜃 (𝑔 𝑋 𝑥) = 𝑔𝑌 N𝜃 (𝑥),

(13)

which is just the usual definition of equivariance. Consequently, for the loss we have ℓ𝑔¯ 𝜃 (𝑥, 𝑦) = ℓ(N𝑔¯ 𝜃 (𝑥), 𝑦) = 𝑙 (𝑔𝑌 N𝜃 (𝑔 𝑋−1 𝑥), 𝑦) = ℓ(N𝜃 (𝑔 𝑋−1 𝑥), 𝑔𝑌−1 𝑦) = ℓ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦) .

(14)

Remark 3.1. All practically used equivariant neural networks are constructed out of equivariant layers, i.e. 𝜃 ∈ E in our language. However in general, that a network is equivariant does not imply that 𝜃 ∈ E. As shown in Shahverdi et al. [17], if the parameters are identifiable (intuitively, this means that the map 𝜃 → 𝑓 𝜃 is injective up to symmetries), there is an equivalence between equivariant layers and equivariant functions. We assume the group 𝐺 to be finite. We will furthermore assume that the representation 𝜌¯ is unitary with respect to the Euclidean inner product. Then, we can consider the orthogonal complement 𝑇 E ⊥ of 𝑇 E. Furthermore, the orthogonal projection of a vector 𝑣 ∈ 𝑇 L to 𝑇 E is given by [39, Prop. 2.8] 𝑝 : 𝑇L → 𝑇E ,

𝑣 ↦→

1 ∑︁ 𝑔𝑣 ¯ . |𝐺 | 𝑔∈𝐺

(15)

Learning Algorithm We train the network with respect to the loss function with stochastic gradient descent (SGD). With learning rate 𝜂 and batch size 𝐵, the training step could be formulated as 𝐵

1 ∑︁ 𝜃 𝑡+1 = 𝜃 𝑡 − 𝜂 ∇ 𝜃 ℓ 𝜃𝑡 (𝑋𝑖 , 𝑌𝑖 ), 𝐵 𝑖=1 6

(16)

where (𝑋𝑖 , 𝑌𝑖 ) are independent random variables identically distributed with (𝑋, 𝑌 ). The sample gradient ∇ 𝜃 𝑙 𝜃𝑡 (𝑋, 𝑌 ) has mean and covariance ∫ 𝜇(𝜃 𝑡 ) = ∇ 𝜃 ℓ 𝜃𝑡 (𝑥, 𝑦)d𝜇 𝑋,𝑌 = ∇ 𝜃 𝐿 (𝜃 (𝑡)) (17) ∫Ω Σ(𝜃 𝑡 ) = ∇ 𝜃 ℓ 𝜃 (𝑡 ) (𝑥, 𝑦) ∇ 𝜃 ℓ 𝜃𝑡 (𝑥, 𝑦) ⊤ d𝜇 𝑋,𝑌 . (18) Ω

Weight averaging The stochasticity of batch sampling provides an opportunity for the training process to explore wide basins in the loss landscape, but it also results in noise that prevents the trajectory from reaching the optimal solution. Weight averaging can be viewed as an estimator to approximate this solution, by taking the mean of the previous weights 𝑇 ∑︁ ¯𝜃 𝑇 = 1 𝜃𝑡 . 𝑇 𝑖=1 𝑖

(19)

Performance boost To quantify the performance boost from stochastic weight averaging, we compare the expected loss before and after averaging. For a given starting iterate 𝜃 0 of the averaging window, define E[𝐿 (𝜃 𝑇 )] 𝑅(𝑇) = , (20) E[𝐿 ( 𝜃¯𝑇 )] where the expectation is over the batch sampling (equivalently, over the SGD trajectory). In particular, 𝑅(𝑇) > 1 indicates that averaging improves the loss, with larger values indicating a greater boost.

4

Theoretical Results

As we outlined above, we will analyse stochastic gradient descent (16). For simplicity, we will ignore effects due to finite step size and instead consider the stochastic differential equation √︃ d𝜃 𝑡 = −∇𝐿 (𝜃 𝑡 )d𝑡 + 𝐵𝜂 𝑃(𝜃 𝑡 )d𝑊𝑡 , (21) where d𝑊 (𝑡) is Brownian motion with covariance 𝐼, and 𝑃(𝜃)𝑃(𝜃) 𝑇 = Σ(𝜃). For small values of 𝜂, (21) is a good approximation of the SGD dynamics. A heuristic argument for this is as follows: The sample gradient ∇ 𝜃 𝑙 𝜃 (𝑡 ) (𝑋, 𝑌 ) has mean ∇ 𝜃 𝐿 (𝜃 (𝑡)) and covariance Σ(𝜃 (𝑡)). By the central limit Í𝐵 theorem, 𝐵1 𝑖=1 ∇ 𝜃 𝑙 𝜃𝑡 (𝑋𝑖 , 𝑌𝑖 ) approximates a Gaussian distribution with mean ∇ 𝜃 𝐿 (𝜃 (𝑡)) and variance √1 Σ(𝜃 (𝑡)). Adding such increments to the variable 𝜃 𝑡 (as we do in SGD) hence corresponds essentially 𝐵 to an Euler-Maruyama scheme to solve (21). The limit can also be justified more rigorously – see [40, Corollary 10]. We will analyze (21) in the following. √︃ Remark 4.1. To simplify the notation, we will normalize 𝐵𝜂 = 1 from here on.

4.1

Equivariance properties of the stochastic gradients

When training on augmented data, the Hessian ∇2 𝐿 and covariance matrix Σ satisfy the following equivariance relations. They will be crucial to our entire analysis. Lemma 4.2. For any 𝜃 ∈ L, the following transformation rules hold, 𝑔Σ(𝜃) ¯ 𝑔¯ ⊤ = Σ( 𝑔𝜃), ¯

𝑔¯ ⊤ ∇2 𝐿 ( 𝑔𝜃) ¯ 𝑔¯ = ∇2 𝐿 (𝜃). 7

(22)

Proof. Applying the chain rule to (14) implies 𝑔¯ 𝑇 ∇ℓ 𝜃 (𝑥, 𝑦)| 𝜃=𝑔¯ 𝜃 = ∇ℓ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦), or equivalently ∇ℓ 𝜃 (𝑥, 𝑦)| 𝜃=𝑔¯ 𝜃 = 𝑔∇ℓ ¯ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦).

(23)

Applying the chain rule one more time yields ∇2 ℓ 𝜃 (𝑥, 𝑦)| 𝜃=𝑔¯ 𝜃 = 𝑔∇ ¯ 2 ℓ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦) 𝑔¯ ⊤ .

(24)

Because the data is augmented, (𝑥, 𝑦) ∼ (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦), and hence, using (23) and (24), we obtain Σ(𝜃)| 𝜃=𝑔¯ 𝜃 = E[∇ℓ 𝜃 (𝑥, 𝑦)∇ℓ 𝜃 (𝑥, 𝑦) ⊤ | 𝜃=𝑔¯ 𝜃 ] = E[ 𝑔∇ℓ ¯ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦)∇ℓ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦) ⊤ 𝑔¯ ⊤ ] ⊤ ⊤ = 𝑔E[∇ℓ ¯ 𝜃 (𝑥, 𝑦)∇ℓ𝑔¯ 𝜃 (𝑥, 𝑦) ] 𝑔¯ ⊤

(25) (26) (27)

= 𝑔Σ(𝜃) ¯ 𝑔¯ ,

(28)

2

(29)

2

∇ 𝐿 (𝜃)| 𝜃=𝑔¯ 𝜃 = E(∇ ℓ 𝜃 (𝑥, 𝑦)| 𝜃=𝑔¯ 𝜃 ) = E[ 𝑔∇ ¯ ℓ 𝜃 (𝑔 𝑋−1 𝑥, 𝑔𝑌−1 𝑦) 𝑔¯ ⊤ ] 2 ⊤ 2

= 𝑔E[∇ ¯ ℓ(𝑥, 𝑦)] 𝑔¯ 2

⊤

= 𝑔∇ ¯ 𝐿 (𝜃) 𝑔¯ ,

(30) (31) (32)

which was the claim. Consider the flow Φ𝑡 associated to (21), that maps the initial probability measure 𝜇0 of 𝜃 to the measure 𝜇𝑡 at time 𝑡, i.e. 𝜇𝑡 = Φ∗𝑡 𝜇0 . A first consequence of Lemma 4.2 is the following equivariance of Φ, Φ∗𝑡 𝑔¯ ∗ = 𝑔¯ ∗ Φ∗𝑡 .

(33)

Here, 𝑓 ∗ denotes the pushforward operation of a map 𝑓 . The discrete version of this equivariance guarantee was in [7] shown to imply a weak equivariance guarantee for an ensemble of neural networks trained with SGD. Concretely, if we initialize an ensemble using an invariant distribution, the same will be true for all times 𝑡, which in turn will imply that the ensemble mean will be equivariant at all times. Proposition 4.3. Let 𝜇0 be an invariant parameter distribution: 𝑔¯ ∗ 𝜇0 = 𝜇0 for all 𝑔 ∈ 𝐺. Let 𝜇𝑡 = Φ∗𝑡 𝜇0 be the distribution obtained by independently training an ensemble of networks using (21) for a time 𝑡. Then, the ensemble-averaged network 𝑁¯ 𝑡 (𝑥) = E 𝜃∼𝜇𝑡 [𝑁 𝜃 (𝑥)] .

(34)

is 𝐺-equivariant for every 𝑡 > 0. Proof. See Appendix A. Note that this result is not very informative for a single model – its initial distribution 𝜇0 is a Dirac delta, and hence only invariant if it is concentrated in a point 𝜃 0 ∈ E (applying the theorem in this context in fact recovers the results from [13] about single networks). Still, it is interesting that (33) persists in our stochastic setting, and we therefore provide a proof – based on analysing the Fokker-Planck-equation – in Appendix A.

8

4.2

Dynamics close to equivariant minima and the Ornstein-Uhlenbeck process

We will now move on to study SWA (19) more specifically. Since stochastic averages are taken along single training runs, we need to analyse the solution paths of (21). We will assume that solutions converge towards equivariant minima since they are trained on augmented data. Put differently, we will study the dynamics of (21) near a local minimum 𝜃 ∗ of the cumulative loss 𝐿, that we furthermore assume is equivariant in the strict sense 𝜃 ∗ ∈ E (cf. Remark 3.1). To simplify the notation, we will shift the coordinates, so that 𝜃 ∗ = 0. Near 𝜃 ∗ , we can approximate 𝐿 with the quadratic function 𝐿 (𝜃) = 12 (𝜃 −𝜃 ∗ )𝐻 (𝜃 −𝜃 ∗ ), with 𝐻 = ∇2 𝐿 (𝜃 ∗ ). We can furthermore make the approximation 𝑃(𝜃) ≈ 𝑃(𝜃 ∗ ) =: 𝑃. The dynamics (21) then become d𝜃 𝑡 = −𝐻𝜃 𝑡 + 𝑃d𝑊𝑡 .

(35)

This is the well-known Ornstein-Uhlenbeck process. Its solution has a closed form: If the process is initalized at 𝜃 0 , we have ∫ 𝑡

𝜃 𝑡 = 𝑒 −𝐻𝑡 𝜃 0 +

𝑒 −𝐻 (𝑡 −𝑠) 𝑃d𝑊𝑠 .

(36)

0

Since we are studying the dynamics at an equivariant point 𝜃 ∗ , 𝐻 and 𝑃 decompose into block matrices with respect to the decomposition E ⊕ E ⊥ , as specified by the following Lemma 4.4. The matrices 𝐻 and 𝑃 commute with 𝑔. ¯ Consequently, they have a block diagonal structure: They map E to E and E ⊥ to E ⊥ . Proof. Since 𝜃 ∗ ∈ E, we have 𝑔𝜃 ¯ ∗ = 𝜃 ∗ . Together with Lemma 4.2, this shows that 𝑔¯ 𝑇 𝐻 𝑔¯ = 𝐻, 𝑇 𝑔¯ Σ(𝜃 ∗ ) 𝑔¯ = Σ(𝜃 ∗ ), which is the first claim. Let 𝑣 ∈ E and 𝑤 ∈ E ⊥ . To prove the second claim, we note that we by definition of E have 𝑔𝑣 ¯ = 𝑣. Furthermore, the projection (15) yields 1 ∑︁ 𝑔𝑤 ¯ = 0. |𝐺 | 𝑔∈𝐺

(37)

We therefore obtain 𝑣 ⊤ 𝐻𝑤 =

1 ∑︁ ⊤ ⊤ 1 ∑︁ 1 ∑︁ 𝑣 𝑔¯ 𝐻 𝑔𝑤 ¯ = ( 𝑔𝑣) ¯ ⊤ 𝐻 𝑔𝑤 ¯ = 𝑣⊤ 𝐻 𝑔𝑤 ¯ = 0, |𝐺 | 𝑔∈𝐺 |𝐺 | 𝑔∈𝐺 |𝐺 | 𝑔∈𝐺

(38)

The argument for 𝑣 ⊤ Σ𝑤 = 0 proceeds in the same manner. Since E and E ⊥ are orthogonal complements, ⊥ ⊥ we may choose orthonormal bases {𝑒 𝑖 } for E and {𝑒 ⊥ 𝑖 } for E , so that {𝑒 𝑖 } ∪ {𝑒 𝑖 } forms an orthonormal basis for the ambient space. In this basis, Σ takes the block-diagonal form Σ = diag{Σ E , Σ E ⊥ }, where Σ E and Σ E ⊥ denote the restrictions of Σ to E and E ⊥ , respectively. We choose 𝑃 to be block-diagonal in the same basis, 𝑃 = diag{𝑃 E , 𝑃 E ⊥ }, with 𝑃 E 𝑃⊤E = Σ E and 𝑃 E ⊥ 𝑃⊤E⊥ = Σ E ⊥ . Then 𝑃𝑃⊤ = Σ, and for any 𝑣 ∈ E, 𝑤 ∈ E ⊥ , we have 𝑣 ⊤ 𝑃𝑤 = 0. The previous lemma, and basic properties of the matrix exponential, immediately show that the solution ⊥ (36) also decomposes: We get that 𝜃 = (𝜃 E , 𝜃 E ), with ∫ 𝑡 E −𝐻E 𝑡 E 𝜃𝑡 = 𝑒 𝜃0 + 𝑒 −𝐻E (𝑡 −𝑠) 𝑃 E d𝑊𝑠E ∫0 𝑡 ⊥ E⊥ −𝐻E ⊥ 𝑡 E 𝜃𝑡 = 𝑒 𝜃0 + 𝑒 −𝐻E ⊥ (𝑡 −𝑠) 𝑃 E ⊥ d𝑊𝑠E , (39) 0

9

⊥

for independent Brownian motions 𝑊𝑠E , 𝑊𝑠E on E and E ⊥ , respectively. With the block structure of the Hessian, we introduce a quantity to measure the equivariance of the network. Remember that N𝜃 is equivariant if 𝜃 ∈ E. Hence, a reasonable measure of the “nonequivariance” of a parameter is the norm of the projection of 𝜃 onto E ⊥ . By Lemma 4.4 (and the fact that 𝐻 ≻ 0 at the strict local minimum 𝜃 ∗ ), this is proportional to the quantity  ⊥ ⊤ ⊥ 𝐿 ⊥ (𝜃 𝑡 ) = 𝜃 𝑡E 𝐻 E ⊥ 𝜃 𝑡E . (40) Here, we measure the norm of the non-equivariant component of 𝜃 in the metric defined by 𝐻 E ⊥ which captures the length scale of the curvature of the loss landscape. I.e. if 𝐿 ⊥ is small, then the non-equivariant component of the parameters is small in comparison to the length scale of the loss and the network is close to equivariance. Analogous to (20), we introduce a quantity to measure the gain in equivariance from performing SWA. Concretely, we define it as the ratio of 𝐿 ⊥ evaluated on averaged weights and on unaveraged weights at the end of training, E[𝐿 ⊥ (𝜃 𝑇 )] 𝑅⊥ (𝑇) = . (41) E[𝐿 ⊥ ( 𝜃¯𝑇 )] We will in the following argue that 𝑅⊥ (𝑇) will be large (corresponding to a large gain in equivariance) in two settings: First, we will look at the case of weights sampled far apart from each other in time. Next, we will consider a more realistic scenario of dense samples, in a specialized setting of invariant classification problems. In the latter case, we will give detailed analysis of the behaviour of 𝑅⊥ (𝑇) in the NTK limit.

4.3

Well-separated samples

Since we are at a strict local minimum, the operators 𝐻 E and 𝐻 E ⊥ are positive definite. Therefore, 𝑒 −𝐻E 𝑡 and 𝑒 −𝐻E ⊥ 𝑡 will for large values of 𝑡 be vanishingly small. This together with the decomposition (39) shows ⊥ that 𝜃 𝑡E and 𝜃 𝑡E will be approximately independent of each other, and, more importantly, approximately 𝑇 in (19) are well separated in independent of the initialization. Hence, if the sampling times 𝑡 𝑖 = (𝑖Δ𝑡 )𝑖=1 ⊥ time (Δ𝑡 sufficiently large), the sequence of samples 𝜃 𝑡E𝑖 will be approximately i.i.d. samples of the limit distribution ∫ Δ𝑡 ⊥ ⊥ 𝜃 ΔE𝑡 = 𝑒 −𝐻E ⊥ (𝑡 −𝑠) 𝑃 E ⊥ d𝑊𝑠E , (42) 0

which is a Gaussian with mean zero and covariance matrix ∫ Δ𝑡 𝜎Δ𝑡 = 𝑒 −𝐻E ⊥ (𝑡 −𝑠) 𝑃 E ⊥ 𝑃𝑇E ⊥ 𝑒 −𝐻⊥ (𝑡 −𝑠) d𝑠 .

(43)

0

The same argument applies to 𝜃 𝑡E𝑖 . As an i.i.d mean of Gaussians, the stochastic weight average 𝜃 𝑡 also becomes Gaussian distributed for large 𝑇, with covariance 𝑇1 𝜎Δ𝑡 . Since the covariance of 𝜃 𝑡 does not scale with 𝑇, we therefore conclude avg.

𝑅(𝑇Δ𝑡 ) ∼ 𝑇 .

(separated samples)

(44)

Note that this result, while it conveys the main intuition behind why SWA improves equivariance, is not readily applicable to practical training runs. For one, it is hard to estimate a priori for how long we need to wait between taking samples in order for our approximations above to be valid. Longer wait times would furthermore increase training time. In the next section, we will therefore give a more detailed analysis that does not rely on the assumption of well-separated samples.

10

4.4

Dense samples in a classification setting

Let us now consider the complete opposite limit of taking samples well separated in time: taking them infinitesimally often. That is, make the approximation 𝑇

1 ∑︁ 1 𝜃𝑇 = 𝜃 𝑡𝑖 ≈ 𝑇 𝑖=1 𝑇

∫ 𝑇 𝜃 𝑡 d𝑡.

(45)

0

Using standard results of stochastic calculus, one sees that 𝜃 𝑇 also here becomes Gaussian distributed, with mean 𝑇, but with a more complicated expression for the covariance: ∫ 𝑇 ∫ 𝑇 ∫ min(𝑠,𝑡 ) 1 cov(𝜃 T ) = 2 𝑒 −𝐻 (𝑡 −𝑟 ) 𝑃𝑃𝑇 𝑒 −𝐻 (𝑠−𝑟 ) d𝑟 d𝑠 d𝑡. (46) 𝑇 0 0 0 To give a direct comparison between 𝐿 (𝜃 𝑡 ) and 𝐿 (𝜃 𝑡 ) is therefore harder than in the last section. It will depend highly on the interplay of the eigendecompositions of 𝐻 and 𝑃, and we will therefore not do it in full generality. We instead again specialize, this time to group-invariant classification problems. That is, the group acts invariantly on the class label 𝑦 ∈ {1, . . . , 𝑛𝑌 }. As usual in such a setting, we use our network outputs to define an a-posteriori probability distribution 𝑝 𝜃 (𝑦 | 𝑥) = softmax(N𝜃 (𝑥))

(47)

over the label space {1, . . . , 𝑛𝑌 }. This distribution could be interpreted as a vector in the Euclidean space R𝑛𝑌 . We train the network using a cross-entropy loss: ℓ 𝜃 (𝑥, 𝑦) = − log( 𝑝 𝜃 (𝑦|𝑥)).

(48)

We further restrict the analysis to the loss landscape near a well-fitted minimum. We claim that in this regime, the matrix Σ = 𝑃𝑃𝑇 is approximately equal to the Hessian 𝐻. By well-fitted, we mean that the Í𝑌 residual error 𝜀(𝑥) = 𝑝 𝜃 (𝑥) − 𝑒 𝑦 ( 𝑥 ) ∈ R𝑛𝑌 between the predicted distribution 𝑝 𝜃 (𝑥) = 𝑛𝑦=1 𝑝 𝜃 (𝑦 |𝑥)𝑒 𝑦 and the data distribution 𝑝 data (𝑦 |𝑥) = 𝑒 𝑦 ( 𝑥 ) is small. Here, 𝑒 𝑘 ∈ R𝑛𝑌 is the unit vector with components (𝑒 𝑘 )𝑖 = 𝛿𝑖𝑘 , and 𝑦(𝑥) is the true label class for the sample 𝑥. To formalize this approximation, let 𝐽 (𝑥) = ∇ 𝜃 N𝜃 (𝑥) be the Jacobian of the network. Using the chain rule, the gradient of the loss is ∫ ∫ ⊤ ∇ 𝜃 𝐿 (𝜃) = 𝐽 (𝑥) ( 𝑝 𝜃 (𝑥) − 𝑒 𝑦 ( 𝑥 ) )d𝜇 𝑋 = 𝐽 (𝑥) ⊤ 𝜀(𝑥)d𝜇 𝑋 . (49) 𝑉𝑋

𝑉𝑋

Consequently, the Hessian takes the form1 ∫ ∫ 2 ⊤ 𝐽 (𝑥) Λ(𝑥)𝐽 (𝑥)d𝜇 𝑋 + 𝐻 = ∇ 𝐿 (𝜃 ∗ ) = 𝑉𝑋

|

∇2𝜃 N𝜃 (𝑥)𝜀(𝑥)d𝜇 𝑋 ,

(50)

𝑉𝑋

{z

}

G

|

{z R

}

where we defined Λ(𝑥) = diag( 𝑝 𝜃 (𝑥)) − 𝑝 𝜃 (𝑥) 𝑝 𝜃 (𝑥) ⊤ . Note that the residual term R is linear in the small error 𝜀. Now consider the covariance matrix Σ. Applying (18) and the chain rule, we obtain ∫ ∇ 𝜃 log 𝑝 𝜃 (𝑦|𝑥) ∇ 𝜃 log 𝑝 𝜃 (𝑦|𝑥) ⊤ d𝜇 𝑋 Σ= ∫𝑉𝑋   = 𝐽 (𝑥) ⊤ 𝜀(𝑥)𝜀(𝑥) ⊤ 𝐽 (𝑥) d𝜇 𝑋 . (51) 𝑉𝑋 1 Here, we use the notation ∇2 N (𝑥) for the tensor ∇2 N (𝜃) 2 𝑖 𝑗 𝑘 = 𝜕 𝜃𝑖 𝜕 𝜃 𝑗 N 𝜃 (𝑥) 𝑘 , and denote by ∇ 𝜃 N 𝜃 (𝑥)𝜀(𝑥) the contraction 𝜃 𝜃 𝜃 with 𝜀 over the 𝑘-dimensional space.

11

By expanding the bracket 𝜀(𝑥)𝜀(𝑥) ⊤ = diag(𝑒 𝑦 ( 𝑥 ) ) − 2𝑒 𝑦 ( 𝑥 ) 𝑝 𝜃 (𝑥) ⊤ + 𝑝 𝜃 (𝑥) 𝑝 𝜃 (𝑥) ⊤ , we see that it equals Λ(𝑥) − Δ(𝑥), where Δ(𝑥) = diag(𝜀(𝑥)) − 2𝜀(𝑥) 𝑝 𝜃 (𝑥) ⊤ . Combining (50) and (51), the difference between the two matrices is ∫ ∑︁ 𝐻−Σ =R+ 𝐽 (𝑥) ⊤ Δ(𝑥)𝐽 (𝑥)d𝜇 𝑋 . (52) 𝑉𝑋

𝑦

Since both terms on the right-hand side are linear in 𝜀, we have 𝐻 ≈ Σ when 𝜀 ≪ 1. Furthermore, this structure allows us to formalize the relation between 𝐿 eq and 𝐿 ⊥ . Another direct outcome is that the two have a common eigenbasis, in which the dynamics decompose further into one-dimensional solutions ∫ 𝑡 √︁ 𝑖 −𝜆𝑖 𝑡 𝑖 𝜃𝑡 = 𝑒 (53) 𝜃0 + 𝑒 −𝜆𝑖 (𝑡 −𝑠) 𝜆 𝑖 d𝐵𝑖𝑠 . 0

Here, 𝐵𝑖𝑠 are independent,√ one-dimensional Brownian motions and −𝜆𝑖 are the eigenvalues of 𝐻 (the eigenvalues of 𝑃 are then 𝜆𝑖 ). This gives the following simplified form of the averaged process (45): ∫ 𝑇 ∫ 𝑇 ∫ 𝑇∫ 𝑡 √︁ 1 1 𝑖 −𝜆𝑖 𝑡 𝑖 ¯𝜃 𝑇𝑖 = 1 𝜃 𝑡 d𝑡 = 𝑒 𝜃 0 d𝑡 + 𝑒 −𝜆𝑖 (𝑡 −𝑠) 𝜆𝑖 d𝐵𝑖𝑠 d𝑡 (54) 𝑇 0 𝑇 0 𝑇 0 0 ∫ 𝑇 1 − 𝑒 −𝜆𝑖 𝑇 𝑖 1 = 𝜃0 + √ (1 − 𝑒 −𝜆𝑖 (𝑇 −𝑠) ) d𝐵𝑖𝑠 . (55) 𝜆𝑖 𝑇 𝜆𝑖 𝑇 0 We will now use (53) and (55) to derive an expression for 𝑅⊥ (𝑇) from (41). Using the Itô isometry, we obtain ∫ 𝑇 E[𝜃 𝑇𝑖 (𝐻𝜃 𝑇 ) 𝑖 ] = 𝜆𝑖 E[(𝜃 𝑇𝑖 ) 2 ] = 𝜆𝑖 𝑒 −2𝜆𝑖 𝑇 (𝜃 0𝑖 ) 2 + 𝜆2𝑖 𝑒 −2𝜆𝑖 (𝑇 −𝑠) d𝑠 (56) 0   1 −2𝜆𝑖 𝑇 −2𝜆𝑖 𝑇 𝑖 2 ) (57) = 𝜆𝑖 𝑒 (𝜃 0 ) + (1 − 𝑒 2 ∫ 𝑇 (1 − 𝑒 −𝜆𝑖 𝑇 ) 2 𝑖 2 1 𝑖 𝑖 𝑖 2 ¯ ¯ ¯ E[ 𝜃 𝑇 (𝐻 𝜃 𝑇 ) ] = 𝜆𝑖 E[( 𝜃 𝑇 ) ] = (𝜃 0 ) + 2 (1 − 𝑒 −2𝜆𝑖 (𝑇 −𝑠) ) 2 d𝑠 (58) 𝜆𝑖 𝑇 2 𝑇 0   1 1 − 𝑒 −2𝜆𝑖 𝑇 1 − 𝑒 −4𝜆𝑖 𝑇 (1 − 𝑒 −𝜆𝑖 𝑇 ) 2 𝑖 2 (𝜃 ) + − + . (59) = 0 𝑇 𝜆𝑖 𝑇 2 𝜆𝑖 𝑇 2 4𝜆 𝑖 𝑇 2 The behavior of the above expressions for large 𝑇 depends on whether 𝜆 𝑖 𝑇 is large or small. For 𝜆𝑖 𝑇 ≫ 1, we can approximate 𝑒 −𝑐𝜆𝑖 𝑇 ≈ 0, so that we get 𝜆𝑖 ∼ 𝜆𝑖 2   1 1 1 1 1 𝑖 2 E[ 𝜃¯𝑇𝑖 (𝐻 𝜃¯𝑇 ) 𝑖 ] ≈ (𝜃 ) + − + ∼ 𝑇 𝜆𝑖 𝑇 2 4𝜆 𝑖 𝑇 2 𝑇 𝜆𝑖 𝑇 2 0

E[𝜃 𝑇𝑖 (𝐻𝜃 𝑇 ) 𝑖 ] ≈

(60) (𝜆𝑖 𝑇 ≫ 1) .

(61)

For 𝜆𝑖 𝑇 ≪ 1, we instead have 1 − 𝑒 −𝑐𝜆𝑖 𝑇 ≈ 𝑐𝜆 𝑖 𝑇, so that E[ 𝜃¯𝑇𝑖 (𝐻 𝜃¯𝑇 ) 𝑖 ] ≈ 𝜆𝑖 ((1 − 2𝜆 𝑖 𝑇) (𝜃 0𝑖 ) 2 + 𝜆 𝑖 𝑇) ∼ 𝜆 𝑖 (𝜃 0𝑖 ) 2   1 2 1 𝑖 𝑖 𝑖 2 ¯ ¯ E[ 𝜃 𝑇 (𝐻 𝜃 𝑇 ) ] ≈ 𝜆𝑖 (𝜃 0 ) + − + ∼ 𝜆𝑖 (𝜃 0𝑖 ) 2 𝑇 𝑇 𝑇

(62) (𝜆𝑖 𝑇 ≪ 1) .

(63)

Summing over the above expressions over all dimensions recovers the expectation of the loss functions 𝐿(𝜃 𝑇 ) and 𝐿( 𝜃¯𝑇 ). For large 𝑇, we obtain ∑︁ ∑︁ ∑︁ E[𝐿 (𝜃 𝑇 )] = E[𝜃 𝑇𝑖 (𝐻𝜃 𝑇 ) 𝑖 ] ∼ 𝜆𝑖 (𝜃 0𝑖 ) 2 + 𝜆𝑖 = Tr(𝐻) + 𝜖𝑇 (64) E[𝐿 ( 𝜃¯𝑇 )] =

𝑖

𝜆𝑖 𝑇 ≪1

∑︁

∑︁

𝑖

E[ 𝜃¯𝑇𝑖 (𝐻 𝜃¯𝑇 ) 𝑖 ] ∼

𝜆𝑖 𝑇 ≫1

𝜆𝑖 (𝜃 0𝑖 ) 2 +

𝜆𝑖 𝑇 ≪1

12

∑︁ 1 1 = 𝑁 (1/𝑇) + 𝜖¯𝑇 , 𝑇 𝑇 𝜆 𝑇 ≫1 𝑖

(65)

where 𝑁 (1/𝑇) denotes the number of large eigenvalues 𝜆 𝑖 with 𝜆 𝑖 ≫ 1/𝑇 and |𝜖𝑇 |, | 𝜖¯𝑇 | ≤ 𝑇 → ∞. Hence, for large 𝑇, 𝑅(𝑇) = 𝑇 ·

Í

𝜆𝑖 𝑇 ≪1 𝜆 𝑖 → 0,

Tr(𝐻) + 𝜖𝑇 . 𝑁 (1/𝑇) + 𝜖¯𝑇

Analogously, defining 𝑁⊥ (1/𝑇) as the count of large eigenvalues in 𝐻 E ⊥ , we have ∑︁ ∑︁ ∑︁ 𝜆𝑖 (𝜃 0𝑖 ) 2 + 𝜆𝑖 = Tr(𝐻 E ⊥ ) + 𝜖𝑇 E[𝜃 𝑇𝑖 (𝐻𝜃 𝑇 ) 𝑖 ] ∼ E[𝐿 ⊥ (𝜃 𝑇 )] = 𝜆𝑖 𝑇 ≪1 𝑖

𝑖

E[𝐿 ⊥ ( 𝜃¯𝑇 )] =

∑︁

∑︁

E[ 𝜃¯𝑇𝑖 (𝐻 𝜃¯𝑇 ) 𝑖 ] ∼

𝑖 ∈ E⊥

(66)

(67)

𝜆𝑖 𝑇 ≫1 𝑖

𝜆𝑖 (𝜃 0𝑖 ) 2 +

𝜆𝑖 𝑇 ≪1 𝑖 ∈ E⊥

∑︁ 1 1 = 𝑁⊥ (1/𝑇) + 𝜖¯𝑇 , 𝑇 𝑇 𝜆 𝑇 ≫1

(68)

𝑖

𝑖 ∈ E⊥

Hence, for large 𝑇, 𝑅⊥ (𝑇) = 𝑇 ·

Tr(𝐻 E ⊥ ) + 𝜖𝑇 . 𝑁⊥ (1/𝑇) + 𝜖¯𝑇

(69)

These already shows that, even with dense samples, both 𝑅(𝑇) and 𝑅⊥ (𝑇) converges to infinity as 𝑇 → ∞. Hence, stochastic weight averaging does encourage equivariance and boost performance at the same time. This begs the question: How much of the equivariance gain from stochastic weight averaging simply comes from a generic improvement in overall performance? To measure this, we should study the ratio between 𝑅⊥ (𝑇) and 𝑅(𝑇), i.e. 𝑁 (1/𝑇) + 𝜖𝑇′ Tr(𝐻 E ⊥ ) + 𝜖¯𝑇 𝑅⊥ (𝑇) = . 𝑅(𝑇) 𝑁⊥ (1/𝑇) + 𝜖¯𝑇 Tr(𝐻) + 𝜖𝑇′

(70)

Asymptotically for 𝑇 → ∞, 𝑁 (1/𝑇) and 𝑁⊥ (1/𝑇) will count all dimensions in L and E ⊥ , respectively, and the 𝜖-terms converge to zero. Consequently, 𝑅⊥ (𝑇) dim L Tr(𝐻 E ⊥ ) = . 𝑇→∞ 𝑅(𝑇) dim E ⊥ Tr(𝐻) lim

(71)

In the next section, we will derive a bound on this ratio in the NTK-, or infinitely wide, limit [33].

4.5

A refined analysis in the NTK limit

We begin by specifying the theoretical setting. We consider a base architecture as above, i.e. with input space 𝑋, intermediate spaces 𝑉𝑙 and output space 𝑌 on which representations 𝜌 𝑋 , 𝜌𝑙 and 𝜌𝑌 act. We increase the network’s width by a factor 𝑁 by taking 𝑁 direct sums of the intermediate spaces, i.e. 𝑉𝑙 → 𝑉𝑙⊕ 𝑁 . The group then acts by the direct sum representation on the enlarged spaces, 𝜌𝑙 → 𝜌𝑙⊕ 𝑁 . In the following, we will consider the limit 𝑁 → ∞. As finite-dimensional representations of a finite group, the actions 𝜌𝑙 on the base intermediate spaces decompose into direct sums of irreducible representations (irreps) 𝜌𝜆 as [39] 𝜌𝑙 =

Ê

⊕𝑚𝜆𝑙

𝜌𝜆

,

(72)

𝜆

where 𝑚 𝑙𝜆 ∈ N0 denotes the multiplicity of the irrep 𝜆 and the sum runs over all irreps of 𝐺. In particular, Í 𝑛𝑙 = dim 𝑉𝑙 = dim 𝜌𝑙 = 𝜆 𝑚 𝑙𝜆 dim 𝜌𝜆 . The extended representations decompose in the same way, the only thing that changes are the multiplicities: 𝑚 𝑙𝜆 → 𝑁𝑚 𝑙𝜆 . 13

𝜌0

𝜌1

𝜌2

𝜌3

𝜌1

𝜌0

𝑉1

𝑋

𝑉1

...

𝑉𝐿

𝑌

.. .

𝑋

𝜌𝐿

...

𝑉𝐿

.. .

...

𝜌1 𝑉1

𝜌3

𝑌

𝜌𝐿

...

𝑉𝐿

Figure 1: Definition of the NTK limit we consider here. On the left is the ’base architecture’, on the right is the extended architecture by a factor 𝑁. Now, the base equivariant parameters 𝜃 ∈ E satisfy 𝑔𝑙+1𝑊 𝑙 = 𝑊 𝑙 𝑔𝑙 . That is, they belong to the space Hom(𝜌𝑙 , 𝜌𝑙+1 ). By Schur’s lemma, that space is Hom(𝜌𝑙 , 𝜌𝑙+1 )  {⊕𝜆 𝐴𝜆𝑙 ⊗ Iddim 𝜌𝜆 : 𝐴𝜆𝑙 ∈ R𝑚𝜆 ×𝑚𝜆 }  ⊕𝜆 R𝑚𝜆 ×𝑚𝜆 . 𝑙+1

𝑙

𝑙+1

𝑙

Hence, the dimension of the space of equivariant parameters at layer 𝑙 is ∑︁ 𝑙 𝑙 dim Ebase = 𝑚 𝑙+1 𝜆 𝑚𝜆 .

(73)

(74)

𝜆

The dimension of the space of all parameters at base layer 𝑙 is on the other hand given by ! ! ∑︁ ∑︁ 𝑙 𝑙+1 𝑙 dim Lbase = 𝑛𝑙+1 𝑛𝑙 = 𝑚 𝜆 dim 𝜌𝜆 𝑚 𝜆′ dim 𝜌𝜆′ . 𝜆

Since 𝐺 is a finite group, Cauchy–Schwarz:

Í

𝜆 (dim 𝜌𝜆 )

(75)

𝜆′

2 = |𝐺 | (see e.g. [39]) and we obtain the following bound by using

𝑙 dim Lbase ≤|𝐺 |

∑︁

! 12 2 (𝑚 𝑙+1 𝜆 )

∑︁

! 12 (𝑚 𝑙𝜆 ) 2

(76)

𝜆

𝜆

We now assume that the intermediate spaces all scale with the same rate in 𝑁 so that the ratios of their 𝑙 widths stay fixed. Formally, 𝑟 −1 ≤ 𝑚 𝑙+1 𝜆 /𝑚 𝜆 ≤ 𝑟 for all 𝜆 and 𝑙 > 1 for some 𝑟 ≥ 1. ! ∑︁ 𝑙 𝑙 dim Lbase ≤|𝐺 |𝑟 (𝑚 𝑙𝜆 ) 2 ≤ |𝐺 |𝑟 2 dim Ebase . (77) 𝜆 𝑙 = 𝑁 2 dim E 𝑙 𝑙 , this inequality remains valid Since dim Eext for 1 < 𝑙 < 𝐿 − 1, and similarly for dim Lext base as we take the limit. The first and last layers however need to be treated separately. This is simple though: We have

𝑛𝑋 𝑛𝑋 1 · (𝑁𝑛1 ) · (𝑁𝑛2 ) = · dim Lext 𝑁𝑛2 𝑁𝑛2 𝐿−1 𝐿−1 𝑛 𝑋 ∑︁ 𝑛 𝑋 |𝐺 |𝑟 2 ∑︁ 𝑛 𝑋 |𝐺 |𝑟 2 𝑙 𝑙 ≤ dim Lext ≤ dim Eext ≤ dim Eext 𝑁𝑛2 𝑙=1 𝑁𝑛2 𝑙=1 𝑁

0 dim Lext = 𝑛 𝑋 · (𝑁𝑛1 ) =

(78) (79)

and similarly 𝑛𝑌 𝑛𝑌 𝐿−1 · (𝑁𝑛 𝐿−1 ) · (𝑁𝑛 𝐿 ) = · dim Lext 𝑁𝑛 𝐿−1 𝑁𝑛 𝐿−1 𝐿−1 𝐿−1 𝑛𝑌 ∑︁ 𝑛 𝑋 |𝐺 |𝑟 2 ∑︁ 𝑛𝑌 |𝐺 |𝑟 2 𝑙 𝑙 ≤ dim Lext ≤ dim Eext ≤ dim Eext 𝑁𝑛 𝐿−1 𝑙=1 𝑁𝑛 𝐿−1 𝑙=1 𝑁

𝐿 dim Lext = (𝑁𝑛 𝐿 )𝑛𝑌 =

14

(80) (81)

We conclude that dim Eext 1 dim Eext 1 = Í 𝐿−1 ≥ ∼ −1 2 0 𝐿 dim Lext (1 + (𝑛 𝑋 + 𝑛𝑌 )𝑁 )|𝐺 |𝑟 |𝐺 |𝑟 2 𝑙=1 dim L𝑙 + dim Lext + dim Lext

(82)

for large 𝑁. Therefore, ⊥ dim Eext 1 ≤1− dim Lext |𝐺 |𝑟 2

⇒

dim Lext |𝐺 |𝑟 2 ≥ ⊥ dim Eext |𝐺 |𝑟 2 − 1

(83)

for large 𝑁, where 𝑟 is a constant that depends on the network structure but does not scale with 𝑁. 4.5.1

Ratio of the traces in the infinite-width limit

From here we consider the network training in the infinite-width limit. Namely, denoting the network width scaling as 𝑁, the learning rate is 𝑂 𝑁 (1), while the weights are initialized as 𝑊𝑖𝑙 𝑗 ∼ N (0, 𝑐2𝑙 ). Under this condition and mild assumptions on the nonlinearity, the pre-activations ℎ𝑙 (𝑥) converge to Gaussian processes with covariance kernel (84) Σ (𝑙) (𝑥, 𝑥 ′ ) = Σ (𝑥, 𝑥 ′ )Id 𝑁 𝑛𝑙 . Furthermore, the neural tangent kernel (NTK), defined as 𝐾 (𝑥, 𝑥 ′ ) = 𝐽 (𝑥)𝐽 (𝑥 ′ ) ⊤ ,

(85)

′ ′ converges in operator norm to 𝐾∞ (𝑥, ∫ 𝑥 ) = K (𝑥, 𝑥 ) 𝐼𝑌 for some symmetric positive semidefinite kernel K : R𝑛𝑋 × R𝑛𝑋 → R with Tr( K) = 𝑋 K (𝑥, 𝑥) 𝑝 data (𝑥)d𝑥 < ∞ [33].

The limiting kernel can be obtained recursively layer-by-layer. For an MLP with nonlinearities 𝜎𝑙 , the recursion relations are given by (0)

′

𝑐20

𝑥⊤𝑥′ , K (0) (𝑥, 𝑥 ′ ) = 0 𝑛𝑋   Σ (𝑙+1) (𝑥, 𝑥 ′ ) = 𝑐 2𝑙+1 EΣ (𝑙) 𝜎𝑙 (𝑢)𝜎𝑙 (𝑢 ′ ) ,   Σ¤ (𝑙+1) (𝑥, 𝑥 ′ ) = 𝑐 2𝑙+1 EΣ (𝑙) 𝜎𝑙′ (𝑢)𝜎𝑙′ (𝑢 ′ ) , K (𝑙+1) (𝑥, 𝑥 ′ ) = K (𝑙) (𝑥, 𝑥 ′ ) Σ¤ (𝑙+1) (𝑥, 𝑥 ′ ) + Σ (𝑙+1) (𝑥, 𝑥 ′ ), Σ

(𝑥, 𝑥 ) =

where, by EΣ (𝑙) we denote the expectation value with respect to    (𝑙) Σ (𝑥, 𝑥) Σ (𝑙) (𝑥, 𝑥 ′ ) ′ (𝑢, 𝑢 ) ∼ N 0, (𝑙) ′ . Σ (𝑥 , 𝑥) Σ (𝑙) (𝑥 ′ , 𝑥 ′ )

(86) (87) (88) (89)

(90)

We introduce the label-smoothed cross entropy loss, introduced by Szegedy et al. [41], to avoid the nondegenerate Hessian of the one-hot cross-entropy (the derivatives from the non ground-truth classes are 0 in this case). The dependency on the smoothing parameter 𝛼 will drop out at the end of our calculation. Lemma 4.5. Consider the label-smoothed cross entropy loss ∫ ∑︁ 1 L 𝛼 (𝜃) =(1 − 𝛼)𝐿 CrossEntropy (𝜃) + 𝛼 − log 𝑝 𝜃 (𝑦|𝑥)d𝑥 𝑛𝑌 𝑋 𝑦

(91)

with smoothing parameter 𝛼 ∈ (0, 1). Define 𝑦 𝑥𝛼 = (1 − 𝛼)𝑒 𝑦 +

𝛼 1, 𝑛𝑌

(92)

where 1 is the vector with all components one. Assume that ∫ the smoothed labels approximate the data 𝛼 well, i.e. 𝑝(𝑥) ≈ 𝑦 𝑥 and so 𝜖 (𝑦, 𝑥) ≈ 0 for all 𝑥, thus 𝐻 ≈ 𝑋 𝐽 ⊤ (𝑥)Λ(𝑥)𝐽 (𝑥)d𝑥. Then: Tr(𝐻) 𝛼(2 − 𝛼) (𝑛𝑌 − 1) = =: 𝜏𝛼 . 𝑚→∞ Tr( K) 𝑛𝑌 lim

15

(93)

Proof. By assumption, 2   2 ∥ 𝑝(𝑥) ∥ 2 ≈ ||𝑦 𝑥𝛼 || 2 = (1 − 𝛼) + 𝑛𝛼𝑌 + (𝑛𝑌 − 1) 𝑛𝛼𝑌

(94)

2

𝛼) = (1 − 𝛼) 2 + 2𝛼(1− + 𝑛𝛼𝑌 𝑛𝑌

= 1 − 2𝛼 + 𝛼2 +

2𝛼−2𝛼2 +𝛼2 𝑛𝑌

= 1 − 𝛼(2− 𝛼)𝑛𝑌(𝑛𝑌 −1) . Therefore Tr(Λ(𝑥)) = 1 − ∥ 𝑝(𝑥) ∥ 2 =

(95) (96) (97)

𝛼(2 − 𝛼) (𝑛𝑌 − 1) = 𝜏𝛼 . 𝑛𝑌

(98)

And ∫ Tr(𝐻) = ∫𝑋 =

Tr(𝐽 (𝑥) ⊤ Λ(𝑥)𝐽 (𝑥))d𝜇 𝑋

(99)

Tr(Λ(𝑥)𝐽 (𝑥)𝐽 (𝑥) ⊤ )d𝜇 𝑋

(100)

Tr(Λ(𝑥)𝐾 (𝑥, 𝑥))d𝜇 𝑋

(101)

K (𝑥, 𝑥) Tr(Λ(𝑥))d𝜇 𝑋

(102)

∫𝑋 = 𝑁 →∞

∫𝑋

−→

𝑋

=𝜏𝛼 Tr( K) We would now like to apply NTK theory to also estimate Tr(𝐻 E ) in the limit. To this end, we consider the equivariant subnetwork of the full network eq

N𝜃 (𝑥) = N𝑃E 𝜃 (𝑥).

(103)

Where 𝑃 E represents the orthogonal projection from L onto E with respect to the standard inner product on the parameter coordinates. It takes a simple calculation to argue that in points 𝜃 ∗ ∈ E, 𝐻 E is the Hessian of the network N eq . In order to repeat the above calculation for 𝐻 E , we hence need to study the NTK function K E of the network N eq , which is the focus of the next section. 4.5.2

NTK for an equivariant subnetwork

We begin with a detailed discussion on the irrep decomposition introduced in (73). Standard basis. We work in the setting discussed above: intermediate pre-activations 𝑥 𝑙 and postactivations ℎ𝑙 live in 𝑉𝑙  R𝑛𝑙 , carrying a representation 𝜌 𝑙 of 𝐺. The network’s actual coordinates — the ones in which the nonlinearity 𝜎 is applied component-wise, and in which the weight matrices 𝑊 𝑙 have i.i.d. Gaussian entries before any equivariance constraint — are the standard basis coordinates, which we write 𝑥 𝑙 ∈ R𝑛𝑙 . This standard basis is generally not adapted to 𝜌 𝑙 . Change of basis. position:

Fix an orthogonal matrix 𝑃𝑙 ∈ 𝑂 (𝑛𝑙 ) that block-diagonalizes 𝜌 𝑙 into its irrep decom𝑃𝑙

⊤

𝜌 𝑙 (𝑔) 𝑃𝑙 =

Ê

𝜌𝜆 (𝑔) ⊕𝑚𝜆 𝑙

𝜆

16

for all 𝑔 ∈ 𝐺,

(104)

É ⊕𝑚𝜆𝑙 so that 𝑉𝑙 = in the basis produced by 𝑃𝑙 . We call this the irrep-adapted basis, and define the 𝜆 𝜌𝜆 irrep coordinates of an arbitrary vector 𝑣 by
𝑣˜ 𝑙𝜆,𝑘,𝑎 = 𝑃𝑙 𝑣 𝑙 𝜆,𝑘,𝑎 , (105) where 𝜆 indexes the irreps, 𝑘 ∈ [𝑚 𝑙𝜆 ] indexes the multiplicities of each irrep 𝜌𝜆 , and 𝑎 ∈ [𝑑𝜆 ], 𝑑𝜆 = dim 𝜌𝜆 , indexes the basis inside the irrep 𝜌𝜆 itself. Every quantity subscripted (𝜆, 𝑘, 𝑎) below lives in this irrep-adapted basis. Since 𝑃𝑙 is orthogonal, this change of basis preserves inner products and hence the (isotropic) Gaussian law of the unconstrained weights. After enlarging by 𝑁 times, we acquire a new index 𝑗, referring to the 𝑁 copies of the base space R𝑛𝑙 . Note that after enlarging by 𝑁 times, 𝑃𝑙 is applied channel-by-channel — i.e. we have
𝑣˜ 𝑙𝜆, 𝑗,𝑘,𝑎 = 𝑃𝑙 𝑣 𝑙𝑗,· 𝜆,𝑘,𝑎 . (106) The nonlinearity on the full space. The activation function
𝜎 : R → R is applied entrywise only in the standard basis: 𝑥 𝑙+1 = 𝜎(ℎ𝑙 ) means (𝑥 𝑙+1 )𝑖 = 𝜎 (ℎ𝑙 )𝑖 for each standard coordinate 𝑖 ∈ [𝑛𝑙+1 ]. Passing to irrep coordinates via 𝑃𝑙 , we define
⊤
𝜍 : 𝑉 𝑙 → 𝑉 𝑙 , 𝑥 𝑙 ↦→ 𝑃𝑙 𝜎 𝑃𝑙 𝑥 𝑙 , (107) i.e. 𝜍 is 𝜎 conjugated by the change of basis. Because 𝑃𝑙 mixes standard coordinates across different (𝜆, 𝑘, 𝑎), 𝜍 is in general not a component-wise map on 𝑉𝑙 . However, by the block structure of 𝑃𝑙 across channels, 𝜍 is still defined channel-by-channel on each copy of the base space 𝑉 𝑙 , even after we extend it by 𝑁 times: 𝜍 acting on the full layer decomposes as 𝑁 independent copies of the same base map, so the 𝑙+1 ˜𝑙 post-activations 𝑥˜𝜆, 𝑗,𝑘,𝑎 = 𝜍 ( ℎ 𝑗,· )𝜆,𝑘,𝑎 depend only on channel 𝑗’s own pre-activation and never on any ′ other channel 𝑗 ≠ 𝑗. This is what licenses treating the 𝑁 channels as conditionally i.i.d. throughout. With this notation in place, the non-equivariant network’s irrep coordinates across consecutive layers are related through 𝑐 1 ∑︁ 1 ℎ˜ 0𝜆, 𝑗,𝑘,𝑎 = √ 𝐴𝜆𝜆′ , 𝑗,𝑘𝑘 ′ ,𝑎𝑎′ 𝑥˜𝜆0 ′ ,𝑘 ′ ,𝑎′ (108) 0 𝑛 𝜆′ ,𝑘 ′ ,𝑎′ ∑︁ 𝑐𝑙 𝑙 ℎ˜ 𝑙𝜆, 𝑗,𝑘,𝑎 = √ 𝐴𝜆𝜆 ˜𝜆𝑙 ′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ (109) ′ , 𝑗 𝑗 ′ ,𝑘𝑘 ′ ,𝑎𝑎 ′ 𝑥 𝑙 𝑛 𝑁 𝑗 ′ ,𝜆′ ,𝑘 ′ ,𝑎′ 𝑙+1 ˜𝑙 𝑥˜𝜆, 𝑗,𝑘,𝑎 = 𝜍 ( ℎ 𝑗,· )𝜆,𝑘,𝑎 ,

(110)

𝑙 𝑙 where 𝐴𝜆𝜆 ′ , 𝑗 𝑗 ′ ,𝑘𝑘 ′ ,𝑎𝑎 ′ are 𝑊 with a change of basis, hence are still standard Gaussian matrices (entries i.i.d. N (0, 1)).

Equivariant weights. For the equivariant subnetwork, the weights 𝐴𝑙 are further restricted to intertwiners. Since the equivariant weights are obtained by orthogonally projecting the isotropically Gaussian weights of the non-equivariant network onto the (Schur-lemma) intertwiner subspace, they are themselves isotropic Gaussians on the equivariant subspace: within each irrep 𝜌𝜆 , Schur’s lemma forces the weight to be shared across the 𝑑𝜆 copies indexed by 𝑎, so the projection reduces to a single scalar 𝑙 ′ ′ 𝐴𝜆, 𝑗 𝑗 ′ ,𝑘𝑘 ′ ∼ 𝑁 (0, 1) per (𝜆, 𝑗, 𝑗 , 𝑘, 𝑘 ), i.i.d. across all indices and layers. Relation (109) becomes, for the equivariant network, 𝑐 1 ∑︁ 0 0 𝐴𝜆, 𝑗,𝑘𝑘 ′ 𝑥˜𝜆,𝑘 (111) ℎ˜ 0𝜆, 𝑗,𝑘,𝑎 = √︁ ′ ,𝑎 , 0 ′ 𝑛 𝑑𝜆 𝑘 ∑︁ 𝑐𝑙 𝑙 𝑙 ℎ˜ 𝑙𝜆, 𝑗,𝑘,𝑎 = √︁ 𝐴𝜆, (112) 𝑗 𝑗 ′ ,𝑘𝑘 ′ 𝑥˜𝜆, 𝑗 ′ ,𝑘 ′ ,𝑎 , 𝑙 𝑛 𝑁 𝑑𝜆 𝑗 ′ ,𝑘 ′

17

iid

𝑙+1 𝑙 with 𝐴𝜆, 𝑗 𝑗 ′ ,𝑘𝑘 ′ ∼ N (0, 1). In particular, conditioned on the previous layer, the 𝑥˜𝜆, 𝑗,𝑘,𝑎 are identically distributed for same 𝜆, 𝑘 and 𝑎 but different 𝑗; and are independent for different 𝑗, 𝜆, and 𝑘, since they are built from disjoint sets of 𝛼’s.

Analogous to the standard MLP case, as 𝑁 → ∞, the pre-activations in irrep-adapted bases ℎ˜ 𝑙 converge to Gaussian processes with kernels   Σ𝑙E (𝑥, 𝑥 ′ ) (𝜆, 𝑗,𝑘,𝑎) (𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ ) = E ℎ˜ 𝑙𝜆, 𝑗,𝑘,𝑎 (𝑥) ℎ˜ 𝑙𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ (𝑥 ′ ) , (113) and the neural tangent kernel of the equivariant subnetwork converge to deterministic limits,  
⊤ 𝐾 𝑙E (𝑥, 𝑥 ′ ) (𝜆, 𝑗,𝑘,𝑎) (𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ ) = E ∇ ℎ˜ 𝑙𝜆, 𝑗,𝑘,𝑎 (𝑥) ∇ ℎ˜ 𝑙𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ (𝑥 ′ ) ,

(114)

It worths clarifying that, Σ𝑙E , 𝐾 𝑙E : R𝑛𝑋 × R𝑛𝑋 → R 𝑁 𝑛𝑙+1 × 𝑁 𝑛𝑙+1 are the matrix-valued equivariant NNGP kernel and NTK. Their ambient dimension 𝑁𝑛𝑙+1 × 𝑁𝑛𝑙+1 grows with 𝑁, so a priori “convergence as 𝑁 → ∞” is not meaningful, since the kernels for different 𝑁 live in different spaces. Lemma 4.6 resolves this. It shows that Σ𝑙E and 𝐾 𝑙E are, for every 𝑁, entirely determined by a fixed collection of 𝑑𝜆 × 𝑑𝜆 blocks Σ𝑙𝜆 , K𝜆𝑙 , which are independent of 𝑁 and 𝑗, 𝑘. It is these 𝑁-independent blocks that converge, and (113) and (114) should be read as shorthand for that block-wise statement. Lemma 4.6. The kernels Σ𝑙E and 𝐾 𝑙E admits the following block structure,

Σ𝑙E (𝜆, 𝑗,𝑘,𝑎) , (𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ ) (𝑥, 𝑥 ′ ) = 𝛿𝜆𝜆′ 𝛿 𝑗 𝑗 ′ 𝛿 𝑘𝑘 ′ Σ𝑙𝜆 𝑎𝑎′ (𝑥, 𝑥 ′ )

𝐾 𝑙E (𝜆, 𝑗,𝑘,𝑎) , (𝜆′ , 𝑗 ′ ,𝑘 ′ ,𝑎′ ) (𝑥, 𝑥 ′ ) = 𝛿𝜆𝜆′ 𝛿 𝑗 𝑗 ′ 𝛿 𝑘𝑘 ′ K𝜆𝑙 𝑎𝑎′ (𝑥, 𝑥 ′ ),

(115) (116)

where each Σ𝑙𝜆 and K𝜆𝑙 are 𝑑𝜆 -dimensional matrices. Moreover, they have the following recursive formulas, Σ0𝜆 𝑎𝑎′ (𝑥, 𝑥 ′ ) =

′ Σ𝑙+1 𝜆 𝑎𝑎′ (𝑥, 𝑥 ) =

𝑐20 ∑︁ 𝑛0 𝑑 𝜆

𝑘¯ 𝑙+1 2 𝑐 𝑙+1 𝑚 𝜆

𝑛𝑙+1 𝑑𝜆 0 ′ K𝜆 (𝑥, 𝑥 ) = Σ0𝜆 (𝑥, 𝑥 ′ ),

′ 𝑥˜𝜆, 𝑘,𝑎 ¯ 𝑥˜𝜆, 𝑘,𝑎 ¯ ′,

(117)

  E𝑢,𝑢′ ∼Σ𝑙 ( 𝑥, 𝑥 ′ ) 𝜍 (𝑢) 𝑎 𝜍 (𝑢 ′ ) 𝑎′ ,

(118)

𝜆

(119)

′ K𝜆𝑙+1 𝑎𝑎′ (𝑥, 𝑥 ′ ) = Σ𝑙+1 𝜆 𝑎𝑎′ (𝑥, 𝑥 ) +

𝑐2𝑙+1 𝑚 𝑙+1 𝜆

  E𝑢,𝑢′ ∼Σ𝑙 ( 𝑥, 𝑥 ′ ) 𝜍 ′ (𝑢) 𝑎 𝜍 ′ (𝑢 ′ ) 𝑎′ K𝜆𝑙 (𝑥, 𝑥 ′ ),

𝜆 𝑛𝑙+1 𝑑𝜆    𝑙 Σ𝜆 (𝑥, 𝑥) Σ𝑙𝜆 (𝑥, 𝑥 ′ ) 𝑙 ′ ′ ′ where 𝑢, 𝑢 ∼ Σ𝜆 (𝑥, 𝑥 ) means (𝑢, 𝑢 ) ∼ N 0, 𝑙 ′ . Σ𝜆 (𝑥 , 𝑥) Σ𝑙𝜆 (𝑥 ′ , 𝑥 ′ )

(120)

Proof. For 𝑙 = 0: by (112) and independence of 𝐴0 from 𝑥˜ 0 , 𝑐2 ∑︁  0 𝑐20 ∑︁    ′ 0 ′ ′ ,𝑎 𝑥 E ℎ˜ 0𝜆, 𝑗,𝑘,𝑎 (𝑥) ℎ˜ 0𝜆, 𝑗,𝑘,𝑎′ (𝑥 ′ ) = 0 0 E 𝐴𝜆, 𝑗,𝑘𝑘 ′ 𝐴𝜆, 𝑥 ˜ ˜ = 𝑥˜𝜆,𝑘 ′ ,𝑎 𝑥˜𝜆,𝑘 ′ ,𝑎 ′ , 𝜆,𝑘 ′ ′ ′ 𝜆, 𝑘¯ ,𝑎 𝑗,𝑘 𝑘¯ 𝑛 𝑑𝜆 ′ ¯ ′ 𝑛0 𝑑 𝜆 𝑘 ′ 𝑘 ,𝑘

(121) 0 0 using E[ 𝐴𝜆, 𝑗,𝑘𝑘 ′ 𝐴𝜆, 𝑗,𝑘 𝑘¯ ′ ] = 𝛿 𝑘 ′ 𝑘¯ ′ . For the induction step, condition on layer 𝑙; by (112),

  ˜ 𝑙+1 ′ (𝑥 ′ ) = E ℎ˜ 𝑙+1 ℎ 𝜆, 𝑗,𝑘,𝑎 𝜆, 𝑗,𝑘,𝑎

𝑐2𝑙+1 𝑁𝑛𝑙+1 𝑑𝜆

∑︁ ∑︁ 𝑗′

𝑙+1 𝑙+1 ′ 𝑥˜𝜆, 𝑗 ′ ,𝑘 ′ ,𝑎 𝑥˜𝜆, 𝑗 ′ ,𝑘 ′ ,𝑎′ (𝑥 ),

(122)

𝑘′

′ 𝑙 𝑙 𝑙 𝑙 ′ after using E[ 𝐴𝜆, 𝑗 𝑗 ′ ,𝑘𝑘 ′ 𝐴𝜆, 𝑗 𝑗¯′ ,𝑘 𝑘¯ ′ ] = 𝛿 𝑗 ′ 𝑗¯′ 𝛿 𝑘 ′ 𝑘¯ ′ . Since the pairs 𝑥˜𝜆, 𝑗 ′ ,𝑘 ′ ,· (𝑥), 𝑥˜𝜆, 𝑗 ′ ,𝑘 ′ ,· (𝑥 ) , 𝑗 ∈ [𝑁], are i.i.d. copies of (𝜍 (𝑢), 𝜍 (𝑢 ′ )) with 𝑢, 𝑢 ′ ∼ Σ𝑙𝜆 (𝑥, 𝑥 ′ ); the average over 𝑗 ′ converges almost surely to Σ𝑙+1 𝜆 18

by the law of large numbers as 𝑁 → ∞, and the sum over 𝑘 ′ collapses into 𝑚 𝑙𝜆 counts of the single expectation, since the 𝑘 ′ -summands are identically distributed. For the NTK, 𝐾 E0 = Σ0E at layer 0, so the block structure is automatically satisfied. For the induction step, split the tangent vector into the last-layer and earlier-layer parts, ′ ˜ 𝑙+1 ∇ 𝐴 ℎ˜ 𝑙+1 𝜆, 𝑗,𝑘,𝑎 (𝑥), ∇ 𝐴 ℎ𝜆, 𝑗,𝑘,𝑎′ (𝑥 )

(123)

𝑙+1 ′ The first term is computed exactly as for Σ𝑙+1 𝑎𝑎′ (𝑥, 𝑥 ). For the second, 𝜆 above and converges to Σ𝜆 the chain rule applied to (112) gives ∑︁ 𝑐 𝑙+1 𝜕𝐴<𝑙 ℎ˜ 𝑙+1 = 𝐴𝑙+1 ′ ′ 𝜍 ′ ( ℎ˜ 𝑙𝑗 ′ )𝜆,𝑘 ′ ,𝑎 𝜕𝐴<𝑙 ℎ˜ 𝑙𝜆, 𝑗 ′ ,𝑘 ′ ,𝑎 , (124) √ 𝜆, 𝑗,𝑘,𝑎 𝑁𝑛𝑙+1 𝑑𝜆 𝑗 ′ ,𝑘 ′ 𝜆, 𝑗 𝑗 ,𝑘𝑘 so that the second term, analogous to the NNGP kernel, converges a.s. by the induction hypothesis as 𝑁 → ∞ to 𝑐2𝑙+1 𝑚 𝑙+1   𝜆 (125) E𝑢,𝑢′ ∼Σ𝑙 ( 𝑥, 𝑥 ′ ) 𝜍 ′ (𝑢) 𝑎 𝐾 𝑙E (𝑥, 𝑥 ′ ) 𝜍 ′ (𝑢 ′ ) 𝑎′ . E 𝑛𝑙+1 𝑑𝜆 Furthermore, since K𝑙 is fixed, we could extract it from the expectation, which results in (120). Remark 4.7. In particular, the NNGP kernels and NTKs we calculate here can be interpreted as the NTK for a layerwise equivariant architecture. Such a kernel has previously only been calculated in the case when all intermediate representations are regular representations of a finite group [38]. We extend these results to a much more general setting. 𝑚𝑙

Note that there is an extra term 𝑛𝑙 𝑑𝜆𝜆 for the recursive formula. These terms accumulate through layers and result in a bound of trace against that of the ambient kernels Σ𝑙 , 𝐾 𝑙 . Lemma 4.8. Assume 𝜎 is the elementwise ReLU. Define 𝑟 𝑙 = max 𝜆

𝑚 𝑙+1 𝜆 (126)

, 𝑑𝜆 𝑛𝑙+1

Then for all 𝑙 ≥ 0, Tr Σ𝑙E (𝑥, 𝑥) ≤

𝑙 Ö

! 𝑟

𝑙′

Tr Σ𝑙 (𝑥, 𝑥),

Tr 𝐾 𝑙E (𝑥, 𝑥) ≤

𝑙 ′ =0

𝑙 Ö

! 𝑟

𝑙′

Tr 𝐾 𝑙 (𝑥, 𝑥),

(127)

𝑙 ′ =0

Proof. GP kernel. By Lemma 4.6, summing the trace of Σ0E over the 𝑁𝑚 𝜆0 copies of each irrep, we have, Tr Σ0E (𝑥, 𝑥) =

𝑐20 ∑︁ 𝑁𝑚 𝜆1 ∑︁ 2 𝑁𝑛1 𝑥˜𝜆,𝑘 ′ ,𝑎 ≤ 𝑐20𝑟 0 ∥ 𝑥∥ ˜ 2 = 𝑟 0 𝑁𝑛1 Σ0 (𝑥, 𝑥) = 𝑟 0 Tr Σ0 (𝑥, 𝑥). 0 𝑛0 𝑛 𝜆 𝑑𝜆 𝑘 ′ ,𝑎

(128)

The second to last equality is because changing of basis preserves norm, and the last equality is because Σ0 (𝑥, 𝑥 ′ ) is a 𝑁𝑛1 × 𝑁𝑛1 dimensional matrix. For the induction step, homogeneity of 𝜎 (ReLU) gives ˆ ∥2] E𝑢∼Σ𝑙 [∥𝜍 (𝑢) ∥ 2 ] =E𝑢∼Σ𝑙 [∥𝜎( 𝑃𝑢) 𝜆

(129)

𝐸


⊤ =E 𝜒∼𝑁 (0,1) [𝜎( 𝜒) 2 ] Tr 𝑃𝑙 Σ𝑙𝜆 (𝑥, 𝑥) 𝑃𝑙

(130)

=E 𝜒∼𝑁 (0,1) [𝜎( 𝜒) 2 ] Tr Σ𝑙𝜆 (𝑥, 𝑥),

(131)

2

2

𝑙

E𝑢∼Σ𝑙 [𝜎(𝑢) ] =𝐸 𝜒∼𝑁 (0,1) [𝜎( 𝜒) ] Σ (𝑥, 𝑥). 19

(132)

For notational convenience, we denote E 𝜒∼𝑁 (0,1) [𝜎( 𝜒) 2 ] as 𝑘 𝜎 , so that Tr Σ𝑙+1 (𝑥, 𝑥) =𝑛𝑙+2 Σ𝑙+1 (𝑥, 𝑥) = 𝑛𝑙+2 𝑐2𝑙+1 E𝑢∼Σ𝑙 [𝜎(𝑢) 2 ] =𝑛𝑙+2 𝑐2𝑙+1 𝑘 𝜎 Σ𝑙 (𝑥, 𝑥) =

𝑛𝑙+2 𝑐2𝑙+1 𝑛𝑙+1

(133)

𝑘 𝜎 Tr Σ𝑙 (𝑥, 𝑥).

(134)

By induction, we have 𝑙+2

Tr Σ𝑙+1 E (𝑥, 𝑥) =

𝜆 ∑︁ ∑︁ 𝑐2 𝑚 𝑙+1 𝑚 𝑙+1 𝜆

𝜆, 𝑗

≤𝑟 𝑙+1

𝑛𝑙+1 𝑑𝜆

E𝑢∼Σ𝑙 [∥𝜍 (𝑢) ∥ 2 ] =

∑︁ 𝑁𝑐2 𝑚 𝑙+2 𝑙+1 𝜆

𝜆

𝜆

𝑘=1

𝑛𝑙+2 𝑐2𝑙+1 ∑︁

𝑛𝑙+1 𝑑𝜆

𝑙 𝑙+1 𝑁𝑚 𝑙+1 𝜆 Tr Σ𝜆 (𝑥, 𝑥)𝑘 𝜎 = 𝑟

2 𝑚 𝑙+1 𝜆 E𝑢∼Σ𝑙 [∥𝜍 (𝑢) ∥ ] 𝜆

(135)

𝑛𝑙+2 𝑐2𝑙+1

𝑘 𝜎 Tr Σ𝑙E (𝑥, 𝑥) 𝑛𝑙+1 𝑛 𝑙+1 𝜆 ! ! 𝑙+1 𝑙+1 2 Ö ′ 𝑐 𝑛𝑙+2 Ö 𝑙+1 𝑙 𝑙′ 𝑙 𝑟 ≤ 𝑟 Tr Σ𝑙+1 (𝑥, 𝑥) 𝑘 𝜎 Tr Σ (𝑥, 𝑥) = 𝑛 𝑙+1 𝑙 ′ =0 𝑙 ′ =0

(136)

(137)

NTK. For ReLU, 𝜎 ′ is the indicator of R>0 , so for 𝑢 ∼ N (0, Σ𝑙E (𝑥, 𝑥)), E[𝜎 ′ (𝑢) 2 ] = P(𝑢 ≥ 0) = 21 , independent of the variance of 𝑢. A calculation analogous to the GP bound gives the desired bound for NTK. Since 𝐺 acts trivially on the output layer, 𝜌𝜆 = 𝜌triv there, 𝑑𝜆 = 1, and only 𝑎 = 𝑎 ′ occurs in (120). Hence K𝜆𝐿 (𝑥, 𝑥 ′ ) is a scalar function, and the matrix-valued 𝐾 E𝐿 (𝑥, 𝑥 ′ ) collapses to K𝜆𝐿 (𝑥, 𝑥 ′ ) times the identity Id𝑛𝑌 , and the bound applies directly to the scalar kernel. Corollary 4.9. 𝐾 E𝐿 (𝑥, 𝑥 ′ ) = K 𝐿E (𝑥, 𝑥 ′ ) 𝐼𝑛𝑌 , and K 𝐿E (𝑥, 𝑥) ≤

𝐿 Ö  𝑟 𝑙 K 𝐿 (𝑥, 𝑥),

(138)

𝑙=1

Note that for the output layer, 𝑟 𝐿+1 = 1, hence do not contribute to the coefficient. Combining Lemma 4.8 with the identity relating tr 𝐻 E to 𝐾 E𝐿 used in the proof of Lemma 4.5 then gives, exactly as before, Lemma 4.10. Under the assumptions above, for the cross-entropy loss and near an equivariant minimizer, Î𝐿 𝑙 the spectral proportion of 𝐻 E relative to 𝐻 is at most 𝑙=1 𝑟 , with 𝑟 𝑙 = max𝜆 𝑚 𝑙𝜆 /(𝑑𝜆 𝑛𝑙 ).

4.6

Main theoretical result

Plugging the results of Lemma 4.10 and (83) into (71) leads to our main theorem. For convenience we list here the assumptions we made during its derivation: • We average over long times, i.e. 𝑇 large • The neural networks are wide, i.e. 𝑁 large • The group 𝐺 is finite • The model is well trained, i.e. the error 𝜖 is small • The non-linearity is the ReLU function. 20

Under these conditions, we have Theorem 4.11. The asymptotic relative ratio of improvement for the equivariance of the solution satisfies the bound ! 𝐿 Ö 𝑚 𝑙𝜆 𝑅⊥ (𝑇) |𝐺 |𝑟 2 lim ≥ 1− max , (139) 𝑇→∞ 𝑅(𝑇) 𝜆 𝑑 𝜆 𝑛𝑙 |𝐺 |𝑟 2 − 1 𝑙=1 In particular, the asymptotic relative ratio is bigger than one if 𝐿 Ö 𝑚 𝑙𝜆 1 > max . |𝐺 |𝑟 2 𝑙=1 𝜆 𝑑𝜆 𝑛𝑙

(140)

As an important special case, consider rectangular networks, i.e. 𝑚 𝑙𝜆 = 𝑚 𝜆 and 𝑛𝑙 = 𝑛 for all 𝜆 and 𝐿 − 1 ≥ 𝑙 > 1. In this case, Theorem 4.11 simplifies to Corollary 4.12. For rectangular networks, the asymptotic relative ratio of improvement for the equivariance of the solution satisfies the bound  𝐿−1 !  𝑅⊥ (𝑇) |𝐺 |𝑟 2 𝑚𝜆 lim ≥ 1 − max , (141) 𝑇→∞ 𝑅(𝑇) 𝜆 𝑑𝜆 𝑛 |𝐺 |𝑟 2 − 1 which is larger than one if  𝐿 1 𝑚𝜆 > max 𝜆 𝑑𝜆 𝑛 |𝐺 |𝑟 2

(142)

Remark 4.13. The bound in Theorem 4.11 explicitly depends on the representation 𝜌¯ on H via the multiplicities 𝑚 𝑙𝜆 . This may seem unintuitive since there is considerable freedom in the choice of the 𝜌. ¯ However, the bounded improvement ratio also depends on 𝜌¯ via the equivariant minimum around which the loss is expanded. Hence, the equivariant minimum to which the network converged at the end of training selects the 𝑚 𝑙𝜆 in (139). Note the significance of the asymptotic ratio being greater than one – in this case 𝑅⊥ > 𝑅 in the limit of 𝑇 → ∞ and hence the expected equivariance improvement due to SWA (as measured by 𝑅⊥ ) outpaces the expected equivariance improvement due to SWA (as measured by 𝑅), under data augmentation. Some remarks are in order. First, note that for a given architecture, the condition (140) is harder to satisfy if the group is larger. This aligns with the intuition that tasks with larger symmetries are harder to learn, see for instance the discussion in [42]. Secondly, for any intermediate representation, we have ! ∑︁ 𝑚 𝑙𝜆 𝑚𝑙 1 𝑙 𝑙 𝑙 2 𝑛 = 𝑚 𝜆 𝑑𝜆 ≥ max 𝑚 𝜆 𝑑𝜆 ≥ (inf 𝑑𝜆 ) max ⇒ max 𝑙 𝜆 ≤ ≤ 1. (143) 2 𝜆 𝜆 𝜆 𝑑𝜆 𝜆 𝑛 𝑑𝜆 inf 𝑑 𝜆 𝜆 𝜆 Hence, as long as sufficiently many intermediate representations are either mixtures of irreps or irreps of non-trivial dimension, the condition (140) will be satisfied. This aligns with the observation that in a manifestly symmetric network, using equivariant features up until the final layer is beneficial even if the 𝑚𝑙

task considered is invariant. Finally, we note that in in many cases, max𝜆 𝑛𝑙 𝑑𝜆 is significantly smaller than 𝜆 1. In particular, if an intermediate representation is regular (as in a group-CNN [10]), we have 𝑚 𝑙𝜆 = 𝑑𝜆 𝑚𝑙

for all 𝜆 [39, Cor. 2.18], so that max𝜆 𝑛𝑙 𝑑𝜆 = 𝑛1𝑙 = |𝐺1 | for that 𝑙. 𝜆

4.7

A proxy for the orthogonal loss

Before moving on to the experiments, we give an auxiliary result that justifies the equivariance metric we use in practice. 21

Recall that we restrict the scope to group invariant classification tasks. Under this setting, a natural metric for the equivariance of the network is the Kullback-Leibler divergence between the network output and its group-averaged version: ! 1 ∑︁ eq (144) 𝑙 𝜃 (𝑥) =𝐷 KL 𝑝 𝜃 (𝑔 𝑋 𝑥) 𝑝 𝜃 (𝑥) . |𝐺 | 𝑔 ∫ eq 𝐿 eq (𝜃) = 𝑙 𝜃 (𝑥)d𝜇 𝑋 (145) 𝑉𝑋

We now show that it is an accurate proxy for the orthogonal loss under our assumptions. Proposition 4.14. Under the assumption that 𝜀(𝑥) = 𝑝 𝜃 (𝑥) − 𝑒 𝑦 ( 𝑥 ) is small, for 𝜃 in a neighbourhood of the equivariant minimizer 𝜃 ∗ = 0,     𝐿 eq (𝜃) = 𝐿 ⊥ (𝜃) + 𝑂 ∥𝜃 ∥ 3 + 𝑂 𝜀∥𝜃 ∥ 2 . (146) In particular, the equivariant component 𝜃 E does not contribute at this order. Proof. Let 𝐽 (𝑥) = ∇ 𝜃 𝑁 𝜃 (𝑥) 𝜃=0 be the Jacobian at the minimizer. Linearizing the network output gives 𝑁 𝜃 (𝑥) = 𝑁0 (𝑥) + 𝐽 (𝑥)𝜃 + 𝑂 (∥𝜃 ∥ 2 ). Since the minimizer 𝜃 ∗ = 0 is equivariant, its Jacobian satisfies 𝐽 (𝑔 𝑋 𝑥) = 𝐽 (𝑥) 𝑔¯ ⊤ . Therefore, the deviation of the perturbed logits from their group average is entirely ⊥ driven by the orthogonal component 𝜃 E :   ⊥ 𝑢 𝑔 (𝑥) = 𝑁 𝜃 (𝑔 𝑋 𝑥) − 𝑁¯ 𝜃 (𝑥) = 𝐽 (𝑥) 𝑔¯ ⊤ 𝜃 E + 𝑂 ∥𝜃 ∥ 2 . (147) Next, we expand the Kullback–Leibler divergence to second order,  
𝐷 KL 𝑝 𝜃 (· | 𝑔 𝑋 𝑥) ∥ 𝑝¯ 𝜃 (· | 𝑥) = 21 𝑢 𝑔 (𝑥) ⊤ Λ(𝑥) 𝑢 𝑔 (𝑥) + 𝑂 ∥𝑢 𝑔 (𝑥) ∥ 3  
⊤ = 21 𝜃 ⊥ 𝑔¯ 𝐽 (𝑥) ⊤ Λ(𝑥)𝐽 (𝑥) 𝑔¯ ⊤ 𝜃 ⊥ + 𝑂 ∥𝜃 ∥ 3 . Taking the expectation over 𝑥 ∼ 𝜇 𝑋 and averaging over 𝑔 ∈ 𝐺 yields 𝐿 eq (𝜃).  
⊤ 𝐿 eq (𝜃) = 12 𝜃 ⊥ G 𝜃 ⊥ + 𝑂 ∥𝜃 ∥ 3 .

(148)

(149)

Where G is the Gauss-Newton matrix that appears in (50). Since the full Hessian is 𝐻 = G + 𝑅 with the residual 𝑅 being linear in 𝜀, we can replace G with 𝐻 at the cost of an 𝑂 (𝜀∥𝜃 ∥ 2 ) error. Due to the block-diagonality of 𝐻 with respect to E ⊕ E ⊥ , we can restrict 𝐻 to 𝐻 E ⊥ , concluding that        
⊤ 𝐿 eq (𝜃) = 12 𝜃 ⊥ 𝐻⊥ 𝜃 ⊥ + 𝑂 ∥𝜃 ∥ 3 + 𝑂 𝜀∥𝜃 ∥ 2 = 𝐿 ⊥ (𝜃) + 𝑂 ∥𝜃 ∥ 3 + 𝑂 𝜀∥𝜃 ∥ 2 . (150) This completes the proof. Corollary 4.15. The ratio of equivariance losses evaluated on unaveraged and averaged weights satisfies   E 𝐿 eq (𝜃 𝑇 )
  = 𝑅⊥ (𝑇) 1 + 𝑜(1) (151) ¯ E 𝐿 eq ( 𝜃 𝑇 ) in the regime in which Proposition 4.14 applies. The arguments above also apply to the following variant. Replacing the orbit average of the logits by Í the orbit average of the predictive distributions, 𝑝¯ 𝜃 (𝑥) = |𝐺1 | 𝑔 𝑝 𝜃 (𝑔 𝑋 𝑥), changes 𝑝¯ 𝜃 by 𝑂 (∥𝜃 ∥ 2 ) and hence 𝐿 eq by 𝑂 (∥𝜃 ∥ 3 ). 22

Remark 4.16. From a practical standpoint, the orthogonal loss 𝐿 ⊥ is a theoretical construct: computing it explicitly requires knowledge of the Hessian and the orthogonal projection in the parameter space, which is typically intractable for deep networks. In contrast, 𝐿 eq is an easily computable, well-motivated metric that directly measures the equivariance of the network’s predictions. Proposition 4.14 essentially bridges this gap, establishing 𝐿 eq as a proxy for 𝐿 ⊥ near a well-fitted minimum. Consequently, any theoretical predictions or dynamics derived for 𝐿 ⊥ , in particular, the ratio 𝑅⊥ , can directly transit to 𝐿 eq . This justifies the use of 𝐿 eq in our experiments to evaluate the learned equivariance of the network, with the confidence that this empirical metric follows the theoretical behavior of the orthogonal loss 𝐿 ⊥ .

5

Experiments

5.1

Synthetic experiments

In this section, we conduct synthetic experiments to validate our theoretical settings, and verify our theoretical results. Throughout the section, we work with multi-layer perceptrons (MLPs) since they have completely unrestricted parameter space, and our theory applies without modification. 5.1.1

Empirical validation of the quadratic approximation

Before applying our theoretical framework to weight averaging, we first conduct an experiment to validate the quadratic approximation of the loss landscape near a well-fitted minimum. We employ a teacherstudent setup where both models are MLPs with identical architectures, and the input data is sampled from a standard Gaussian distribution. In this controlled setting, the global minimum is explicitly known, allowing us to compute the exact empirical Hessian at the minimizer. To simulate the dynamics of late-stage training, we perturb the parameters away from the minimum using two sampling strategies: (1) stochastic gradient ascent, which explicitly pushes the parameters along the noise distribution of SGD, and (2) Gaussian diffusion, which adds naive, homogeneous isotropic noise to the weights. We then perform standard stochastic gradient descent from these perturbed points, tracking both the true loss and its theoretically estimated quadratic approximation. We first analyze the spectral properties of the Hessian. As illustrated in Figure 2a, the effective rank of the Hessian monotonically increases with the number of data samples used for its computation, eventually plateauing at a sample size of 45. Furthermore, the resulting matrix is nearly full rank; the difference is primarily due to the positive homogeneity of the ReLU activation functions, which causes scale invariances in the parameter space. Despite this slight rank deficiency, the quadratic approximation remains accurate under the training dynamics. Figure 2b compares the true loss against the quadratic estimation for the optimization trajectory originating from the gradient ascent perturbations. The approximation aligns closely with the full loss, and the approximation error strictly decreases and stabilizes as the parameters converge back toward the minimizer. Crucially, this high accuracy holds because the parameter trajectory is not exploring the landscape arbitrarily. As is shown with the calculation, the covariance matrix of the SGD noise is asymptotically the Hessian, causing the optimization path to be heavily biased toward the top eigenspaces where the Hessian dominates the local geometry. By naturally avoiding flat or highly non-linear directions, the SGD trajectory remains confined to a region where the loss is strictly dominated by the quadratic term, rendering our approximation highly accurate. Similar results from the Gaussian diffusion sampling method can be found in Section X of the Appendix, which demonstrates the robustness and ability of self-correction of SGD.

23

10−1

loss

rank

10−3

Approximation error

mean actual loss mean quadratic estimate

4000

2000 Hessian rank Hessian dimension

10−5 10−7

0 0

50 samples

100

0

1000 SGD step

(a)

2000

10−1 10−2 10−3 10−4 0 −10−4 −10−3

0.1 0.2 distance from minimizer

(b)

Figure 2: Empirical validation of the quadratic approximation in a teacher-student MLP setting. (a) Effective rank of the empirical Hessian evaluated at the global minimum, plotted against the number of data samples used. (b) Evaluation of the quadratic approximation along an optimization trajectory (sampled via gradient ascent). Left: The true loss strictly follows the quadratic estimation. Right: The approximation error between full- and quadratic loss stabilizes and strictly decreases as the parameter distance to the minimizer diminishes. 5.1.2

Verification of the theorem

With the quadratic approximation for the loss function justified, we now extend the teacher-student setup of Section 5.1.1, and verify our main theorem (Theorem 4.11). Setup. In this experiment, the teacher and student share the same expressivity in terms of equivariant functions, but they are no longer identical: the teacher is a multilayer perceptron whose hidden layers carry the regular representation of 𝐺 = 𝐶𝑛 , so that the network is exactly 𝐺-equivariant by construction. On the other hand, the student is an ordinary MLP of matching width and depth, with no constraint tying its weights together. Working with MLPs rather than convolutional architectures is deliberate: the student’s parameter space is then the full space with no affine restriction L ⫋ H , which is exactly the setting in which Theorem 4.11 was derived and the theorem applies without modification. The teacher’s weights are embedded into the student’s parameter space. Therefore, the minimizer is accessible to the student without relying on the convergence of SGD to recover it. Because the layerwise representations are explicit, the group action 𝜌¯ on the parameter space is a permutation of coordinates, the decomposition E ⊕ E ⊥ is explicit, and the structural factors 𝑟 and 𝑚 𝜆 are directly accessible. As a result, the trace ratio Tr(𝐻 E ⊥ )/Tr(𝐻) and the asymptotic ratios are evaluated exactly. Starting from 𝜃 ∗ , we perturb the student away from 𝜃 ∗ by entrywise independent Gaussian diffusion added directly to the weights, with the standard deviation proportional to that of the each weight. From each sampled parameter, we train the network with SGD, evaluating 𝐿 (𝜃 𝑇 ), 𝐿 ( 𝜃¯𝑇 ), 𝐿 eq (𝜃 𝑇 ) and 𝐿 eq ( 𝜃¯𝑇 ) along the trajectory, from which we obtain empirical estimates of 𝑅(𝑇), 𝑅⊥ (𝑇) and their ratio. Dependence on the averaging time. For a fixed cyclic group 𝐺 = 𝐶4 acting on the input, we track 𝑅⊥ (𝑇)/𝑅(𝑇) as a function of the number of SWA epochs 𝑇; the result is shown in Figure 3 (left). The ratio increases monotonically from values below the asymptotic bound early in averaging, before gradually decreasing and plateauing above the theoretical bound for large 𝑇. Dependence on the group size. We repeat the same procedure for teachers equivariant under cyclic groups 𝐶𝑛 of increasing order, holding the student architecture and training budget fixed, and record the plateaued value of 𝑅⊥ (𝑇)/𝑅(𝑇) for each 𝑛; see Figure 3 (right). The resulting values decrease with 24

G = C4

2.5

Empirical R⊥ /R Theoretical bound

2.25

R⊥ (T )/R(T ) Bound: 1.31

2.00

2.0 R⊥ /R

R⊥ /R

1.75 1.50

1.5

1.25

1.0

1.00

10k 20k 30k SWA averaging steps T

40k

2

(a)

6

10 group order |G|

16

(b)

Figure 3: Synthetic verification of Theorem 4.11 in the teacher-student setting. (a) The empirical ratio 𝑅⊥ (𝑇)/𝑅(𝑇) for 𝐺 = 𝐶4 as a function of the SWA averaging time 𝑇. The curve gradually rises above the asymptotic bound |𝐺 |𝑟 2 /(|𝐺 |𝑟 2 − 1) = 4/3 (dashed line) and plateaus after mild decrease. (b) Plateaued value of 𝑅⊥ /𝑅 against the group order |𝐺 | for cyclic groups 𝐶𝑛 . |𝐺 |, in qualitative agreement with the trend implied by Theorem 4.11 and the discussion following it: larger symmetry groups lead to lower dimension ratio, since |𝐺 |𝑟 2 /(|𝐺 |𝑟 2 − 1) → 1 as |𝐺 | grows, so the guaranteed excess of 𝑅⊥ over 𝑅 shrinks correspondingly. Within the range of group sizes tested, the empirical ratio nonetheless remains bounded above 1 throughout, indicating that the equivariance-specific benefit of SWA persists even as it weakens for larger groups. Taken together, these two experiments substantiate both qualitative predictions of Theorem 4.11: convergence of 𝑅⊥ /𝑅 toward its asymptotic value as the averaging window grows, and a decreasing dependence of that asymptotic value on the size of the augmentation group.

5.2

Practical tasks

We evaluate the performance of SWA on a range of image and graph classification tasks, chosen to cover diverse input modalities and architectural families while keeping the augmentation group fixed across each experimental block. 5.2.1

Datasets and architectures

Image classification. We use [MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, and ImageNet-100], augmented under the action of discrete cyclic subgroups of 𝑆𝑂 (2). For the primary experiments, we use the group of 90-degree rotations 𝐶4 . This range spans low-resolution grayscale images, mediumresolution natural images with varying class counts (CIFAR-10/100), and higher-resolution natural images (ImageNet-100). To explore the generality in architecture of the equivariance gain from SWA, we evaluate four families on this image-classification task: • Multilayer perceptrons (MLPs). These serve as a baseline and provide the closest match to the theoretical setting of Section 4, where the parameter space is unrestricted, and Theorem 4.11 applies without modification. • Convolutional networks (CNNs): VGG-16, PreResNet-164, and WideResNet-28×10. CNNs are 25

the standard architecture for image classification and admit a natural lift to their group-equivariant counterparts (G-CNNs), which are exactly 𝐺-equivariant by construction. Comparing standard CNNs trained with 𝐺-augmentation against this architectural baseline isolates the contribution of SWA-induced statistical equivariance from architecturally-enforced equivariance. • Vision transformers: ViT-S. ViTs process inputs as patch sequences, so any equivariance present is learned from data rather than architecturally imposed. This makes ViTs a test of whether SWA improves equivariance as a generic sampling mechanism, rather than a consequence of the architectural inductive biases of CNNs. • Graph neural networks on superpixels: ViG-Ti applied to superpixel graphs extracted from the image inputs. This architecture lies between CNNs and pure GNNs. The input geometry is derived from the image, but the computation over graphs ignores the pixel grid information. It tests whether the SWA equivariance gain persists when the data geometry is diluted into node and edge features from exact pixel values. Graph classification. To verify that the framework extends beyond image data, we additionally evaluate on three graph classification benchmarks: [DHFR, BZR, COX2], using graph isomorphism networks. Here the relevant symmetry is 3D rotation, and we augment with discrete subgroups of 𝑆𝑂 (3). The model we apply is GIN introduced by [43], it takes in positional features of the atoms, therefore is not manifestly equivariant. The graph classification experiments serve as a confirmation of consistency of the SWA effect (Δ𝐿 eq < 0). 5.2.2

Training and evaluation

For training, we use cross entropy loss and the stochastic gradient descent (SGD) optimizer, along with canonical choices of hyperparameters for each architecture. For each model, we train for a number of epochs until convergence, and then train for an additional number of epochs while taking the average of the weights at the end of each epoch. We evaluate the models before and after applying SWA, and the gain from applying SWA. Besides the cross entropy loss and accuracy on the test set, we use the following metrics for evaluation of equivariance: • Orbit Same Prediction (OSP): the average proportion of the prediction within the same orbit that matches the mode prediction of the orbit. This metric is a direct measure of the equivariance of the model, as it quantifies how consistent the predictions are within the same orbit. Mathematically, the OSP can be defined as: OSP =

𝑁 ∑︁ 1 ∑︁ 1 max 1{ 𝑦ˆ 𝑔𝑋 𝑥𝑖 = 𝑗 } 𝑁 𝑖=1 |𝐺 | 1≤ 𝑗 ≤𝑛𝑌 𝑔∈𝐺

(152)

Where 𝑁 is the number of test samples, and 𝑦ˆ 𝑥 denotes the prediction of sample 𝑥. Note that the OSP is lower bounded by the accuracy, and for a perfectly equivariant model, the OSP would be 100%. • Equivariance Loss: the KL-divergence between the prediction of a sample and the average prediction of its orbit. This serves as a continuous measure that captures how much the predictions deviate from being equivariant. The quantity is defined in Eq. (145).

26

Accuracy (%) Before

OSP (%) Δ

After

Before

After

Δ

Dataset

Model

MNIST

MLP CNN

94.55 ±0.20 94.68 ±0.15 0.13 ±0.11 97.73 ±0.12 97.93 ±0.09 74.23 ±6.59 71.74 ±3.91 −2.49 ±6.70 94.56 ±0.57 96.85 ±0.24

0.20 ±0.07 2.30 ±0.56

FashionMNIST

MLP CNN

85.28 ±0.09 85.98 ±0.07 75.43 ±2.79 85.21 ±0.08

0.71 ±0.11 95.62 ±0.21 96.71 ±0.16 9.77 ±2.74 96.08 ±0.33 98.45 ±0.06

1.09 ±0.15 2.37 ±0.30

CIFAR-10

VGG-16 87.83 ±0.16 89.83 ±0.09 PRN 87.24 ±0.44 92.14 ±0.23 WRN 90.09 ±0.41 94.72 ±0.20

2.00 ±0.20 92.97 ±0.15 95.62 ±0.06 4.89 ±0.29 94.70 ±0.20 98.07 ±0.09 4.63 ±0.27 95.40 ±0.13 98.90 ±0.07

2.65 ±0.13 3.37 ±0.15 3.50 ±0.17

CIFAR-100

VGG-16 54.76 ±0.66 60.17 ±0.63 PRN 58.51 ±0.46 67.38 ±0.79 WRN 65.89 ±0.39 74.93 ±0.20

5.41 ±0.33 81.43 ±0.53 87.53 ±0.26 8.87 ±0.38 86.87 ±0.10 93.63 ±0.08 9.04 ±0.46 87.91 ±0.24 96.05 ±0.10

6.10 ±0.37 6.76 ±0.15 8.14 ±0.18

ImageNet-100

ViG-Ti ViT-S

72.36 ±0.54 72.19 ±0.97 −0.17 ±0.46 91.10 ±0.14 93.00 ±0.29 61.59 ±1.04 76.27 ±0.08 14.68 ±1.04 88.71 ±0.32 95.95 ±0.22

1.90 ±0.24 7.25 ±0.45

Table 1: Accuracy and OSP results under 𝐶4 rotation augmentation, showcasing the impact of stochastic weight averaging (SWA). Positive Δ values indicate an improvement after applying SWA. Model abbreviations: PRN (PreResNet-164), WRN (WideResNet-28×10). 5.2.3

Results

We begin with image classification results, reporting the absolute gains in accuracy and equivariance from SWA, and show that the relative improvement ratio 𝑅⊥ /𝑅 is consistently larger than 1. That is, SWA improves equivariance more than it improves general performance. Boost of accuracy and equivariance Table 1 summarizes the empirical impact of SWA on test accuracy and OSP across various datasets and architectures. A primary observation is the positive correlation between the gains in standard generalization (ΔAcc) and the improvements in equivariance (ΔOSP). Across different models, particularly on more complex tasks such as CIFAR-100 and CIFAR-100, boosts in test accuracy are consistently accompanied by proportionally large increases in OSP. For instance, the models that benefit the most in accuracy (e.g., PreResNet and WideResNet on CIFAR-100) concurrently achieve the highest absolute gains in OSP (up to 8.14%). This strong correlation suggests that a significant portion of SWA’s performance benefit stems directly from the network acquiring better transformation invariance. Furthermore, the improvement in equivariance via SWA appears as a robust and universal phenomenon. The improvement in OSP is strictly positive across all tested configurations, even in regimes where standard accuracy gains are negative (the ViG-Tiny setting on ImageNet-100). This persistent enhancement in equivariance, irrespective of the fluctuations in overall task accuracy, agrees with our theoretical framework: weight averaging cancels out the weights on the non-equivariant subspace, with an improvement rate generally higher than the total space. Relative Improvement Ratio In Section 5.1.2, we empirically verified our main theorem (Theorem 4.11) on synthetic tasks tailored to our theoretical setting. On practical datasets, however, several assumptions underlying the theorem no longer hold (e.g., the unrestricted parameter space), and the architectures involve structural information — such as the exact layerwise irrep decomposition — that is intractable to compute. We therefore restrict ourselves to demonstrating the qualitative claim that SWA improves equivariance more than it improves predictive performance. We revisit the image classification tasks and compare the gains in loss and in equivariance loss, listed in Table 2.

27

Loss (𝐿)

Equivariance Loss (𝐿 eq )

Ratio

Dataset

Model

Before

After

Before

After

𝑅⊥ /𝑅

MNIST

MLP CNN

0.187 0.791

0.183 ±0.003 0.795 ±0.099

0.021 0.019

0.018 ±0.001 0.005 ±0.000

1.14 3.93

FashionMNIST

MLP CNN

0.422 0.720

0.406 ±0.002 0.455 ±0.002

0.021 0.011

0.014 ±0.001 0.003 ±0.000

1.44 2.52

CIFAR-10

VGG-16 PRN WRN×10

0.430 0.499 0.407

0.311 ±0.004 0.300 ±0.009 0.220 ±0.009

0.125 0.079 0.083

0.041 ±0.001 0.017 ±0.001 0.009 ±0.001

2.21 2.79 4.98

CIFAR-100

VGG-16 PRN WRN

2.627 2.050 1.652

1.495 ±0.032 1.608 ±0.042 1.359 ±0.010

0.327 0.155 0.161

0.056 ±0.005 0.045 ±0.001 0.026 ±0.001

3.32 2.70 5.09

ImageNet-100

ViG-T ViT-S

1.178 1.907

1.352 ±0.058 1.154 ±0.011

0.088 0.118

0.066 ±0.004 0.027 ±0.001

1.53 2.60

Table 2: Loss and Equivariance loss (𝐿 eq ) results under 𝐶4 rotation augmentation. 𝑅⊥ /𝑅, measures how much more efficiently SWA minimizes the equivariance loss compared to the overall cross-entropy loss. Model abbreviations: PRN (PreResNet-164), WRN (WideResNet-28×10). Across all architectures and datasets we observe 𝑅⊥ /𝑅 > 1, i.e. SWA reduces the equivariance loss 𝐿 eq strictly more efficiently than it reduces the overall loss 𝐿. The magnitudes admit a more quantitative reading. For the augmentation group 𝐶4 used throughout Table 2, the right-hand side of (139) is at most |𝐺 |𝑟 2 /(|𝐺 |𝑟 2 −1) = 4/3, attained at 𝑟 = 1 and in the limit of vanishing multiplicity factor max𝜆 𝑚 𝜆 /(𝑑𝜆 𝑛); larger 𝑟 or larger multiplicities only decrease it. The MLPs, which are the architectures closest to the setting of Theorem 4.11 in that the parameter space is unrestricted (L = H ), give ratios of 1.14 (MNIST) and 1.44 (FashionMNIST), i.e. they lie closely to the theoretical value 4/3. In particular, the MNIST value falling slightly below 4/3 is not a contradiction, since we disregard the structural facts that lead to actual values for the factors 𝑟 and max𝜆 𝑚 𝜆 /(𝑑𝜆 𝑛). In summary, this reads as evidence that the bound is close to tight in the regime where its assumptions hold. The deep convolutional models behave differently, reaching ratios between 2.21 and 5.09, with bound 4/3 being merely conservative for these settings. Specifically, the ratio shows a correlation with task complexity and model capacity.The deviation from perfect prediction is potentially the root cause of the looseness of the bound. Furthermore, the dimension count (83) assumes an unrestricted parameter space, whereas for residual networks, their admissible parameters form a proper affine subspace L ⊂ H whose equivariant part E is a smaller fraction of L than the count suggests. Both effects push the true ratio upwards relative to the bound, and the trend in Table 2 broadly increases with model depth within each dataset, from VGG-16 to WideResNet-28×10. Generalization to graph classification tasks The results above establish that SWA robustly improves equivariance on image classification tasks, occasionally independently of whether accuracy improves at all. However, it is natural to ask whether this pattern is specific to convolutional architectures and image data, or reflects a more general property of SWA. To test this, we repeat the same analysis on graph classification tasks, where both the underlying group action and the network architecture differ substantially from the image setting. Table 3 reports accuracy and OSP for the three benchmarks under augmentation by the tetrahedral, octahedral and icosahedral rotation groups. The pattern differs from the image experiments. The effect of weight averaging on the test accuracy is essentially random: five of the nine configurations show a negative ΔAcc, the magnitudes range from −2.81% to +2.48%, and in every setting the error bar

28

over seeds is comparable to or larger than the mean. Equivariance, in contrast, improves in all nine configurations, with ΔOSP between 0.68% and 2.66% and a clearly resolved effect in the majority of them. Accuracy (%)

OSP (%)

Dataset

Group

Before

After

Δ

Before

After

Δ

DHFR

Tetra Octa Icosa

79.90 ±2.07 82.28 ±2.44 80.23 ±2.44

81.45 ±1.57 79.47 ±2.15 78.42 ±4.09

1.55 ±1.82 −2.81 ±3.58 −1.81 ±2.23

96.36 ±0.18 97.11 ±0.27 96.96 ±0.40

98.13 ±1.15 97.79 ±0.76 99.42 ±0.51

1.77 ±1.13 0.68 ±0.61 2.46 ±0.47

BZR

Tetra Octa Icosa

85.83 ±3.15 83.13 ±2.51 86.46 ±2.47

84.15 ±3.34 85.61 ±3.04 83.66 ±2.38

−1.68 ±2.99 2.48 ±1.07 −2.80 ±2.01

98.22 ±0.67 97.68 ±0.57 98.82 ±0.40

98.98 ±1.22 99.29 ±0.71 99.76 ±0.22

0.76 ±1.14 1.61 ±0.49 0.94 ±0.29

COX2

Tetra Octa Icosa

78.54 ±2.05 78.91 ±1.88 78.68 ±3.13

77.87 ±0.89 80.64 ±1.17 80.00 ±1.17

−0.67 ±1.27 1.73 ±2.43 1.32 ±3.99

96.75 ±0.45 97.30 ±0.36 97.66 ±0.18

98.62 ±0.95 99.96 ±0.10 99.65 ±0.45

1.87 ±0.98 2.66 ±0.40 2.00 ±0.52

Table 3: Accuracy and OSP on the graph classification benchmarks under augmentation by discrete rotation subgroups of 𝑆𝑂 (3) (tetrahedral, octahedral and icosahedral), before and after stochastic weight averaging. Positive Δ values indicate improvement.

Acknowledgements This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. The computations were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS), partially funded by the Swedish Research Council through grant agreement no. 2025/22-1341.

References [1]

Josh Abramson et al. “Accurate Structure Prediction of Biomolecular Interactions with AlphaFold 3”. In: Nature 630.8016 (2024), pp. 493–500. doi: 10.1038/s41586-024-07487-w.

[2]

Erik J. Bekkers et al. “Roto-Translation Covariant Convolutional Networks for Medical Image Analysis”. In: Medical Image Computing and Computer Assisted Intervention – MICCAI 2018. Lecture Notes in Computer Science. Cham: Springer International Publishing, 2018, pp. 440–448. doi: 10.1007/978-3-030-00928-1_50.

[3]

Michael M. Bronstein et al. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges. 2021. arXiv: 2104.13478.

[4]

Jan E. Gerken et al. “Geometric Deep Learning and Equivariant Neural Networks”. In: Artificial Intelligence Review (2023). doi: 10.1007/s10462-023-10502-7. arXiv: 2105.13926.

[5]

Jan E. Gerken et al. “Equivariance versus Augmentation for Spherical Images”. In: Proceedings of the 39th International Conference on Machine Learning. PMLR, 2022, pp. 7404–7421. arXiv: 2202.03990.

[6]

Jan E. Gerken and Pan Kessel. “Emergent Equivariance in Deep Ensembles”. In: Proceedings of the 41st International Conference on Machine Learning. PMLR, 2024, pp. 15438–15465. arXiv: 2403.03103.

[7]

Oskar Nordenfors and Axel Flinth. Ensembles provably learn equivariance through data augmentation. 2024. arXiv: 2410.01452.

[8]

Pavel Izmailov et al. “Averaging Weights Leads to Wider Optima and Better Generalization”. In: Proceedings of the Thirty-Fourth Conference on Uncertainty in Artificial Intelligence (2018). Monterey, California, USA: AUAI Press, 2018. arXiv: 1803.05407.

[9]

Taco Cohen and Max Welling. “Group equivariant convolutional networks”. In: International conference on machine learning. PMLR. 2016, pp. 852–860. arXiv: 1602.07576.

[10]

Taco Cohen and Max Welling. “Group equivariant convolutional networks”. In: International conference on machine learning. PMLR. 2016, pp. 852–860. arXiv: 1602.07576.

29

[11]

Taco S. Cohen et al. Spherical CNNs. 2018. arXiv: 1801.10130.

[12]

Risi Kondor and Shubhendu Trivedi. On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups. 2018. arXiv: 1802.03690.

[13]

Oskar Nordenfors, Fredrik Ohlsson, and Axel Flinth. “Optimization dynamics of equivariant and augmented neural networks”. In: Transactions on Machine Learning Research (2025).

[14]

Bryn Elesedy and Sheheryar Zaidi. Provably Strict Generalisation Benefit for Equivariant Models. 2021. arXiv: 2102. 10333.

[15]

Akiyoshi Sannai, Masaaki Imaizumi, and Makoto Kawano. “Improved generalization bounds of group invariant / equivariant deep networks via quotient feature spaces”. In: Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence. Vol. 161. Proceedings of Machine Learning Research. PMLR, 2021, pp. 771–780. url: https: //proceedings.mlr.press/v161/sannai21a.html.

[16]

Georg Bökman and Fredrik Kahl. “Investigating How ReLU-networks Encode Symmetries”. In: Advances in Neural Information Processing Systems. Vol. 36. Curran Associates, Inc., 2023, pp. 13720–13744. doi: 10.52202/0752800605.

[17]

Vahid Shahverdi et al. “Identifiable Equivariant Networks are Layerwise Equivariant”. In: Proceedings of the 43th International Conference on Machine Learning. PMLR, 2026. arXiv: 2601.21645.

[18]

Yuyang Wang et al. Swallowing the Bitter Pill: Simplified Scalable Conformer Generation. 2024. arXiv: 2311.17932. url: https://arxiv.org/abs/2311.17932.

[19]

Johann Brehmer et al. Does Equivariance Matter at Scale? 2025. arXiv: 2410.23179.

[20]

Shuxiao Chen, Edgar Dobriban, and Jane H. Lee. A Group-Theoretic Framework for Data Augmentation. 2020. arXiv: 1907.10905.

[21]

M.A. Ganaie et al. “Ensemble Deep Learning: A Review”. In: Engineering Applications of Artificial Intelligence 115 (2022), p. 105151. doi: 10.1016/j.engappai.2022.105151.

[22]

Gao Huang et al. Snapshot Ensembles: Train 1, Get M for Free. 2017. arXiv: 1704.00109.

[23]

Hugh Chen, Scott Lundberg, and Su-In Lee. Checkpoint Ensembles: Ensemble Methods from a Single Training Process. 2017. arXiv: 1710.03282.

[24]

Timur Garipov et al. Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs. 2018. arXiv: 1802.10026.

[25]

Boris Polyak and Anatoli Juditsky. “Acceleration of Stochastic Approximation by Averaging”. In: SIAM Journal on Control and Optimization 30 (1992), pp. 838–855. doi: 10.1137/0330046.

[26]

Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. 2018. arXiv: 1703.01780.

[27]

Zeke Xie, Issei Sato, and Masashi Sugiyama. A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima. 2021. arXiv: 2002.03495.

[28]

Alex Damian, Tengyu Ma, and Jason D. Lee. Label Noise SGD Provably Prefers Flat Global Minimizers. 2021. arXiv: 2106.06530.

[29]

Hao Guo, Jiyong Jin, and Bin Liu. Stochastic Weight Averaging Revisited. 2022. arXiv: 2201.00519.

[30]

Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic Gradient Descent as Approximate Bayesian Inference. 2018. arXiv: 1704.04289.

[31]

Alexandru Meterez et al. A Defense of the Quadratic Model. 2026. arXiv: 2607.21716. url: https://arxiv.org/ abs/2607.21716.

[32]

Gerard Ben Arous et al. Spectral Alignment of Stochastic Gradient Descent for High-Dimensional Classification Tasks. https://arxiv.org/abs/2310.03010v2. 2023.

[33]

Arthur Jacot, Franck Gabriel, and Clement Hongler. “Neural Tangent Kernel: Convergence and Generalization in Neural Networks”. In: Advances in Neural Information Processing Systems. Vol. 31. Curran Associates, Inc., 2018. arXiv: 1806.07572.

[34]

Jaehoon Lee et al. Deep Neural Networks as Gaussian Processes. 2018. arXiv: 1711.00165. url: https://arxiv. org/abs/1711.00165.

[35]

Arthur Jacot, Franck Gabriel, and Clément Hongler. The Asymptotic Spectrum of the Hessian of DNN throughout Training. 2020. doi: 10.48550/arXiv.1910.02875. arXiv: 1910.02875.

[36]

Ronen Basri et al. Frequency Bias in Neural Networks for Input of Non-Uniform Density. 2020. doi: 10.48550/arXiv. 2003.04560. arXiv: 2003.04560.

[37]

Mariia Seleznova et al. “GradPCA: Leveraging NTK Alignment for Reliable Out-of-Distribution Detection”. In: The Fourteenth International Conference on Learning Representations. 2025.

30

[38]

Philipp Misof, Pan Kessel, and Jan E Gerken. “Equivariant Neural Tangent Kernels”. In: Proceedings of the 42nd International Conference on Machine Learning. Vol. 267. Proceedings of Machine Learning Research. PMLR, 2025, pp. 44470–44503. url: https://proceedings.mlr.press/v267/misof25a.html.

[39]

William Fulton and Joe Harris. Representation Theory: A First Course. Springer Science & Business Media, 1991. isbn: 978-0-387-97495-8.

[40]

Qianxiao Li, Cheng Tai, et al. “Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations”. In: Journal of Machine Learning Research 20.40 (2019), pp. 1–47.

[41]

Christian Szegedy et al. Rethinking the Inception Architecture for Computer Vision. 2015. arXiv: 1512.00567. url: https://arxiv.org/abs/1512.00567.

[42]

Longde Huang et al. “Learning Chern Numbers of Multiband Topological Insulators with Gauge Equivariant Neural Networks”. In: Advances in Neural Information Processing Systems. Vol. 38. Curran Associates, Inc., 2025, pp. 147997– 148026. arXiv: 2502.15376.

[43]

Keyulu Xu et al. How Powerful are Graph Neural Networks? 2019. arXiv: 1810.00826. url: https://arxiv.org/ abs/1810.00826.

A

Proof of Proposition 4.3

It is well known that the evolution of (21) is described by the Fokker-Planck equation: If 𝑝 𝑡 is the density of the distribution at time 𝑡, we have 𝜕𝑡 𝑝 𝑡 (𝜃) = − ∇ · ( 𝑝 𝑡 ∇𝐿) + ∇2 · (𝐷 𝑝 𝑡 )

(153)

Here, 𝐷 is the diffusion tensor: 𝐷 (𝜃) =

1 1 𝑃(𝜃)𝑃(𝜃) 𝑇 = Σ(𝜃) 2 2

(154)

We can express (153) in operator form 𝜕𝑡 𝑝 𝑡 = A 𝑝 𝑡 , with A 𝑝 = ∇ · ( 𝑝 𝑡 ∇𝐿) + ∇2 · (𝐷 𝑝 𝑡 ).

(155)

It will be convenient to introduce also the push-forward operation as an operator: (S𝑔 𝑓 ) (𝜃, 𝑡) = 𝑓 ( 𝑔𝜃, ¯ 𝑡). This for instance makes the difference between ∇𝐿 (𝑔𝜃), which is S𝑔 ∇𝐿, and ∇[𝐿 ( 𝑔𝜃)], ¯ which is ∇(S𝑔 𝐿). The time-derivative 𝜕𝑡 trivially commutes with S𝑔 . Its commutation relations with the divergence and Laplace operators are as follows. Lemma A.1. For all 𝑔 ∈ 𝐺, ∇ · [ 𝑔¯ 𝑇 S𝑔 𝑓 ] = S𝑔 [∇ · 𝑓 ] 2

𝑇

2

∇ · [ 𝑔S ¯ 𝑔 𝐷 𝑔¯ ] = S𝑔 [∇ · 𝐷]

(156) (157)

Proof. Using the Einstein summation convention, the chain rule implies for any vector field ℎ ∇ · [S𝑔 ℎ] = 𝜕𝑖 S𝑔 ℎ𝑖 = 𝑔 𝑗𝑖 𝜕 𝑗 ℎ𝑖

(158)

where 𝑔𝑖 𝑗 denote the entries of the matrix corresponding to the action 𝑔. ¯ Since ( 𝑔¯ 𝑇 S𝑔 𝑓 )𝑖 = 𝑔 𝑘𝑖 S𝑔 𝑓 𝑘 , we get ∇ · ( 𝑔¯ 𝑇 S𝑔 𝑓 ) = 𝑔 𝑘𝑖 𝜕𝑖 (S𝑔 𝑓 𝑘 ) = 𝑔 𝑘𝑖 𝑔 𝑗𝑖 S𝑔 𝜕 𝑗 𝑓 𝑘 .

31

(159)

Now, since 𝑔¯ is orthogonal, we have 𝑔 𝑘𝑖 𝑔 𝑗𝑖 = 𝑔 𝑘𝑖 𝑔𝑖𝑇𝑗 = 𝛿 𝑘 𝑗 , so 𝑔 𝑘𝑖 𝑔 𝑗𝑖 S𝑔 𝜕 𝑗 𝑓 𝑘 = 𝛿 𝑘 𝑗 S𝑔 𝜕 𝑗 𝑓 𝑘 = S𝑔 𝜕𝑘 𝑓 𝑘 = S𝑔 [∇ · 𝑓 ],

(160)

which is the first claim. Similarly, the chain rule implies for every tensor field 𝐸 ∇2 · [S𝑔 𝐸] = 𝜕𝑖 𝜕 𝑗 S𝑔 𝐷 𝑖 𝑗 = 𝑔 𝑘𝑖 𝑔ℓ 𝑗 𝜕𝑘 𝜕ℓ 𝐷 𝑖 𝑗 .

(161)

Since ( 𝑔¯ 𝑇 𝐷 𝑔) ¯ 𝑖 𝑗 = 𝑔𝑚𝑖 𝐷 𝑚𝑛 𝑔𝑛 𝑗 , we get ∇2 · [ 𝑔S ¯ 𝑔 𝐷 𝑔¯ 𝑇 ] = 𝑔𝑚𝑖 𝑔𝑛 𝑗 𝜕𝑖 𝜕 𝑗 [S𝑔 𝐷 𝑚𝑛 ] = 𝑔𝑚𝑖 𝑔𝑛 𝑗 𝑔 𝑘𝑖 𝑔ℓ 𝑗 S𝑔 𝜕𝑘 𝜕ℓ 𝐷 𝑚𝑛 = 𝛿 𝑘𝑚 𝛿 𝑛ℓ S𝑔 𝜕𝑘 𝜕ℓ 𝐷 𝑚𝑛 2

= S𝑔 𝜕𝑘 𝜕ℓ 𝐷 𝑘ℓ = S𝑔 [∇ · 𝐷],

(162) (163)

where we again applied 𝑔 𝑘𝑖 𝑔 𝑗𝑖 = 𝛿 𝑘 𝑗 . Using the above formulas, it is now not hard to prove that A commutes with S𝑔 . Lemma A.2. We have AS𝑔 = S𝑔 A for all 𝑔 ∈ 𝐺. Proof. Lemma A.1 implies S𝑔 A 𝑝 = S𝑔 [∇ · ( 𝑝 𝑡 𝐿)] + S𝑔 [∇2 · (𝐷 𝑝 𝑡 )] = ∇ · ( 𝑔¯ 𝑇 S𝑔 [ 𝑝 𝑡 ∇𝐿]) + ∇2 · ( 𝑔S ¯ 𝑔 [𝐷 𝑝 𝑡 ] 𝑔¯ 𝑇 ). Now, we note that the chain rule and the invariance S𝑔 𝐿 = 𝐿 of the cumulative loss, and Lemma 4.2 implies that S𝑔 ∇𝐿 = 𝑔∇S ¯ 𝑔 𝐿 = 𝑔∇𝐿, ¯

S𝑔 𝐷 = 𝑔¯ 𝑇 𝐷 𝑔, ¯

so that ∇ · ( 𝑔¯ 𝑇 S𝑔 [ 𝑝 𝑡 ∇𝐿]) + ∇2 · ( 𝑔S ¯ 𝑔 [𝐷 𝑝 𝑡 ] 𝑔¯ 𝑇 ) = ∇ · ( 𝑔¯ 𝑇 𝑔S ¯ 𝑔 𝑝∇𝐿) + ∇2 · (S𝑔 𝑝 𝑡 𝑔¯ 𝑔¯ 𝑇 𝐷 𝑔¯ 𝑔¯ 𝑇 ) = ∇ · (S𝑔 𝑝 𝑡 ∇𝐿) + ∇2 · (S𝑔 𝑝 𝑡 𝐷) = AS𝑔 𝑝 𝑡 .

With the commutation AS𝑔 = S𝑔 A, it is not hard to prove Proposition 4.3. Proof of Proposition 4.3. Since both A and 𝜕𝑡 commute with S𝑔 , so does the flow – if 𝑝 𝑡 solves the equation initialized at 𝑝 0 , S𝑔 𝑝 𝑡 will do it for the one initialized at S𝑔 𝑝 0 . In particular, if we initialize at an invariant distribution, S𝑔 𝑝 0 = 𝑝 0 , we remain invariant, S𝑔 𝑝 𝑡 = 𝑝 𝑡 . Consequently, 𝜃 ∼ Θ𝑡 ⇒ 𝑔¯ −1 𝜃 ∼ Θ𝑡 , and therefore (12)

𝑁 𝑡 (𝑔 𝑋 𝑥) = E 𝜃∼Θ𝑡 (𝑁 𝜃 (𝑔 𝑋−1 𝑥)) = E 𝜃∼Θ𝑡 (𝑁𝑔¯ −1 𝜃 (𝑔 𝑋 𝑥)) = E 𝜃∼Θ𝑡 (𝑔𝑌 𝑁 𝜃 (𝑥)) = 𝑔𝑌 𝑁 𝑡 (𝑥), i.e., 𝑁 𝑡 is equivariant.

32

10−4

mean actual loss mean quadratic estimate Approximation error

loss

10−5

10−6

10

−7

10−5

10−6 0 −10−6

−10−5 0

500

1000 SGD step

1500

2000

0.148

0.150 0.152 0.154 Distance from minimizer

0.156

Figure 4: Empirical validation of the quadratic approximation in a teacher-student MLP setting, with Gaussian diffusion sampling. Left: The true loss strictly follows the quadratic estimation. Right: The approximation error between full- and quadratic loss stabilizes and strictly decreases as the parameter distance to the minimizer diminishes.

B

Quadratic approximation with Gaussian diffusion sampling

We repeat the experiment in Section 5.1.1, replacing the gradient-ascent sampler with Gaussian diffusion; see Figure 4. The results show comparable approximation accuracy, demonstrating the robustness of the quadratic approximation to the choice of perturbation. The loss values are, however, systematically lower under Gaussian diffusion than under gradient ascent at matched perturbation scale, and the approximation error stabilizes considerably faster. We attribute this to the mismatch between the isotropic noise of the Gaussian sampler and the anisotropy of the loss landscape, in contrast to the gradient ascent case, whose covariance is shaped by the Hessian and can concentrate the perturbation along directions where the loss is more sensitive to gradient steps and the quadratic approximation is least accurate.

33

Related documents

• VITA: A Multi-Source Vicinal Transfer Augmentation Method for Out-of-Distribution Generalization#1154998
• An optimized gradient boosting framework for IoT intrusion detection: a comprehensive evaluation on the CICIoT2023 dataset.#258279
• Data Augmentation and Generation for Improving Learning-based Autonomous Driving Models#634906
• Aggregating many estimators using estimated weights#1014009
• Boosting targeted adversarial transferability via fine-grained feature mixup perturbation and reference-based gradient refinement.#404527