Title: Symmetries in Overparametrized Neural Networks: A Mean Field View
Venue: NeurIPS 2024
Accepted: yes
Keywords: Overparametrized Neural Networks

Mean Field Limit of Neural Networks

Symmetries in Neural Networks

Wasserstein Gradient Flow

Data Augmentation

Feature Averaging

Equivariant Architechtures.
Abstract: We develop a Mean-Field (MF) view of the learning dynamics of overparametrized Artificial Neural Networks (NN) under distributional symmetries of the data w.r.t. the action of a general compact group $G$. We consider for this a class of  generalized shallow NNs given by an ensemble of  $N$ multi-layer units, jointly trained  using stochastic gradient descent (SGD) and possibly symmetry-leveraging (SL) techniques, such as Data Augmentation (DA), Feature Averaging  (FA) or Equivariant Architectures (EA).  We introduce the notions of weakly and strongly invariant  laws (WI and SI) on the parameter space of each single unit, corresponding, respectively, to $G$-invariant distributions, and to distributions supported on parameters fixed by the group action (which encode EA). This allows us to define symmetric models compatible with taking $N\to\infty$ and  give an interpretation of the asymptotic dynamics of DA, FA and EA in terms of Wasserstein Gradient Flows describing their MF limits. When activations respect the group action, we show that, for  symmetric data, DA, FA and freely-trained models obey the exact same MF  dynamic, which stays in the space of WI parameter laws and attains therein the population risk's minimizer. We also provide a counterexample to the general attainability of such an optimum over SI laws.
Despite this,  and quite remarkably, we show that the space of SI laws  is also preserved by these MF distributional dynamics even when freely trained. This sharply contrasts the finite-$N$ setting, in which EAs are generally not preserved by unconstrained SGD. We illustrate the validity of our findings as $N$ gets larger,  in a teacher-student experimental setting, training a student NN to learn from a WI, SI  or arbitrary teacher model through various SL schemes. We lastly deduce a data-driven heuristic to discover the largest subspace of parameters supporting SI distributions for a problem, that could be used for designing EA with minimal generalization error.

Symmetries in Overparametrized Neural Networks: A Mean-Field View

Abstract

We develop a Mean-Field (MF) view of the learning dynamics of overparametrized Artificial Neural Networks (NN) under distributional symmetries of the data w.r.t. the action of a general compact group G . We consider for this a class of generalized shallow NNs given by an ensemble of N multi-layer units, jointly trained using stochastic gradient descent (SGD) and possibly symmetry-leveraging (SL) techniques, such as Data Augmentation (DA), Feature Averaging (FA) or Equivariant Architectures (EA). We introduce the notions of weakly and strongly invariant laws (WI and SI) on the parameter space of each single unit, corresponding, respectively, to G -invariant distributions, and to distributions supported on parameters fixed by the group action (which encode EA). This allows us to define symmetric models compatible with taking N →∞ and give an interpretation of the asymptotic dynamics of DA, FA and EA in terms of Wasserstein Gradient Flows describing their MF limits. When activations respect the group action, we show that, for symmetric data, DA, FA and freely-trained models obey the exact same MF dynamic, which stays in the space of WI parameter laws and attains therein the population risk's minimizer. We also provide a counterexample to the general attainability of such an optimum over SI laws. Despite this, and quite remarkably, we show that the space of SI laws is also preserved by these MF distributional dynamics even when freely trained. This sharply contrasts the finiteN setting, in which EAs are generally not preserved by unconstrained SGD. We illustrate the validity of our findings as N gets larger, in a teacher-student experimental setting, training a student NN to learn from a WI, SI or arbitrary teacher model through various SL schemes. We lastly deduce a data-driven heuristic to discover the largest subspace of parameters supporting SI distributions for a problem, that could be used for designing EA with minimal generalization error.

1 Introduction

Learning in complex tasks, employing ever larger datasets, has strongly benefited from the implementation and training of Artificial Neural Networks (NN) with a huge number of parameters; as well as from training schemes or architectures that can leverage underlying symmetries of the data in order to reduce the problem's complexity (see [33, 34] for general reference). This raises questions, on one hand, of understanding the puzzling generalizability in overparametrized NN; and on the other, of when and how symmetry-leveraging (SL) techniques (such as Data Augmentation, Feature Averaging or Equivariant Architectures), can induce useful biases towards learning with symmetries, without hindering approximation and generalization properties. The recent Mean-Field ( MF ) theory of NN (see [16] and further references below) provides a partial, yet promissory, viewpoint to address the first question for shallow NN: in the Mean-Field Limit ( MFL ) of an infinitely wide hidden layer, stochastic gradient descent (SGD) training dynamics approximates the Wasserstein Gradient Flow

( WGF ) of certain convex population risk on the space of distributions on parameters. Confluently, the incorporation of combined algebraic and probabilistic viewpoints have yielded a more complete view of the benefits of SL techniques under symmetry (see e.g. [13, 27, 59] and further references below); however, it is not clear if and how those findings can scale to overparametrized NN and their MFL .

In this work we develop a systematic MF analysis of the limiting learning dynamics of a class of generalized shallow NNs, under distributional symmetries of the data w.r.t. the action of a compact group, and including the possible effects of employing some of the most popular SL techniques. The effect of symmetries on the WGF dynamics was already studied in [35], in the particular case of two-layer ReLU networks, under data generated by a function symmetric w.r.t. a single orthogonal transformation. We consider our (independent 1 ) work to largely broaden the scope and applicability of such initial contributions, as it provides a unified MF interpretation for both the use of SL techniques under general distributional invariances, and the interplay of such symmetries at the levels of data, architectures and training dynamics. The paper unfolds as follows:

In Section 2 we introduce a class of generalized shallow models with multi-layer units on which we will focus, we recall WGF s and their role in the MFL of NN training dynamics, and review the SL techniques to be studied. Section 3 contains the bulk of our contributions, as we study how SL techniques applied on these models can be interpreted in terms of their limiting WGF s, how they relate to each other in terms of the optima of their corresponding population risks, and how their limiting MF training dynamics behave with or without symmetric data. Finally, Section 4 presents the empirical validation of our main theoretical results through some numerical simulations; it also suggests a potential heuristic for discovering data-driven parameter-sharing schemes that lead to optimal equivariant architectures in ML problems. Proofs and complements to our results can be found in the Supplementary Material (henceforth SuppMat for short), together with a discussion of the scope and limitations of our results, as well as a summary of the notation and abbreviations used.

2 Preliminaries

2.1 Supervised learning with generalized shallow neural networks

Let X , Y and Z be separable Hilbert Spaces, termed as the feature , label and parameter spaces respectively. Typically, these are finite-dimensional, e.g. X = R d and Y = R c (for c, d ∈ N ∗ ) with Z the space of affine transformations between hidden layers. We write P ( · ) for the space of Borel probability measure on a metric space ( · ) . Let π ∈ P ( X × Y ) denote the data distribution from which i.i.d. samples ( X,Y ) ∈ X × Y will be drawn, and ℓ : Y × Y → R be a convex loss function. Consider also an activation function σ ∗ : X × Z → Y . We introduce a general class of shallow NN: Definition 1. A shallow neural network model of parameter θ := ( θ i ) N i =1 ∈ Z N is the function Φ N θ : X → Y given by Φ N θ ( x ) := 1 N ∑ N i =1 σ ∗ ( x ; θ i ) , ∀ x ∈ X . Equivalently, if ν N θ := 1 N ∑ N i =1 δ θ i is the empirical measure associated with θ ∈ Z N , we can write ∀ x ∈ X , Φ N θ ( x ) = ⟨ σ ∗ ( x ; · ) , ν N θ ⟩ or, abusing notation, simply Φ N θ = ⟨ σ ∗ , ν N θ ⟩ .

In the setting where X = R d , Y = R c and Z = R c × b × R d × b × R b (for b ∈ N ∗ ), if we consider, for z = ( W,A,B ) ∈ Z and σ : R b → R b , σ ∗ ( x, z ) := Wσ ( A T x + B ) ; then Φ N θ (with N ∈ N , θ ∈ Z N ) corresponds exactly to a single-hidden-layer neural network with N hidden units. Depending on σ ∗ , however, these shallow NN models can represent settings that go far beyond this first example. In fact, σ ∗ can be taken to be an entire Multi-Layer NN model, in which case Φ N θ will represent an ensemble of N such units trained simultaneously (see SuppMat-C.1). As we will also shortly see, for suitable subspaces of Z , this modelling extends to renowned equivariant architectures such as CNNs, DeepSets and GNNs. Beyond NNs, this setting can also model the deconvolution of sparse spikes, RBF networks, density estimation via MMD minimization, among many others (see [16, 62, 69]).

This class thus allows for non-trivial internal units, while enabling the width N →∞ consistently, and regardless of the possible underlying structure of the (fixed size) units represented by σ ∗ . Inspired by this possibility, and by our writing of shallow NN models , we define a more general notion:

Definition 2 ( Shallow Model ) . A shallow model is any function of the form Φ µ ( x ) := ⟨ σ ∗ ( x ; · ) , µ ⟩ for some µ ∈ P ( Z ) (whenever the integral makes sense for all x ∈ X ). We write Φ µ := ⟨ σ ∗ , µ ⟩ and denote the space of such models as F σ ∗ ( P ( Z )) .

Classically, we want to find a NN model that performs well with respect to π and ℓ . More precisely, having fixed an architecture (given here by N and σ ∗ ), we consider the generalization error or population risk given by R ( θ ) = E π [ ℓ (Φ N θ ( X ) , Y ) ] , and look for a vector of parameters θ ∈ Z N attaining min θ ∈Z N R ( θ ) . However, not only is this function highly non-convex and hard to optimize; but in practice we generally don't have access to π (and thus R ) and we have to solve this problem only with a set of i.i.d. data samples { ( X k , Y k ) } k ∈ N drawn from π . Thus, the usual approach to minimizing this population risk is to train a NN model Φ N θ , through an SGD scheme (see e.g. [7]):

First, initialize θ 0 i , ∀ i ∈ { 1 , . . . , N } , i.i.d. from a fixed distribution µ 0 ∈ P ( Z ) .

Iterate, for k ∈ N , defining ∀ i ∈ { 1 , . . . , N } :

Here, s N k = ε N ς ( kε N ) is the step-size (or learning rate ), parametrized in terms of ς : R + → R + a regular function and ε N > 0 . Also, we have a penalization function r : Z → R , regularizing Gaussian noise ξ k i i.i.d. ∼ N (0 , Id Z ) independent from the initialization and data, and τ, β ≥ 0 . When τ, β > 0 , the method is called stochastic gradient Langevin dynamics, noted SGLD ([74]), or simply noisy SGD. An infinite i.i.d. sample from π will be needed when letting later N →∞ . When π is the empirical measure of a finite dataset, we are performing empirical-risk minimization (which of course is not the same as minimizing generalization error, but follows the same mathematical formulation).

In principle, there are no guarantees that this training procedure will be truly optimizing R ( θ ) let alone approaching its minimum. However, by extending the definition of the generalization error to models in F σ ∗ ( P ( Z )) , one gets the convex functional R : P ( Z ) → R given by R ( µ ) := E π [ ℓ (Φ µ ( X ) , Y ) ] . The problem on Z N is thus lifted to the convex optimization problem on P ( Z ) :

Accordingly, this motivates looking at the evolution of empirical measures ( ν N k ) k ∈ N := ( ν N θ k ) k ∈ N ⊆ P ( Z ) instead of that of the specific parameters ( θ k ) k ∈ N ⊆ Z N . The MF approach to NNs (see [16, 53, 62, 67]) aims at providing theoretical guarantees for problem (2), justifying that a global optimum of the population risk can be approximated by training a NN with SGD for large N . We next provide some necessary background on WGF s and on the MF theory of shallow NN models.

2.2 Wasserstein gradient flow and mean-field limit of shallow models

We briefly recall some elements of Optimal Transport and Wasserstein Gradient Flows, referring to [1, 63, 71] for further background. Let Z be a Hilbert space with norm ∥ · ∥ and, for p ∈ [1 , ∞ ) , let P p ( Z ) := { µ ∈ P ( Z ) : ∫ Z ∥ θ ∥ p µ ( dθ ) < + ∞} be the space of probability measures on Z with finite p -th moment. We endow this space with the p -th Wasserstein distance , defined as: W p ( µ, ν ) := [ inf γ ∈ Π( µ,ν ) E γ [ ∥ X -Y ∥ p ] ] 1 p , ∀ µ, ν ∈ P p ( Z ) with Π( µ, ν ) being the set of couplings between µ and ν (the infimum is always attained). The metric space ( P p ( Z ) , W p ) is Polish and called the p -th Wasserstein Space. In the remainder of this section we consider p = 2 and Z = R D .

We recall central objects for the sequel, including Lions' derivative [9, 47], popularized in mean-field games (see e.g. [10, 12, 15, 38]) and shown (in [32]) to coincide with the Wasserstein gradient ([1]):

Definition 3 ( Linear Functional Derivative and Intrinsic Derivative ) . Given F : P 2 ( Z ) → R , its linear functional derivative is the function (if it exists) ∂F ∂µ : Dom ( F ) ×Z → R such that ∀ µ, ν ∈ Dom ( F ) , lim h → 0 F ((1 -h ) µ + hν ) -F ( µ ) h = ∫ Z ∂F ∂µ ( µ, z ) d ( ν -µ )( z ) and ∫ Z ∂F ∂µ ( µ, z ) dµ ( z ) = 0 . The function F ′ : µ ∈ P 2 ( Z ) ↦→ ∂F ∂µ ( µ, · ) is known as the first variation of F at µ . Moreover, if ∂F ∂µ exists and is differentiable in its second argument, we define the intrinsic derivative of F at µ to be: D µ F ( µ, z ) = ∇ z ( ∂F ∂µ ( µ, z ) ) . Abusing notation, we will write ∂F ∂µ : P 2 ( Z ) × Z → R and D µ F : P 2 ( Z ) ×Z → Z , even if they are only partially defined.

This allows us to define next a Wasserstein Gradient Flow (following e.g. [1, 16]):

Definition 4 ( Wasserstein Gradient Flow ) . Let ς : R + → R + be a regular scalar function and F : P 2 ( Z ) → R be a convex functional for which the intrinsic derivative D µ F is defined. We define a Wasserstein Gradient Flow (WGF) for F (shortened WGF ( F )) as any absolutely continuous trajectory ( µ t ) t ≥ 0 in P 2 ( Z ) that satisfies, distributionally on [0 , ∞ ) ×Z :

Several authors ([1, 16, 63, 71], among others) have proven under various sets of assumptions that, given an initial condition µ 0 ∈ P 2 ( Z ) , the WGF ( F ) admits a unique (weak) solution, ( µ t ) t ≥ 0 . In a sense, WGF ( F ) 'follows the negative gradient' of F . Unfortunately, even for convex F , stationary points of WGF ( F ) need not be global minima, see [16].

We are interested in the case where F is the following convex, entropy-regularized population risk : R τ, β ( µ ) := R ( µ ) + τ ∫ rdµ + βH λ ( µ ) , where τ, β ≥ 0 , λ is the Lebesgue Measure on Z , r : Z → R + is a penalization , and H λ defined as H λ ( µ ) := ∫ log( dµ dλ ( z )) dµ ( z ) if µ ≪ λ or + ∞ otherwise, is the Boltzmann entropy of µ . In this case, WGF ( R τ, β ) reads as the PDE:

known as McKean-Vlasov equation in the probability and PDE communities (see the classic references [54, 70], and the recent review [11]) and popularized as 'distributional dynamics' in NN literature (e.g. [53]). When β > 0 , a solution to (4) has a density w.r.t. λ and is actually strong. Under rather simple technical assumptions (see SuppMat-D.3, or [12, 15, 38, 57, 69]), when τ, β > 0 it is known that the WGF ( R τ,β ) W 2 -converges to a (unique) minimizer. When τ, β = 0 a sort of converse holds (see [16]): if WGF ( R ) converges in W 2 , then the limit minimizes R.

Proven by [16, 53, 62, 67] and later refined e.g. by [14, 22, 23, 51, 66, 69], the main result in the MF Theory of overparametrized shallow NNs states that SGD training for a shallow NN, in the right scaling limit as N →∞ , approximates a WGF :

Theorem 1 ( Mean-Field limit, sketch ) . For each T > 0 , under relevant technical assumptions including regularity of σ ∗ and a proper asymptotic behaviour of ε N → 0 as N →∞ , the rescaled empirical process given by µ N := ( ν N ⌊ t/ε N ⌋ ) t ∈ [0 ,T ] converges in law (in the Skorokhod space D P ( Z ) ([0 , T ]) ) to µ := ( µ t ) t ∈ [0 ,T ] given by the unique WGF ( R τ,β ) starting at µ 0 .

Despite the MF limit of NNs being a theoretical approximation, the behavior it predicts can effectively be observed in practice, even for finite, not too large N (see the numerical experiments in many of the aforementioned works and below). Moreover, it is the asymptotic regime that most closely describes the actual feature-learning behavior observed in large, overparametrized NNs during training (as compared e.g. to the lazy-training regime described by the Neural Tangent Kernel approximation [17, 41]). Note that, for β > 0 , the entropy term H λ in WGF ( R τ,β ) (as well as the Laplace operator in equation (4)) is approximated, in practice, by the Gaussian noise term in the SGLD (1), as N →∞ .

2.3 Symmetry-leveraging techniques

We next discuss mathematical formulations of the main techniques to leverage posited distributional symmetries of the data at the training or architecture levels. We henceforth fix a compact group G of normalized Haar measure λ G , acting on X and Y , which we denote G ⟳ ρ X , G ⟳ ˆ ρ Y . 2 A function f : X → Y is termed equivariant if ∀ g ∈ G , ˆ ρ g -1 .f ( ρ g .x ) = f ( x ) dπ X ( x ) -a.s. We further say that the data ( X,Y ) ∼ π is equivariant , and write π ∈ P G ( X ×Y ) , if ∀ g ∈ G , ( ρ g .X, ˆ ρ g .Y ) ∼ π (this is not enforced unless stated). The space of functions f : X → Y square-integrable (in Bochner sense) w.r.t π X = Law ( X ) is called L 2 ( X , Y ; π X ) . Further relevant concepts are introduced as needed.

Data Augmentation (DA): This training scheme considers { g k } k ∈ N i.i.d. ∼ λ G independent from the { ( X k , Y k ) } k ∈ N in (1), and carries out SGD on samples { ( ρ g k .X k , ˆ ρ g k .Y k ) } k ∈ N . DA and the vanilla training scheme would thus be equivalent if π ∈ P G ( X × Y ) . One can show (see [13, 48]) that, performing SGD with DA , results in an optimization scheme for the symmetrized population risk , R DA ( θ ) := E π [∫ G ℓ ( Φ N θ ( ρ g .X ) , ˆ ρ g .Y ) dλ G ( g ) ] . Despite being effective in practice, DA gives no guarantee that the resulting model will be equivariant. For deeper insights, see [13, 21, 39, 46, 48, 52].

Feature Averaging (FA): Instead of focusing on the data, FA works with symmetrized versions of the vanilla NN models Φ N θ at hand, averaging model copies over all possible translations through the group action. This amounts to constructing (or approximating) the symmetrization operator over L 2 ( X , Y ; π X ) defined as ( Q G .f )( x ) := ∫ G ˆ ρ g -1 .f ( ρ g .x ) dλ G ( g ) (see [27]), and trying to minimize R FA ( θ ) := E π [ ℓ ( ( Q G . Φ N θ )( X ) , Y )] (see [13, 21, 46, 48]). The resulting model will be equivariant, however, FA is inefficient, as ≈ | G | times more evaluations are needed for training and inference.

Equivariant Architectures (EA): Following [8], EA in multilayer NNs are configurations yielding models equivariant between each of the hidden layers (where G is assumed to act). As stated in [29, 30, 61, 64, 65, 75] and [2, 18, 44, 45, 73], once the (equivariant) activation functions between the different layers have been fixed, EA s are plainly parameter-sharing schemes (determined by the space of intertwiners/group convolutions between layers). In our context, assuming that G ⟳ M Z is some group action, we require that σ ∗ : X × Z → Y is jointly equivariant , namely, ∀ ( g, x, z ) ∈ G ×X ×Z , σ ∗ ( ρ g .x, M g .z ) = ˆ ρ g σ ∗ ( x, z ) ; to ensure G -actions over different spaces are properly related. Introducing the set of fixed points for G ⟳ M Z , E G := { z ∈ Z : ∀ g ∈ G, M g .z = z } , a shallow NN model Φ N θ thus has an EA if θ ∈ ( E G ) N . Under the right choices of σ ∗ and M , the obtained EA s can encode interesting and widely applied architectures, such as CNNs [19] and DeepSets [77] (see SuppMat-C.1 for further discussion). We call E G the subspace of invariant parameters , which is a closed linear subspace of Z , with unique orthogonal projection P E G : Z → E G , explicitly given by P E G .z := ∫ G M g .z dλ G ( g ) for z ∈ Z (see [30]). We are thus led to solve: min θ ∈ ( E G ) N R ( θ ) or, equivalently, to find the best projected model Φ N,EA θ := ⟨ σ ∗ , P E G # ν N θ ⟩ , by minimizing R EA ( θ ) := E π [ ℓ ( Φ N,EA θ ( X ) , Y )] . This can considerably reduce the parameter space dimension; however EA might generally yield a decreased expressivity or approximation capacity.

3 Symmetries in overparametrized neural networks: main results

3.1 Two notions of symmetries for parameter distributions

The following notions regarding distributions from P ( Z ) are central to our work:

Definition 5. Given µ ∈ P ( Z ) , we respectively define its symmetrized and projected versions as µ G := ∫ G ( M g # µ ) dλ G and µ E G := P E G # µ . Moreover, we introduce two subspaces of P ( Z ) : P G ( Z ) := { µ ∈ P ( Z ) : ∀ g ∈ G, M g # µ = µ } and P ( E G ) := { µ ∈ P ( Z ) : µ ( E G ) = 1 } .

Example. For G = {± 1 } acting multiplicatively on Z = R , one has E G = { 0 } , hence P ( E G ) = { δ 0 } , while P G ( Z ) = { 1 2 ( ν + ν ( -· )) : ν ∈ P ( R + ) } . In particular, for z ∈ Z , ( δ z ) G = 1 2 ( δ z + δ -z ) . Definition 6 ( Invariant Probability Measures ) . We say that µ ∈ P ( Z ) is:

We notice that: P ( E G ) ⊆ P G ( Z ) , µ G ∈ P G ( Z ) and µ E G ∈ P ( E G ) . Thus, SI implies WI . Next result relates the symmetrization operation on P ( Z ) with the one on shallow models F σ ∗ ( P ( Z )) :

Proposition 1. Let Φ µ ∈ F σ ∗ ( P ( Z )) with σ ∗ : X × Z → Y jointly equivariant. Then:

That is to say, the closest equivariant function (in L 2 ( X , Y ; π X ) ) to Φ µ is given by the shallow model associated to the symmetrized version of µ .

Remark. In particular, Φ µ is equivariant as soon as µ is WI only. Conversely, if Φ µ : X → Y is an equivariant function, then Φ µ = Φ µ G , i.e. it can be expressed in terms of a WI distribution. This highlights a priority role of WI distributions on Z in representing invariant shallow models.

The alternative, 'projected model' Φ µ E G , in turn, is never the closest equivariant shallow model, to Φ µ in L 2 ( X , Y , π X ) , unless equal to Φ µ G . The latter rarely is the case (unlike commonly implied in the literature). In fact, the symmetrized version Q G Φ N θ of a shallow NN model Φ N θ involves ( ν N θ ) G = 1 N ∑ N i =1 φ θ i , where ∀ z ∈ Z , φ z is the orbit measure of the action, 3 and has N · | G |

Z -valued parameters (possibly with | G | = ∞ ). This sharply contrasts ( ν N θ ) E G = 1 N ∑ N i =1 δ P E G .θ i , which has ≤ N distinct parameters, all living in E G . So, in general, depending on σ ∗ and G , one might have ⟨ σ ∗ , ( ν N θ ) E G ⟩ ̸ = Q G Φ N θ . A notable case in which the equality holds is the class of linear models, which is discussed in SuppMat-E.1.

̸

Example. In the previous example, for µ = δ z and σ ∗ jointly equivariant, we have Φ µ = σ ∗ ( · , z ) , Φ µ G = Q G Φ µ = 1 2 ( σ ∗ ( · , z ) + σ ∗ ( · , -z )) and Φ µ E G = σ ∗ ( · , 0) which are generally distinct if z = 0 . Notice that Φ µ G is an equivariant function without any of its 'parameters' living in E G .

3.2 Invariant functionals on P ( Z ) and their optima

In the same spirit as when defining the population risk R : P ( Z ) → R in (2), the risk functions associated with SL-techniques from Section 2.3 can be lifted to functionals over P ( Z ) , namely to: R DA ( µ ) := E π [∫ G ℓ ( Φ µ ( ρ g .X ) , ˆ ρ g .Y ) dλ G ( g ) ] , R FA ( µ ) := E π [ ℓ ( Q G (Φ µ )( X ) , Y )] and R EA ( µ ) := E π [ ℓ ( Φ µ E G ( X ) , Y )] , respectively. This will allow us to study these SL-techniques, in the overparametrized regime, under a common MF framework. We need the following assumption: Assumption 1. π ∈ P 2 ( X × Y ) ; ℓ : Y × Y → R is convex, jointly invariant and differentiable with ∇ 1 ℓ linearly growing; and σ ∗ : X × Z → Y is bounded, jointly equivariant and differentiable.

The quadratic loss ℓ ( y, ˆ y ) = 1 2 || y -ˆ y || 2 is an example of such ℓ . Having σ ∗ bounded and differentiable is a simplifying assumption, usually made in the MF literature, when establishing key results such as global convergence of NN (see e.g. [12, 38, 53]); relaxing this condition to include further commonlyused functions σ ∗ seems feasible, up to some additional technicalities (see SuppMat-A.2 for further discussion). Finally, having σ ∗ be jointly equivariant (as defined in section 2.3) isn't a truly restrictive assumption: under the right choice of σ ∗ and M , any usual single-hidden-layer NN architecture can be made to satisfy it (see SuppMat-C.1 for a deeper discussion). We also need:

Definition 7. A functional F : P ( Z ) → R is invariant if F ( M g # µ ) = F ( µ ) ∀ ( g, µ ) ∈ G ×P ( Z ) ; equivalently, if it equals its symmetrized version F G ( µ ) := ∫ G F ( M g # µ ) dλ G ( g ) .

Proposition 2. Under Assumption 1, R DA , R FA and R EA are invariant (and convex) and we have:

In particular, R = R DA if R is invariant. Moreover, ∀ µ ∈ P G ( Z ) , R ( µ ) = R DA ( µ ) = R FA ( µ ) . Last, if π ∈ P G ( X × Y ) (the data distribution is equivariant), then R is invariant.

The proof relies on Proposition 1 and calculations as in [30], see SuppMat-E.2. Next result is a general property of functionals over P ( Z ) , which is key for the forthcoming analysis:

Proposition 3 ( Optimality for Invariant Functionals ) . Let F : P ( Z ) → R be convex, C 1 and invariant. Then: ∀ µ ∈ P ( Z ) , F ( µ G ) ≤ F ( µ ) ; and so, inf µ ∈P G ( Z ) F ( µ ) = inf µ ∈P ( Z ) F ( µ ) . In particular, if F has a unique minimizer over P ( Z ) , it must be WI .

The proof relies on an ad-hoc version of Jensen's inequality. Next, we state that optimizing under DA and FA is essentially equivalent, and corresponds to optimizing R exclusively over WI measures:

Theorem 2 ( Equivalence of DA and FA ) . Under assumption 1, we have:

Note that, on the other hand, R EA only satisfies: inf µ ∈P ( Z ) R EA ( µ ) = inf µ ∈P ( E G ) R ( µ ) . In the case of the quadratic loss, Theorem 2 can be made more explicit:

Corollary 1. Under Assumption 1, when the loss is quadratic and π X is invariant, we have:

where f ∗ = E π [ Y | X = · ] , and R ∗ , ˜ R ∗ are constants only depending on π and f ∗ . That is, optimizing under DA and FA corresponds to approximating the symmetrized version of f ∗ .

Under equivariance of the data distribution π , the following general result also holds:

Corollary 2. Let Assumption 1 hold and suppose π ∈ P G ( X × Y ) . Then, R is invariant and therefore: inf µ ∈P ( Z ) R ( µ ) = inf µ ∈P G ( Z ) R ( µ ) = inf µ ∈P ( Z ) R DA ( µ ) = inf µ ∈P ( Z ) R FA ( µ ) .

Remark. Consequently, equivariant data allow us to globally optimize the population risk by only considering WI measures. It also shows that DA and FA provide no advantage for this optimization.

The same unfortunately is not true for SI measures (answering a question in [27]), as shown by the following result, which constructs a simple example in which E G is trivial:

Proposition 4. Even with a finite group G acting orthogonally on X = R d , Y = R and Z = R ( d +2) ; with π being compactly-supported and equivariant; with ℓ being quadratic; and with σ ∗ being C ∞ , bounded and jointly equivariant; we can have: inf µ ∈P ( Z ) R ( µ ) < inf ν ∈P ( E G ) R ( ν ) .

In fact, even if R is invariant, when E G is too restrictive , it might become impossible to globally optimize R over SI measures (which amounts to using R EA as a proxy for R ). This subtlety has to be considered when deciding to use EA s on problems where symmetries exist. Nevertheless, if E G has good universality properties, a true SI solution to the learning problem can be sought for:

Proposition 5. Let Assumption 1 hold, ℓ be quadratic and π ∈ P G 2 ( X ×Y ) . If F σ ∗ ( P ( E G )) is dense in L 2 G ( X , Y ; π X ) := Q G ( L 2 ( X , Y ; π X )) , then: inf µ ∈P ( Z ) R ( µ ) = inf ν ∈P ( E G ) R ( ν ) = R ∗ .

Remark. See e.g. [50, 60, 76, 77] for conditions on E G and σ ∗ guaranteeing this 'restricted' universality on L 2 G ( X , Y ; π X ) . These allow for effectively solving the problem in fewer dimensions, which is key in successful EA like CNNs and DeepSets. See SuppMat-E.2.5 for a deeper discussion.

3.3 Symmetries and SL training dynamics in the overparametrized regime

We now study the MFL of the various training dynamics when Z = R D . We begin with the general:

Theorem 3. Let F : P ( Z ) → R be an invariant functional such that WGF ( F ) is well defined and has a unique (weak) solution ( µ t ) t ≥ 0 . If µ 0 ∈ P G 2 ( Z ) , then, for dt -a.e. t ≥ 0 we have µ t ∈ P G 2 ( Z ) .

The proof of Theorem 3 relies on D µ F being equivariant (in a suitable sense) and ( M g µ t ) t ≥ 0 satisfying also, as a consequence, WGF ( F ) (See SuppMat-E.3 for the details). Note that µ 0 ∈ P G 2 ( Z ) is simply verified, e.g. by a standard Gaussian in Z . Specializing this result, we get:

Corollary 3. Let Assumption 1 and technical assumptions (as in [16]) hold. Then, if R and r are invariant, WGF ( R τ,β ) starting from µ 0 ∈ P G 2 ( Z ) satisfies: for dt -a.e. t ≥ 0 , µ t ∈ P G 2 ( Z ) . If moreover β > 0 , each µ t has a density function invariant with respect to G ⟳ M Z .

Remark. If π is equivariant, R is invariant, and this result is valid for a freely-trained NN, without employing SL-techniques . In a way, MFL incorporates these symmetries from infinite SGD iterations.

Theorem 3 and Corollary 3 can thus be seen as significant generalizations of Proposition 2.1 from [35], which addresses the case of wide 2-layer ReLU networks with a target function that's symmetric w.r.t. a single orthogonal transformation. The fact that strong solutions to WGF ( R τ,β ) can be sought among invariant functions to reduce the complexity when π is equivariant, was also first hinted in [53]. The natural domain of invariant functions is in fact the quotient space of G ⟳ M Z (and not E G , which is strictly embedded in it).

Comparing the different training dynamics at the MF level and applying Proposition 2, we also get:

Theorem 4. Under assumptions of Corollary 3, if µ 0 ∈ P G 2 ( Z ) , WGF ( R DA ) and WGF ( R FA ) solutions are equal. If further R is invariant, the WGF ( R ) solution coincides with them too.

Remark. In particular, with equivariant data (i.e. invariant R ), training with DA or FA is essentially the same, at least at the MF level, as using no SL-technique whatsoever . Hence, a relevant, practical open question, is: how do the convergence rates to the MFL compare in all three cases, as N →∞ ?

We will now see that similar results hold for P ( E G ) instead of P G ( Z ) . Notice that the entropyregularized risk forces each µ t to have a density w.r.t. λ in Z if β > 0 . Therefore, if G ⟳ M Z is non-trivial (thus E G is a strict subspace), we always have µ t ̸∈ P ( E G ) . It thus seems natural to restrain the noise in equation (1) to stay in E G ; namely, to consider the projected noisy SGD dynamic:

Note that projecting only the noise in (5) doesn't force θ k +1 i to be in E G , even if θ k i was. Introducing the related projected-regularized functional : R τ, β E G ( µ ) := R ( µ ) + τ ∫ rdµ + βH λ E G ( µ E G ) , with λ E G the Lebesgue Measure over E G , we get the following analogue of Corollary 3:

Theorem 5. Let Assumption 1 and technical assumptions on R hold. Suppose that R and r are invariant. Then, if µ 0 ∈ P 2 ( E G ) , ( µ t ) t ≥ 0 solution of WGF ( R τ,β E G ) satisfies ∀ t ≥ 0 , µ t ∈ P 2 ( E G ) .

The result holds for β ≥ 0 . Its proof is based on pathwise properties of the McKean-Vlasov stochastic differential equation (studied e.g. in [22]) associated with the WGF ( R τ,β E G ), see SuppMat-D.2.

Remark. Theorem 5 is significantly stronger than Corollary 3: it implies that, for equivariant π , the flow will remain in the set of SI distributions all throughout its evolution, despite there being no explicit constraint on the network parameters during training (they can all be freely updated), nor any SL-technique being used. This is a highly non-intuitive fact, and a large N exclusive phenomenon, as our numerical experiments will show. See SuppMat-D.2 for a deeper discussion.

Remark. Notice that, despite the computation of E G being generally hard (see [29]), µ 0 ∈ P 2 ( E G ) can be achieved by simply setting µ 0 = δ ⃗ 0 . Moreover, since one can also take β = 0 , 'having access' to the noise projection P E G is never explicitly required, allowing for a broader applicability of the result. In particular, as we'll show in our experiments, usual shallow NNs with all parameters initialized at { 0 } , freely trained with 'noiseless' SGD, will satisfy Theorem 5 in the MFL .

Remark. Theorem 5 holds too for the invariant functionals R DA , R FA and R EA in the role of R, even if π is not equivariant . Notably, DA , FA and EA procedures starting on a SI distribution, despite being free to involve all parameters, will keep the distribution SI all throughout training.

Last, we also have:

Theorem 6. Let the conditions for Theorem 5 hold. If µ 0 ∈ P 2 ( E G ) , the WGF ( R FA ), WGF ( R DA ) and WGF ( R EA ) solutions coincide . If R is invariant, WGF ( R ) solution coincides with them too.

4 Numerical experiments and architecture-discovery heuristic

To empirically illustrate some of our results from the previous section, we consider synthetic data produced in a teacher-student setting (see e.g. [14, 16]). Code necessary for replicating the obtained results, as well as a detailed description of our experimental setting, can be sought in the SuppMat.

We study a simple setting with: X = Y = R 2 , Z = R 2 × 2 ∼ = R 4 , and G = C 2 acting on X and Y by permuting the coordinates ; and on Z via the natural intertwining action (for which E G is explicit). We take the jointly equivariant activation σ ∗ ( x, z ) = σ ( z · x ) , ∀ ( x, z ) ∈ X × Z with σ : R → R a sigmoidal function applied pointwise; and consider normally distributed features, and labels produced from a teacher model f ∗ . This teacher is given by a shallow NN model, either f ∗ = Φ N ∗ θ ∗ with N ∗ = 5 arbitrary particles θ ∗ ∈ Z N ∗ , or its symmetrized version f ∗ = Q G . Φ N ∗ θ ∗ (referred to as WI ), with 10 particles. 4 Notice that the data distribution π will be equivariant only if the teacher is. We try to mimic such teacher with a student model , Φ N θ , with the same σ ∗ , but different particles θ ∈ Z N that will be trained to minimize the regularized population risk R τ,β (with quadratic loss and penalization). For this we employ the SGD dynamic given by Equation (1) (or projected, if required, as in Equation (5)), possibly involving DA , FA or EA techniques. We refer to free training , with no SL-techniques involved, as vanilla training. Each experiment was repeated N r = 10 times to ensure consistency. Explicit values of the fixed training parameters are found in SuppMat-F.

4.1 Study for varying N

We demonstrate how properties of WGF ( R τ,β ) stated in Section 3.3 become apparent as N grows. We consider a teacher with ν N ∗ θ ∗ either arbitrary or WI , and different training schemes performed over N e epochs, all initialized with the same particles drawn from given µ 0 ∈ P 2 ( Z ) that is either SI or WI . Figure 1 displays the behavior of the student's particle distribution, ν N N e , after training, in terms of certain 'normalized version' of the W 2 -distance, which we call simply Relative Measure Distance, or RMD for short. 5 We refer to SuppMat-F for further insights and, additionally: a deeper

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

7

10

5

10

3

10

1

Relative Measure Distance (RMD)

Distance to Projected Version

Model

V vs. P(V)

DA vs. P(DA)

FA vs. P(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

(a)

SI

initialization and

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

4

10

3

10

2

10

1

10 0

Relative Measure Distance (RMD)

Distance to Symmetrized Version

Model

V vs. G(V)

DA vs. G(DA)

FA vs. G(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

(b)

WI

initialization and

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

6

10

4

10

2

Relative Measure Distance (RMD)

Distance to Projected Version

Model

V vs. P(V)

DA vs. P(DA)

FA vs. P(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

(c)

SI

initialization and

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

6

10

5

10

4

10

3

10

2

10

1

10 0

Relative Measure Distance (RMD)

Distance to Symmetrized Version

Model

V vs. G(V)

DA vs. G(DA)

FA vs. G(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

(d)

WI

initialization and

WI

teacher

analysis of the case of ν N ∗ θ ∗ being SI , comparisons between different techniques and EA , and L 2 comparisons between Φ N θ and both f ∗ and Q G .f ∗ (to illustrate Corollary 1).

We first look at the SI -initialized training. Though from [30] we know that (exact) DA or FA during training will respect E G without needing to pass to the MFL . This is certainly not true in general for the vanilla scheme, where the symmetry is never explicit for the model. We notice in Figure 1, however, that, as N grows big, the SI -initialized vanilla training scheme, under only a WI teacher, does remain SI throughout training, as predicted in Theorem 5. This is absolutely remarkable, since there is no intuitive reason why the vanilla scheme (were parameters can be updated freely ) shouldn't escape E G to better approximate f ∗ . For instance, for an arbitrary teacher (with the same particles, but un-symmetrized ) vanilla training readily leaves E G to better approximate f ∗ . Though this isn't a predicted behaviour from our theory, it motivates a heuristic we present in the upcoming section. On the other hand, and as expected, both DA and FA consistently remain within E G almost independently of N , and even if f ∗ isn't equivariant. Finally, as N grows bigger, the end-of-training distribution of the vanilla scheme approaches that of DA and FA (as expected from Theorem 4).

Regarding the WI -initialized training, unlike the SI case, particles sampled i.i.d. from a WI distribution don't necessarily yield a WI empirical distribution ν N 0 . On the one hand, this means we require large N to see ν N N e being (approx.) WI ; and on the other hand, it means we have no guarantee that DA and FA will be close unless we look at larger N (where Theorem 4 applies). The second column of Figure 1 precisely shows these behaviours as N grows: both a trend of ν N N e towards becoming WI , and a clear coincidence between the DA , FA and vanilla schemes (the latter only for equivariant f ∗ ).

4.2 Heuristic algorithm for discovering EA parameter spaces

From these experiments, for non-equivariant f ∗ , the SI -initialized WGF is seen to eventually escape E G . In turn, for equivariant f ∗ , Figure 2 shows that a training scheme initialized at E ⊊ E G (i.e. ν N θ 0 ∈ P ( E ) ), eventually leaves E , but stays within E G (as expected from Theorem 5). These empirical observations hint to an heuristic for discovering the 'good' EAs for the task at hand .

Assume indeed π equivariant w.r.t. G . We want to find the unknown, largest (i.e. most expressive) subspace of Z supporting SI measures. We hence consider some (large) N , a shallow NN model with e.g. σ ∗ ( x, z ) = σ ( z.x ) , and numerical thresholds ( δ j ) j ∈ N ⊆ R + . We define E 0 = { 0 } ≤ E G as an initial subspace and initialize ν N θ 0 = ν N ⃗ 0 ∈ P ( E 0 ) . Then, we iteratively proceed as follows:

Teacher Particles

Student Particles

Origin

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions

End of Step 0

z4

Teacher Particles

Student Particles

Origin

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions

Start of Step 1

z4

Teacher Particles

Student Particles

Origin

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions

End of Step 1

z4

0.6

0.4

0.2

z3

0

-0.2

0.4

0.6

0.6

0.4

-0.2

0.d

z1

0

0.2

0.2

0

22

0).4

0.6

0.4

0.2

z3

0

-0.2

-0.4

0.6

0.A

-02

0.2

0.4

0.6

z2

z1

0.6

0.4

0.2

z3

0

0.2

0.4

0.6

0.4

0.4

0.6

-0.2

0.2

z1

0

0.2

0.2

0

22

0.4

0.6

0.4

0.2

z3

0

0.2

-0A

-0.6

0.A

0.4

0.6

02

0.2

z2

z1

0.6

0.4

0.2

z3

0

0.2

0.4

0.6

0.4

Z0-

0.4

0.6

z1

0

0.2

z2

0.4

0.6

0.4

0.2

z3

0

-02

-0.4

-0.6

0.A

0.6

02

Q

22

z1

̸

For j = 0 , 1 , . . . , initialize a model at ν N θ 0 ∈ P ( E j ) , train it for N e epochs, and check whether RMD 2 ( ν N N e , P E j # ν N N e ) ≤ δ j . If that is the case, the training didn't escape E j , and one could suppose E G := E j . Otherwise, it left E j (so E G = E j ) and one can set e.g. E j +1 := E j ⊕ v E j , with v E j := 1 N ∑ N i =1 ( θ N e i -P E j .θ N e i ) . Allegedly, this scheme would eventually leave all strict subspaces to reach the 'right' E G . Figure 2 indeed illustrates this behaviour in our simple teacherstudent example (see SuppMat-F.2 for further details). Notice that we start knowing close to nothing about data symmetries ( E 0 = { 0 } ), and end up 'discovering' a data-based parameter-sharing scheme ( E ∗ = E G ) that allows for building SI NNs. This idea might have potential for real world applications, yet a larger scale experimental analysis and rigorous theoretical guarantees need to be provided.

We refer to [72] for a different approach to this idea of 'discovering the real symmetries of the data'. Their work uses relaxed group convolution layers to discover 'data-driven symmetry-breaking' in ML problems. A deeper comparison between both approaches shall be found in SuppMat-F.2.

5 Conclusion

In the light of theoretical guarantees given by the MF theory of overparametrized shallow NN, we explored their training dynamics when data is possibly equivariant for some group action and/or SL techniques are employed. We thus described how DA , FA and EA schemes can be understood in the limit of infinite internal units, and studied in that setting the qualitative advantages that can be attired from their use. In this MFL , DA and FA are essentially equivalent in terms of the optimization problem they solve and the trajectory of their associated WGF s. Moreover, for equivariant data, freely -trained NN, in the MFL , obey the same WGF as DA / FA . They also 'respect' symmetries during training, as WI and SI initializations (corresponding to symmetric parameter distributions and EA configurations) are preserved throughout, even if potentially all NN parameters can be updated. We also highlighted the prominent role of WI laws in representing equivariant models. We illustrated our results with appropriate numerical experiments, which in turn suggested a data-driven heuristic to find appropriate parameter subspaces for EA s in a given task. Providing theoretical guarantees for this heuristic is an interesting problem left for future work. A further relevant question to address, is to quantify and compare the speeds at which all studied training schemes approach the MFL , as this would a provide a full comparative picture of their performances. Extending the MF analysis of symmetries to NNs with more complex inner structures is another interesting line of work.

References

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows in Metric Spaces and in the Space of Probability Measures . Lectures in Mathematics ETH Zürich. Birkhäuser, 2. ed edition, 2008. OCLC: 254181287.

Jimmy Aronsson. Homogeneous vector bundles and g-equivariant convolutional neural networks. Sampling Theory, Signal Processing, and Data Analysis , 20(2), jul 2022.

Dominique Bakry, Ivan Gentil, Michel Ledoux, et al. Analysis and geometry of Markov diffusion operators , volume 103. Springer, 2014.

Raphaël Barboni, Gabriel Peyré, and François-Xavier Vialard. Understanding the training of infinitely deep and wide resnets with conditional optimal transport. arXiv preprint arXiv:2403.12887 , 2024.

Andrew Barron. Universal approximation bounds for superpositions of a sigmoidal function. ieee trans. on information theory 39, 930-945. Information Theory, IEEE Transactions on , 39:930 - 945, 06 1993.

Benjamin Bloem-Reddy, Yee Whye, et al. Probabilistic symmetries and invariant neural networks. Journal of Machine Learning Research , 21(90):1-61, 2020.

Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review , 60(2):223-311, 2018.

Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veliˇ ckovi´ c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 , 2021.

Pierre Cardaliaguet. Notes on mean-field games (from P.-L. Lions lectures at Collège de France). 2013. Available at: https://www.ceremade.dauphine.fr/~cardaliaguet/ MFG20130420.pdf .

R. Carmona and F. Delarue. Probabilistic Theory of Mean Field Games with Applications I: Mean Field FBSDEs, Control, and Games . Probability Theory and Stochastic Modelling. Springer International Publishing, 2018.

Louis-Pierre Chaintron and Antoine Diez. Propagation of chaos: A review of models, methods and applications. I. models and methods. Kinetic and Related Models , 15(6):895-1015, 2022.

Fan Chen, Yiqing Lin, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for kinetic mean field langevin dynamics. Electronic Journal of Probability , 29:1-43, 2024.

Shuxiao Chen, Edgar Dobriban, and Jane H Lee. A group-theoretic framework for data augmentation. Journal of Machine Learning Research , 21(245):1-71, 2020.

Zhengdao Chen, Grant Rotskoff, Joan Bruna, and Eric Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. Advances in Neural Information Processing Systems , 33, 2020.

Lénaïc Chizat. Mean-field langevin dynamics : Exponential convergence and annealing. Transactions on Machine Learning Research , 2022.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in Neural Information Processing Systems , 31, 2018.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems , 32, 2019.

Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant CNNs on homogeneous spaces. Advances in Neural Information Processing Systems , 32, 2019.

Y. Le Cun, B. Boser, J. S. Denker, R. E. Howard, W. Habbard, L. D. Jackel, and D. Henderson. Handwritten digit recognition with a back-propagation network , page 396-404. Morgan Kaufmann Publishers Inc., San Francisco, CA, USA, 1990.

George V. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems , 2:303-314, 1989.

Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data augmentation. In International conference on machine learning , pages 1528-1537. PMLR, 2019.

Valentin De Bortoli, Alain Durmus, Xavier Fontaine, and Umut Simsekli. Quantitative propagation of chaos for SGD in wide neural networks. Advances in Neural Information Processing Systems , 33, 2020.

Arnaud Descours, Arnaud Guillin, Manon Michel, and Boris Nectoux. Law of large numbers and central limit theorem for wide two-layer neural networks: the mini-batch and noisy case. arXiv preprint arXiv:2207.12734 , 2022.

J. Diestel and J.J. Uhl. Vector Measures . Mathematical surveys and monographs. American Mathematical Society, 1977.

C. Dru¸ tu and M. Kapovich. Geometric Group Theory . Colloquium Publications. American Mathematical Society, 2018.

Paul Dupuis and Richard S Ellis. A weak convergence approach to the theory of large deviations . John Wiley & Sons, 2011.

Bryn Elesedy and Sheheryar Zaidi. Provably strict generalisation benefit for equivariant models. In International conference on machine learning , pages 2959-2969. PMLR, 2021.

Stewart N Ethier and Thomas G Kurtz. Markov processes: characterization and convergence . John Wiley & Sons, 2009.

Marc Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International conference on machine learning , pages 3318-3328. PMLR, 2021.

Axel Flinth and Fredrik Ohlsson. Optimization dynamics of equivariant and augmented neural networks. arXiv preprint arXiv:2303.13458 , 2023.

Qiang Fu and Ashia Wilson. Mean-field underdamped langevin dynamics and its space-time discretization. arXiv preprint arXiv:2312.16360 , 2023.

Wilfrid Gangbo and Adrian Tudorascu. On differentiability in the wasserstein space and wellposedness for hamilton-jacobi equations. Journal de Mathématiques Pures et Appliquées , 125:119-174, 2019.

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, 2016. http://www.deeplearningbook.org .

Philipp Grohs and Gitta Kutyniok. Mathematical aspects of deep learning . Cambridge University Press, 2022.

Karl Hajjar and Lénaïc Chizat. On the symmetries in the dynamics of wide two-layer neural networks. Electron. Res. Arch. , 31(4):2175-2212, 2023.

B.C. Hall. Lie Groups, Lie Algebras, and Representations: An Elementary Introduction . Graduate Texts in Mathematics. Springer, 2003.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks , 2(5):359-366, 1989.

Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. In Annales de l'Institut Henri Poincare (B) Probabilites et statistiques , volume 57, pages 2043-2065. Institut Henri Poincaré, 2021.

Kevin H Huang, Peter Orbanz, and Morgane Austern. Quantifying the effects of data augmentation. arXiv preprint arXiv:2202.09134 , 2022.

Ningyuan Huang, Ron Levie, and Soledad Villar. Approximately equivariant graph networks. Advances in Neural Information Processing Systems , 36, 2024.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems , 31, 2018.

Olav Kallenberg. Probabilistic Symmetries and Invariance Principles . Springer, Dordrecht, 2005.

Olav Kallenberg. Random Measures, Theory and Applications , volume 77. Springer, 01 2017.

Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International conference on machine learning , pages 2747-2755. PMLR, 2018.

Leon Lang and Maurice Weiler. A wigner-eckart theorem for group equivariant convolution kernels. In International Conference on Learning Representations , 2020.

Zhiyuan Li, Ruosong Wang, Dingli Yu, Simon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora. Enhanced convolutional neural tangent kernels. arXiv preprint arXiv:1911.00809 , 2019.

Pierre Luis Lions. Cours au College de France . 2008.

Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178 , 2020.

Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902 , 2018.

Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant networks. In International conference on machine learning , pages 4363-4371. PMLR, 2019.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory , pages 3351-3418. PMLR, 2021.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665E7671, 2018.

Sylvie Méléard. Asymptotic behaviour of some interacting particle systems; McKean-Vlasov and Boltzmann models , pages 42-95. Springer Berlin Heidelberg, Berlin, Heidelberg, 1996.

Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. Mathematical Statistics and Learning , 6(3):201-357, 2023.

Tomohiro Nishiyama. Convex optimization on functionals of probability densities. arXiv preprint arXiv:2002.06488 , 2020.

Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics , pages 9741-9757. PMLR, 2022.

F. Otto and C. Villani. Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis , 173(2):361-400, 2000.

Mircea Petrache and Shubhendu Trivedi. Approximation-generalization trade-offs under (approximate) group equivariance. Advances in Neural Information Processing Systems , 36, 2024.

Siamak Ravanbakhsh. Universal equivariant multilayer perceptrons. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 7996-8006. PMLR, 13-18 Jul 2020.

Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parametersharing. In International conference on machine learning , pages 2892-2901. PMLR, 2017.

Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics , 75(9):1889-1935, jul 2022.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

John Shawe-Taylor. Threshold network learning in the presence of equivalences. In J. Moody, S. Hanson, and R.P. Lippmann, editors, Advances in Neural Information Processing Systems , volume 4. Morgan-Kaufmann, 1991.

John Shawe-Taylor. Sample sizes for threshold networks with equivalences. Information and Computation , 118(1):65-72, 1995.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics , 80(2):725-752, 2020.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. Mathematics of Operations Research , 47(1):120-152, 2022.

Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Alain-Sol Sznitman. Topics in propagation of chaos. In Paul-Louis Hennequin, editor, Ecole d'Eté de Probabilités de Saint-Flour XIX - 1989 , pages 165-251, Berlin, Heidelberg, 1991. Springer Berlin Heidelberg.

Cédric Villani. Optimal transport. Old and new , volume 338 of Grundlehren der mathematischen Wissenschaften . Springer-Verlag, Berlin, 2009.

Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, and Tess E. Smidt. Discovering symmetry breaking in physical systems with relaxed group convolution. In Proceedings of the 41st international conference on machine learning (ICML2024) , 2024.

Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable CNNs. Advances in neural information processing systems , 32, 2019.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML2011) , 2011.

Jeffrey Wood and John Shawe-Taylor. Representation theory and invariant neural networks. Discrete Applied Mathematics , 69(1):33-60, 1996.

Dmitry Yarotsky. Universal approximations of invariant maps by neural networks. Constructive Approximation , 55(1):407-474, 2022.

Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in Neural Information Processing Systems , 30, 2017.

A General Considerations for the Reader

This section presents a summary of recurrent notation, abbreviations and key concepts used in our work, as well as a discussion on its limitations, scope and possible extensions. We thank anonymous Reviewers for suggesting us to add this section to the original manuscript.

A.1 Summary of recurrent notation, abbreviations and key concepts

In this section we summarize the main abbreviations and notation employed throughout the body of the paper, as well as simple definitions of fundamental concepts from probability theory and algebra that might be useful for understanding our work

Table 1 serves as a glossary for the most used abbreviations in the body of the paper. Table 2 contains a summary of the notation that we recurrently use in our definitions, statements and proofs. It also provides some simple references to mathematical concepts that are key in our work.

For clarity, beyond the contents of both tables, we here also explain a few relevant concepts to the unfamiliar reader:

Bochner Integrals: These correspond to the right generalization of Lebesgue integrals for vector-valued functions (see [24] for further reference).

Say we have a function f : X → Y between Hilbert spaces, and such that π X is some measure on X , then we say f is square-integrable (in Bochner sense) if it satisfies: ∫ X ∥ f ( x ) ∥ 2 Y dπ X ( x ) < ∞ .

This integral also respects closed linear operators, as shown by Hille's theorem (see Theorem II.2.6 in [24]). Namely, if L : Y → Y is a closed linear operator over Y , we have: ∫ L.fdπ = L. ∫ fdπ . In particular, this also holds for bounded linear operators.

X X X X

In general, most of the basic and most intuitive properties of traditional integrals can also be expressed for Bochner integrals.

Compact Groups, Group Actions and the Haar measure: We recurrently talk about group actions via orthogonal representations throughout our work, so a due clarification may be required. We assume G to be a compact group. Namely, G has a topology that makes it compact, and so that the multiplication and inversion operations are continuous . We say that G acts on Z , which we denote G ⟳ Z , whenever there exists a map:

that satisfies T ( e G , z ) = z and T ( g 1 , T ( g 2 , z )) = T ( g 1 .g 2 , z ) , ∀ g 1 , g 2 ∈ G, ∀ z ∈ Z . We always assume the actions to be continuous ; namely, T is continuous with respect to the product topology. Alternatively, by denoting T g := T ( g, · ) , T · : g ∈ G ↦→ T g ∈ Sym ( Z ) is a group homomorphism and, if the action is continuous, T g is an homeomorphism ∀ g ∈ G . If we further assume that ∀ g ∈ G, T g is linear, we call T · a group representation . Further, if { T g } g ∈ G are orthogonal transformations, we call T · an orthogonal group representation

and denote the group action by G ⟳ T Z . This is the case for all of the G -actions considered throughout this work.

Working only with group representations is common-practice in the context of symmetries for NNs (see e.g. [13, 48]). Beyond NNs, a whole field of mathematics is of course dedicated to the study of group representations (see e.g. [36] for reference). In this work we borrow some of this theory's terminology, as we refer to equivariant linear maps also as intertwining maps or intertwiners . This is also why we refer to the layer-by-layer action in EA s for shallow NNs as an intertwining action (see e.g. SuppMat-C.1).

Finally, it is well known that a compact group G admits a unique normalized Haar measure λ G ∈ P ( G ) (see [25]), which is left and right invariant, finite on every compact set, outer regular on Borel sets and inner regular on open sets. It can be interpreted as the uniform distribution on G , and it is extensively used throughout this work.

For further references in the topic, the curious reader might be interested in [25, 42, 43], among many others.

Weak convergence of measures: Wealso recall one of the most used notions of convergence in a space of probability measures. Given a sequence ( µ n ) n ∈ N ⊆ P ( Z ) , we say it weakly converges to some point µ ∈ P ( Z ) if, for any continuous and bounded function f : Z → R , we have ⟨ f, µ n ⟩ - - - - → n →∞ ⟨ f, µ ⟩ . Notice that this type of convergence is weaker than W p convergence for p ≥ 1 (which additionally also requires the convergence of p -th order moments of the involved measures).

Notice how we write ⟨ f, µ ⟩ to denote ∫ Z fdµ . This notation is heavily used (and abused) throughout the core of our work.

A.2 Limitations, scope and possible extensions of our work

For convenience of the reader we here present a discussion of some limitations and of the scope of application of the present work. This subsection does not provide any mathematical results required for the sequel.

Some of these limitations are of technical nature, and regard specific assumptions made in order for the specific proofs of our theoretical results to hold. In absence of these conditions, or under less restrictive ones, some of our results (or weaker forms of them) might still hold true, but further research is needed in order to properly establish their validity. Other limitations are the object of more general research questions in this area.

On the infinite i.i.d. data sample: The MF theory of NN makes the assumption that it is possible to take an infinite i.i.d. sample from the data distribution π . This may appear as a limitation of our results, since real-world datasets are naturally finite. We acknowledge that this is in fact an abstraction, but it can nevertheless be a potentially good approximation of the behavior of the SGD algorithm when large datasets are available. Indeed, Theorem 1, which holds for ε N = o (1 /N ) , ensures that a sample size of the order of tN is required in order to approximate the WGF up to a ('macroscopic') time t . When the long-time convergence of the WGF can be effectively quantified (an active research question today, see e.g. [57] and references therein), we can estimate how large a data sample will be required in order to attain a predetermined generalization error level. On the other hand, even for π with finite support, the MF approximation will end up minimizing the empiricalrisk. This, as pointed out in Section 2.1, is not the same as minimizing generalization error w.r.t. some underlying data distribution, but it is widely interpreted (in the application of most NN-based machine learning algorithms) as a proxy of doing so. Once again, this approximation can be reasonably good when the data sample size is large enough.

On the infinite width of NN models: Despite the MFL being a theoretical tool (which requires that the width N goes to + ∞ ), it is still the asymptotic regime that most closely describes the feature-learning behavior of large NNs during training (see the discussion at the end of Section 2.2). We believe that truly useful insights can be obtained from it for real, moderately wide NNs. For instance, our experiments show that, in quite reasonable practical settings of shallow NNs (with standard pointwise sigmoid activation and objax' default SGD training), the predicted MF behavior is seen to emerge in practice, already for finite, not too large N (1000 was generally enough). Actually, most of the insights obtained from the

MF analysis of NN are, in fact, unaccessible from a fixed N standpoint (as is the case for Theorem 5). Further non-asymptotic and quantitative answers to practical questions (e.g. how many neurons are needed in order to attain a given level of generalization/population error at given computational cost?) could also be obtained from MF theory, via quantitative propagation of chaos results (see e.g. [12]).

On Assumption 1: This is the main assumption underlying many of our core results. Although it is generally simple and not excessively hard to satisfy in practice, it might seem constraining in the context of neural networks. In particular, the technical condition on the gradient of ℓ , as well as the boundedness of σ ∗ , could seem to limit the applicability of our results. However, these conditions can be replaced by alternative properties of the data distribution (e.g. that π has compact support, or finite moments up to a given degree). Similarly, lifting some of the rest of the technical assumptions required for establishing the MFL , is part of the ongoing research work in the field (see e.g. [12] or [38] for some alternative conditions).

Regarding the assumption that ℓ is G -invariant, it is known that traditional loss functions naturally satisfy this condition. The joint-equivariance of σ ∗ , on the other hand, is much less constraining than it may initially seem. In practice, depending on the choice of σ ∗ and M , it might even end up being a trivial constraint. A deeper discussion on this very assumption shall be found in SuppMat-C.1.

On the generality of shallow models Φ µ ∈ F σ ∗ ( P ( Z )) : These allow for modeling a wide range of situations (including some variants of multi-layer models). In fact, σ ∗ can by itself encode a complex deep architecture (see SuppMat-C.1) and the resulting shallow model can represent way more interesting structures than single-hidden-layer NNs (e.g. ensembles of such multilayer 'units' trained in parallel). Nevertheless, these shallow models are still far from including all possibilities in the context of NNs. In fact, a fully unified, satisfactory MF theory for general deep NNs is still an open, actively tackled question (for advancements on it see e.g. [4, 55, 68]). We believe that some of our key results (e.g. Theorems 3 and 4) can be extended to some of those settings, which is a question we will leave for future work.

Universality guarantees for shallow models Φ µ ∈ F σ ∗ ( P ( Z )) : In this work, we have only explored the simple setting of tensor-of-order-1 NNs as modelled through σ ∗ , see SuppMatC.1. In particular, it is known from [49] that the desired universality on equivariant models is not always possible with these kinds of networks (which is, of course, a limitation to the applicability of Proposition 5). Exploring other interesting situations that could be modeled with our current framework, or plainly reformulating it to account for new variants of NNs (e.g. high-order tensor NNs thay might allow for easier EA universality) is undoubtedly part of our future work. See SuppMat- E.2.5 for a related discussion.

On G -equivariant data distributions: The assumption of π being G -equivariant implies that π X is G -invariant as well [6]. This implication can seem a bit too restrictive in some settings: e.g. for image classification, it amounts to assuming that 'images can arrive with any possible orientation' at training time, which is not necessarily the case. However, as extensively discussed in the literature (see [13, 27, 30, 48]), assuming π X to be G -invariant is usually reasonable when the aim is to 'exploit symmetries' of the problem. Not having such assumption means that there are little to no properties from the data that can be exploited in a proof. On the other hand, as mentioned in remarks after Theorems 3 and 5, our proofs don't explicitly rely on having π ∈ P G ( X × Y ) , but rather on the risk functional R : P ( Z ) → R being G -invariant. In consequence, one could simply neglect the G -equivariance of π altogether, and introduce the implied inductive bias by applying DA (which forces the marginal on X to be G -invariant), or any other SL technique that achieves the same result.

On the numerical experiments: The suite of experiments that were presented in Section 4 and SuppMat-F, though quite insightful for our theoretical results, are quite limited in reach. In particular, finding ways to illustrate our theoretical results on less controlled settings , such as a real-world equivariant datasets, will be a key part of our future work.

The presented experiments come from a controlled setting, mainly looking to avoid heavy practical constraints (e.g. visualizing the huge parameter space of a NN trained over an image dataset can be exceedingly hard), as well as an increased complexity of the involved objects (e.g. even for finite groups, group actions can get severely more complicated than what we experienced). Due to our currently limited computational resources, we

leave these inquiries for future work. Similarly, taking our experiments to a larger scale (e.g. N = 100 . 000 ) is also a challenge left for future inquiries (our current results, with N ≤ 5000 , might still be considered as small N ). A significant problem for scaling our experimental verifications is the fact that we relied on calculating Wasserstein distances (which is usually really computationally-expensive) to provide rigorous numerical evidence of our theoretical findings.

Furthermore, we are yet to experiment with architectures that are compatible with infinite compact groups G ; namely for examples coming from physics and NeuralODEs. Different variants of the activation function should also be tested out (e.g. a ReLU, tanh, and many others), as well as variants of the optimization algorithm (e.g. Adam, RMSProp, etc.).

Finally, our heuristic needs to be tested on a larger scale, with more complex datasets and architectures. Also, theoretical guarantees to sustain it shall be provided in future work. We believe that similar arguments as in [72] could be developed for our case; and alternative approaches could be based on understanding the support properties of the McKean-Vlasov SDE studied in Appendix E.3.2.

On possible variants of the training dynamic: In this work we focus mostly on the 'traditional' WGF learning dynamics, without delving much into other interesting possible variants of the training process. This might be seen as a possible limitation to the applicability of our work.

Firstly, the decision to work with the usual WGF follows from the standards set by [16, 53, 62, 67] among many others. Beyond this fact, we believe that results such as Theorems 3 and 4 are somewhat 'natural', and that they should hold regardless of small differences in the training dynamics.

For instance, we know that our proofs for Theorems 3 and 4 work straightforwardly in the setting of [22], where the MFL is established even with a fixed learning rate that does not necessarily decrease with N . For large fixed learning rates, this might shed some light onto the MF behaviour of symmetries under 'Edge of Stability' (EoS) dynamics. Similarly, we believe some of our more 'natural' results to be applicable as well for Wasserstein sub-Gradient Flows [16], Underdamped Dynamics [31], Annealed Dynamics [15], among many others. In contrast, Theorems 5 and 6, which involve stronger notions of symmetry, don't seem immediate to generalize to many other kinds of asymptotic dynamics. Studying the applicability of our results under different variants of the training process, is surely an interesting question to be tackled as part of our future work.

B Symmetries in functions, measures, data and shallow models

In this section we state some useful, basic results on the effect of symmetries acting on measures, functions and data, that will be used in the sequel. We also explain how symmetries of interest can be incorporated into the generalized shallow NN setting from Definition 1, complementing also the discussions given in Section 2.1 and Section 2.3 in that regard.

Recall X , Y and Z are (separable) Hilbert Spaces and G a compact group with Haar measure λ G , such that G ⟳ ρ X , G ⟳ ˆ ρ Y and G ⟳ M Z . Also, let E G ⊆ Z be the linear subspace of parameters that are fixed points of the action of G over Z , and P E G the orthogonal projection onto it.

B.1 Differentials and integrals of equivariant functions

The following lemma characterizes the differential of jointly equivariant functions with respect to the action of some group G . Here we can assume G be a lcsH group w.l.o.g. and consider representations that aren't necessarily orthogonal (we denote them differently to avoid confusion).

Proposition 6. Let G ⟳ χ X , G ⟳ ˜ χ Z , G ⟳ ˇ χ Y via some representations χ, ˜ χ and ˇ χ respectively (not necessarily orthogonal). Let f : X × Z → Y be jointly G -equivariant with respect to these actions (i.e. ∀ g ∈ G, ∀ x ∈ X , ∀ z ∈ Z , ˇ χ g .f ( x, z ) = f ( χ g .x, ˜ χ g .z ) ) and Fréchet-differentiable on its first argument. Then:

Proof of Proposition 6. Indeed, we know that ∀ z ∈ Z D x f ( · , z ) is the unique function that satisfies, ∀ ˜ x ∈ X :

Since we want to prove that ∀ ˜ x ∈ X , ∀ z ∈ Z , ∀ g ∈ G : D x f ( χ g . ˜ x, ˜ χ g .z ) = ˇ χ g D x f (˜ x, z ) χ -1 g , it will be enough to check that:

̸

which by uniqueness implies the result. Thanks to the joint equivariance of f , we have ∀ h = 0 :

Now, recall that for every g ∈ G , the operator ˇ χ g is bounded, i.e. it has finite operator norm 0 < ∥ ˇ χ g ∥ < ∞ (non-zero as ˇ χ g is invertible). By defining ˜ h := χ -1 g .h , we have:

Multiplying by 1 = ∥ χ -1 g χ g ˜ h ∥ ∥ ˜ h ∥ the last term is seen to be bounded by

Since χ g and χ -1 g are bounded operators, we have that: h → 0 ⇐⇒ ˜ h = χ -1 g h → 0 . Thus

This, in particular, allows us to characterize the differential of equivariant functions.

Corollary 4. If G ⟳ χ X , G ⟳ ˜ χ Y , and f : X → Y is a G -equivariant and Fréchet-differentiable function, then:

Proof of Corollary 4.

We can also get some interesting integral properties of jointly equivariant functions.

Proposition 7. Let X , Y and Z be (separable) Hilbert Spaces and G be a lcsH group. Let G ⟳ χ X , G ⟳ ˜ χ Z , G ⟳ ˇ χ Y via some representations χ, ˜ χ and ˇ χ respectively (not necessarily orthogonal).

Let f : X × Z → Y be a jointly G -equivariant function (with respect to these actions). Consider a measure µ ∈ P ( Z ) and let f be Bochner integrable on its second argument with respect to µ . Then:

Proof of Proposition 7. Let µ ∈ P ( Z ) and f be as stated. Notice that, ∀ x ∈ X , ∀ g ∈ G , we have:

where we've used the linearity of the Bochner integral under continuous linear operators (as is ˇ χ g ). It follows, from the joint G -equivariance of f and the definition of the pushforward measure , that:

We conclude the desired result.

B.2 Properties of symmetric measures

Consider again compact G acting orthogonally over the different spaces. Recall that, given µ ∈ P ( Z ) , we defined µ G := ∫ G ( M g # µ ) dλ G as its symmetrized version and µ E G := P E G # µ as its projected version. The following two results assumed in the core of the paper are elementary, but we provide their detailed proofs for completeness:

Lemma 1. We have the following inclusion: P ( E G ) ⊆ P G ( Z ) . Also, for any µ ∈ P ( Z ) the following equalities hold ∀ g ∈ G : µ E G = ( M g # µ ) E G = ( µ G ) E G = ( µ E G ) G and ( M g # µ ) G = µ G .

Proof of Lemma 1. Let µ ∈ P ( E G ) (i.e. µ ( E G ) = 1 ), g ∈ G and consider any positive measurable f : Z → R . We can see:

So, that ∀ g ∈ G , µ = M g # µ , and thus µ ∈ P G ( Z ) . Regarding the equalities, consider µ ∈ P ( Z ) and A ∈ B Z some borel set. Since the λ G is right-invariant, we have ∀ g ∈ G, ∀ z ∈ Z :

Then, for g ∈ G , M -1 g P -1 E G ( A ) = ( P E G M g ) -1 ( A ) = ( P E G ) -1 ( A ) and so:

and:

On the other hand, since µ E G ∈ P E G ( Z ) ⊆ P G ( Z ) , ( · ) G leaves it unchanged: ( µ E G ) G = µ E G . For the last equality, let g ∈ G and f : Z → R + be measurable. We have:

once again by the right-invariance of λ G . Namely, ( M g # µ ) G = µ G .

Proposition 8. For µ ∈ P ( Z ) , we have: µ G ∈ P G ( Z ) and µ E G ∈ P ( E G ) .

Proof of Proposition 8. Indeed, let h ∈ G and B ∈ B Z (borel set of Z ), using the properties of M and the left-invariance of λ G , we get that:

So, ∀ g ∈ G, µ G = M g # µ G , implying that µ G ∈ P G ( Z ) . On the other hand, by definition we have (as the projection is surjective) µ E G ( E G ) = µ ( P -1 E G ( E G )) = µ ( Z ) = 1 , so that µ E G ∈ P ( E G ) .

Remark. It's not hard to notice that, on Z = R D and with λ being the lebesgue measure, if µ ∈ P ( Z ) is such that µ ≪ λ and has density u : Z → R + , then: µ G ∈ P G ( Z ) has density u G := ∫ G u ◦ M g dλ G ( g ) w.r.t. λ (whereas µ E G doesn't have a density w.r.t. λ unless the action is trivial). This follows from the O ( D ) -invariance of λ and some standard calculations.

Since we will be working on P 2 ( Z ) on Section 3.3, it's useful to also notice that:

Remark. If µ ∈ P p ( Z ) , then µ G , µ E G ∈ P p ( Z ) . Indeed, it follows from the fact that ∥ M g ∥ ≤ 1 ∀ g ∈ G (since the representation is orthogonal) and ∥ P E G ∥ ≤ 1 (since P E G is a projection).

Also, we have that:

Proposition 9. P p ( E G ) := P ( E G ) ∩P p ( Z ) and P G p ( Z ) := P G ( Z ) ∩P p ( Z ) are closed and convex subspaces of P p ( Z ) (under the topology induced by W p ).

Proof of Proposition 9. Convexity is direct by definition of the involved spaces. For closedness under the Wasserstein topology, recall (see e.g. [1]) that W p ( µ n , µ ) - - - - → n →∞ 0 is equivalent to: µ n - - - -⇀ n →∞ µ (weak convergence) and ∫ Z ∥ θ ∥ p dµ n ( θ ) - - - - → n →∞ ∫ Z ∥ θ ∥ p dµ ( θ ) . Since for f ∈ C b ( Z ) , f ◦ M g (for g ∈ G ) is continuous and bounded, ( µ n ) n ∈ N ⊆ P G ( Z ) implies ∀ g ∈ G , M g # µ = µ (namely, µ ∈ P G ( Z ) ). Similarly, f ◦ P E G is continuous and bounded, and so if ( µ n ) n ∈ N ⊆ P ( E G ) , then P E G # µ = µ (i.e. µ ∈ P ( E G ) ).

B.3 Properties of equivariant data

We are representing the idea of 'data being symmetric' with respect to the action of G , by assuming the data distribution π to be equivariant. We will next see that this is a natural generalization of more intuitive, though restrictive, notions of data being symmetric, for instance when X is a r.v. on X with invariant law π X and Y = f ∗ ( X ) + ξ for some equivariant function f ∗ : X → Y and some centered independent noise ξ . Indeed the following results tell us that, assuming π ∈ P G ( X × Y ) , implies such a structure of the data, but with a more general ξ = ξ X , possibly correlated to X but still 'conditionally' centered given X . This result will be required in proving Proposition 5.

Proposition 10. Let π ∈ P ( X × Y ) be an equivariant data distribution such that E π [ ∥ Y ∥ 2 ] < ∞ . Then f ∗ = E π [ Y | X = · ] is an equivariant function.

Proof of Proposition 10. Indeed, as E [ ∥ Y ∥ 2 ] < ∞ , we know the conditional expectation E [ Y | X ] is well defined and there exists a measurable f ∗ : X → Y s.t. f ∗ ( X ) = E [ Y | X ] . Now, by properties of the conditional expectation): Given any h : X → R square integrable, we will show that: E π [ Y h ( X )] = E π [ ∫ G ˆ ρ -1 g .f ∗ ( ρ g .X ) dλ G ( g ) h ( X )] . Indeed, notice that by Fubini's theorem (as f ∗ ∈ L 2 ( X , Y ; π X ) and h square integrable), linearity of the integral and G -invariance of π :

By uniqueness of the conditional expectation, we know therefore that f ∗ ( X ) a.s. = ∫ G ˆ ρ -1 g .f ∗ ( ρ g .X ) λ G ( g ) . In particular, π X -a.e. f ∗ = Q G ( f ∗ ) , making f ∗ G -equivariant.

As a particular example, if we assumed that the data distribution was given by some function f : X → Y , i.e. Y = f ( X ) , taking π to be equivariant would be equivalent to assuming that π X is invariant and f is an equivariant function (which is the setting of the data simulated in our numerical experiments; see Section 4).

We notice that Proposition 10 can also be used to recover a celebrated result from [27] (later generalized by [40]), in the general setting where data symmetry is encoded by the condition that π ∈ P 2 ( X × Y ) . Define the symmetrization gap of a learning problem with quadratic loss as:

The following extension of mentioned statements from [27, 40] is not needed for our results, but we provide a proof of it in SuppMat-G, in view of its potential, independent interest:

Lemma 2 (Symmetrization Gap Characterization) . Consider the quadratic loss and π ∈ P ( X × Y ) such that E π [ ∥ Y ∥ 2 ] < ∞ . Also, assume that π | X is G -invariant, but π is only H -invariant with respect to some H ≤ G (closed). Then, the generalization gap satisfies:

where f ∗ ( x ) = E π [ Y | X = x ] is the conditional expectation function, and f ⊥ G := f -Q G f .

In particular, if π is G -invariant as well, we get ∆( f, Q G f ) = ∥ f ⊥ G ∥ 2 L 2 ( X , Y ; π X )

C Concrete realizations of shallow models

The class of models we have introduced in Definition 1 allows for taking an arbitrary, common parameter space Z for all hidden units, as well as an arbitrary function σ ∗ : X × Z → Y . As noted in the description of EA s in Section 2.3, by requiring only that σ ∗ is jointly-equivariant, we moreover ensure that G ⟳ M Z is, in some sense, properly related to the actions G ⟳ ρ X and G ⟳ ˆ ρ Y . This abstract property of σ ∗ allows us, in fact, to encode a wide range of situations and interesting results, without delving into the specifics of a particular choice of architecture. We next analyze this notion in the concrete example of the setting of a traditional single-hidden-layer shallow NN.

C.1 Traditional single layer neural networks and large ensembles of multi-layer units

Recall the finite-dimensional setting of single-hidden-layer neural networks. We considered X = R d , Y = R c and Z = R c × b × R d × b × R b , and defined, for z = ( W,A,B ) ∈ Z and σ : R b → R b , σ ∗ ( x, z ) := Wσ ( A T x + B ) . This allows us to express a shallow NN with N hidden units, of parameters θ = ( θ i ) N i =1 ∈ Z N , with θ i = ( W i , A i , B i ) , as the function Φ N θ : X → Y given by:

where we write the expression by blocks , considering

This corresponds exactly to the usual single-hidden-layer setting (see e.g. [23, 53, 62, 67]), only that we allow for the structure to involve the use of block matrices . This allows us to translate many relevant EA s, such as CNNs (see [19]) or DeepSets (see [77]) into the shallow NN framework that we propose (see Appendix C.2 for fully developed examples).

We now also consider a G -action on the intermediate layer , G ⟳ η R b (so that G ⟳ η ⊗ Id N ( R b ) N ). This allows us to define the natural intertwinning action 6 of G on Z , which is given by:

for any g ∈ G and z = ( W,A,B ) ∈ Z . This is exactly the action under which the fixed points (i.e. E G ) will correspond exactly to EA s in the traditional sense for this architecture (i.e. each layer being an equivariant function). In particular, we get the following straightforward result:

Proposition 11 ( Joint equivariance of σ ∗ for single-hidden-layer NNs ) . In the setting of singlehidden-layer NNs described above, if σ : R b → R b is G -equivariant (with respect to the action given by η ), then σ ∗ is jointly equivariant.

Proof. This fact follows directly from the specific definition of σ ∗ , the linearity and the orthogonality of all the relevant group representations. Indeed, for any equivariant σ : R b → R b , any g ∈ G and any z = ( W,A,B ) ∈ Z , we have:

Thus, in this particular case, in which σ ∗ represents the unit of a single hidden layer neural network, we require only for σ : R b → R b to be G -equivariant (with respect to the action given by η on R b ) in order for the joint equivariance of σ ∗ to hold (and, in consequence, most subsequent results from our work).

In particular, if η is chosen to be trivial (i.e. η ≡ id R b ), any single-hidden-layer NN is jointlyequivariant, 7 regardless of σ : R b → R b . Namely, all of the results contained in the core of the paper can be applied to an arbitrary single-hidden-layer NN; most importantly, those relating vanilla , DA and FA training. The disadvantage of having a trivial η is that the space E G might not encode very interesting EA s. This might make Theorem 5 lose part of its impressiveness, but it takes no credit off the rest of our results (such as Corollary 3 and Theorem 4).

Analogously, if η acts via permutations of the coordinates in R b , it is enough to consider a σ that is the pointwise application of a scalar function. In practice, this usually isn't a restrictive condition, since most commonly used NN architectures are naturally built following this pattern. Therefore, almost any single-hidden-layer architecture can yield a jointly-equivariant σ ∗ for some of the most common and interesting finite groups (e.g. S n , C n , among many others; see [30] for further discussion).

For a more complex (possibly infinite) compact group G acting orthogonally on the data and parameters (with non-trivial η ), for σ ∗ to be jointly-equivariant we have to start properly restraining σ : R b → R b . For instance, choosing an O ( b ) -equivariant σ (e.g. a Norm-ReLU) would grant Proposition 11 to always hold; but such a restraining choice could potentially harm the model's expressiveness and applicability. We leave a deeper exploration of this more challenging case as future work.

Finally, all of the above discussion (including Proposition 11) readily generalizes to the multilayer case. Namely, if σ ∗ encodes a Multi-Layer Perceptron (MLP) with multiple hidden layers whose parameters live in some linear space Z , we can define G ⟳ M Z as the intertwining action between each successive layer (similar to the previous example). σ ∗ can be made jointly-equivariant by making every activation function on each hidden layer equivariant (see [30]). Then, E G corresponds exactly to the parameters that make the entire MLP an equivariant architecture (in the sense that every layer is an equivariant function). With this, Φ N θ is an ensemble of N such MLPs trained in parallel, to which our results would also apply. Also as before, if the G -action on all of the hidden layers (but not on input/output) is made trivial, then any σ ∗ representing a multilayer architecture can be jointly-equivariant.

In the upcoming sections we will further develop these ideas, to show that some of the most relevant and widely applied equivariant architectures can be realized as part of our setting.

C.2 Shallow DeepSet models, shallow GNNs and shallow CNNs

C.2.1 DeepSets

An emblematic example of neural networks with equivariant architecture are the Deep Sets , introduced by [77]; which correspond, in practice, to NN architectures designed to be invariant/equivariant to the action of the group of permutations G = S n .

Consider that our NN processes sets of size n , which contain real-valued vectors of dimension ˜ d ∈ N . We can represent this input space simply as X := R n × ˜ d . Say we wanted to build a wide single-hidden-layer network that is invariant/equivariant to the action of S n , and returns a new set of n vectors, but now of dimension ˜ c ∈ N (i.e. Y := R n × ˜ c ). For this, let's consider the same architecture as in the previous section (replace d = n ˜ d, c = n ˜ c ), where we'll let our intermediate layer be simply another set of n vectors, now of dimension ˜ b with ˜ b ∈ N (such that b = n ˜ b above), and repeated N ∈ N times (as in the multiple hidden units that we want to achieve). With this, our network flows as follows: Φ N θ : R n × ˜ d → R ( n × ˜ b ) × N → R n × ˜ c . Notice how S n acts naturally on each hidden space by simply permuting the vectors of the set (i.e. ρ , ˆ ρ and η are defined in this way).

Following the same structure as in the previous section, we will have a parameter space given by: Z := R ( n × ˜ c ) × ( n × ˜ b ) × R ( n × ˜ d ) × ( n × ˜ b ) × R n × ˜ b , and a unit that acts on z = ( W,A,B ) ∈ Z as σ ∗ ( x, z ) := Wσ ( A T x + B ) , with σ : R n × ˜ b → R n × ˜ b some usual activation function applied pointwise (which, as mentioned above, will be immediately equivariant to the defined action of G via η ). Notice how building this architecture requires no hard a priori knowledge of the underlying symmetry of the problem (beyond the fact that the inputs and outputs are sets).

Under the natural action from this context (i.e. ρ , ˆ ρ and η acting on the sets by permuting their elements, and M defined as in SuppMat-C.1), E G corresponds exactly to the set of parameters that yield a S n -equivariant shallow neural network (which can be interpreted as a DeepSet). For the interested reader, we will make this connection explicit in the rest of this section.

As shown in [77], the only way to have an S n -equivariant affine transformation between two spaces R n × d 1 and R n × d 2 is if the parameters ˜ A ∈ R ( n × d 1 ) × ( n × d 2 ) and ˜ B ∈ R n × d 2 (from the definition x ↦→ ˜ A T .x + ˜ B ) are of the form:

Where α, β ∈ R d 1 × d 2 , γ ∈ R d 2 are the truly trainable parameters of the layer; I = Id n × n and J = ⃗ 1 n ⃗ 1 T n are two n × n matrices; and ⊗ is the usual tensor product. More explicitly, writing the matrices by blocks, this is:

In particular, Equation (8) gives us an explicit expression for our space E G . Indeed, an element z = ( W,A,B ) ∈ Z will live in E G if and only if there exists w 1 , w 2 ∈ R ˜ c × ˜ b , a 1 , a 2 ∈ R ˜ d × ˜ b and b 1 ∈ R ˜ b such that:

In a sense, we can think of E G simply as being equivalent to R 2(˜ c × ˜ b ) × R 2( ˜ d × ˜ b ) × R ˜ b . In particular, we went from having D = dim( Z ) = n 2 · ˜ b · (˜ c + ˜ d ) + n · ˜ b free parameters on each unit, to simply ˜ D = dim( E G ) = 2 · ˜ b · (˜ c + ˜ d ) + ˜ b , which should be easier to manage in general.

Now, recall our construction from SuppMat-C.1, and consider the matrices W,A and B from Equation (6). We notice that, in this example, they are of the form:

Namely, sensibly replacing d 1 = ˜ d , d 2 = ˜ bN for A and B ; and d 1 = ˜ c , d 2 = ˜ bN for W in Equation (8); we also get an explicit condition under which W,A and B yield a globally S n -equivariant architecture. By properly writing these matrices by blocks (as in SuppMat-C.1), and separating the action of each of the N units, one shall notice that the condition for W,A and B to be S n -equivariant corresponds exactly to every θ i = ( W i , A i , B i ) ∈ Z , i = 1 , . . . , N , being of the form given in Equation (9). That is to say, Φ N θ has an S n -equivariant architecture (a DeepSet from [77]) if and only if ∀ i ∈ { 1 , . . . , N } , θ i ∈ E G , which is exactly the condition stated in Section 2.3.

Remark. The reader could notice that our previous construction is not truly involving the complete richness of DeepSets. Namely, these architectures often involve using some "pooling" mechanisms, such as a global average pooling (GAP) operation to force invariance into the network (see [8]). Namely, if A : R n × ˜ b → R ˜ b is the usual linear GAP operation, we might want a NN that flows as Φ N θ : R n × ˜ d → R ( n × ˜ b ) × N A - → R ˜ bN → R ˜ c . This is no trouble within our framework, since we can simply consider the 'adequate' unit: σ ∗ ( x, z ) := W A ( σ ( A T x + B )) for z = ( W,A,B ) ∈ Z := R ˜ c × ˜ b × R ( n × ˜ d ) × ( n × ˜ b ) × R n × ˜ b , and all of our results would be applicable.

The main disadvantage of doing this is that we are forced to encode some a priori knowledge of the symmetries of the problem into our choice of architecture. While this isn't useful for our heuristic from Section 4.2, all other results relating the DA , FA and EA training dynamics still apply.

Remark. Similar to the last remark, more complex equivariant NN, with multiple layers and various inner operations involved, can be modeled by choosing σ ∗ properly. Namely, set σ ∗ to be a whole multi-layer structure, with parameters in Z , and E G being the subspace of those that make the architecture equivariant (see e.g. Equation (9) ).

In such case, the shallow model Φ N θ , with θ ∈ Z N , will represent an ensemble of N multi-layer units, each one given, for every i ∈ { 1 , . . . , N } , by σ ∗ ( · , θ i ) : X → Y . The output of the ensemble is simply the average of the outputs of each one of the units.

C.2.2 GNNs

Generalizing the ideas from DeepSets to GNNs is fairly straightforward. Namely, consider the input of a layer to be a graph, represented by a set of features, coupled with an adjacency matrix that also contains relevant edge-features; and the output to be analogous. Namely, let's say the input space is X 1 = R ( n × d 1 ) × R ( n × n ) × d 2 and the output space is X 2 = R ( n × d 3 ) × R ( n × n ) × d 4 . Consider also the natural S n action acting on these spaces, i.e. permuting the vertices of the graph, acting jointly between vertex features and adjacency matrix. With this in place, one can find an analogous characterization to Equation (8), but for affine layers that are S n -equivariant between graph spaces (see e.g. [29, 49] for a more explicit construction). From there, it is not hard to construct, following the same steps as for DeepSets, an explicit characterization of E G for single-hidden-layer GNNs, analogous to Equation (9). Also as before, more complex GNN structures (with multiple layers, pooling operations, etc.) can be encoded in this setting through the unit σ ∗ (with the possibly same drawbacks as in DeepSets). For brevity, we don't delve into GNNs in their full complexity and leave that to the curious reader.

C.2.3 CNNs

Finally, following the same logic as before, we can model another one of the most emblematic traditional equivariant models: CNNs. We here consider only 1D-CNNs for simplicity, but all arguments can be readily generalized to 2D or 3D CNNs.

In this setting, take an array of n vectors of dimension ˜ d ∈ N as input ( X := R n × ˜ d ) and, analogous to DeepSets, say that the output and hidden layers have the same structure, so that Y := R n × ˜ c ; and Z := R ( n × ˜ c ) × ( n × ˜ b ) × R ( n × ˜ d ) × ( n × ˜ b ) × R n × ˜ b , for ˜ b, ˜ c ∈ N . We consider σ ∗ simply as before.

The single main difference with the study of DeepSets is that, in this case, we consider the action of C n (i.e. the cyclic group of order n , also denoted Z n ) instead of S n . In particular, the natural action

of C n on each space (through ρ , ˆ ρ and η ) consists simply of shifting the array's coordinates in a given direction (modulo n ). As before, this makes the set E G correspond exactly to the set of parameters that yield each layer C n -equivariant (as in a 1D-CNN).

More specifically, we can characterize (analogous to Equation (8)) how a single C n -equivariant layer looks like between two spaces X 1 = R ( n × d 1 ) and X 2 = R ( n × d 2 ) under the natural action. It is well known (see [8]) that a C n -equivariant affine layer, x ↦→ ˜ A T .x + ˜ B , between X 1 and X 2 , of parameters ˜ A ∈ R ( n × d 1 ) × ( n × d 2 ) and ˜ B ∈ R n × d 2 , must be of the form:

with β ∈ R d 2 , α i ∈ R d 2 × d 1 , ∀ i ∈ { 0 , . . . , n -1 } , and the associated circulant matrix being defined (by blocks, and considering the indices modulo n ) as:

As a consequence, we get an analog of Equation (9) to explicitly describe E G : we have that z = ( W,A,B ) ∈ E G if and only if ∃ ( w i ) n -1 i =0 ⊆ R ˜ c × ˜ b , ( a i ) n -1 i =0 ⊆ R ˜ d × ˜ b and β ∈ R ˜ b such that:

Thus, the parameter space goes from having dim( Z ) = n 2 · ˜ b · (˜ c + ˜ d ) + n · ˜ b free parameters on each unit, to simply dim( E G ) = n · ˜ b · (˜ c + ˜ d ) + ˜ b . One might also notice that the obtained parameter space E C n contains E S n (from DeepSets), which is expected from the fact that C n ≤ S n . In global terms, the shallow CNNs we have modeled here, correspond to models that will grow asymptotically in terms of the number of different convolutional filters (encoded by N ) that are being used.

Finally, as it was also mentioned in previous examples, more complex CNN structures (with multiple layers, pooling operations, etc.) can be encoded in this setting through modifications to the unit σ ∗ .

D Further elements from the MF theory of shallow neural networks

In this section we study several theoretical notions required in the MF approach to overparametrized shallow NN. For the purpose of our results, Subsections D.1 and D.2 are the most relevant ones, as we establish therein some useful properties or formula that will be explicitly required. Subsections D.3 and D.4 review some results and recent literature regarding well-posedness and long-time convergence of WGFs, which are relevant to the MF interpretation of the training dynamics of shallow models.

D.1 Linear functional derivatives and intrinsic derivatives

Let X , Y and Z be separable Hilbert spaces. Recall the definitions of the linear functional derivative and intrinsic derivatives:

Definition 8 (Linear Functional Derivative (First Variation)) . For a functional F : P 2 ( Z ) → R , its linear functional derivative (lfd) is a function: ∂F ∂µ : P 2 ( Z ) ×Z → R such that ∀ µ, ν ∈ P 2 ( Z ) :

The function F ′ : µ ∈ P 2 ( Z ) ↦→ ∂F ∂µ ( µ, · ) is also known as the first variation of R at µ .

Definition 9 (Intrinsic Derivative) . Whenever ∂F ∂µ : P 2 ( Z ) ×Z → R exists and is differentiable on its second argument, the intrinsic derivative of F is defined as: D µ F ( µ, z ) = ∇ z ( ∂F ∂µ ( µ, z ) ) .

Example. To better illustrate the notion of the linear functional derivative and the intrinsic derivative, consider the following usual examples:

In the important case of the Boltzmann entropy F = H λ , defined for µ ≪ λ by H λ := ∫ log( dµ dλ ( z )) dµ ( z ) (and + ∞ otherwise), we have that (modulo an additive constant that doesn't depend on z , see [56]):

Whenever F ( µ ) := ∫ Z ϕ ( z ) dµ ( z ) for some bounded continuously differentiable function ϕ : Z → R , it is well known that :

The most relevant example, in our case, regards the linear functional and intrinsic derivatives of the population risk functional, R ( µ ) = E π [ ℓ (Φ µ ( X ) , Y )] . We can consider the general setting presented in [16], where some Hilbert Space H is considered, and it is assumed that F : P ( Z ) → R can be written as F ( µ ) = L ( ⟨ Φ , µ ⟩ ) , where Φ : Z → H is a parametrization of elements in H ; L : H → R is some loss functional, and the integral ⟨ Φ , µ ⟩ is a Bochner integral on H . This generalizes the shallow NN setting, as one might consider H = L 2 ( X , Y , π X ) , L : H → R given by L ( f ) = E π [ ℓ ( f ( X ) , Y )] and Φ : Z → H defined as ∀ z ∈ Z , Φ( z ) = σ ∗ ( · ; z ) ; so that R ( µ ) = L ( ⟨ Φ , µ ⟩ ) . In this setting, we can prove the following result: 8

Proposition 12. Let H be a separable Hilbert Space and F ( µ ) := L ( ⟨ Φ , µ ⟩ ) , for some function that's Gateaux-differentiable L : H → R on every direction and of continuous differential; and Φ : Z → H such that ∀ µ ∈ P ( Z ) , ∥⟨ Φ , µ ⟩∥ H < ∞ .

Then ∀ z ∈ Z , ∀ µ ∈ P ( Z ) :

Here, C F,µ := ⟨∇ h L ( ⟨ Φ , µ ⟩ ) , ⟨ Φ , µ ⟩⟩ H is exactly the constant needed to avoid ambiguity in the definition; ( · ) ∗ denotes the adjoint operator and, in particular, ∇ z Φ( z ) = ( D z Φ( z )) ∗ : H → Z . When Z = R D this corresponds to the usual definition of the gradient.

Proof of Proposition 12. We know that, ∀ µ, ν ∈ P ( Z ) , h ∈ [0 , 1] :

Let's denote by q µ := ⟨ Φ , µ ⟩ (analogously q ν -µ := ⟨ Φ , ν -µ ⟩ ) and s µ,ν := hq ν -µ , so we can write:

As L is Gateaux differentiable, we have the following first order Taylor expansion ∀ x, s ∈ H , ∀ t ∈ R :

which allows us to write:

As ∥ q ν -µ ∥ < ∞ by hypothesis, we can say that: o ( | h |∥ q ν -µ ∥ ) = o ( h ) . Therefore, taking the limit with h → 0 , we get that:

Now, developping this last term (using, for instance, the linearity of the Bochner integral, as we know D h L ( x, · ) . ( · ) to be linear and bounded), we get that:

and so by definition of the gradient of L :

From here we deduce that:

where C F,µ is a fixed constant, given by:

On the other hand, for the intrinsic derivative, notice that D z ( ∂F ∂µ ( µ, z )) : Z → R is a bounded linear functional over Z , so (by Riesz Representation) ∃ D µ F ( µ, z ) := ∇ z ( ∂F ∂µ ( µ, z )) ∈ Z such that:

However, we can develop the RHS, and as the constant C F,µ doesn't depend on z , we get that:

Now, by the chain rule and the definition of the adjoint operator of D z Φ( z ) :

So, as they coincide for every z ∈ Z , we conclude that:

Proposition 12 applies directly to our population risk functional R ( µ ) := E π [ ℓ (Φ µ ( X ) , Y ) ] , by considering:

The Hilbert space: H = L 2 ( X , Y , π X )

L : H → R given by L ( f ) = E π [ ℓ ( f ( X ) , Y )] , which is Gateaux-differentiable on every direction in H if we assume ℓ : Y × Y → R to be continuously differentiable on its first argument, with ∇ 1 ℓ linearly growing. The differential (which is continuous) can be explicitly computed to be:

Φ : Z → H defined as ∀ z ∈ Z , Φ( z ) = σ ∗ ( · ; z ) , which satisfies ∀ µ ∈ P ( Z ) , ∥⟨ Φ , µ ⟩∥ H < ∞ under the assumption of σ ∗ being bounded and continuous.

Corollary 5. We can explicitly compute the linear functional derivative and the intrinsic derivative for the learning problem's population risk:

Beyond this particular example, the linear functional derivative and the intrinsic derivative behave well when the underlying functional is invariant, as shown by the following result:

Proposition 13. Let F : P ( Z ) -→ R be invariant and of class C 1 . Then: ∀ z ∈ Z , ∀ µ ∈ P ( Z ) , ∀ g ∈ G :

i.e. ∂F ∂µ is jointly invariant and D µ F jointly equivariant.

Proof of Proposition 13. To prove this, recall that the linear functional derivative of F is the only function ∂F ∂µ : P ( Z ) ×Z → R satisfying ∀ µ, ν ∈ P ( Z ) :

̸

In particular, as F is G -invariant (and M g linear), we can write ∀ µ, ν ∈ P ( Z ) , ∀ h = 0 and g ∈ G :

Taking the limit as h → 0 on both sides, we get:

and also:

So, by uniqueness, we get ∀ g ∈ G :

and, from proposition 6 (since ∂F ∂µ is jointly invariant), we get ∀ g ∈ G :

Finally, for convenience, let us denote by H E G the function defined on P ( Z ) by H E G ( µ ) := H λ E G ( µ E G ) = H λ E G ◦ P E G #( µ ) (presented in Section 3.3). We can straightforwardly compute its linear functional derivative and its intrinsic derivative in that space:

Example. (LFD and intrinsic derivative for H E G ) By definition of the linear derivative ∂H λ E G ∂η of H λ E G on P ( E G ) and the form we know it takes (see the examples from Appendix D.1), we see that, whenever µ E G , ν E G ≪ λ E G , we have:

which yields that ∂H E G ∂µ ( µ, z ) = log ( dµ E G dλ E G ( P E G .z ) ) + C (for C an appropriate constant). A formal expression for the intrinsic derivative follows, which is given by :

D.2 Expression for the WGF of the regularized population risk

In the case of the regularized population risk functional R τ,β : P ( Z ) → R , we can explicitly write its intrinsic derivative. Consider a slightly more general functional, denoted by R τ,β ν , where the entropy is calculated against a Gibbs measure ν ≪ λ such that ν ( dz ) = e -U ( z ) λ ( dz ) for some function U : Z → R (as in [38]). We have ∀ µ ∈ P ( Z ) s.t. µ ≪ ν , ∀ z ∈ Z :

so that WGF ( R τ,β ν ) as in definition 4 reads:

We recover the expression for WGF ( R τ,β ) in Equation (4) by considering U ≡ 0 :

We can see that Equation (4) corresponds to a Fokker-Planck equation, which can be interpreted in terms of a non-linear SDE system, representing the behaviour of the type parameter: the McKeanVlasov SDE [11, 54, 70] (also known as the Mean Field Langevin Dynamics (MFLD) in the NN literature). In the case of R τ,β it reads:

where ( B t ) t ≥ 0 is a D -dimensional standard Brownian Motion. It is indeed standard to check (by applying Itô's formula to φ ( Z t , t ) for φ a smooth function, and taking expectation) that µ t = Law ( Z t ) is a weak solution to (10). Under mild regularity conditions, both formulations are equivalent. See [11, 54, 70] for details. The previous correspondence also holds true when β = 0 (in which case (11) is an ODE). The process (11) or variants of it will prove useful to establish some of the relevant results of the paper.

D.3 Global convergence in the regularized case

For the example we just presented of the entropy-regularized population risk, multiple authors (see [12, 15, 38, 57, 69] among many others) have studied the properties of WGF ( R τ,β ), particularly, the global convergence results that can be obtained. For instance, consider the following results from [38] (where we look at R τ,β ν for generality). First, define:

Definition 10. We say that a functional R : P p ( Z ) → R is of class C 1 if ∂R ∂µ ( µ, · ) is well defined and bounded for every µ ∈ P p ( Z ) , and the function ( µ, z ) ∈ P p ( Z ) ×Z ↦→ ∂R ∂µ ( µ, z ) is continuous .

Now, from [15, 38], we get the following key result. We include the proof for completeness:

Lemma 3 (as in [38, 15]) . Assume that R : P p ( Z ) → R is convex and of class C 1 . Then, for any µ, µ ′ ∈ P p ( Z ) , we have:

Proof of Lemma 3 (from [38]). Define µ ϵ := (1 -ϵ ) µ + ϵµ ′ . Since R is convex, we have

Since the map s ∈ [0 , 1] ↦→ µ s is continuous, it is of compact image (denoted [ µ, µ ′ ] ). In particular, as ∂R ∂µ is continuous and bounded on its second argument, we get that, it is bounded on [ µ, µ ′ ] ×Z . The dominated convergence and Lebesgue differentiation theorems (as ε → 0 ) allow us to conclude.

Consider now the following assumption:

Assumption 2 (As in [38]) . U : Z → R is assumed to be C ∞ , with ∇ U Lipschitz continuous, and such that ∃ C U > 0 , ∃ C ′ U ∈ R such that ∀ x ∈ Z : ∇ U ( x ) · x ≥ C U ∥ x ∥ 2 + C ′ U . When required, we will also assume that r : Z → R satisfies these conditions.

Notice that these conditions imply that ∃ 0 ≤ C ′ ≤ C s.t. ∀ x ∈ Z , C ′ ∥ x ∥ 2 -C ≤ U ( x ) ≤ C (1 + ∥ x ∥ 2 ) (i.e. U has quadratic growth) and | ∆ U ( x ) | ≤ C

Since R τ,β ν includes an entropy term , it guarantees strict convexity, weak lower semicontinuity and compact sublevel sets (see e.g. [38] or [26]). On the other hand, assumption 2 implies that U (or r ) will have quadratic growth. Namely, we get (see [38] for a detailed proof):

Proposition 14 (Existence and Uniqueness of the minimizer (regularized case)) . Let R be convex , of class C 1 and bounded from below. Let ν be the Gibbs measure with potential U . Then, R τ, β ν has a unique minimizer , µ ∗ , τ, β, ν ∈ P ( Z ) , absolutely continuous with respect to Lebesgue measure λ . When either U or r satisfies assumption 2, this minimizer also belongs to P 2 ( Z ) .

For establishing global convergence results further assumptions are requred

Assumption 3 (Assumptions for well definedness (from [12] and [38])) . Assume that the intrinsic derivative D µ R : P ( Z ) ×Z → Z of the functional R : P ( Z ) → R exists and satisfies either one of the following:

1. (From [38]). Assume:

D µ R is bounded and Lipschitz continuous, i.e. ∃ C R > 0 s.t. ∀ z, z ′ ∈ Z , ∀ µ, µ ′ ∈ P 2 ( Z ) ,

∀ µ ∈ P ( Z ) , D µ R ( µ, · ) ∈ C ∞ ( Z ) .

∇ D µ R : P ( Z ) ×Z → Z × Z is jointly continuous.

(From [12], who relax some differentiability conditions at the cost of boundedness assumptions; this allows them to avoid altogether the coercivity condition from assumption 2, which is used in [38]):

∀ x ∈ Z , ∀ m,m ′ ∈ P 2 ( Z ) , | D µ R ( m,x ) -D µ R ( m ′ , x ) | ≤ M R mm W 1 ( m,m ′ ) for some constant M R mm ≥ 0 (i.e. it is lipschitz on the measure argument).

Suppose that

for some constant M R mx ≥ 0 i.e. ∇ D µ R ( µ, x ) is uniformly bounded.

This allows to establish a traditional global convergence result from the MF Theory of NNs:

Theorem 7 (from [12] and [38]) . Let µ 0 ∈ P 2 ( Z ) , and let assumption 2 and 3 hold; then:

where u t denotes the density of µ t := Law ( X t ) , the solution to equation (4) . i.e. following the WGF makes the regularized risk decrease at a known rate. This is known as the energy dissipation equation .

Remark. Notice that this equation can be rewritten using the Fisher divergence (or relative Fisher Information) between two measures. This quantity is defined as:

Then, almost by definition, we get:

This allows us to characterize the stationary points of the dynamic explicitly, as done in [12, 38, 57, 69].

Theorem 7 implies that the WGF converges to the unique global optimizer of the regularized problem:

Theorem 8 (from [38]) . Let R be convex , bounded from below and C 1 ; also assume that assumption 2 and 3 hold. Consider µ 0 ∈ ∪ p> 2 P p ( Z ) and let ( µ t ) t ≥ 0 be the WGF ( R τ,β ν ) starting from µ 0 . Then, the equation has a stationary distribution, µ ∞ , that satisfies:

Remark. Global Convergence Results such as Theorem 7 or Theorem 8 have been established as early as in [53] (for the quadratic loss). However, settings such as those of [12, 38, 57, 69, 15] are of notorious interest to establish essentially the same results under fundamentally more general assumptions.

Making further technical assumptions on our regularized functionals leads to better convergence results . Namely, consider the following definition:

Definition 11. We say µ ∈ P ( Z ) satisfies the Log-Sobolev Inequality with constant ϑ > 0 (in short, LSI( ϑ )), if for any ν ∈ P ( Z ) such that ν ≪ µ , we have:

where D ( ν || µ ) is the KL divergence and I ( µ || ν ) is the Fisher divergence .

(see [3, 58] for background on functional inequalities and [15, 12, 57, 69] for applications of it to the NN context). In our setting, as done by most authors in recent years to achieve the desired global convergence results, the following 'uniform-LSI' on the functional R : P ( Z ) → R is assumed to hold:

Assumption 4 ( Uniform LSI from [15, 12, 57, 69]) . There exists ϑ > 0 such that ∀ µ ∈ P 2 ( Z ) , ˆ µ satisfies LSI( ϑ ). Here ˆ µ is the probability measure with density w.r.t. λ given by (slightly abusing notation, and considering U the potential of a Gibbs measure ν used in the entropy, which is 0 if ν = λ ):

Remark. This LSI is a recurrent element in the literature of WGF and Optimal Transport in general. In particular, it implies (see [3]) Poincaré Inequality: ∀ ϕ ∈ C 1 b ( Z ) , Var ˆ µ ( ϕ ) ≤ 1 2 ϑ E ˆ µ [ |∇ ϕ | 2 ] , and ([58]) the Talagrand's T 2 -transport inequality as well: ∀ ν ∈ P 2 ( Z ) , ϑW 2 2 ( ν, ˆ µ ) ≤ D ( ν || ˆ µ )

Beyond the characterization of the decay (from [38]), we have the following guarantee:

Theorem 9 (from [12, 15]) . Let R be convex, C 1 and bounded from below, and let assumptions 3 and 4 hold. Then, if for some t 0 ≥ 0 , µ t 0 has finite entropy and finite second moment ; then ∀ t ≥ t 0 ,

where µ ∞ = µ τ, β, ν = arg min µ ∈P ( Z ) R τ β ν ( µ ) . That is, the value function following the WGF thus converges exponentially fast to the optimum value of the problem, and this implies an exponential convergence in relative entropy.

One thus recovers, under the right technical assumptions, a version of Theorem 4 from [53] and actually a quantitative improvement of it. By Talagrand's inequality this also implies exponential W 2 convergence of µ t to µ ∞ . We note that the result in [12] is established in the setting with τ = 0 , β = 1 and ς ≡ 1 ; however, one can show that the result holds as stated by using standard arguments.

D.4 Conditions for well-posedness of WGF

Most of the technical conditions that will be here presented are directly taken from both [16] and [22]. We only adapt them slightly to fit into our notation.

Regarding the existence of weak solutions to the WGF presented in equation (4), [16] are able to guarantee it under the following assumptions (more general assumptions might be sought in [1, 63], but these are relatively standard in the MF context):

Assumption 5 (Assumptions for existence and uniqueness of the WGF solutions (taken from [16])) . Consider a setting as described in proposition 12, with R ( µ ) := L ( ⟨ Φ , µ ⟩ ) + V ( µ ) , with V ( µ ) = τ ∫ Z rdµ .

Let Z to be the closure of a convex open set within some finite-dimensional euclidean space.

Let L : H → R + be differentiable, with a differential dL that is Lipschitz on bounded sets and bounded on sublevel sets .

Let Φ : Z → H be differentiable and V : Z → R + be semiconvex (i.e. ∃ λ ∈ R : V + λ | · | 2 is convex).

There exists a family ( Q r ) r> 0 of nested nonempty closed convex subsets of Z such that:

(a) { u ∈ Ω; dist ( u, Q r ′ ) ≤ r } ⊂ Q r + r ′ for all r, r ′ > 0 ,

(b) Φ and V are bounded, and d Φ is Lipschitz on each Q r

(c) ∃ C 1 , C 2 > 0 such that sup u ∈ Q r ( ∥ d Φ( u ) ∥ + ∥ ∂V ( u ) ∥ ) ≤ C 1 + C 2 r for all r > 0 , where ∥ ∂V ( u ) ∥ stands for the maximal norm of an element in ∂V ( u ) .

On the other hand, [22] are able to prove (based on Theorem 1.1. from [70]) the existence of strong solutions with pathwise uniqueness for McKean-Vlasov SDE given by

where b and Σ satisfies the conditions of B2 (presented below) and for all t ≥ 0 , µ t = Law ( Z t ) ∈ P 2 ( R D ) , ( B t ) t ≥ 0 is an r -dimensional Brownian motion (with r ∈ N ∗ potentially different from D ∈ N ∗ ), and Z 0 has the (fixed) law µ 0 ∈ P 2 ( R D ) . For this, consider the following technical assumptions ( B1 and B2 ) which have been taken directly from [22]:

Assumption 6 (Assumptions for the existence and uniqueness of solutions in [22]) . Consider:

B1. There exist a measurable function g : R D ×W → R , M 1 ≥ 0 and µ 0 ∈ P 2 ( R D ) such that for any N ∈ N , the following hold.

(a) For any w 1 , w 2 ∈ R D and z ∈ W we have

with ∫ W ζ 2 ( z ) dπ W ( z ) < + ∞

(b) b N ∈ C ( R + × R D ×P 2 ( R D ) , R D ) and Σ N ∈ C ( R + × R D ×P 2 ( R D ) , R D × r ) .

(c) For any w 1 , w 2 ∈ R D and µ 1 , µ 2 ∈ P 2 ( R D )

B2. There exist M 2 ≥ 0 , κ > 0 , b ∈ C ( R + × R D ×P 2 ( R D ) , R D ) and Σ ∈ C ( R + × R D × P 2 ( R D ) , R D × r ) such that

Proposition 15 (Proposition 11 in [22]) . Assuming B1 and B2 . Given µ 0 ∈ P 2 ( R D ) as a fixed initial condition; then, there exists an ( F t ) t ≥ 0 -adapted process ( Z t ) t ≥ 0 that is the unique (pathwise) strong solution of the McKean-Vlasov SDE:

Additionally, it satisfies for each T ≥ 0 : sup t ∈ [0 ,T ] E [ ∥ Z t ∥ 2 ] < ∞

E Proofs and discussions of main results

E.1 Proofs of Section 3.1

Proof of Proposition 1. By definition of the symmetrization operator, we know that ∀ x ∈ X :

For g ∈ G , since σ ∗ is equivariant and M g is invertible, we can write:

As mentioned in Section 3.1, a simple case where we will have Φ µ G = Φ µ E G is when σ ∗ is linear : Proposition 16. If σ ∗ : X × Z → Y is jointly equivariant and π X -a.s. ∀ x ∈ X , [ z ↦→ σ ∗ ( x ; z )] is a bounded linear operator, then, for any µ ∈ P ( Z ) : Φ µ G = ⟨ σ ∗ , µ G ⟩ = ⟨ σ ∗ , µ E G ⟩ = Φ µ E G .

Proof of Proposition 16. A straightforward computation yields (using Fubini's theorem and the linearity of integrals and σ ∗ ), ∀ µ ∈ P ( Z ) , ∀ x ∈ X ( π X -a.s.):

Example. Any usual linear model written in terms of some feature function ϑ : X → Z (where Z is possibly an Reproducing Kernel Hilbert Space) enters this framework, by defining: σ ∗ ( x, z ) = ⟨ z, ϑ ( x ) ⟩ . This won't satisfy Assumption 1, since it is not bounded; but it still serves as an illustration.

E.2 Proofs of results in Section 3.2

E.2.1 Proof of Proposition 2

Now consider, as a shorthand notation, ∀ x ∈ X , ∀ y ∈ Y the functional L x,y : P ( Z ) → R given by ∀ µ ∈ P ( Z ) : L x,y ( µ ) = ℓ ( Φ µ ( x ) , y ) . The following lemma that shall be useful for later stages. Lemma 4. Let σ ∗ be jointly equivariant and ℓ be invariant. Then, ∀ g ∈ G, ∀ x ∈ X , ∀ y ∈ Y , ∀ µ ∈ P ( Z ) ,

L ρ g .x, ˆ ρ g .y ( M g # µ ) = L x,y ( µ ) Equivalently, the map L : P ( Z ) → L 2 ( X × Y , π ) given by L ( µ ) ↦→ [( x, y ) ↦→ L x,y ( µ )] is equivariant (under the appropiate 9 G -actions).

Proof of Lemma 4. Using the joint equivariance of σ ∗ (via proposition 7) and the invariance of ℓ , a straightforward computation yields, for all x ∈ X , y ∈ Y , and g ∈ G :

With this we can prove Proposition 2. Notice that we will basically utilize the equivariance properties of σ ∗ and ℓ in Assumption 1 (the convexity of the functions comes directly from the convexity of R when ℓ is convex, together with the linearity of ( · ) G , ( · ) E G and ∫ G ( · ) dλ G .).

Proof of Proposition 2. We can readily see that: R EA ( µ ) = E π [ ℓ ( Φ µ E G ( X ) , Y )] = R ( µ E G ) . On the other hand, as σ ∗ is jointly equivariant, from Proposition 1, we have:

Next, using Lemma 4, Fubini's theorem and the inversion-invariance of λ G , we get:

From these expressions we can quickly verify that R DA , R FA and R EA are invariant. Namely, by Lemma1, for g ∈ G and µ ∈ P ( Z ) we have: R FA ( M g # µ ) = R (( M g # µ ) G ) = R ( µ G ) = R FA ( µ ) , and: R EA ( M g # µ ) = R (( M g # µ ) E G ) = R ( µ E G ) = R EA ( µ ) . On the other hand, the rightinvariance of λ G implies:

We can also see that whenever R is invariant, we have that, for µ ∈ P ( Z ) and g ∈ G,R ( M g # µ ) = R ( µ ) , so that: R DA ( µ ) = R G ( µ ) = ∫ G R ( M g # µ ) dλ G ( g ) = ∫ G R ( µ ) dλ G ( g ) = R ( µ ) .

Also, if µ ∈ P G ( Z ) , we have, for all g ∈ G , µ = µ G = M g # µ , so that: R DA ( µ ) = ∫ G R ( M g # µ ) dλ G ( g ) = ∫ G R ( µ ) dλ G ( g ) = R ( µ ) = R ( µ G ) = R FA ( µ ) .

Finally, we verify that our population risk R : P ( Z ) → R is invariant whenever π ∈ P G ( X × Y ) . Indeed, ∀ g ∈ G and ∀ µ ∈ P ( Z ) , by the invariance of ℓ and π , together with Proposition 7 from the equivariance of σ ∗ , we get:

That is, R is invariant.

Notice that the equivariance of the data distribution π can also make the regularized population risk be invariant, under the right choice of r . Namely:

Corollary 6. If R : P ( Z ) → R and r : Z → R are invariant (in their respective sense), then R τ,β is invariant.

The result can be proven for R τ,β ν with ν some G -invariant measure (such as λ for orthogonal representations). Notice that r ( θ ) = ∥ θ ∥ 2 is an example of invariant function for orthogonal representations .

Proof of Corollary 6. It is enough to notice that V ( µ ) = ∫ Z r ( θ ) dµ ( θ ) and H ν ( µ ) = ∫ Z log( dµ dν ) dµ (with µ ≪ ν ) are invariant when r : Z → R and ν ∈ P ( Z ) are invariant (in their respective sense):

For V , notice that for g ∈ G :

thanks to the invariance of r . i.e. V is invariant.

For H ν , notice that, for g ∈ G , as ν is invariant, we know that d ( M g # µ ) dν ( x ) = dµ dν ( M -1 g x ) . Therefore:

Which proves that H ν is invariant.

We can readily conclude, since: R τ,β ( µ ) = R ( µ ) + τ ∫ rdµ + βH λ ( µ ) , for all µ ∈ P ( Z ) .

E.2.2 Proof of Proposition 3

In order to prove Proposition 3, we require a version of Jensen's inequality that's suited for our context. Such a result might exist in the literature, but since we couldn't find a complete proof under our assumptions, we provide our own.

Proposition 17 ( Jensen's Inequality ) . Let F : P ( Z ) -→ R be such that Lemma 3 holds. Let S be some measurable space, λ ∈ P ( S ) and s ∈ S ↦→ µ s ∈ P ( Z ) a measurable function. Define ˜ µ ∈ P ( Z ) as the intensity measure: ˜ µ = ∫ S µ s dλ ( s ) ∈ P ( Z ) . Then, Jensen's inequality holds:

Proof of Proposition 17. Since Lemma 3 holds, we have that ∀ µ 1 , µ 2 ∈ P ( Z ) :

Let ˜ s ∈ S be arbitrary and consider µ 2 = ˜ µ := ∫ µ s dλ ( s ) ; and µ 1 = µ ˜ s . Then:

Integrating the inequality with respect to λ (on ˜ s ):

We notice that the LHS doesn't depend on ˜ s , so that ∫ S F (∫ µ s dλ ( s ) ) dλ (˜ s ) = F (∫ µ s dλ ( s ) ) . On the other hand, the right-most term in △ can be developed as:

Notice that the linear functional derivative is chosen in such a way so that it satisfies ∀ ν ∈ P ( Z ) , ∫ Z ∂F ∂µ ( ν, z ) dν ( z ) = 0 . In particular, the second term of the previous expression vanishes. We get that

But, by definition: ∀ f : Z → R integrable,

So this is, by definition, and applying the same convention on the definition of the linear functional derivative 10 :

With this, we conclude that △ = ∫ S F ( µ s ) dλ ( s ) , and so, we get that:

which corresponds to Jensen's inequality.

Remark. We believe that the C 1 hypothesis can be lifted. Understanding what happens when equality holds should be of interest both in our context and in more general scenarios.

Thanks to this Jensen inequality , we readily get the following result:

Corollary 7. If F : P ( Z ) -→ R is convex, C 1 and invariant, then ∀ µ ∈ P ( Z ) : F ( µ G ) ≤ F ( µ )

Proof. Direct from the definition of ( · ) G and Proposition 17.

With these results in place, we are ready to prove Proposition 3:

Proof of Proposition 3. Evidently, since P G ( Z ) ⊆ P ( Z ) , we have:

For the other inequality, take ( µ n ) n ∈ N ⊆ P ( Z ) to be an infimizing sequence for F ; i.e. such that F ( µ n ) ≥ F ( µ n +1 ) and F ( µ n ) - - - - → n →∞ inf µ ∈P ( Z ) F ( µ ) ). Such a sequence always exists. By Corollary 7, we have ∀ n ∈ N , F ( µ G n ) ≤ F ( µ n ) ; thus, ∀ n ∈ N :

Which allows us to infer, by taking n →∞ , that: inf µ ∈ P G ( Z ) F ( µ ) ≤ inf µ ∈ P ( Z ) F ( µ ) . In turn, we can conclude that:

Notice that if there was some minimizer µ ∗ ∈ arg min µ ∈P ( Z ) F ( µ ) , then by Corollary 7 we would also have µ G ∗ ∈ arg min µ ∈P ( Z ) R ( µ ) . Namely, if such a minimizer was unique , then it would satisfy: µ ∗ = µ G ∗ ∈ P G ( Z ) . That is, the unique solution would be WI .

E.2.3 Proof of Theorem 2 and Corollary 1

Notice that, under Assumption 1, from Corollary 5 we know R is of class C 1 (as well as convex). This properties actually transfers to the functionals R DA , R FA and R EA , as shown by the following result:

Proposition 18. If R : P ( Z ) → R is a convex and C 1 functional, then R DA , R FA and R EA are convex and C 1 as well, with linear functional derivatives given by:

And intrinsic derivatives given by (when well defined):

In particular, from Proposition 17, we have that: ∀ µ ∈ P ( Z ) , R FA ( µ ) ≤ R DA ( µ )

Proof of Proposition 18. We can calculate the linear functional derivatives (l.d.f. for short) as follows. Let µ, ν ∈ P ( Z ) , and consider:

We have used Fubini's theorem, which is applicable 11 thanks to the fact that R is of class C 1 , and we've used the definition of the linear functional derivative for R . Also, we see that (using Fubini's theorem once again, as well as the definition of the linear functional derivative of R ):

We can then identify:

and, by taking the gradient:

We analogously calculate the expression for the l.f.d. of R FA ; let µ, ν ∈ P ( Z ) :

and also:

So, by the definition of the lineal functional derivative, we identify:

and taking the gradient we get:

Lastly, the l.f.d. of R EA is calculated similarly, noticing that:

and that:

We can thus identify:

The last remark is direct from proposition 17.

With all of these different elements in place, we are ready to prove Theorem 2.

Proof of Theorem 2. Under Assumption 1, R is convex and of class C 1 from Corollary 5; and so are R G , R FA and R EA , from Proposition 18. Since Proposition 2 ensures that the latter are always invariant, Proposition 3 implies that R DA , R FA and R EA can all be optimized by only considering weakly equivariant models (explaining the first and last equalities). The two middle equalities follow directly from Proposition 2, since R , R DA and R FA coincide over P G ( Z ) .

In the case of the quadratic loss, one can employ the properties of Q G from [27] to show Corollary 1.

Proof of Corollary 1. Notice that, for µ ∈ P G ( Z ) , Φ µ is a G -invariant function (i.e. Φ µ ∈ L 2 G ( X , Y ; π X ) ). Also, a simple calculation (see e.g. [5, 51]) allows us to write: R ( µ ) = E π [ ∥ Φ µ ( X ) -Y ∥ 2 Y ] = R ∗ + E π X [ ∥ Φ µ ( X ) -f ∗ ( X ) ∥ 2 Y ] = R ∗ + ∥ Φ µ -f ∗ ∥ 2 L 2 ( X , Y ; π X ) with

R ∗ = E π [ ∥ Y -f ∗ ( X ) ∥ 2 Y ] being independent of µ . We can thus write (simplifying subscripts for simplicity):

where ( f ∗ ) ⊥ G := f ∗ - Q G .f ∗ . We notice that, since Φ µ and Q G .f ∗ are G -equivariant functions, we have that ⟨ (Φ µ - Q G .f ∗ ) , ( f ∗ ) ⊥ G ⟩ L 2 = 0 . That is, for any µ ∈ P G ( Z ) , we have R ( µ ) = ˜ R ∗ + ∥ Φ µ -Q G .f ∗ ∥ 2 L 2 , where ˜ R ∗ := R ∗ + ∥ ( f ∗ ) ⊥ G ∥ 2 L 2 ( X , Y ; π X ) is independent of µ (and doesn't intervene in the optimization). Finally, we get:

E.2.4 Proof of Corollary 2

When π is assumed to be equivariant, we can summon our previous results to prove Corollary 2.

Proof of Corollary 2. From Proposition 2 we know that equivariant data implies R : P ( Z ) → R is invariant, and also that this makes R = R DA . We conclude using Theorem 2.

We can readily extend this to the regularized case by recalling Corollary 6:

Corollary 8. When R : P ( Z ) → R and r : Z → R are G -invariant, a minimum for the regularized population risk R τ,β can be found within P G ( Z ) . When β > 0 such WI minimum is unique .

Proof of Corollary 8. Direct from Corollary 6, together with Proposition 3 (as in Corollary 2). The uniqueness comes from the strict convexity of the entropy term (see proposition 14).

E.2.5 Proof of Proposition 4 and Proposition 5

Proof of Proposition 4. Consider the group G = C 4 acting on R 2 via 90 ◦ rotations. Let K = B (0 , 1) ⊆ R 2 be a compact set. Consider a random variable X ∼ N (0 , Id 2 ) | K (i.e. given by X = Z ✶ Z ∈ K for Z ∼ N (0 , Id 2 ) ) and set Y = ∥ X ∥ 2 . Notice that π defined this way is compactly supported. Clearly G is finite (thus compact) and it can be seen as its ortogonal representation:

By the definition of our r.v.s, it is clear that:

X ( d ) = ρ g X ∀ g ∈ G because X ∼ N (0 , Id 2 )

∀ g ∈ G, ( X,Y ) ( d ) = ( ρ g X, ˆ ρ g Y ) (since ˆ ρ is the trivial representation, it is enough to notice that ∥ ρ g .X ∥ 2 = ∥ X ∥ 2 for all g ∈ G .

Therefore, π = Law ( X,Y ) is G -invariant (and compactly supported). Consider a shallow NN given by: Φ N θ : R 2 -→ R N × b -→ R (with b ∈ N and some action G ⟳ η R b ) as: Φ N θ ( x ) = 1 N ∑ N i =1 W i σ ( A T i x + B i ) , ∀ x ∈ R d ; where θ i = ( W i , A i , B i ) ∈ Z := R 1 × b × R 2 × b × R b ∼ = R D . We let G ⟳ M Z as described in appendix C.1:

We can assume, for instance, that b = 1 and η is the trivial representation (so that no condition is required for σ ∗ to be jointly G -equivariant) and recall that: θ i ∈ E G ⇐⇒ ∀ g ∈ G, M g θ i = θ i .

However, if we assume that: ∀ g ∈ G, ρ g A i = A i , then, in particular: ( -1 0 0 -1 )( A 1 i A 2 i ) = ( A 1 i A 2 i ) . This in turn implies, as A 1 i = -A 1 i and A 2 i = -A 2 i that A 1 i = A 2 i = 0 . i.e. A i ≡ 0 . Thus, any θ i = ( w i A i ) ∈ E G has A i = 0 . Therefore, if we choose any activation σ (e.g the sigmoid activation or σ = tanh , both C ∞ and bounded) and we choose N ∈ N ∗ and θ i ∈ E G ∀ i = 1 , ..., N ; then:

which is a constant independent of x . i.e. any equivariant architecture in this context is a constant function (whereas Y = ∥ X ∥ 2 is not). Notice, in particular, that any shallow model Φ ν with ν ∈ P ( E G ) will also be a constant function. In particular, notice that we will never do 'better' than minimizing over all posible constants:

The problem on the right has a known answer, which is C ∗ = E π [ ∥ X ∥ 2 ] > 0 . On the other hand, consider a fully conected neuronal network. By the universal approximation theorem (which applies for the chosen σ , as in [37, 20, 5]), as π is compactly supported (in particular, π X ( K ) = 1 ); we consider the parameters that approximate the function f ( x ) = ∥ x ∥ 2 in K = B (0 , 1) to precision ε > 0 . i.e. For ε ∈ (0 , √ C ∗ ) , we know: = ⇒∃ N ∈ N , ∃ a 1 , ..., a N ∈ R 2 , ∃ w 1 , ..., w N ∈ R 1 such that: √

Then:

But, in particular, ∃ ν N θ ∈ P ( Z ) such that:

and so:

In particular, we can't expect an optimum of the learning problem to be achieved within P ( E G ) .

To overcome situations as in Proposition 4, the usual setting is to assume some universality condition. This leads to Proposition 5, which we will now prove:

Proof of Proposition 5. A standard calculation from the quadratic loss case (see the proof of Corollary 1 for details) yields that, for any µ ∈ P ( Z ) :

where f ∗ ( x ) := E π [ Y | X = x ] and R ∗ is the Bayes risk of the problem. From Proposition 10, we know that f ∗ ∈ L 2 G ( X , Y ; π X ) , and so by universality of F σ ∗ ( P ( E G )) onto that space (as well as that of F σ ∗ ( P ( Z )) ), we conclude directly that: inf ν ∈P ( E G ) R ( ν ) = R ∗ = inf µ ∈P ( Z ) R ( µ ) .

Remark. Works such as [50, 60, 76, 77] precisely provide conditions under which universality of F σ ∗ ( P ( E G )) can be guaranteed (modulo some adaptations from their setting to ours).

Particularly, our single-hidden-layer NNs, with the width N →∞ , correspond to what is referred to as 'networks of tensor order 1' in the literature. As noted in [50], such kind of equivariant NNs are unable to achieve universality for certain types of group actions (see Theorem 2 from [50]). Despite this, 'first order universality' has already been established for some of the most important examples of equivariant architectures, such as Deep Sets ([76, 77]) and CNNs ([76]).

Adapting our setting, in order to eventually allow for arbitrary order tensors in the MF formulation, is part of the future challenges to make our work more broadly applicable.

E.3 Proofs of results in Section 3.3

Having laid out all the different relevant elements for our work, we can now procede with the proofs of some of our main results.

E.3.1 Proof of Theorem 3, Corollary 3 and Theorem 4

We start by proving Theorem 3 on the general case.

Proof of Theorem 3. We know that a family ( µ t ) t ≥ 0 ⊆ P 2 ( Z ) satisfies WGF ( F ) in the weak sense if ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

Now, profiting from the uniqueness of the solutions of this equation, it will be enough to show that, given a solution ( µ t ) t ≥ 0 ⊆ P 2 ( Z ) of WGF ( F ), then ( µ G t ) t ≥ 0 ⊆ P G 2 ( Z ) is also a solution . Indeed, consider, for g ∈ G , ˜ µ t = M g # µ t , and notice that for φ ∈ C ∞ c ( Z × (0 , T )) :

Now, we can define φ g ∈ C ∞ c ( Z × (0 , T )) given by ∀ ( z, t ) ∈ Z × (0 , T ) φ g ( z, t ) = φ ( M g .z, t ) , which satisfies:

So that, by also using proposition 13 and the orthogonality of the group action, we get:

Where the last equality comes from the fact that ( µ t ) t ≥ 0 is a solution to the WGF .

In particular, as we also have that ˜ µ 0 = M g # µ 0 = µ 0 (because µ 0 ∈ P G ( Z ) ), by uniqueness we can conclude that this means that ∀ g ∈ G, ∀ t ∈ (0 , T ) λ -a.e. , µ t = M g # µ t .

This may seem weaker that what we want to prove. Nevertheless, as our group is compact and has a unique normalized Haar measure, we can proceed as follows: let f : [0 , T ] ×Z → R + be any positive and measurable function. Given g ∈ G , take Ω g ⊆ [0 , T ] a full measure set where it holds that µ t = M g # µ t . In particular, f t = f ( t, · ) : Z → R is positive and measurable, so that:

and we can integrate this equality to get: ∫ T 0 ⟨ f t , µ t ⟩ dt = ∫ T 0 ⟨ f t ◦ M g , µ t ⟩ dt Now, by integrating both sides with respect to the Haar measure, and applying Fubini's theorem (because everything is positive) we get:

Implying (by a standard argument) that ∀ t ∈ [0 , T ] a.e. µ t = µ G t , and therefore: ∀ t ∈ [0 , T ] a.e. µ t ∈ P G ( Z ) .

τ,β

Proof of Corollary 3. From Corollary 2, Corollary 6 and Corollary 8, we know that R is invariant. On the other hand, from [38] (or [70]) we know that, under our assumptions, a unique weak solution to the Fokker-Planck equation exists. Furthermore, this solution is known to be strong if β > 0 .

In particular, theorem 3 applies and allows us to conclude that if µ 0 ∈ P G 2 ( Z ) , then ∀ t ≥ 0 (a.e.) µ t ∈ P G 2 ( Z ) .

When β > 0 , since solutions are strong , we conclude that the densities ( u t ) t ≥ 0 are all G -invariant functions ( λ -a.e.). This follows from the remark about densities of invariant measures provided in SuppMat-B.2.

Remark. When β > 0 , we have a unique weakly-invariant minimizer (from proposition 3 and/or corollary 8); and also, under mild assumptions, a global convergence result. That is, independently of the network's initialization, we will converge to the G -invariant solution. An interesting question in this setting is then: At which point does the WGF enter the space P G 2 ( Z ) ?

We can also prove Theorem 4:

Proof of Theorem 4. This proof follows from the fact that ∀ z ∈ Z , ∀ µ ∈ P G ( Z ) :

Indeed, notice that, from proposition 18:

while also:

Now, let ( µ FA t ) t ≥ 0 and ( µ DA t ) t ≥ 0 be the WGF solutions starting from µ 0 for R FA and R DA respectively. As R FA is G -invariant, by corollary 3, ( a.e. ) ∀ t ≥ 0 , µ FA t ∈ P G ( Z ) . Now, let's see that this process actually also satisfies WGF ( R DA ), forcing both processes to coincide by uniqueness.

Indeed, we know that ( µ FA t ) t ≥ 0 satisfies: ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

Now, as ( a.e. ) ∀ t ≥ 0 , µ FA t ∈ P G ( Z ) , we have ∀ z ∈ Z : D µ R FA ( µ FA t , z ) = D µ R DA ( µ FA t , z ) . In particular, ( µ FA t ) t ≥ 0 satisfies ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

Implying that ( µ FA t ) t ≥ 0 solves WGF ( R DA ) starting from µ 0 ; thus by uniqueness: ( µ FA t ) t ≥ 0 = ( µ DA t ) t ≥ 0 .

The last part of the theorem comes from Proposition 2, since if R is invariant, its WGF will exactly coincide with that of R DA (they are the same functional).

Remark. This results tells us the ultimate bottom line: at the MF level, training with DA or FA results in the exact same dynamic. Furthermore, whenever data is equivariant, they are both essentially equivalent to applying no technique whatsoever . Despite this result concerning infinitely wide NNs, it provides meaningful practical insights (as shown Appendix F) for large enough NNs, that could be used in applications.

Remark. Since R DA , R FA and R all coincide on P G ( Z ) , one could expect Theorem 4 to hold for R even without assuming π to be equivariant. However, the invariance of R is crucial for such a result: if R isn't invariant, nothing guarantees that its WGF process will stay within P G ( Z ) , whereas WGF ( R DA ) and WGF ( R FA ) always do so.

E.3.2 Proof of Theorem 5 and Theorem 6

We can now provide the proof for the stronger result , Theorem 5, stating that the WGF of an invariant functional will respect E G all along training. This proof uses the McKean-Vlasov non-linear SDE (Equation (11)) presented in Appendix D.2 (see [22] for a reference). Namely, we will consider the following projected McKean-Vlasov SDE (with ( B t ) t ≥ 0 a BM on Z ), given by:

which corresponds to the MF limit dynamics arising from performing the projected noisy SGD scheme from Equation (5). This can be shown by adapting relatively standard arguments from the MF literature on NN, see e.g. [22, 23, 53].

Proof of Theorem 5. The proof has two steps. The first will consist in showing that the process (12) satisfies ∀ t ≥ 0 , µ t ∈ P ( E G ) . In the second step we will check that ( µ t ) t ≥ 0 = ( Law ( Z t )) t ≥ 0 is a solution to the WGF ( R τ,β E G ) presented in Section 3.3.

Step 1: The ( pathwise unique ) solution of the projected McKean-Vlasov SDE (12), Z = ( Z t ) t ≥ 0 satisfies a.s. for all t ≥ 0 :

Here, ξ 0 is such that Law ( ξ 0 ) = µ 0 , and R τ := R + ⟨ r, ·⟩ is being used as shorthand notation. We first let g ∈ G be an arbitrary group element, and we study how the process ˜ Z = ( ˜ Z t ) t ≥ 0 := ( M g Z t ) t ≥ 0 satisfies this same equation (13).

Denote ν s := M g # µ s as the law of ˜ Z s , we want to show that for all t ≥ 0 :

Indeed, first notice that:

Let Ω be the full measure set where ξ 0 ∈ E G (which we can do since µ 0 ∈ P ( E G ) , or, equivalently: P ( ξ 0 ∈ E G ) = 1 ). Then, ∀ ω ∈ Ω , Z 0 ( ω ) = ξ 0 ( ω ) ∈ E G . In particular, ∀ ω ∈ Ω , ∀ g ∈ G, ˜ Z 0 ( ω ) = M g Z 0 ( ω ) = M g ξ 0 ( ω ) = ξ 0 ( ω ) = Z 0 ( ω ) . That is, ˜ Z 0 a.s. = Z 0 .

Now, the equation is satisfied by ( Z t ) t ≥ 0 and therefore, for t ≥ 0 , we have:

Here, we used the linearity of the integral (and the stochastic integral), the fact that ∀ g ∈ G, M g P E G = P E G , and Proposition 13, which holds for ∀ θ ∈ Z , ∀ µ ∈ P ( Z ) (in particular for θ = Z s ( ω ) , ∀ ω ∈ Ω and µ s = Law ( Z s ) ). Thus, ∀ g ∈ G , (14) holds.

By the pathwise uniqueness of the solution ( Z t ) t ≥ 0 , we have (following, for instance, [28]):

In particular, as g ∈ G was arbitrary, we have that:

We now want to be able to interchange the ∀ g ∈ G with the probability measure. Fortunately, we are dealing with a compact group with a normalized Haar measure λ G . Indeed, from equation (15) we deduce that ∀ g ∈ G, ∀ t ≥ 0 , P ( ∥ Z t -M g Z t ∥ = 0) = 1 .

Now, notice that, for any t ≥ 0 and ω ∈ Ω :

We can integrate both sides by P to get (using Fubini as functions are positive and measurable):

where in the last step we have used the fact that ∀ g ∈ G, ∀ t ≥ 0 , P ( ∥ Z t -M g Z t ∥ = 0) = 1 , so that ∀ t ≥ 0 , ∀ g ∈ G, ∫ Ω ∥ Z t ( ω ) -M g .Z t ( ω ) ∥ d P ( ω ) = 0 .

This implies that ∀ t ≥ 0 P -a.s. Z t = P E G Z t , i.e. P ( Z t ∈ E G ) = µ t ( E G ) = 1 , or, in other words, ∀ t ≥ 0 , µ t ∈ P ( E G ) as required. Note that all arguments work as well in the case that β = 0

Step 2: We now prove that ( µ t ) t ≥ 0 , studied in the previous step, is a weak solution to equation (3); that is, to the WGF ( F ) , with (using the previously introduced notation R τ ):

Also recall the notation H E G ( µ ) := H λ E G ( µ E G ) = H λ E G ◦ P E G #( µ ) presented in the end of Appendix D.1, as well as the calculations for its intrinsic derivative.

It is standard to check, applying Itô's formula and taking expectation, that the family ( µ t ) t ≥ 0 = ( Law ( Z t )) t ≥ 0 satisfies, ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

with tr denoting the trace of a square matrix and D 2 z the Hessian matrix acting on the z variable. Notice that the process P E G B is classically a Brownian motion in E G . Together with the fact that µ t is supported on E G , this implies that, for β > 0 , µ t has a density w.r.t. λ E G .

It is clear from the case β = 0 that the terms involving first order spatial derivatives of φ exactly give rise to the terms associated with functional R τ in the definition (3) of the WGF ( R τ,β E G ) . Therefore, we just need to check that for all t ≥ 0 , the distribution defined for every ϕ ∈ C ∞ c ( Z ) by ϕ ↦→

∫ Z tr [ P E G D 2 z ϕ ( z )] dµ t ( z ) and the distribution div ( D µ H E G ( µ t , · ) µ t ) are equal. In fact, for all such ϕ we have:

Since µ t is concentrated in E G , we have P E G # µ t = µ t . This and the equality P 2 E G = P E G imply the previous expression is equal to:

Noticing that P T E G ∇ z is the gradient calculated on E G , integrating by parts with respect to λ E G , and using again the fact that µ t is concentrated in E G , a straightforward calculation yields the desired expression ∫ Z tr [ P E G D 2 z ϕ ( z )] dµ t ( z ) . This concludes the proof.

Remark. Notice that, by a.s. continuity of the McKean-Vlasov diffusion (12) and the fact that E G is closed, Step 1 of the previous proof actually shows that P ( Z t ∈ E G , ∀ t ≥ 0) = 1 .

Notice also that Theorem 5 bears some resemblance to Corollary 1 in [30], which states that E G is stable under the traditional gradient flow of the augmented risk ( [ θ ∈ Z N ↦→ R DA ( θ ) ∈ R ] , as in Section 2.3). Our result shares a similar flavor, but for the MF dynamics of freely-trained NNs with equivariant data.

Remark. Unlike with WI distributions, initializing a shallow NN with µ 0 ∈ P ( E G ) isn't as straightforward as using a normal distribution. Effectively (and efficiently) computing the space E G is actually quite challenging (as noted in [29]).

A natural way to ensure that µ 0 ∈ P ( E G ) , independently of the form of E G , is to initialize all parameters to be 0 . The question of whether under such initialization the parameters will eventually exit { 0 } (or some larger, strict subspace E ⊊ E G ) and find values over the entire space E G is left for future work. Some insights on this behaviour can be sought in our experimental results, see Section 4. If true, this behavior could point towards some type of underlying hypoellipticity of the McKean-Vlasov dynamics (12) (or variants) on E G , which would be interesting to analyze, in particular in view of potential theoretical guarantees for architecture-discovering heuristics as suggested in Section 4.2.

Note that there is no need for seeing E G as a subspace of an ambient space Z . When training with EA s, we simply force our parameters to live on E G , since we fix the architechture beforehand. Namely, our 'whole space' is ˜ Z = E G (regarded directly as a vector space E G ∼ = R ˜ D ) rather than Z . Thus, the relevant population risk is the restricted version of the original: ˜ R := R | P ( ˜ Z ) : P ( ˜ Z ) → R ; and we can apply the usual results from the MF Theory when the relevant hypothesis are satisfied by ˜ Z and ˜ R . Notably, we can have global convergence of ˜ R τ,β to inf µ ∈P ( ˜ Z ) ˜ R τ,β ( µ ) = inf µ ∈P ( E G ) R τ,β E G ( µ )

Remark. As shown in [38] (see Proposition 19 in Appendix G for the details) the regularized versions of the involved functionals (i.e. R τ,β and R τ,β E G ) Γ -converge to the original R as τ, β → 0 ; meaning that, for small values of the regularization parameters, we should expect the achieved optima to ressemble inf µ ∈P G ( Z ) R ( µ ) and inf ν ∈P ( E G ) R ( ν ) respectively (or, under Proposition 5, both to R ∗ ). We will also leave the exploration of how this approximation behaves as future work.

Finally, we provide a proof for Theorem 6

Proof of Theorem 6. The proof structure is very similar to that of Theorem 4. Namely, it comes from noticing that for µ ∈ P ( E G ) and z ∈ E G , we have:

We already know the first equality, as seen in Theorem 4 (since µ ∈ P G ( Z ) from lemma 1). We only need to show the last equality. Indeed, notice that:

while also, from Proposition 18:

Knowing this, the rest of the proof is analogous to that of Theorem 4. Let ( µ FA t ) t ≥ 0 , ( µ DA t ) t ≥ 0 and ( µ EA t ) t ≥ 0 be the WGF solutions starting from µ 0 for R FA , R DA and R EA respectively. Since µ 0 ∈ P ( E G ) ⊆ P G ( Z ) , from Theorem 4 we know that ( µ FA t ) t ≥ 0 and ( µ DA t ) t ≥ 0 coincide. Let's see that, w.l.o.g., ( µ FA t ) t ≥ 0 coincides with ( µ EA t ) t ≥ 0 .

As R FA is G -invariant, by theorem 5, ∀ t ≥ 0 , µ FA t ∈ P ( E G ) . Now, let's see that this process also satisfies WGF ( R EA ), forcing both processes to coincide by uniqueness.

As before, we know that ( µ FA t ) t ≥ 0 satisfies: ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

Now, as ∀ t ≥ 0 , µ FA t ∈ P ( E G ) , we can restrict our integral to E G . Also, we have ∀ z ∈ E G : D µ R FA ( µ FA t , z ) = D µ R EA ( µ FA t , z ) . With these properties, ( µ FA t ) t ≥ 0 satisfies ∀ φ ∈ C ∞ c ( Z × (0 , T )) :

Making the integral over Z once again (we can since µ FA t ∈ P ( E G ) for all t ≥ 0 ), we get that ( µ FA t ) t ≥ 0 solves WGF ( R EA ) starting from µ 0 ; thus by uniqueness: ( µ FA t ) t ≥ 0 = ( µ EA t ) t ≥ 0 .

The last part of the theorem comes, once again, from Proposition 2, since if R is invariant, its WGF exactly coincides with that of R DA .

F Experimental setting and further experiments

All the different experiments were run on Python 3.10, on a Google Colab session consisting (by default) of 2 Intel Xeon virtual CPUs (2.20GHz) and with 13GB of RAM.

In order to obtain results that can be visualized, we consider a simple setting where X = Y = R 2 and Z = R 2 × 2 ∼ = R 4 . We let G = C acting on X and Y via the coordinate transposition

2 action (i.e. the group generated by the orthogonal matrix ( 0 1 1 0 ) ); and on Z via the natural intertwining action (i.e. M g .z = ˆ ρ g .z.ρ T g ). We also consider the jointly equivariant activation given by σ ∗ ( x, z ) = σ ( z · x ) ∀ x ∈ R 2 , ∀ z ∈ R 2 × 2 with σ : R → R a sigmoidal activation function (which is C ∞ and bounded) applied pointwise. Under this setting, E G can be explicitly computed

as E G = 〈( 1 √ 2 0 0 1 √ 2 ) , ( 0 1 √ 2 1 √ 2 0 )〉 , which is a 2 -dimensional subspace of the ambient 4

dimensional space. It's projection operator P E G is also explicitly known.

We consider a teacher model f ∗ = Φ N ∗ θ ∗ with N ∗ fixed particles, such that ν N ∗ θ ∗ is either arbitrary , WI or SI . Let ϑ = 0 . 5 be a scale parameter. The arbitrary particles were chosen to be, for N ∗ = 5 :

This was fixed in order to make the task non-trivial and interesting . The WI teacher distribution was simply chosen to be ( ν N ∗ θ ∗ ) G , with θ ∗ as just described, so that the corresponding teacher function resulted to be f ∗ = Q G . Φ N ∗ θ ∗ . In other words, the WI distribution has 10 particles, corresponding to each of those of θ ∗ , together with their image under the G -action. The SI particles were also fixed, but their chosen coordinates had to be expressed in terms of the basis vector of E G (i.e. only providing 2 parameters). Particularly, they were fixed to be N ∗ = 5 and, denoting them by a ∗ = ( a ∗ i ) N ∗ i =1 to avoid confusion, explicitly described as:

As seen on Section 3.1, the teacher f ∗ : X → Y will be an equivariant function as soon as its parameter distribution is chosen either WI or SI . Our data distribution π will be such that ( X,Y ) ∼ π will satisfy X ∼ N (0 , σ 2 π . Id 2 ) (with σ π = 4 ), and Y = f ∗ ( X ) . Namely, π X will always be G -invariant, whereas π will only be G -invariant if f ∗ is. This setting allows for testing the different results provided, without losing the properties of Q G as a projection (which require π X to be G -invariant, as shown in [27]).

-

We will try to mimic the teacher network by using student networks , which will be given by Φ N θ ; namely, with the same σ ∗ , but varying values of N and θ ∈ Z N . We will train them to minimize the regularized population risk R τ,β given by a quadratic loss , ℓ ( y, ˆ y ) = ∥ y -ˆ y ∥ 2 Y , and a quadratic penalization , r ( z ) = ∥ z ∥ 2 Z . For this purpose, we employ a minibatch variant of the SGD training scheme provided in Equation (1) (possibly projected , as in Equation (5), when required). We will also employ the different symmetry-leveraging techniques presented in Section 2.3, such as DA , FA and EA . We refer to the free training with no SL-techniques whatsoever as the vanilla training.

The training parameters were fixed to be (unless explicitly stated otherwise):

Step Size: ς ≡ α > 0 (with α = 50 in most experiments), ε N = 1 N , so that s N k = α N . This was convenient, since it corresponds to the usual implementation of SGD on most common NN frameworks in Python (namely, pytorch and jax ).

Regularization parameters: τ = 10 -4 and β = 10 -6 .

Batch Size: It was chosen to be B = 20 .

Number of Training Epochs: In line with the statement of Theorem 1, to observe phenomena at a MF scale, we need an amount of iterations (commonly known as epochs in the ML literature) that is proportional to the number of particles. For this purpose, we fix an observation time horizon of T = 20 . All training schemes were performed for a total of N e = N · T epochs (iterations). An additional 'granularity' parameter (usually set to be gr = 5 ) is introduced to determine how often in the dynamic we will observe and save the training losses and particle positions : we do so every ⌊ N e gr ⌋ steps. Notice that N e depends on N , and so models with different values of N were trained for a different amount of epochs.

Student Initialization: The student's particles, θ ∈ Z N , are initialized i.i.d. from some µ 0 that is chosen to be either WI or SI . When WI -initialized, they are sampled from a random gaussian Z ∼ N (0 , 1 16 ) . When SI -initialized, particles are taken to be P E G .Z with Z as before.

Number of Repetitions: Each experiment was repeated a total of N r = 10 times to ensure consistency. Each repetition, a different random seed was employed to: initialize the student's particles, generate the training data, and generate the noise for the SGD iteration. In particular, on a fixed repetition, all models were trained with the same data and the same noise being applied on SGD updates.

Remark. As N e is chosen to be proportional to the number of particles, N , computational burden and memory requirements quickly became heavy for the simple machines we employed (which didn't even have a dedicated GPU). This is the reason why we don't scale our experiments beyond the N = 5000 case. As reference, for N = 5000 a single training (with the above hyperparameters) of a single model (either of vanilla , DA , FA or EA ) took ≈ 15 minutes (which quickly amounts to large amounts of running time for the N r = 10 repetitions, the 4 different training schemes and the 6 possible settings with WI or SI initialization and arbitrary , WI or SI teacher). This is a clear point of improvement and shall be tackled in future work.

Remark. As here mentioned, on every fixed 'repetition' of the experiments, the same noise was used during the SGD training iterations for the vanilla , DA and FA schemes. The EA scheme, despite using the same seeds for the data and student initializations, didn't have the same noise applied during SGD. This was because, despite using the same seed for the noise generator, noise for our EAs was only 2 -dimensional (since EA s are parametrized by E G ), while it was 4 -dimensional for the other schemes. This made the resulting training schemes have an additional layer of noise separating them; and so solving this issue, in order to properly visualize Theorem 6, becomes fundamental.

Remark. Notice that the N r = 10 performed repetitions were largely enough to allow for plots with error bars (actually, we do boxplots which encode the variability of the different quantities more precisely) that allow for significant analysis of the observed phenomena. We do not go beyond N r = 10 due to the low computational capabilities of our machines, and the already high computational cost of running the experiments (for thousands of hidden units and epochs, in many different settings, and involving the calculation of Wasserstein Distances, as we'll comment below).

To facilitate the implementation of the ideas behind our EA models, we use the group and representations tools from the emlp repository provided as part of [29]. This code is openly available and

has an MIT License for unrestricted access. We employ it to numerically (and efficiently) determine the space E G (namely, its basis), as well as P E G . We do remark that these calculations were correct only up to a precision of 10 -8 , which results in a slight burden for our empirical results. On the other hand, regarding the implementation of EMLP s provided in the package ( EA s in our setting), some slight modifications to the source code had to be performed in order to correctly represent our setting.

On a similar note, we can numerically compute the squared Wasserstein2 distance between two empirical distributions of particles by employing the pyot library. This allows us to evaluate to what extent our resulting models are close to each other in terms of their particle distribution ν N θ (which is what the MF approach suggests). In order to fix a common scale in which the experiments can be compared, for different values of N , and mitigate the effects of fluctuating empirical estimates (mainly for small values of N , and in the low dimensions considered), we consider a natural normalization of the Wasserstein2 distance, which we refer to as the RMD (Relative-Measure-Distance). This is defined as: RMD 2 ( µ, ν ) = W 2 2 ( µ,ν ) M 2 µ + M 2 ν where M 2 µ = 2 E [ ∥ Z ∥ 2 ] for Z ∼ µ (so that 0 ≤ RMD ≤ 1 ).

The RMD provided a good metric for the experiments here presented. Notice that, as N →∞ , by the L.L.N. for empirical distribution following from the MFL convergence, the RMD is expected to stabilize at the corresponding value of the limiting distributions. Therefore, up to a multiplicative quantity approaching a (finite, non null, in our case) constant, we are observing the behavior of the Wasserstein2 distance. A drawback from using the Wasserstein-2 metric, is that calculating them can be very expensive computationally. This is another one of the reasons why our experiments only get to N = 5000 particles.

F.1 Study for varying N

Beyond the analyisis already provided in Section 4, we here provide some meaningful insights. We want to observe to what extent the properties proved in Section 3.3 for the WGF of R τ,β can be observed in practice. For this purpose, we observe, for N ∈ { 5 , 10 , 50 , 100 , 500 , 1000 , 5000 } , different relevant quantities to evaluate the different combinations of teachers and students .

For this set of experiments, WI -initialized students were trained with the usual SGD scheme from Equation (1); while, SI -initialized students, were trained with the projected SGD dynamics from Equation (5). Models for the different schemes to be compared, are all initialized with the exact same (random) particles.

Figure 3 displays some comparisons between particle distributions in terms of RMD at the end of training, knowing that they were all initialized with the same particles drawn from µ 0 ∈ P ( E G ) . From this, we can visually see that, when the teacher distribution is either WI of SI , the resulting distribution from vanilla training stays on E G (since RMD 2 ( ν N NT , ( ν N NT ) E G ) is small) increasingly more as N becomes large. This fact is absolutely remarkable, since, for a WI teacher there should be no reason why the vanilla training (that's completely free , in principle) shouldn't escape E G to achieve a better approximation of f ∗ . On the other hand, in every single teacher setting, almost independently of N , both DA and FA consistently remain within E G (as expected, even if f ∗ isn't equivariant). For an arbitrary teacher , we see that the vanilla training distribution readily leaves E G to better approximate f ∗ , which isn't a predicted behaviour from our theory (we have no guarantees of 'leaving E G when data isn't equivariant'), but motivates the heuristic defined in Section 4.

Still from Figure 3, we can see that, as N grows bigger, the end-of-training distribution of the vanilla scheme becomes closer and closer to that of DA and FA (from Theorem 4 we actually expect them to be equal in the limit). A similar result is obtained relating vanilla , DA and FA to the EA scheme; the values are however larger than before and less significantly close in general. This is possibly due to the different noises employed during training (as mentioned in a remark above). We do however notice that for increasing values of N , the EA , DA , FA and vanilla schemes (the latter only under equivariant f ∗ ) tend towards coinciding, which serves to illustrate the constatations from Theorem 6. Finally, notice that the results obtained for WI and SI teachers present almost no quantitative differences whatsoever between them.

In Figure 4 we present a visualization of the final particle distribution, after an SI -initialized training under a WI teacher f ∗ , of the vanilla , DA , FA and EA schemes (on a single realization of the experiment). At least visually (and macroscopically), it seems like all these regimes followed (approximately) the same flow, as they end up with an approximately equal particle distribution. This isn't a rigorous comparison at all, and providing better quantitative comparisons between the methods

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

7

10

5

10

3

10

1

Relative Measure Distance

Distance to Projected Version

Model

V vs. P(V)

DA vs. P(DA)

FA vs. P(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

6

10

4

10

2

Relative Measure Distance

Distance to Projected Version

Model

V vs. P(V)

DA vs. P(DA)

FA vs. P(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

10

1

Relative Measure Distance

Distance to Projected Version

Model

V vs. P(V)

DA vs. P(DA)

FA vs. P(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

6

10

4

10

2

10 0

Relative Measure Distance

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

10

1

Relative Measure Distance

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

10

1

Relative Measure Distance

Comparison between vanilla, DA and FA schemes

Model

V vs. DA

V vs. FA

DA vs. FA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

10

1

10 0

Relative Measure Distance

Comparison between EA and other techniques

Model

V vs. EA

DA vs. EA

FA vs. EA

(a)

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

10

1

Relative Measure Distance

Comparison between EA and other techniques

Model

V vs. EA

DA vs. EA

FA vs. EA

(b)

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

10

1

Relative Measure Distance

Comparison between EA and other techniques

Model

V vs. EA

DA vs. EA

FA vs. EA

(c)

SI

teacher

is to be considered for future work. As a counterfactual, we provide in Figure 5 the results of an SI -initialized training that is performed under a non-equivariant f ∗ . We can see that the vanilla model readily leaves E G to achieve a better approximation of f ∗ , while the DA , FA and EA schemes 'stay inside' (roughly coinciding between them as well).

Figure 6 also displays RMD comparisons between particle distributions at the end of training, but for an (identic) initialization with particles drawn from µ 0 ∈ P G ( Z ) . Now, unlike the SI case, with particles sampled i.i.d. from a WI distribution, nothing ensures that the resulting ν N 0 will be WI as well. Namely, in this case the limit as N →∞ becomes significantly more important to visualize the theoretical results. Indeed, since ν N 0 isn't necessarily WI , we no longer have a guarantee that the finiteN networks trained with DA , FA or vanilla methods will be close to each other in any sense. We do however notice on Figure 6 that, for increasing N , the end-of-training distributions of DA , FA and vanilla schemes (the latter only when f ∗ is equivariant) become increasingly closer to their symmetrized versions (namely, RMD 2 ( ν N NT , ( ν N NT ) G ) becomes smaller, though never as small as in the SI -initialized experiments). Also, as guaranteed by Theorem 4, for large N we see that DA and FA become indistinguishably close, no matter the teacher's properties; also, when f ∗ is equivariant, they both 'coincide' with the vanilla scheme. A comparison between the EA scheme and the result from projecting the FA scheme on the last step is also presented. It is used simply to illustrate that, in principle, directly training on E G isn't necessarily comparable to performing 'free-training' and projecting the resulting particle distribution only on the last step (even when an SL technique such as FA is used).

0.4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (vanilla)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (DA)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (FA)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (EA)

z4

0.6

0.4

0.2

z3

0

-0.2

-0.4

-0.6

-0.4

0.4

0.6

-0.2

0.2

0

z1

0.2

0

0.4

9'0

0.4

0.2

z3

0

0.2

0.4

0.6

0.4

0.2

b0

0.6

0.2

z2

0

z1

0.6

0.4

0.2

z3

0

-0.2

-0.4

-0.6

-0.4

0.4

0.6

-0.2

0.2

0

z1

0.2

-0.2

0

z2

9'0

0.4

0.2

z3

0

0.2

0.4

0.6

0.4

0.2

b0

0.6

0.2

z2

0

z1

0.6

0.4

0.2

z3

0

-0.2

-0.4

-0.6

-0.4

0.6

-0.2

0.2

0

z1

0.2

-0.2

0

z2

0.4

0.6

0.4

0.2

z3

0

0.2

0.4

0.6

-0.4

0.6

0.4

0.2

0.2

0

0

70

TZ

0.6

0.4

0.2

Z3

0

-0.2

-0.4

-0.6

-0.4

0.6

-0.2

0.2

0

z1

-0.2

0

z2

0.4

0.6

0.4

0.2

z3

0

0.2

0.4

0.6

-0.4

0.6

-0.4

0.2

0.2

z2

0.2

z1

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (vanilla)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (DA)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (FA)

z4

Teacher Particles

Student Particles

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

-0.6

-0.4

-0.2

0

0.2

0.4

0.6

Student Particles' Positions (EA)

z4

0.6

0.4

0.2

Z3

0

-0.2

-0.4

-0.6

0.6

0.6

0.4

-0.4

0.2

-0.2

0

z1

0

-0.2

z2

0.2

0.6

0.4

0.2

z30

-0.2

0.4

0.6

-0.6

0.4

0.2

0.2

b0

0

z1

>

0.6

0.4

0.2

0

Z3

-0.2

-0.4

-0.6

-0.4

0.4

0.6

-0.2

0.2

0

z1

0.2

-0.2

0

z2

0.4

0.6

0.4

0.2

z3

0

0.2

0.4

0.6

-0.4

0.6

0.4

0.2

0.2

z2

z1

0.6

0.4

0.2

z3

0

-0.2

-0.4

-0.6

-0.4

0.4

0.6

-0.2

0.2

0

z1

-0.2

0

z2

0.4

0.6

0.4

0.2

Z3

0

0.2

0.4

9'0-

-0.4

0.6

0.4

0.2

0.2

0

z2

z1

0.6

0.4

0.2

z3

0

-0.2

-0.4

-0.6

-0.4

0.4

0.6

-0.2

0.2

0

z1

0.2

0

z2

0.4

O

0.6

0.4

0.2

Z3

0

0.2

0.4

0.6

-0.4

0.6

0.4

0.2

0.2

0

z2

Z1

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

1

10 0

Relative Measure Distance

Distance to Symmetrized Version

Model

V vs. G(V)

DA vs. G(DA)

FA vs. G(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

1

Relative Measure Distance

Distance to Symmetrized Version

Model

V vs. G(V)

DA vs. G(DA)

FA vs. G(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

2

10

1

10 0

Relative Measure Distance

Distance to Symmetrized Version

Model

V vs. G(V)

DA vs. G(DA)

FA vs. G(FA)

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

4

10

3

10

2

10

1

10 0

Relative Measure Distance

Distance between SL techniques after training

Model

V vs. DA

V vs. FA

DA vs. FA

P(FA) vs. EA

(a)

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

6

10

5

10

4

10

3

10

2

10

1

10 0

Relative Measure Distance

Distance between SL techniques after training

Model

V vs. DA

V vs. FA

DA vs. FA

P(FA) vs. EA

(b)

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

6

10

5

10

4

10

3

10

2

10

1

10 0

Relative Measure Distance

Distance between SL techniques after training

Model

V vs. DA

V vs. FA

DA vs. FA

P(FA) vs. EA

(c)

SI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

10

1

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

10

1

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

(a)

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

(b)

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

8

10

7

10

6

10

5

10

4

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Model

Model

vanilla

DA

FA

EA

(c)

SI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

2

10

1

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

2

10

1

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

(a)

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

(b)

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to T eacher

Model

vanilla

DA

FA

EA

(c)

SI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

2

10

1

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

10

1

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

(a)

arbitrary

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

(b)

WI

teacher

5

10

50

100

500

1000

5000

Number of Particles (

N

)

10

3

10

2

L

2  Distance (estimated)

Distance to Symmetrized Teacher

Model

vanilla

DA

FA

EA

(c)

SI

teacher

Now, beyond the analysis of the underlying particle distributions after training, we turn our focus to comparisons of the resulting models . Wemeasure some distances in L 2 ( X , Y ; π X ) , by approximating ∥ · ∥ L 2 ( X , Y ; π X ) with a Monte-Carlo sample of 100 random data points drawn from π .

Figure 7 shows that the observed behaviour for the underlying particles , ν N θ , of each model, is consistent with the behaviour of the obtained model Φ N θ . That is, as particles become close to being symmetric in some sense, the resulting shallow model also becomes increasingly equivariant as well.

Finally, Figure 8 and Figure 9 illustrate quite well the observations from Corollary 1. When our teacher isn't equivariant , models trained using any kind of SL technique end up suffering from the inductive

bias introduced by the symmetric assumption (something that's hinted by the symmetrization gap characterization from Lemma 2 presented in SuppMat-B.3). On the other hand, the vanilla model thrives in approximating f ∗ as it is capable of breaking the assumed symmetry, unlike DA , FA and EA . The training regimes that use SL techniques are effectively approximating Q G .f ∗ , as shown in Figure 9 (and proven in Corollary 1). We also notice that, for the WI -initialized experiments, EA s end up suffering from their constraint of staying within E G , as they can't approximate f ∗ (or Q G .f ∗ ) as well as DA or FA (even when the teacher is SI ). This isn't the case for the SI -initialized experiments, where the performance of DA , FA and EA (and vanilla only for equivariant f ∗ ) is quite closely comparable (once again, hinting at Theorem 6). We also notice a general trend showing that, for bigger N , the approximations of f ∗ (or, eventually, Q G .f ∗ ) become increasingly better (specially for the WI -initialization).

F.2 Heuristic algorithm for discovering EA parameter spaces

The proposed heuristic that we infer from the results on the previous experimental setting is quite thoroughly described in Section 4.2. We only notice that, for this particular setting, the learning rate was chosen to be α = 20 (to better approximate the MFL conditions). Beyond the description of the proposed heuristic and the simple example visualized in Figure 2, we also provide Figure 10 here, illustrating a possible threshold choice in that setting.

Considering E j for j = 0 , 1 , . . . as the spaces that are discovered on each step of the heuristic, Figure 10 displays the values of: RMD 2 ( ν N θ , P E j # ν N θ ) and RMD 2 ( ν N θ , P E G # ν N θ ) ; both before and after training on a given heuristic step j . The red line simbolizes a possible value of δ that could be fixed to detect whenever the obtained particle distribution after training stayed on E j . In the case of this example, on steps 0 and 1 we would decide that the training left the original space E j , but we wouldn't do so on step 2 , allowing us to fix E G := E 2 . As shown by the values of RMD 2 ( ν N θ , P E G # ν N θ ) , we wouldn't be too far off with our prediction.

Despite this proposed heuristic being potentially interesting for real-world applications, we acknowledge that the setting where it is applied here might be too simple , synthetic and idealized . On one hand, this provided a clean-enough framework, where the underlying phenoma could be easily observed. However, in order to properly validate our heuristic approach, experiments with more complex settings (and with larger and more intricate datasets) need to be performed. These should also be coupled with sound theoretical guarantees, whose exploration we leave for future work.

Finally, we here provide some further details on the possible connections of our proposed heuristic, to the 'symmetry-discovery' method presented in [72]. In their work, they employ an architecture based on relaxed group convolution layers, which allows to 'detect' breaks of the supposed data symmetry, by observing the un-alignment of the layer weights .

As in our work, their method starts with the most possibly constrained architecture: the null space in our case, and the 'perfectly aligned weights' in theirs. This ensures that the model will start respecting symmetry with respect to the largest possible group; only for the training on data to cause these symmetries to 'break' overtime. Their method seemingly works in a single training iteration, while ours iteratively constructs the invariant linear subspace E G by adding one new 'symmetry-breaking dimension' at a time. Our Theorem 5 guarantees that, in the MF scale, our method won't leave E G ; but we have yet to establish symmetry-breaking guarantees for our heuristic, comparable to Proposition 3.1 in [72].

Finally, it's important to note that neither one of the methods is truly discovering the 'underlying symmetry' of the data. They are both closer to simply 'optimizing architectures' which are compatible with data symmetries: either finding the 'right' subspace E G , or the 'right' weights for each group convolution filter. Identifying the true underlying structure of data symmetries is a much harder problem that is yet to be tackled in both cases.

G Further theoretical insights

The following result provides consistency guarantees when the regularization parameters τ and β are small, and is a slight extension of a result in [38].

fase0_e

fase1_s

fase1_e

fase2_s

fase2_e

Heuristic Step

10

7

10

5

10

3

10

1

Relative Measure Distance (RMD)

Distance to Original Subspace when applying the Heuristic

RMD Comparison

V vs. P_0(V)

V vs. P_E(V)

Proposition 19 ( Γ -convergence , as in ) . Let Z = R D . If R is W p -continuous, ν is a Gibbs measure of potential U , and both U and r satisfy assumption 2 (or alternatively, are equal to 0), then R τ, β ν Γ -converges to R when τ, β ↓ 0 . Particularly, given µ ∗ , τ, β, ν the minimizer of R τ, β ν , we have

In particular, every cluster point of ( µ ∗ , τ, β, ν ) τ,β is a minimizer of R .

Proof of Proposition 19. We follow the exact same proof structure as [38], employing essentially their same techniques. However, we do adapt it to the case of taking the simultaneous limit of τ, β → 0 , and so we do include it for completeness.

Let ( τ n ) n ∈ N and ( β n ) n ∈ N be two positive sequences decreasing to 0 . On the one hand, since R is continuous (weakly if p = 0 or in W p for other p ≥ 1 ) and H ν ( µ ) = D ( µ || ν ) ≥ 0 , for all µ n → µ (in the appropiate sense), we have

On the other hand, given µ ∈ P p ( Z ) , consider ρ to be the heat kernel in Z = R D and ρ n ( x ) := β -n D ν ( x/β n ) . In particular, from [1] (as the heat kernel has finite p -th moments) we know that µ n := µ ∗ ρ n - - - - → n →∞ µ in W p (or weakly if it is the case).

Now, since the function h ( x ) := x log( x ) is convex, from Jensen's inequality we get that

Besides, we have (denoting here g ( x ) = e -U ( x ) ):

The last inequality is due to the quadratic growth of U ; and by the same argument on r :

Notice that whenever U ≡ 0 or r ≡ 0 , despite them not satisfying assumption 2, we still get the same inequalities (since the leftmost term would be 0).

Now, as R is W p -continuous, R ( µ n ) - - - - → n →∞ R ( µ ) , and:

And, as lim n →∞ β n log( √ 2 β n ) = 0 and the rest of the terms are bounded, we conclude that:

In particular, denoting by µ τ,β,ν ∗ the unique minimizer of R τ,β ν , then from the previous expressions we get ∀ n ∈ N and ∀ µ ∈ P p ( Z ) :

So that,

Finally, we provide, for completeness, a proof of Lemma 2 :

Proof of Lemma 2 (based on [27] and [40]). As H ≤ G is a compact group and π is H -invariant, from proposition 10 we know that f ∗ = E π [ Y | X = · ] lives in f ∗ ∈ L 2 H ( X , Y ; π | X ) . Consider any f ∈ L 2 ( X , Y ; π | X ) , by Lemma 1 from [27] (which applies since π | X is G -invariant), we can decompose f as f = f G + f ⊥ G , where f G = Q G f is its symmetric part and f ⊥ G = f -Q G f its antisymmetric part. A standard calculation of the population risk under the quadratic loss setting (see the proof of Corollary 1 for further insight) gives: R ( f ) = R ∗ + ∥ f -f ∗ ∥ 2 L 2 ( X , Y ; π X ) , and so:

Where we used that ⟨ f G , f ⊥ G ⟩ L 2 ( X , Y ; π X ) = 0 . The first term on the right hand side, -2 ⟨ f ∗ , f ⊥ G ⟩ L 2 ( X , Y ; π X ) , is what [40] call the mismatch between the real underlying model (which is only H -equivariant) and the symmetrized version of our model (which is made entirely G -equivariant).

Now, when π is G -equivariant, by proposition 10, Q G f ∗ = f ∗ , and so: -2 ⟨ f ∗ , f ⊥ G ⟩ L 2 ( X , Y ; π X ) = 0 , giving us the desired result:

Lemma 2 essentially says that if we try to symmetrize a model with respect to a group that has 'more symmetries' than what are actually observable in our data (i.e. π in itself is only H -invariant, but we symmetrize with respect to G ≥ H ); we can either win or lose generalization power according to the interplay between the two presented terms. In particular, if π is G equivariant, there's a strict generalization benefit from choosing a symmetric model to tackle our learning problem (which gives the name to the paper [27]). In particular, whenever f ⊥ G is non-zero (on a strictly positive π | X -measure set) there's a strict gain in generalization power from using the symmetrized version of the model.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Within the limited space we could consecrate to those sections, both the abstract and the introduction give quite accurate descriptions of the context of the work (Mean-Field approach to overparametrized NN), the questions and problems addressed (role of symmetries and impact of symmetry-leveraging techniques, as seen from the MeanField limit), our setting (the type of neural network models considered), the specific results obtained (as precisely as it was possible in words, namely the precise description of the Wasserstein gradient flows describing the limiting training dynamics in different settings), the novelty of the paper (what can be learned from this viewpoint), and the contents of our numerical experiments. The main results are clearly understandable from the abstract, and the introduction gives a rather general description of the paper's goals, since it also leaves room for the context description and the outline of the paper. See Section 1 for details.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We acknowledge and discuss many of the key assumptions required for both our theoretical and experimental results. The need for these assumption limits the applicability and practical impact of our work. We explicitly provide relevant insights in this regard, in a dedicated section, SuppMat-A.2. Also, and without including every single instance, we also discuss our limitations in: assumption 1, after presenting equation (1), section 4, and in many remarks in the SuppMat (e.g. in E.2.5,E.3 or F).

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All mathematical statements presented are given complete, rigorous mathematical proofs in the Supplementary Material (mainly in SuppMat-E). We also provide complete sets of assumptions for each result, either in the statement itself or (for some results where the precise assumptions are very technical and lengthy to state) previous to the corresponding proof in the Supplementary Material. Even proofs of some elementary facts or statements are provided for completeness. Furthermore, nearly every mathematical result stated in the main paper is given a comment or discussion regarding its context, its interpretation (beyond the pure mathematical value) and/or an idea of what tools or arguments its proof will rely upon.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Together with our submission, we include a ZIP file with the necessary code (and instructions, following the guidelines provided by the conference) for replicating all of our experimental results. We also provide the necessary details to understand and reproduce our work both in Section 4 and (more extensively) in SuppMat-F.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often

one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: As mentioned in our previous answer, we include a ZIP file with the necessary code and instructions to faithfully reproduce the main experimental results. We closely follow the NeurIPS guidelines for our code submission, providing the exact commands and environments required to run our experiments. Also both Section 4 and SuppMat-F provide a thorough landscape of the necessary assumptions and parameters to consider when replicating our results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: In Section 4 we provide the necessary level of detail for understanding our main experimental results. We provide a thorough and detailed description of the exact parameters (and how they were chosen) employed in our different experiments in SuppMat-F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The numerical experiments considered in this work, involved the application of a noisy SGD dynamic (as in equation (1), or equation (5)) that relied on various levels of randomness: the initial parameter configurations, the simulated training data and the gaussian noise in SGD iterations. These factors are taken into account (see SuppMat-F) and thus the different experiments were repeated a total of 10 times (under different random seeds), which was largely enough to include interpretable error bars in our plots, and to correctly illustrate the true trends of the underlying phenomena (see, again, SuppMat-F).

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In SuppMat-F we indicate both the computer resources employed to run the different experiments (see the first paragraph) as well as an estimate of the total computer time for the most complex simulation considered (namely, the N = 5000 case). We did not make related experiments prior to the preparation of this submission.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This paper rigorously conforms to the NeurIPS Code of Ethics. By its mathematical/theoretical nature, this paper did not require the interaction with humans (besides the authors), and did not expose any living being to any type of harm. Also, no direct or indirect social impact is expected from this work in the short-to-mid term. Since our numerical results come from simulations, there are no major concerns to be had related to the nature of the employed dataset (in terms of privacy, consent, fair use or representative evaluation). Also, Intellectual property in all its forms was thoroughly respected. Finally, all the elements required for reproducibility of our results are disclosed in the material joint to the submission, including software employed and the data-simulation mechanisms.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [No]

Justification: The research object of this paper is mathematical, as it studies abstract mathematical models of artificial neural networks, and their properties and behaviors in specific, idealized situations and contexts. Therefore, no direct or indirect social impact is expected in the short-to-mid term. It is, however, expected that the theoretical results presented could have a positive impact in the long term, at the levels of development or practical use of safer, more transparent or interpretable machine learning algorithms; or at least in enlarging the corpus of conceptual tools available to researchers and practitioners to better understand the advantages, drawbacks, and potential risks or impacts of their professional activities. Though it might be useful eventually, our work is mostly theoretical and its impact remains limited to academia and pure ML practice. By all these reasons, no potential, positive nor negative, societal impacts of the work performed are addressed in the paper.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate

to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: As previously mentioned, this work is mainly theoretical, and thus poses no major threats that could require the implementation of such safeguards.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We employ mostly public ML libraries and packages (e.g. objax , pyot , among others). We do, however, also involve the use of the emlp package provided by [29]. These assets are all properly credited in our work (see SuppMat-F), and the corresponding use licenses (particularly for the use of emlp ) are properly mentioned and respected.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Beyond the provided code to reproduce the experiments presented in the paper (which is accompanied with the details and instructions needed to execute it), no assets are introduced in this publication, hence, no specific documentation needs to be provided in that regard.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the Mean-Field limit of generalized shallow neural networks after learning with Wasserstein gradient flow under data augmentation, feature averaging or equivariant architectures as well as standard training. The provided results provide insights into learning with symmetries by covering two notions of invariant laws on the parameter space, varying assumptions on the data distribution and preserved symmetry over the course of training.\n\nSTRENGTHS:\nThe paper provides several deep and quite general theoretical results on learning with symmetries in the mean-field limit. Although the results are quite technical, the authors manage to present them in good clarity and embedded in the related literature of mean field limits of shallow models. Considering the examples of Data Augmentation, Feature Averaging and Equivariant Architectures seems natural and provides several interesting insights, for example the equivalence of DA, FA and even free training under symmetric data.\n\nWEAKNESSES:\nThe paper already seems quite complete and well written and I could not identify significant weaknesses. The following weaknesses are only minor points:\n\n- The description of Figure 1 is hard to parse. Markers like (a)-(d) could reduce ambiguity of what is meant by a column.\n- The markers in Figure 2 should be larger and it is unclear to me which additional insight to draw from the second row.\n\nAs I am not an expert in the mean-field and Wasserstein gradient flow literature, I did not verify the correctness of all proofs, and did not double-check the claimed novelty.\n\n**Typos:** 22: lastly, 65 and 66: b=N?, 75: Definition name in bold for consistency?, 120: […], 275: will, 329: a heuristic\n\nQUESTIONS:\n- How close is your setting to practical neural networks trained with SGD? In future work, are there hopes to extend your mean-field theory results to deep networks or to SGD with large learning rates and edge of stability dynamics? Similarly is the entropy regularization merely theoretically pleasing or also practically feasible?\n- Can GNNs or Transformers also be encoded in your class of equivariant architectures? If so, these example would significantly broaden the target audience.\n- In future work, in which ways could Assumption 1 be relaxed?\n\nLIMITATIONS:\nThe authors adequately acknowledge the main limitations of the paper. These limitations lie beyond the scope of this paper, but pose interesting directions for future work:\n\n- The provided results are only asymptotic and only consider gradient flow on shallow models as opposed to edge of stability dynamics in finite deep models.\n- The practical relevance of the findings and proposed algorithm at finite width for practical architectures, datasets and optimization procedures could be empirically verified. In particular, does vanilla training remain SI throughout training beyond the simple student-teacher setting?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates the learning dynamics of neural networks with several symmetry-leveraging techniques from a mean-field perspective. The main result indicates that optimizations with data augmentation and feature averaging (and the corresponding Wasserstein gradient flows) are equivalent in the mean-field limit under mild assumptions, while equivariant architecture may suffer from limited expressivity. Moreover, for equivariant data, freely-trained models also have the mean-field flow remaining in the set of strongly invariant distributions if initialized there, which is supported by the experiments. Finally, the authors propose a heuristic to discover equivariant architectures by gradually expanding the equivariant subspace.\n\nSTRENGTHS:\n* The math in the paper is rigorous and detailed. Complete proofs and abundant examples and remarks are provided. As more of a practitioner myself, I can understand the main ideas of the paper without much difficulty, though I did not check the proofs carefully.\n* The experiments are compact and well-designed, clearly supporting the theoretical results.\n\nWEAKNESSES:\n* This is not essentially a weakness, but the theoretical results don't bring many interesting insights to me. They mostly match my existing intuitions. E.g. for equivariant data, freely-trained models tend to have the same training dynamics as symmetry-aware models. I will leave it to other reviewers to comment on the theoretical significance of these results.\n\n* I suggest the authors provide more detailed explanations about some mathematical terminology, e.g. the pushforward of a probability measure; dt-a.e. standing for almost everywhere wrt dt. This could make the paper more accessible to the general audience.\n\nQUESTIONS:\n* Can you explain what is the connection between jointly equivariant and equivariant? The way I understand it is that we have some group actions on $\\mathcal Z$ induced from the definition of joint equivariance, i.e. $\\sigma_* (\\rho_g x, M_g z) = \\hat\\rho_g \\sigma_*(x,z)$. In other words, $M_g$ is determined by the group action on data and the model architecture. Then, if $M_g z = z$, then a NN model parametrized by $z$ is equivariant. Is this correct?\n* Proposition 5 states that equivariant architectures can lead to optimal solutions provided some universality conditions. How restrictive are these conditions?\n* Is the property of data equivariance (defined in L150) equivalent to the marginal distribution on $\\mathcal X$ being invariant and the function $f: \\mathcal X \\to \\mathcal Y$ being equivariant, given that the mapping from $\\mathcal X$ to $\\mathcal Y$ is deterministic? If so, would it be too restrictive to assume the invariance of $\\pi_\\mathcal X$?\n* In Section 4.1, the authors show that SI-initialized but freely-updated models remain in $\\mathcal E^G$ for large $N$s. This is indeed an interesting phenomenon. I wonder if this is related to the exact WI-initialized teachers. Can you provide, in addition to the arbitrary particles in Appendix D, the WI-initialized particles? Most likely, sampling from the pushed-forward empirical measure would result in a non-symmetric teacher, but there's a chance that the WI teacher might contain two pairs of symmetric parameters (e.g. $\\theta_1^*, g\\theta_1^*, \\theta_2^*, g\\theta_2^*$) and another singled out, making it mostly equivariant. After all, if the teacher is not equivariant, then the data is not equivariant, and the assumptions in Theorem 5 do not hold.\n* Section 4.2 reminds me of another work on discovering symmetry [1], which initially specifies a large candidate group and discovers the subgroup symmetry by inspecting the weights of relaxed group convolution filters. Can you comment on the connection between your method and theirs?\n\n## References\n[1] Wang, Rui, et al. \"Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution.\" Forty-first International Conference on Machine Learning.\n\nLIMITATIONS:\nThe authors have discussed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the symmetric structure of the model and data structures with respect to the action of $G$ and studies the Wasserstein gradient flow (WGF) for learning overparameterized neural networks under the mean-field (MF) regime. In particular, the authors consider data augmentation (DA), feature averaging (FA), and equivariant architecture (EA) as symmetry-leveraging techniques. They show that for symmetric data, DA and FA follow the same dynamics by exploiting weakly and strongly invariant laws (WI and SI) on the parameter space. In particular, the following statements are given.\n\n- DA and FA are equivalent in terms of optimal values and also equivalent to standard training (free training) for WI measures (Theorem 2).\n- Furthermore, this equivalence is shown over the space of general measures when the data distribution is equivalent (Corollary 2).\n- WGF for an invariant functional preserves weak invariance of the measure (Theorem 3 and Corollary 3). Moreover, there is an equivalence of WGF for DA, FA, and free training under suitable conditions.\n\nSTRENGTHS:\nThe paper proposes the notion of WI and SI and proves several equivalences between DA, FA, and free training. These results indicate when these techniques (DA and FA) are meaningful.\n\nWEAKNESSES:\n- This work mainly focuses on two-layer neural networks; hence, it is unclear if the results can be applied to deep neural networks.\n- The results, especially for WGF, seem to fully rely on the mean-field limit $N\\to\\infty$. Hence, it’s unclear whether the theory has some implications for finite-neuron settings.\n- (Minor) The paper is somewhat hard to follow because of the numerous notations. I suggest including a table that summarizes these notations.\n\nQUESTIONS:\n- In Theorem 4, the noisy SGD (4) involves the projection step of the noise. I think such a projection is non-trivial in general. Are there any practical examples?\n- Is WGF essential for the results in Section 3.3? Do similar results hold for other dynamics (e.g., underdamped version)?\n\nLIMITATIONS:\nLimitations are addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a mean-field analysis on a class of overparameterized neural networks that are expressed as an ensemble of $N$ units, trained with SGD under symmetries in data distribution and possible use of symmetry-leveraging techniques (data augmentation, feature averaging, and equivariant architectures). While the analysis is based on prior work on mean-field analysis on the same class of neural networks e.g. based on Wasserstein gradient flows [1], the contribution of this work is in extending to learning under symmetries. The setup is as follows:\n\n- In considered neural networks, a unit computes $\\sigma_*:(\\mathbf{x},\\theta_i)\\mapsto\\hat{\\mathbf{y}}$ and the network computes $\\Phi_{\\theta}^N:(\\mathbf{x},\\theta)\\mapsto \\frac{1}{N}\\sum_{i=1}^N\\sigma_*(\\mathbf{x},\\theta_i)$.\n- Overparameterization means the width $N\\to\\infty$, yielding a limiting ensemble $\\Phi_\\mu$ where $\\mu$ is understood a probability measure on the space of parameters $\\theta_i$.\n- The parameters $\\theta$ are optimized using noisy SGD.\n\nThe mean-field theory of the neural networks considered above [1] aims to analyze the training dynamics of the limiting ensemble $\\Phi_\\mu$, given that its optimization using a convex loss function becomes a convex optimization problem (unlike $\\Phi_{\\theta}^N$ which induces non-convex optimization) and that its global optimum may be approximated by training $\\Phi_{\\theta}^N$ with a large width $N$. Specifically, it is known that SGD on $\\Phi_{\\theta}^N$ approximates Wasserstein gradient flow on $\\Phi_\\mu$ in the scaling limit $N\\to\\infty$ under certain assumptions, which then, under certain assumptions, converges to the global optima of the convex optimization problem.\n\nGiven the background, the paper presents a mean-field analysis given that the data generating distribution is possibly under a symmetry described by a compact group $G$, and symmetry-leveraging techniques are possibly used. The authors consider data augmentation, feature averaging (i.e. symmetrization; also called the Reynolds operator), and equivariant architectures, in accordance to literature [2]. The idea for the analysis is to distinguish between two types of symmetries for the measure $\\mu$ on parameters: weakly-invariant and strongly-invariant, where weakly-invariant refers to an invariant measure on the subspace containing possibly non-invariant elements, and strongly-invariant refers to a measure on invariant subspace. The authors first show in Section 3.1 that, given that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant, $\\Phi_\\mu$ is equivariant if and only if $\\mu$ is weakly-invariant. Then, feature averaging (and data augmentation) can be understood as being related to weakly-invariant distributions obtained through symmetrization, and equivariant architecture is related to strongly-invariant distributions.\n\nThe authors then, in Section 3.2, begin mean-field analysis on optimality under a set of assumptions including that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant (Assumption 1), and show that:\n\n- (Proposition 3) Optimizing with respect to an invariant loss function yields a weakly-invariant optimum if it has a unique minimizer,\n- (Corollary 1) If input data distribution is invariant, optimizing under data augmentation and feature averaging are equivalent and corresponds to approximating the symmetrized version of mean of the labeling function,\n- (Corollary 2) If data distribution is equivariant, optimizing under data augmentation and feature averaging provide no advantage,\n- (Proposition 4 and 5) But even when data distribution is equivariant, equivariant architectures (i.e. strongly invariant measures $\\mu$) may not achieve the optimum with respect to a given invariant loss function, unless universality properties that are \"good enough\" for the given data distribution are assumed.\n\nThen, in Section 3.3, the authors provide mean-field analysis on training dynamics, and show that:\n\n- (Theorem 3, Corollary 3, and Theorem 4) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\\mu_{0:t}$ following the Wasserstein gradient flows yields a unique weakly-invariant solution $\\mu_t$. If the data distribution is equivariant, from Corollary 2 this result applies to freely-trained neural networks without symmetry leveraging techniques, and data augmentation and feature augmentation again provide no advantages.\n- (Theorem 5 and 6) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\\mu_{0:t}$ starting at a strongly-invariant initial condition $\\mu_0$ and following the Wasserstein gradient flows yields strongly-invariant solutions $\\mu_t$. Importantly, this means that if the data distribution is equivariant, $\\mu_t$ stays strongly-invariant throughout training even though it is not explicitly forced or constrained to be invariant. Furthermore, equivariance of data distribution can be dropped if one of the symmetry leveraging techniques are used, which all yield coinciding solutions.\n\nIn Section 4, the authors provide synthetic experiments focusing on validating Theorem 3-6, and then propose and demonstrate an empirical method for finding good equivariant architecture for the task at hand (i.e. the most expressive among invariant subspace) by training according to Theorem 5 and 6 until convergence.\n\n[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.\n\nSTRENGTHS:\n- S1. The paper targets an important and challenging problem of understanding training dynamics of overparameterized neural networks under symmetries and their leveraging. The technical approach that applies mean-field analysis of generalized shallow neural networks [1] to learning under symmetries is original as far as I am aware.\n- S2. The paper considers three types of symmetry-leveraging techniques: data augmentation, feature averaging, and equivariant architectures. These three are sufficiently general to represent currently used techniques [2]. The descriptions as well as strengths and weaknesses of each technique discussed in Section 2.3 are correct.\n- S3. The authors' approach of imposing distributional symmetries on the measure $\\mu$ to describe equivariant architectures as strongly-invariant (and feature averaging as weakly-invariant given that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant) is interesting and could be potentially useful in future work.\n- S4. The results presented in Section 3 and 4 are interesting, in particular Theorems 5 and 6 which shows that overparameterization and equivariant units, combined with equivariant data distribution or symmetry leveraging, leads to the property that strongly-invariant parameter initialization stays strongly-invariant throughout the training even without explicit constraints to do so. (But also see W1.)\n\n[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.\n\nWEAKNESSES:\n- W1. I have concerns regarding Assumption 1, which is a key assumption underlying Section 3 and assumes that the unit $\\sigma_*:(\\mathbf{x},\\theta_i)\\mapsto\\hat{\\mathbf{y}}$ is jointly equivariant with respect to the input $\\mathbf{x}$, the parameters $\\theta_i$, and the output $\\hat{\\mathbf{y}}$. This assumption states that we have already (partially) baked in the symmetry constraint to our model class. This may be natural for analyzing equivariant architectures, but does not coincide with the practical uses of data augmentation and feature averaging (e.g. please see [2]), as they are generally applied to an arbitrary, non-equivariant function, while in the setup of this work they are applied on already equivariant units. This limits the usefulness of the results regarding data augmentation and feature averaging. Perhaps the setup of this work could be relevant to symmetry breaking in already equivariant architectures [3], but I am not entirely sure.\n- W2. In Line 92-93, the authors claim that \"Note that requiring an infinite and i.i.d. sample from $\\pi$ is key when letting later $N\\to\\infty$. This is not truly a restriction, since one can always sample from the empirical distribution of finite data points\". This seems problematic since (1) if we can infinitely sample from $\\pi$, then we have access to it, which violates basic assumptions in machine learning (Line 82), and (2) if we infinitely sample from an empirical distribution supported on finite set of data points, the resulting distribution would not be $\\pi$, and consequently the SGD scheme in Equation (1) would not be optimizing the objective in Equation (2). (Please correct me if I am wrong.)\n- W3. The writing of the paper could be in general improved. For example, the paper uses an excessive number of abbreviations (MF, MFL, NN, SL, DA, FA, EA, WI, SI, WGF, RMD, V), which is hurting readability. It would be also helpful to provide an example of a measure and its (non-equivalent) invariant symmetrization and projection in Definition 5; these are important concepts that are used in all later sections, but I was not able to directly grasp them from reading the text. For example, considering a uniform distribution on $[0, 1]\\subset\\mathbb{R}$, and a multiplicative group {1, -1} acting on $\\mathbb{R}$ by multiplication, may be sufficiently informative.\n- W4. The architecture discovery algorithm proposed in Section 4.2 was a bit questionable to me, since the architecture of the ensemble is already specified (at least up to a significant portion) by the jointly equivariant unit(s) $\\sigma_*(\\cdot,\\cdot)$, and the algorithm searches for their parameters within a strongly-invariant regime. I am not sure why this algorithm can be understood as doing architecture discovery, instead of doing the usual parameter optimization.\n- W5. Currently, it is not immediate how the study of the ensemble considered in the paper would benefit the development of practical neural networks that leverage symmetry. Showing how some of the currently used symmetry-leveraging deep neural networks can be viewed as such ensembles could be informative in this directions.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.\n\n[3] Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, and Tess E. Smidt, Discovering symmetry breaking in physical systems with\nrelaxed group convolution, arXiv preprint arXiv:2310.02299, 2023.\n\nQUESTIONS:\nPlease see the Weaknesses section.\n\nLIMITATIONS:\nThe authors have partially addressed limitations of their work, for example in Section 5. Since the discussions are currently scattered across the main text, I encourage the authors to gather them in a dedicated limitations section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}