Symmetries in Overparametrized Neural Networks: A Mean Field View

Authors: javier  maass martínez
joaquín  fontbona
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:04:22

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.

Main paper: 
 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 WGFs 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 WGFs, 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. 
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 whereX = 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 }:EQUATIONHere, 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):EQUATIONAccordingly, 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 WGFs and on the MF theory of shallow NN 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 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 :EQUATIONSeveral 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:EQUATIONknown 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 → ∞.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 [27] ), and trying to minimize [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.L 2 (X , Y; π X ) defined as (Q G .f )(x) := G ρg -1 .f (ρ g .x)dλ G (g) (seeR F A (θ) := E π ℓ (Q G .Φ N θ )(X), Y (seeEquivariant 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, EAs 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 EAs 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 
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 hasE 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:Weakly-Invariant (WI) if µ = µ G and Strongly-Invariant (SI) if µ = µ E G .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: (Q G Φ µ ) = Φ µ G .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 versionQ 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 isan equivariant function without any of its 'parameters' living in E G .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 F A (µ) := 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, ŷ) = 1 2 ||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 F A and R EA are invariant (and convex) and we have:R DA (µ) = R G (µ), R F A (µ) = R(µ G ) and R EA (µ) = R(µ E G ).In particular,R = R DA if R is invariant. Moreover, ∀µ ∈ P G (Z), R(µ) = R DA (µ) = R F A (µ). 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:inf µ∈P(Z) R DA (µ) = inf µ∈P G (Z) R DA (µ) = inf µ∈P G (Z) R(µ) = inf µ∈P G (Z) R F A (µ) = inf µ∈P(Z) R F A (µ).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:inf µ∈P G (Z) R(µ) = R * + inf µ∈P G (Z) ∥Φ µ -f * ∥ 2 L 2 (X ,Y;π X ) = R * + inf µ∈P G (Z) ∥Φ µ -Q G .f * ∥ 2 L 2 (X ,Y;π X ) .wheref * = 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 F A (µ).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 EAs 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 [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.π ∈ 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.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 F A ) 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 Gis 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:EQUATIONNote 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 F A 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 F A ), WGF(R DA ) and WGF(R EA ) solutions coincide. If R is invariant, WGF(R) solution coincides with them too. 
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: 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.X = Y = R 2 , Z = R 2×2 ∼ = RWe 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 Ne , 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 (ν N Ne , (ν N Ne ) E G ) or RMD 2 (ν NNe , (ν N Ne ) G ) depending on initialization (either SI or WI); and the second shows the RMD between DA, FA and vanilla schemes. 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 Ne 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 Ne towards becoming WI, and a clear coincidence between the DA, FA and vanilla schemes (the latter only for equivariant f * ).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: 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 Ne , P Ej #ν N Ne ) ≤ δ 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 Ej , with v Ej := 1 N N i=1 (θ Ne i -P Ej .θ Ne 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. 
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 WGFs. 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 EAs 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. 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. 
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. 
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. 
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 
 
 Appendix: 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.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: X L.f dπ X = L. X f dπ X . In particular, this also holds for bounded linear operators. 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:T : G×Z → Z (g, z) → T (g, z)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)E π ℓ(Φ N θ (X), Y ) for θ ∈ Z N or E π ℓ(Φ µ (X), Y ) for µ ∈ P(Z) F σ * (P(Z))Space of shallow models with unit σ * and measures in P(Z) P p (Z)Probability measures µ ∈ P(Z) with finite p-th moment: Z ∥θ∥ p µ(dθ) < +∞.W p p-th Wasserstein distance, defined as W p (µ, ν) := inf γ∈Π(µ,ν) E γ [∥X -Y ∥ p ] 1 p , ∀µ, ν ∈ P p (Z), where Π(µ, ν) is the set of couplings between µ and ν.Linear Functional Derivative of F :P 2 (Z) → R as in Definition 3 D µ F Intrinsic Derivative of F : P 2 (Z) → R, as in Definition 3 WGF(F ) Wasserstein Gradient Flow for F : P 2 (Z) → R, as in Definition 4 r Some penalization function r : Z → R + λ Lebesgue Measure on Z µ ≪ λ µ ∈ P(Z) is absolutely continuous w.r.t. λ. H λ Boltzmann entropy of µ; H λ (µ) := log( dµ dλ (z))dµ(z) if µ ≪ λ or +∞ otherwise. R τ, βPenalized population risk, given by Rτ, β (µ) = R(µ) + τ rdµ + βH λ (µ), with τ, β ≥ 0 the regularization parameters. L 2 (X , Y; π X ) Space of functions f : X → Y square-integrable (in Bochner sense) w.r.t π X s N kStep-size (or learning rate) used during training. Parametrized as s N k = ε N ς(kε N ), with a regular function ς : R + → R + and some ε N > 0; with k, N ∈ N.Compact group and its normalized Haar measure, which is the unique right and left invariant measure w.r.t. the group multiplication, see [43] . ρ, ρ, M Orthogonal representations of the action of G over X , Y and Z respectively. We also denote these actions asG ⟳ ρ X , G ⟳ ρ Y and G ⟳ M Z respectively. Q G Symmetrization operator over L 2 (X , Y; π X ); (Q G .f )(x) := G ρg -1 .f (ρ g .x)dλ G (g) for any f ∈ L 2 (X , Y; π X ). E G Space of fixed points for G ⟳ M Z (i.e. z ∈ Z s.t. ∀g ∈ G M g .z = z) P E G Orthogonal projection from Z onto E G ; P E G .z := G M g .z dλ G (g) for all z ∈ Z. f #µ Pushforward of measure µ ∈ P(Z) via f : Z → Z, given by f #µ(•) = µ(f -1 (•)) µ G , µ E GSymmetrized and projected versions of µ ∈ P(Z):µ G := G (M g #µ)dλ G and µ E G := P E G #µ. P G (⋆) Space of G-invariant measures over space ⋆. e.g. for G ⟳ M Z, the space of all µ ∈ P(Z) s.t. M g #µ = µ for all g ∈ G. P(E G ) Space of measures concentrated on E G , i.e. µ ∈ P(Z) s.t. µ(E G ) = 1. R DA , R F A , R EAPopulation risks that are optimized when applying different SL techniques; DA, FA and EA respectively.L 2 G (X , Y; π X ) Space of functions f : X → Y square-integrable (in Bochner sense) w.r.t π X that are also G-equivariant. In other words, shorthand for Q G (L 2 (X , Y; π| X )) dt-a.e.Almost everywhere w.r.t. the Lebesgue measuredt on R + λ E G Shorthand for the Lebesgue Measure over E G . R τ, β E G Projected-regularized functional given by R τ, β E G (µ) := R(µ)+τ rdµ+βH λ E G (µ E G ) RMD Relative-Measure-Distance; RMD 2 (µ, ν) = W 2 2 (µ,ν) M 2 µ +M 2 ν where M 2 µ = 2 Z ∥z∥ 2 dµ(z). D P(Z) ([0, T ]) Skorokhod space of càdlàg functions from [0, T ] to P(Z).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 EAs 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: We also 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 f dµ. This notation is heavily used (and abused) throughout the core 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 SuppMat-C.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.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.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 ⟳ χ Yvia 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:∀g ∈ G, ∀x ∈ X , ∀z ∈ Z, D x f (χ g .x, χg .z) = χg .D x f (x, z)χ -1 gProof of Proposition 6. Indeed, we know that ∀z ∈ Z D x f (•, z) is the unique function that satisfies, ∀x ∈ X :lim h→0 ∥f (x + h, z) -f (x) -D x f (x, z)h∥ ∥h∥ = 0Since we want to prove that ∀x ∈ X , ∀z ∈ Z, ∀g ∈ G :D x f (χ g .x, χg .z) = χg D x f (x, z)χ -1g , it will be enough to check that:lim h→0 ∥f (χ g .x + h, χg z) -f (χ g .x, χg .z) -χg D x f (x, z)χ -1 g h∥ ∥h∥ = 0which by uniqueness implies the result. Thanks to the joint equivariance of f , we have ∀h ̸ = 0:∥f (χ g .x + h, χg z) -f (χ g .x, χg .z) -χg D x f (x, z)χ -1 g h∥ ∥h∥ = ∥f (χ g .(x + χ -1 g .h), χg z) -f (χ g .x, χg .z) -χg D x f (x, z)χ -1 g h∥ ∥h∥ = ∥ χg .f (x + χ -1 g .h, z) -χg .f (x, z) -χg D x f (x, z)χ -1 g h∥ ∥h∥ = ∥ χg . f (x + χ -1 g .h, z) -f (x, z) -D x f (x, z)(χ -1 g h) ∥ ∥χ g .χ -1 g .h∥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:∥ χg . f (x + h, z) -f (x, z) -D x f (x, z) h ∥ ∥χ g h∥ ≤ ∥ χg ∥∥f (x + h, z) -f (x, z) -D x f (x, z) h∥ ∥χ g . h∥ Multiplying by 1 = ∥χ -1 g χg h∥ ∥ h∥the last term is seen to be bounded by∥ χg ∥∥χ -1 g ∥ • ∥f (x + h, z) -f (x, z) -D x f (x, z) h∥ ∥ h∥Since χ g and χ -1 g are bounded operators, we have that:h → 0 ⇐⇒ h = χ -1 g h → 0 . Thus lim h→0 ∥f (χ g .x + h, χg z) -f (χ g .x, χg .z) -χg D x f (x, z)χ -1 g h∥ ∥h∥ ≤ lim h→0 ∥ χg ∥∥χ -1 g ∥ • ∥f (x + h, z) -f (x, z) -D x f (x, z) h∥ ∥ h∥ = 0This, 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: ∀g ∈ G, ∀x ∈ X , D x f (χ g .x) = χg D x f (x)χ T g Proof of Corollary 4. Direct.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 ⟳ χ Yvia 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:∀x ∈ X , ∀g ∈ G, χg ⟨f (x; •), µ⟩ = ⟨f (χ g x; •), χg #µ⟩Proof of Proposition 7. Let µ ∈ P(Z) and f be as stated. Notice that, ∀x ∈ X , ∀g ∈ G, we have:χg ⟨f (x, •), µ⟩ = χg Z f (x, θ)dµ(θ) = Z χg f (x, θ)dµ(θ),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:Z χg f (x, θ)dµ(θ) = Z f (χ g x, χg θ)dµ(θ) = Z f (χ g x, θ)d( χg #µ)( θ)We conclude the desired result.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: Z f (M g z)µ(dz) = E G f (M g z)µ(dz) = E G f (z)µ(dz) = Z f (z)µ(dz).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 :P E G M g z = G M h (M g z)dλ G (h) = G (M h M g z)dλ G (h) = G M hz)dλ G ( h) = P E G 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:(M g #µ) E G (A) = µ(M -1 g P -1 E G (A)) = µ(P -1 E G (A)) = µ E G (A),and:(µ G ) E G (A) = G µ(M -1 g P -1 E G (A))dλ G (g) = G µ(P -1 E G (A))dλ G (g) = µ(P -1 E G (A)) = µ E G (A).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:⟨f, (M g #µ) G ⟩ = G ⟨f, M h #(M g #µ)⟩dλ G (h) = G ⟨f, (M h)#µ⟩dλ G ( h) = ⟨f, µ G ⟩, 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:µ G (M -1 h (B)) = G µ(M -1 g (M -1 h (B)))dλ G (g) = G µ(M -1 g (B))dλ G (g) = µ G (B) 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 )).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 thatE π [∥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 measurablef * : 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 π:E π [(Q G .f * )(X)h(X)] = E π G ρ-1 g .f * (ρ g .X)dλ G (g)h(X) = G ρ-1 g .E π [f * (X)h(ρ g . -1 .X)]dλ G (g) = G E π [ρ -1 g .Y h(ρ g . -1 .X)]dλ G (g) = G E π [Y h(X)]dλ G (g) = E π [Y h(X)]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:∆(f, Q G f ) := E π [∥Y -f (X)∥ 2 Y ] -E π [∥Y -(Q G f )(X)∥ 2 Y ]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:∆(f, Q G f ) = -2⟨f * , f ⊥ G ⟩ L 2 (X ,Y;π X ) + ∥f ⊥ G ∥ 2 L 2 (X ,Y;π X ) where f * (x) = E π [Y |X = x]is the conditional expectation function, andf ⊥ 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 )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 EAs 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.Recall the finite-dimensional setting of single-hidden-layer neural networks. We consideredX = 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:EQUATIONwhere we write the expression by blocks, consideringW = (W 1 , . . . , W N ) ∈ R c×bN , A = (A 1 , . . . , A N ) ∈ R d×bN , and B =    B 1 . . . B N    ∈ R bN .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 EAs, 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 action6 of G on Z, which is given by:EQUATIONfor 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 EAs 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:σ * (ρ g .x, M g .z) = σ * (ρ g .x, (ρ g .W.η T g , ρ g .A.η T g , η g .B)) = (ρ g .W.η T g ).σ((ρ g .A.η T g ) T .(ρ g .x) + η g .B) = ρg .(W.η T g .σ(η g .A T .(ρ T g .ρ g ).x + η g .B)) = ρg .(W.η T g .σ(η g .(A T .x + B))) = ρg .(W.η T g .η g .σ(A T .x + B)) = ρg .(W.σ(A T .x + B)) = ρg .σ * (x, z)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 EAs. 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.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 = nc), 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×d1 and R n×d2 is if the parameters Ã ∈ R (n×d1)×(n×d2) and B ∈ R n×d2 (from the definition x → ÃT .x + B) are of the form:EQUATIONWhere α, β ∈ R d1×d2 , γ ∈ R d2 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:α ⊗ I =     α 0 . . . 0 0 α . . . 0 . . . . . . . . . 0 0 0 . . . α     ∈ R (n×d1)×(n×d2) , β ⊗ J =     β β . . . β β β . . . β . . . . . . . . . β β β . . . β     ∈ R (n×d1)×(n×d2) γ ⊗ (1, . . . , 1) = (γ, . . . , γ) ∈ R (n×d2)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 existsEQUATIONIn a sense, we can think of E G simply as being equivalent to R 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:W ∈ R (n×c)×(n× bN ) , A ∈ R (n× d)×(n× bN ), and B ∈ R (n× bN ) .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 nequivariant 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.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×d1) × R (n×n)×d2 and the output space is X 2 = R (n×d3) × R (n×n)×d4 . 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.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×d1) and X 2 = R (n×d2) under the natural action. It is well known (see [8] ) that a C n -equivariant affine layer, x → ÃT .x + B, between X 1 and X 2 , of parameters Ã ∈ R (n×d1)×(n×d2) and B ∈ R n×d2 , must be of the form: Ã = C(α 0 , . . . , α n-1 ) and B = β ⊗ (1, . . . , 1), with β ∈ R d2 , α i ∈ R d2×d1 , ∀i ∈ {0, . . . , n -1}, and the associated circulant matrix being defined (by blocks, and considering the indices modulo n) as:C(α 0 , . . . , α n-1 ) =       α 0 α 1 ... ... α -1 α -1 α 0 α 1 ... ... . . . . . . . . . . . . . . . ... ... α -1 α 0 α 1 α 1 ... ... α -1 α 0       ∈ R (n×d2)×(n×d1) .As a consequence, we get an analog of Equation ( 9) to explicitly describeE 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: W = C(w 0 , . . . , w n-1 ), A = C(a 0 , . . . , a n-1 ), and B = β ⊗ (1, . . . , 1)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 Cn contains E Sn (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 σ * .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.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 ∂µ :EQUATIONZ ∂F ∂µ (µ, z)dµ(z) = 0The 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:1. 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] ):∂F ∂µ (µ, z) = log dµ dλ (z) + 1 and D µ F (µ, z) = 1 dµ dλ (z) ∇ z dµ dλ (z).2. Whenever F (µ) := Z ϕ(z)dµ(z) for some bounded continuously differentiable function ϕ : Z → R, it is well known that :∂F ∂µ (µ, z) = ϕ(z) -ϕdµ and D µ F (µ, z) = ∇ z ϕ(z)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 considerH = 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):∂F ∂µ (µ, z) = D h L(⟨Φ, µ⟩)(Φ(z)) = ⟨∇ h L(⟨Φ, µ⟩), Φ(z)⟩ H -C F,µ D µ F (µ, z) = (D h L(⟨Φ, µ⟩)(D z Φ(z))) * = ∇ z Φ(z)(∇ h L(⟨Φ, µ⟩)).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]:F ((1 -h)µ + hν) -F (µ) h = L(⟨Φ, (1 -h)µ + hν⟩) -L(⟨Φ, µ⟩) h = L(⟨Φ, µ⟩ + h⟨Φ, ν -µ⟩) -L(⟨Φ, µ⟩) h .Let's denote by q µ := ⟨Φ, µ⟩ (analogously q ν-µ := ⟨Φ, ν -µ⟩) and s µ,ν := hq ν-µ , so we can write:F ((1 -h)µ + hν) -F (µ) h = L(q µ + s µ,ν ) -L(q µ ) h .As L is Gateaux differentiable, we have the following first order Taylor expansion ∀x, s ∈ H, ∀t ∈ R:L(x + t s) = L(x) + t D h L(x).s + o(|t|∥s∥),which allows us to write:F ((1 -h)µ + hν) -F (µ) h = L(q µ + h q ν-µ ) -L(q µ ) h = h.D h L(q µ ).q ν-µ + o(|h|∥q ν-µ ∥) h .As ∥q ν-µ ∥ < ∞ by hypothesis, we can say that: o(|h|∥q ν-µ ∥) = o(h). Therefore, taking the limit with h → 0, we get that:lim h→0 F ((1 -h)µ + hν) -F (µ) h = D h L(q µ ).q ν-µ + lim h→0 o(h) h = D h L(q µ ).q ν-µ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:D h L(q µ ).q ν-µ = D h L(q µ ).⟨Φ, ν -µ⟩ = ⟨D h L(q µ )(Φ), ν -µ⟩,and so by definition of the gradient of L:lim h→0 F ((1 -h)µ + hν) -F (µ) h = Z ⟨∇ h L(⟨Φ, µ⟩), Φ(z)⟩ H d(ν -µ)(z).From here we deduce that:∂F ∂µ (µ, z) = ⟨∇ h L(⟨Φ, µ⟩), Φ(z)⟩ H -C F,µ ,where C F,µ is a fixed constant, given by:C F,µ = Z ∂F ∂µ (µ, z)dµ(z) = Z ⟨∇ h L(⟨Φ, µ⟩), Φ(z)⟩ H d(µ)(z) = ⟨∇ h L(⟨Φ, µ⟩), ⟨Φ, µ⟩⟩ HOn 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:∀z ∈ Z, ∇ z ∂F ∂µ (µ, z) , z Z = D z ∂F ∂µ (µ, z) (z)However, we can develop the RHS, and as the constant C F,µ doesn't depend on z, we get that:D z ∂F ∂µ (µ, z) (z) = D z (⟨∇ h L(⟨Φ, µ⟩), Φ(z)⟩ H ) (z)Now, by the chain rule and the definition of the adjoint operator of D z Φ(z):D z ∂F ∂µ (µ, z) (z) = ⟨∇ h L(⟨Φ, µ⟩), D z Φ(z)(z)⟩ H = (D z Φ(z)) * (∇ h L(⟨Φ, µ⟩)) , z ZSo, as they coincide for every z ∈ Z, we conclude that:D µ F (µ, z) = (D z Φ(z)) * (∇ h L(⟨Φ, µ⟩))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:D h L(f )(h) = E π ⟨∇ 1 ℓ ((f (X), Y ) , h(X)⟩ Y• Φ : 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:∂R ∂µ (µ, z) = E π ⟨∇ 1 ℓ (⟨σ * (X; •), µ⟩, Y ) , σ * (X; z)⟩ Y + (constant not depending on z) D µ R(µ, z) = E π [∇ z σ * (X; z).∇ 1 ℓ(⟨σ * (X; •), µ⟩, Y )]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 :∂F ∂µ (M g #µ, M g .z) = ∂F ∂µ (µ, z) and D µ F (M g #µ, M g .z) = M g .D µ F (µ, z)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):lim h→0 F ((1 -h)µ + hν) -F (µ) h = Z ∂F ∂µ (µ, z)d(ν -µ)(z) and Z ∂F ∂µ (µ, z)dµ(z) = 0.In particular, as F is G-invariant (and M g linear), we can write ∀µ, ν ∈ P(Z), ∀h ̸ = 0 and g ∈ G:F ((1 -h)µ + hν) -F (µ) h = F ((1 -h)(M g #µ) + h(M g #ν)) -F (M g #µ) h .Taking the limit as h → 0 on both sides, we get:lim h→0 F ((1 -h)µ + hν) -F (µ) h = Z ∂F ∂µ (M g #µ, z)d(M g #ν -M g #µ)(z) = Z ∂F ∂µ (M g #µ, M g .z)d(ν -µ)(z),and also:Z ∂F ∂µ (M g #µ, M g .z)dµ(z) = Z ∂F ∂µ (M g #µ, z)dM g #µ(z) = 0.So, by uniqueness, we get ∀g ∈ G:∂F ∂µ (µ, z) = ∂F ∂µ (M g #µ, M g .z),and, from proposition 6 (since ∂F ∂µ is jointly invariant), we get ∀g ∈ G:D µ F (M g #µ, M g .z) = ∇ z ∂F ∂µ (M g .µ, M g .z) = M g .∇ z ∂F ∂µ (µ, z) = M g .D µ F (µ, z)Finally, for convenience, let us denote by H E G the function defined on P(Z) by Section 3.3) . We can straightforwardly compute its linear functional derivative and its intrinsic derivative in that space:H E G (µ) := H λ E G (µ E G ) = H λ E G • P E G #(µ) (presented inExample. (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: lim h→0 H λ E G ((1 -h)µ E G + hν E G ) -H λ E G (µ E G ) h = E G ∂H λ E G ∂η (µ E G , x)d(ν E G -µ E G )(x) = E G log dµ E G dλ E G (x) + C d(ν E G -µ E G )(x) = Z log dµ E G dλ E G (P E G .z) + C d(ν -µ)(z), 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 µ H E G (µ, z) = 1 dµ E G dλ E G (z) P T E G ∇ z dµ E G dλ E G (P E G .z)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:D µ R τ,β ν (µ, z) = D µ R(µ, z) + τ ∇ z r(z) + β∇ z U (z) + β 1 µ(z) ∇ z µ(z) ,so that WGF(R τ,β ν ) as in definition 4 reads:∂ t µ t = ς(t) div D µ R τ,β ν (µ t , •) µ t = ς(t) [div ((D µ R(µ t , •) + τ ∇ z r + β∇ z U ) µ t ) + β∆µ t ] .We recover the expression for WGF(R τ,β ) in Equation ( 4) by considering U ≡ 0:EQUATIONWe 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 McKean-Vlasov SDE [11, 54, 70] (also known as the Mean Field Langevin Dynamics (MFLD) in the NN literature). In the case of R τ,β it reads:EQUATIONwhere (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.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:R(µ ′ ) -R(µ) ≥ Z ∂R ∂µ (µ, z)d(µ ′ -µ)(z)Proof of Lemma 3 (from [38] ). Define µ ϵ := (1 -ϵ)µ + ϵµ ′ . Since R is convex, we haveϵ (R(µ ′ ) -R(µ)) ≥ R(µ ϵ ) -R(µ) = ϵ 0 Z ∂R ∂µ (µ s , z)d(µ ′ -µ)(dz) ds. 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.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:• µ R is bounded and Lipschitz continuous, i.e. ∃C R > 0 s.t. ∀z, z ′ ∈ Z, ∀µ, µ ′ ∈ P 2 (Z), |D µ R(µ, z) -D µ R(µ ′ , z ′ )| ≤ C R [|z -z ′ | + W 2 (µ, µ ′ )] • ∀µ ∈ P(Z), D µ R(µ, •) ∈ C ∞ (Z). • ∇D µ R : P(Z) × Z → Z × Z is jointly continuous.2. (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 sup µ∈P2(Z) sup x∈Z |∇D µ R(µ, x)| ≤ M R mxfor 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:∀t > 0, d dt (R τ,β ν (µ t )) = -ς(t) Z D µ R(µ t , z) + τ ∇r(z) + β ∇u t u t (z) + β∇U (z) 2 dµ t (z)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:I(µ||ν) := Z ∇ log( dµ dν (z)) 2 dµ(z)Then, almost by definition, we get:d dt (R τ,β ν (µ t )) = -β 2 ς(t)I(µ t ||μ t )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:µ ∞ := arg min µ∈P(Z) R τ, β ν (µ) and lim t→∞ W 2 (µ t , µ ∞ ) = 0Remark. 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:D(ν||µ) := Z log( dν dµ (z))dν(z) ≤ 1 2ϑ Z ∇ log( dν dµ (z)) 2 dν(z) =: 1 2ϑ I(ν||µ)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 ν = λ):μ(z) ∝ exp - 1 β ∂R ∂µ (µ, z) - τ β r(z) -U (z)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, µ t0 has finite entropy and finite second moment; then ∀t ≥ t 0 ,D(µ t ||µ ∞ ) ≤ R τ, β ν (µ t ) -R τ, β ν (µ ∞ ) ≤ (R τ, β ν (µ t0 ) -R τ, β ν (µ ∞ ))e -2βϑ t t 0 ς(s)dswhere µ ∞ = µ τ, β, ν = 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.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µ.1. Let Z to be the closure of a convex open set within some finite-dimensional euclidean space.2. Let L : H → R + be differentiable, with a differential dL that is Lipschitz on bounded sets and bounded on sublevel sets.3. Let Φ : Z → H be differentiable and V : Z → R + be semiconvex (i.e. ∃λ ∈ R :V + λ| • | 2 is convex).4. 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∈Qr (∥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 bydZ t = b(t, Z t , µ t )dt + Σ(t, Z t , µ t )dB twhere 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. 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 ) sup t≥0 ∥b N (t, w1, µ1) -b N (t, w2, µ2)∥ + ∥Σ N (t, w 1 , µ 1 ) -Σ N (t, w 2 , µ 2 )∥ ≤ M 1 ∥w 1 -w 2 ∥ + W R D |⟨g(•, z), µ 1 ⟩ -⟨g(•, z), µ 2 ⟩| 2 dπ W (z) 1/2 , sup t≥0 {∥b N (t, 0, µ 0 )∥ + ∥Σ N (t, 0, µ 0 )∥} ≤ M 1 . 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 sup t≥0,w∈R D ,µ∈P2(R D ) {∥b N (t, w, µ) -b(t, w, µ)∥ + ∥Σ N (t, w, µ) -Σ(t, w, µ)∥} ≤ M 2 N -κ .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:dZ t = b(t, Z t , µ t )dt + Σ(t, Z t , µ t )dB tAdditionally, it satisfies for each T ≥ 0:sup t∈[0,T ] E[∥Z t ∥ 2 ] < ∞E.1 Proofs of Section 3.1Proof of Proposition 1. By definition of the symmetrization operator, we know that ∀x ∈ X :(Q G Φ µ )(x) = G ρg -1 Φ µ (ρ g x)dλ G (g)For g ∈ G, since σ * is equivariant and M g is invertible, we can write:Φ µ (ρ g x) = ⟨σ * (ρ g x, •), µ⟩ = ⟨σ * (ρ g x, •), M g #(M -1 g #µ)⟩ = ρg ⟨σ * (x, •), M -1g #µ⟩ where we've used Proposition 7 in the last equality. In turn, we can write (via the inversion-invariance of λ G ):(Q G Φ µ )(x) = G ρg -1 ρg ⟨σ * (x, •), M -1 g #µ⟩dλ G (g) = G ⟨σ * (x, •), M g -1 #µ⟩dλ G (g) = G ⟨σ * (x, •), M g #µ⟩dλ G (g) = ⟨σ * (x, •), µ G ⟩ = Φ µ G (x)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.):⟨σ * (x, •), µ G ⟩ = G Z σ * (x, M g .z)dµ(z)dλ G (g) = Z G σ * (x, M g .z)dλ G (g)dµ(z) = Z σ * (x, G M g .z dλ G (g))dµ(z) = Z σ * (x, P E G .z)dµ(z) = ⟨σ * (x, •), µ E G ⟩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.2Now 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:L ρg.x, ρg.y (M g #µ) = ℓ (⟨σ * (ρ g .x, •), M g #µ⟩, ρg .y) = ℓ (ρ g .⟨σ * (x, •), µ⟩, ρg .y) = ℓ (⟨σ * (x, •), µ⟩, y) = L x,y (µ)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:R F A (µ) = E π ℓ Q G (Φ µ )(X), Y = E π ℓ Φ µ G (X), Y = R(µ G )Next, using Lemma 4, Fubini's theorem and the inversion-invariance of λ G , we get:R DA (µ) = E π G ℓ Φ µ (ρ g .X), ρg .Y dλ G (g) = E π G L ρg.X, ρg.Y (µ)dλ G (g) = G E π L ρg.X, ρg.Y (µ) dλ G (g) = G E π L ρg.X, ρg.Y (M g #M -1 g #µ) dλ G (g) = G E π L X,Y (M g -1 #µ) dλ G (g) = G E π [L X,Y (M g #µ)] dλ G (g) = G R(M g #µ)dλ G (g) = R G (µ)From these expressions we can quickly verify that R DA , R F A and R EA are invariant. Namely, by Lemma 1, for g ∈ G and µ ∈ P(Z) we have:R F A (M g #µ) = R((M g #µ) G ) = R(µ G ) = R F A (µ), and: R EA (M g #µ) = R((M g #µ) E G ) = R(µ E G ) = R EA (µ).On the other hand, the rightinvariance of λ G implies:R DA (M g #µ) = G R((M h #(M g #µ))dλ G (h) = G R((M h#µ))dλ G ( h) = R DA (µ).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 F A (µ). 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:R(M g #µ) = E π [ℓ(⟨σ * (X; •), M g #µ⟩, Y )] = E π ℓ(ρ g ⟨σ * (ρ -1 g X; •), µ⟩, ρg ρ-1 g Y ) = E π ℓ(⟨σ * (ρ -1 g X; •), µ⟩, ρ-1 g Y ) = E π [ℓ(⟨σ * (X; •), µ⟩, Y )] = R(µ) 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):1. For V , notice that for g ∈ G:V (M g #µ) = Z r(θ)d(M g #µ)(θ) = Z r(M g θ)dµ(θ) = r(θ)dµ(θ) = V (µ)thanks to the invariance of r. i.e. V is invariant.2. For H ν , notice that, for g ∈ G, as ν is invariant, we know thatd(Mg#µ) dν (x) = dµ dν (M -1 g x).Therefore:H ν (M g #µ) = log d(M g #µ) dν (θ) d(M g #µ)(θ) = log dµ dν (M -1 g θ) d(M g #µ)(θ) = log dµ dν (M -1 g M g θ) dµ(θ) = H ν (µ)Which proves that H ν is invariant.We can readily conclude, since: R τ,β (µ) = R(µ) + τ rdµ + βH λ (µ), for all µ ∈ P(Z).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:F (μ) ≤ S F (µ s )dλ(s)Proof of Proposition 17. Since Lemma 3 holds, we have that ∀µ 1 , µ 2 ∈ P(Z):F (µ 1 ) ≥ F (µ 2 ) + ∂F ∂µ (µ 2 , z)d(µ 1 -µ 2 )(z)Let s ∈ S be arbitrary and consider µ 2 = μ := µ s dλ(s); and µ 1 = µ s. Then:F ( µ s dλ(s) ≤ F (µ s) - ∂F ∂µ µ s dλ(s), z d µ s -µ s dλ(s) (z)Integrating the inequality with respect to λ (on s):S F µ s dλ(s) dλ(s) ≤ △ := S F (µ s)dλ(s) - S Z ∂F ∂µ S µ s dλ(s), • d µ s - S µ s dλ(s) dλ(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:⋆ := S Z ∂F ∂µ S µ s dλ(s), • d µ s - S µ s dλ(s) dλ(s) = S Z ∂F ∂µ S µ s dλ(s), • dµ s - Z ∂F ∂µ S µ s dλ(s), • d S µ s dλ(s) dλ(s) = S Z ∂F ∂µ µ s dλ(s), • d(µ s) dλ(s) - S Z ∂F ∂µ S µ s dλ(s), • d S µ s dλ(s) dλ(s)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⋆ = S Z ∂F ∂µ S µ s dλ(s), z dµ s(z) dλ(s)But, by definition:∀f : Z → R integrable, ⟨f, S µ s dλ(s)⟩ = S ⟨f, µ s ⟩dλ(s) = S Z f (z)dµ s (z) dλ(s)So this is, by definition, and applying the same convention on the definition of the linear functional derivative 10 :⋆ = Z ∂F ∂µ µ s , z d µ sdλ(s) (z) = 0With this, we conclude that △ = S F (µ s )dλ(s), and so, we get that:F (μ) = F S µ s dλ(s) ≤ S F (µ s )dλ(s)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:inf µ∈P G (Z) F (µ) ≥ inf µ∈P(Z) F (µ)For the other inequality, take (µ n ) n∈N ⊆ P(Z) to be an infimizing sequence for F ; i.e. such thatF (µ n ) ≥ F (µ n+1 ) and F (µ n ) ----→ n→∞ inf µ∈P(Z) F (µ)). Such a sequence always exists. ByCorollary 7, we have ∀n ∈ N,F (µ G n ) ≤ F (µ n ); thus, ∀n ∈ N: inf µ∈P G (Z) F (µ) ≤ F (µ G n ) ≤ F (µ 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:inf µ∈P G (Z) F (µ) = inf µ∈P(Z) F (µ)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.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 F A 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 F A and R EA are convex and C 1 as well, with linear functional derivatives given by:∂R DA ∂µ (µ, z) = G ∂R ∂µ (M g #µ, M g .z)dλ G (g), ∂R F A ∂µ (µ, z) = G ∂R ∂µ (µ G , M g .z)dλ G (g) and ∂R EA ∂µ (µ, z) = ∂R ∂µ (µ E G , P E G .z)And intrinsic derivatives given by (when well defined):D µ R DA (µ, z) = G M T g .D µ R(M g #µ, M g .z)dλ G (g) D µ R F A (µ, z) = G M T g .D µ R(µ G , M g .z)dλ G (g) and D µ R EA (µ, z) = P T E G .D µ R(µ E G , P E G .z)In particular, from Proposition 17, we have that:∀µ ∈ P(Z), R F A (µ) ≤ 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:lim h→0 R DA ((1 -h)µ + hν) -R DA (µ) h = lim h→0 G R(M g #((1 -h)µ + hν))dλ G (g) -G R(M g #µ)dλ G (g) h = lim h→0 G R((1 -h)M g #µ + hM g #ν) -R(M g #µ) h dλ G (g) = G lim h→0 R((1 -h)M g #µ + hM g #ν) -R(M g #µ) h dλ G (g) = G Z ∂R ∂µ (M g #µ, z)d(M g #ν -M g #µ)(z)dλ G (g) = G Z ∂R ∂µ (M g #µ, M g .z)d(ν -µ)(z)dλ G (g) = Z G ∂R ∂µ (M g #µ, M g .z)dλ G (g)d(ν -µ)(z).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):Z G ∂R ∂µ (M g #µ, M g .z)dλ G (g)dµ(z) = G Z ∂R ∂µ (M g #µ, z)d(M g #µ)(z)dλ G (g) = 0.We can then identify:∂R DA ∂µ (µ, z) = G ∂R ∂µ (M g #µ, M g .z)dλ G (g),and, by taking the gradient:D µ R DA (µ, z) = G M T g .D µ R(M g #µ, M g .z)dλ G (g).We analogously calculate the expression for the l.f.d. of R F A ; let µ, ν ∈ P(Z):lim h→0 R F A ((1 -h)µ + hν) -R F A (µ) h = lim h→0 R((1 -h)µ G + hν G ) -R(µ G ) h = Z ∂R ∂µ (µ G , z)d(ν G -µ G )(z) = Z G ∂R ∂µ (µ G , M g .z)dλ G (g)d(ν -µ)(z),and also:Z G ∂R ∂µ (µ G , M g .z)dλ G (g)dµ(z) = G Z ∂R ∂µ (µ G , z)d(M g #µ)(z)dλ G (g) = Z ∂R ∂µ (µ G , z)dµ G (z) = 0.So, by the definition of the lineal functional derivative, we identify:∂R F A ∂µ (µ, z) = G ∂R ∂µ (µ G , M g .z)dλ G (g),and taking the gradient we get:D µ R F A (µ, z) = G M T g .D µ R(µ G , M g .z)dλ G (g)Lastly, the l.f.d. of R EA is calculated similarly, noticing that:lim h→0 R EA ((1 -h)µ + hν) -R EA (µ) h = lim h→0 R(P E G #((1 -h)µ + hν)) -R(P E G µ) h = Z ∂R ∂µ (µ E G , z)d(P E G #ν -P E G #µ)(z) = Z ∂R ∂µ (µ E G , P E G z)d(ν -µ)(z),and that:Z ∂R ∂µ (µ E G , P E G .z)dµ(z) = Z ∂R ∂µ (µ E G , z)d(µ P E G )(z) = 0.We can thus identify:∂R EA ∂µ (µ, z) = ∂R ∂µ (µ E G , P E G .z) and D µ R EA (µ, z) = P T E G .D µ R(µ E G , P E G .z)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 F A and R EA , from Proposition 18. Since Proposition 2 ensures that the latter are always invariant, Proposition 3 implies that R DA , R F A 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 F A 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 ) withHowever, if we assume that: ∀g ∈ G, ρ g A i = A i , then, in particular:-1 0 0 -1A 1 i A 2 i = A 1 i A 2 i .This in turn implies, asA 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:∀x ∈ R 2 , Φ N,E G θ (x) = 1 N N 1=1 W i σ(0 T • x + B i ) = 1 N N i=1 W i σ(B i ),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:inf θi∈E G i=1...N N ∈N R(Φ E G θ ) ≥ inf ν∈P(E G ) R(ν) ≥ inf C∈R E[|Y -C| 2 ] = inf C∈R E[|∥X∥ 2 -C| 2 ].The problem on the right has a known answer, which isC * = 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 functionf (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: ∥Φ N θ -f ∥ ∞,K = sup x∈K |Φ N θ (x) -f (x)| < ε < √ C * Then: E[|Y -Φ N θ (x)| 2 ] ≤ E[( sup x∈K |Φ N θ (x) -f (x)|) 2 ] < E[C * ] = C *But, in particular, ∃ν N θ ∈ P(Z) such that:E[|Y -Φ N θ (x)| 2 ] < C * and so: inf µ∈P(Z) R(µ) ≤ inf θ∈Z N E[|Y -Φ N θ (x)| 2 ] < C * ≤ inf ν∈P(E G ) R(ν)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):R(µ) = E π [∥Y -Φ µ (X)∥ 2 Y ] = R * + E π [∥f * (X) -Φ µ (X)∥ 2 Y ] 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.Having laid out all the different relevant elements for our work, we can now procede with the proofs of some of our main results.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 )):T 0 Z (∂ t φ(z, t) -⟨ς(t)D µ F (µ t , z), ∇ z φ(z, t)⟩) dµ t (z) dt = 0Now, 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 )):T 0 Z (∂ t φ(z, t) -⟨ς(t)D µ F (μ t , z), ∇ z φ(z, t)⟩)dμ t (z) dt = T 0 Z (∂ t φ(M g .z, t) -⟨ς(t)D µ F (M g #µ t , M g .z), ∇ z φ(M g .z, t)⟩) dµ t (z) dt =: ⋆ 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:∂ t φ g (z, t) = ∂ t φ(M g .z, t) and ∇ z φ g (z, t) = M T g ∇ z φ(M g .z, t)So that, by also using proposition 13 and the orthogonality of the group action, we get:⋆ = T 0 Z (∂ t φ g (z, t) -⟨M g .ς(t)D µ F (µ t , z), M g ∇ z φ g (z, t)⟩) dµ t (z) dt = T 0 Z (∂ t φ g (z, t) -⟨ς(t)D µ F (µ t , z), ∇ z φ g (z, t)⟩) dµ t (z) dt = 0Where 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:∀t ∈ Ω g , ⟨f t , µ t ⟩ = ⟨f t , M g #µ t ⟩ = ⟨f t • M g , µ t ⟩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:T 0 ⟨f t , µ t ⟩dt = G T 0 ⟨f t , µ t ⟩dtdλ G (g) = G T 0 ⟨f t • M g , µ t ⟩dtdλ G (g) = T 0 ⟨f t , µ G t ⟩dtImplying (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):D µ R DA (µ, z) = D µ R F A (µ, z).Indeed, notice that, from proposition 18:D µ R F A (µ, z) = G M T g .D µ R(µ G , M g .z)dλ G (g) = G M T g .D µ R(µ, M g .z)dλ G (g),while also:D µ R DA (µ, z) = G M T g .D µ R(M g #µ, M g .z)dλ G (g) = G M T g .D µ R(µ, M g .z)dλ G (g)Now, let (µ F A t ) t≥0 and (µ DA t ) t≥0 be the WGF solutions starting from µ 0 for R F A and R DA respectively. As R F A is G-invariant, by corollary 3, (a.e.)∀t ≥ 0, µ F A 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 (µF A t ) t≥0 satisfies: ∀φ ∈ C ∞ c (Z × (0, T )): T 0 Z ∂ t φ(z, t) -⟨ς(t)D µ R F A (µ F A t , z), ∇ z φ(z, t)⟩ dµ F A t (z) dt = 0Now, as (a.e.)∀t ≥ 0, µ F A t ∈ P G (Z), we have ∀z ∈ Z:D µ R F A (µ F A t , z) = D µ R DA (µ F A t , z). In particular, (µ F A t ) t≥0 satisfies ∀φ ∈ C ∞ c (Z × (0, T )): T 0 Z ∂ t φ(z, t) -⟨ς(t)D µ R DA (µ F A t , z), ∇ z φ(z, t)⟩ dµ F A t (z) dt = 0Implying that (µ F A t ) t≥0 solves WGF(R DA ) starting from µ 0 ; thus by uniqueness: (µ F A 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 F A 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 F A ) always do so.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:EQUATIONwhich 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:Z t = Z 0 - t 0 ς(s)D µ R τ (µ s , Z s )ds + 2β t 0 ς(s)P E G dB s , and Z 0 = ξ 0 (initial condition) (13)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 = ( Zt ) t≥0 := (M g Z t ) t≥0 satisfies this same equation ( 13).Denote ν s := M g #µ s as the law of Zs , we want to show that for all t ≥ 0: ZtEQUATIONIndeed, first notice that:1. 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, Z0 (ω) = M g Z 0 (ω) = M g ξ 0 (ω) = ξ 0 (ω) = Z 0 (ω). That is, Z0 a.s. = Z 0 .2. Now, the equation is satisfied by (Z t ) t≥0 and therefore, for t ≥ 0 , we have:Zt = M g Z t = M g Z 0 -M g t 0 ς(s)D µ R τ (µ s , Z s )ds + 2βM g t 0 ς(s)P E G dB s = Z0 - t 0 ς(s)M g D µ R τ (µ s , Z s ))ds + 2β t 0 ς(s)M g .P E G dB s = Z0 - t 0 ς(s)D µ R τ (M g #µ s , M g .Z s )ds + 2β t 0 ς(s)P E G dB s = Z0 - t 0 ς(s)D µ R τ (ν s , Zs )ds + 2β t 0 ς(s)P E G dB sHere, 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, (By the pathwise uniqueness of the solution (Z t ) t≥0 , we have (following, for instance, [28] ):P sup t≥0 ∥Z t -Zt ∥ = 0 = 1In particular, as g ∈ G was arbitrary, we have that:EQUATIONWe 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 ω ∈ Ω:∥Z t (ω) -P E G Z t (ω)∥ = Z t (ω) - G M g .Z t (ω)dλ G (g) ≤ G ∥Z t (ω) -M g .Z t (ω)∥ dλ G (g)We can integrate both sides by P to get (using Fubini as functions are positive and measurable):0 ≤ Ω ∥Z t (ω) -P E G Z t (ω)∥dP(ω) ≤ Ω G ∥Z t (ω) -M g .Z t (ω)∥ dλ G (g)dP(ω) ≤ G Ω ∥Z t (ω) -M g .Z t (ω)∥ dP(ω)dλ G (g) = 0where 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 (ω)∥ dP(ω) = 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 β = 0Step 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 τ ):F (µ) = R τ,β E G (µ) = R τ (µ) + βH λ E G (µ E G ).Also recall the notationH 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 )): T 0 Z ∂ t φ(z, t) -⟨ς(t)D µ R τ (µ t , z), ∇ z φ(z, t)⟩ + β tr[P E G D 2 z φ(z, t)] dµ t (z) dt = 0with 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:- Z ⟨∇ z ϕ(z), D µ H E G (µ t , z)⟩ dµ t (z) = - Z ⟨∇ z ϕ(z), P T E G ∇ z dµ E G t dλ E G (P E G .z)⟩ dµ E G t dλ E G (P E G .z) -1 dµ t (z)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:- Z ⟨∇ z ϕ(P E G .z),P T E G ∇ z dµ E G t dλ E G (P E G .z)⟩ dµ E G t dλ E G (P E G .z) -1 dµ t (z) = - E G ⟨∇ z ϕ(x), P T E G ∇ z dµ E G t dλ E G (x)⟩ dµ E G t dλ E G (x) -1 dµ t (x) = - E G ⟨P T E G ∇ z ϕ(x), P T E G ∇ z dµ E G t dλ E G (x)⟩dλ E G (x)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 EAs, 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 6Proof 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:D µ R DA (µ, z) = D µ R F A (µ, z) = D µ R EA (µ, z).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:D µ R F A (µ, z) = G M T g .D µ R(µ, z)dλ G (g) = P E G .D µ R(µ, z),while also, from Proposition 18:D µ R EA (µ, z) = P T E G .D µ R(µ E G , P E G .z) = P E G .D µ R(µ, z).Knowing this, the rest of the proof is analogous to that of Theorem 4. Let (µ F A t ) t≥0 , (µ DA t ) t≥0 and (µ EA t ) t≥0 be the WGF solutions starting from µ 0 for R F A , R DA and R EA respectively. Since µ 0 ∈ P(E G ) ⊆ P G (Z), from Theorem 4 we know that (µ F A t ) t≥0 and (µ DA t ) t≥0 coincide. Let's see that, w.l.o.g., (µ F A t ) t≥0 coincides with (µ EA t ) t≥0 . As R F A is G-invariant, by theorem 5, ∀t ≥ 0, µ F A 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 (µ F A t ) t≥0 satisfies: ∀φ ∈ C ∞ c (Z × (0, T )):T 0 Z ∂ t φ(z, t) -⟨ς(t)D µ R F A (µ F A t , z), ∇ z φ(z, t)⟩ dµ F A t (z) dt = 0Now, as ∀t ≥ 0, µ F A t ∈ P(E G ), we can restrict our integral to E G . Also, we have ∀z ∈ E G :D µ R F A (µ F A t , z) = D µ R EA (µ F A t , z). With these properties, (µ F A t ) t≥0 satisfies ∀φ ∈ C ∞ c (Z × (0, T )): T 0 E G ∂ t φ(z, t) -⟨ς(t)D µ R EA (µ F A t , z), ∇ z φ(z, t)⟩ dµ F A t (z) dt = 0Making the integral over Z once again (we can since µ F A t ∈ P(E G ) for all t ≥ 0), we get that (µ F A t ) t≥0 solves WGF(R EA ) starting from µ 0 ; thus by uniqueness: (µ F A 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 .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 2 acting on X and Y via the coordinate transposition 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 computedas E G = 1 √ 2 0 0 1 √ 2 , 0 1 √ 2 1 √ 2 0, which is a 2-dimensional subspace of the ambient 4dimensional 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:θ * 1 = ϑ.(-1, 0, 0, 0.5) T θ * 2 = ϑ.(0.5, 1, 0, 1) T θ * 3 = ϑ.(-0.5, 0.3, 1, 0) T θ * 4 = ϑ.(0, -1, -0.5, 1) T θ * 5 = ϑ.(0.7, -0.7, 0.5, 0.7) T . 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] ).a * 1 = ϑ.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 -ŷ∥ 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.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 ⌊ Ne 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 EAs 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 EMLPs provided in the package (EAs 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 Wasserstein-2 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 Wasserstein-2 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 µ = 2E[∥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 Wasserstein-2 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.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 N T , (ν N N T ) 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 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, EAs 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).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 Ej #ν 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.The following result provides consistency guarantees when the regularization parameters τ and β are small, and is a slight extension of a result in [38] . Figure 10 : RMD comparison between the empirical student particle distribution, ν N θ , to both P Ej #ν N θ and (ν N θ ) E G (where j is the heuristic step). These are performed at the beginning and the end of training on every fixed heuristic step. The red line is placed at the value 10 -2 and represents a possible threshold δ, to be used in the heuristic to determine whether training left E j or not.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 limτ,β→0 R(µ * , τ, β, ν ) = inf µ∈P2(Z) R(µ).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 havelim inf n→+∞ R τn,βn ν (µ n ) ≥ lim n→+∞ R(µ n ) = R(µ).On the other hand, given µ ∈ P p (Z), consider ρ to be the heat kernel in Z = R D and ρ n (x) := β -D n ν(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 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). 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 thatf * = E π [Y |X = •] lives in f * ∈ L 2H (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 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:∆(f, Q G .f ) = R(f ) -R(Q G .f ) = E ∥f * (X) -f (X)∥ 2 Y -E ∥f * (X) -f G (X)∥ 2Y , which can be written as:∆(f, Q G f ) = E π ∥f * (X) -f G (X)∥ 2 Y -2⟨f * (X) -f G (X), f ⊥ G (X)⟩ Y + ∥f ⊥ G (X)∥ 2 Y -E π ∥f * (X) -f G (X)∥ 2 Y = -2⟨f * -f G , f ⊥ G ⟩ L 2 (X ,Y;π X ) + ∥f ⊥ G ∥ 2 L 2 (X ,Y;π X ) = -2⟨f * , f ⊥ G ⟩ L 2 (X ,Y;π X ) + ∥f ⊥ G ∥ 2 L 2 (X ,Y;π X )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:∆(f, Q G f ) = ∥f ⊥ G ∥ 2 L 2 (X ,Y;π X )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. 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. , 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.We became aware of the work[35] after a first version of this paper was posted.w.l.o.g. by compactness, all G-actions are assumed to be via orthogonal representationsdefined as: φz = Tz#λG, where Tz := [g ∈ G → Mg.z ∈ Z]An additional variant, with f * = Φ N * θ * having SI particles, is given in SuppMat-F.It is roughly equivalent to W2 when far from the δ0 measure, see SuppMat-F for detailsThe name is inspired from the traditional definition of 'intertwinning linear maps' in representation theory, see SuppMat-A.1 for an explanation.Notice that this fact is explicitly used in the proof of Proposition 4Where the gradient ∇xf (x) is the unique vector in H representing the action of Dxf (x) : H → RIn this case, ∀g ∈ G, let g.µ = Mg#µ and g.f = f g given by ∀x ∈ X , ∀y ∈ Y, f g (x, y) = f (ρ -1 g .x, ρ-1 g .y)Notice that the convention on the linear functional derivative's definition isn't truly important for the proof, since in the end ⋆ simply corresponds to the term Z ∂F ∂µ µs, z d µsdλ(s) (z) being substracted to itself.In particular, as for any fixed µ ∈ P(Z) the function g ∈ G → Mg#µ is continuous (thus, of compact image), then the function (g, z) ∈ G × Z → ∂R ∂µ (Mg#µ, Mg.z) is bounded 


Decision: A fine work investigating the Mean Field limit of overparametrized shallow models under symmetric data and/or symmetry-leveraging techniques. The reviewers have unequivocally voted for acceptance.

Reviews: review: paper_summary: paper_summary: This 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:

- 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)$.
- 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$.
- The parameters $\theta$ are optimized using noisy SGD.

The 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.

Given 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.

The 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:

- (Proposition 3) Optimizing with respect to an invariant loss function yields a weakly-invariant optimum if it has a unique minimizer,
- (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,
- (Corollary 2) If data distribution is equivariant, optimizing under data augmentation and feature averaging provide no advantage,
- (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.

Then, in Section 3.3, the authors provide mean-field analysis on training dynamics, and show that:

- (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.
- (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.

In 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.

[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.

[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.
strength_weakness: strength_weakness: - 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.
- 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.)
- 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.
- 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.
- 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.

[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.

[3] Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, and Tess E. Smidt, Discovering symmetry breaking in physical systems with
relaxed group convolution, arXiv preprint arXiv:2310.02299, 2023.
questions: questions: Please see the Weaknesses section.
limitations: limitations: The 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.
paper_summary: paper_summary: This 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.

- DA and FA are equivalent in terms of optimal values and also equivalent to standard training (free training) for WI measures (Theorem 2).
- Furthermore, this equivalence is shown over the space of general measures when the data distribution is equivalent (Corollary 2).
- 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.
strength_weakness: strength_weakness: - This work mainly focuses on two-layer neural networks; hence, it is unclear if the results can be applied to deep neural networks.
- 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.
- (Minor) The paper is somewhat hard to follow because of the numerous notations. I suggest including a table that summarizes these notations.
questions: questions: - 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?
- Is WGF essential for the results in Section 3.3? Do similar results hold for other dynamics (e.g., underdamped version)?
limitations: limitations: Limitations are addressed in the paper.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: * 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.

* 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.
questions: questions: * 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?
* Proposition 5 states that equivariant architectures can lead to optimal solutions provided some universality conditions. How restrictive are these conditions?
* 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$?
* 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.
* 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?

## References
[1] Wang, Rui, et al. "Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution." Forty-first International Conference on Machine Learning.
limitations: limitations: The authors have discussed the limitations.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: The paper already seems quite complete and well written and I could not identify significant weaknesses. The following weaknesses are only minor points:

- The description of Figure 1 is hard to parse. Markers like (a)-(d) could reduce ambiguity of what is meant by a column.
- The markers in Figure 2 should be larger and it is unclear to me which additional insight to draw from the second row.

As 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.

**Typos:** 22: lastly, 65 and 66: b=N?, 75: Definition name in bold for consistency?, 120: […], 275: will, 329: a heuristic
questions: questions: - 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?
- Can GNNs or Transformers also be encoded in your class of equivariant architectures? If so, these example would significantly broaden the target audience.
- In future work, in which ways could Assumption 1 be relaxed?
limitations: limitations: The 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:

- 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.
- 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?

Comments: title: Thank you
Thank you
Thank you
comment: We thank the reviewer very much for their positive comment and for adjustment of their score.
We thank you for pointing out your concerns, as they have undoubtedly helped us provide a clearer exposition of our central ideas and contributions in the final version of our paper. We are also thankful for your positive comments and the evaluation adjustment.
Thank you for the detailed response. The authors have adequately addressed my concerns. Given the strength of the paper, some limitations seem negligible. I will raise my evaluation. The discussion provided in the general comments will significantly help readers understand the importance and applicability of the theory.
Thank you for the thorough explanation. Together with the original response, my major concerns have been addressed. I will adjust my score accordingly.
We thank the reviewer for their comment, and are very happy for their positive evaluation. Thanks again!
Thank you very much for your response, and we are happy as well for your positive recommendation. Regarding your follow-up questions:

Regarding L.1585, we agree that the current phrasing can easily lead to confusion, and we will be sure to change it in the final version. We only wanted to refer to the arbitrary teacher having 5 particles and, indeed, as you have noticed, the WI teacher has exactly 10 particles (which are the ones we provide in our original rebuttal). Namely, it is exactly $G$-equivariant (and not just an empirical sample from a $G$-equivariant measure). As we mention in our original rebuttal, the current version “implicitly” mentioned both of these facts, as from our remarks in lines 196-203 and the phrase in line 1587-1588; but we will surely make this clear and explicit in our final version.  

We do want to stress that the exact equivariance of the WI teacher doesn’t mean the problem is any easier. Since the WI-teacher’s particles don’t live in ${\cal E}^G$, one intuitively would believe that the freely-trained particles of a NN (initialized within ${\cal E}^G$) would prefer to “escape” ${\cal E}^G$ in order to better approximate the teacher. However, as shown by our numerical experiments, this doesn’t happen, and the particles remain (up to numerical error) within ${\cal E}^G$.

Regarding Q5, we deeply thank you for the provided references and ideas. We will be sure to consider them for our upcoming works regarding the topic.
We thank you very much for your response, and we will take the opportunity to provide a clarifying example for the question at hand.

As noted in  Lines 1012-1016, under the same example you are considering, the “natural” action of $G$ on the parameter space ${\cal Z} = \mathbb{R}^{c \times b} \times \mathbb{R}^{d \times b}\times \mathbb{R}^b$ corresponds to taking, for $g \in G$ and $(z = W, A, B) \in \cal Z$: $$M_g.z = M_g.(W, A, B) := (\hat{\rho}_g.W.\eta_g^T, \rho_g.A.\eta_g^T, \eta_g.B). $$ 

This is what we refer to as the “intertwining action” on the layer parameters, as it recalls the traditional definition of an intertwinning linear map from representation theory (also, see the discussion on lines 165-180). Notice that, indeed, this involves the action of the group on the corresponding input/output spaces of each layer (e.g. from $\mathbb{R}^d$ to $\mathbb{R}^b$ and from $\mathbb{R}^b$ to $\mathbb{R}^c$). 

In our original rebuttal, notice that we referred to “a trivial action on the intermediate layer”, and not plainly “a trivial action”. Indeed, we referred to taking $\eta$ to be trivial, rather than $M$ itself; the latter, as you have correctly noticed, doesn’t work in general.

Thanks to the orthogonality (and linearity) of our group action, a similar development to yours yields, for any $\sigma: \mathbb{R}^b \to \mathbb{R}^b$, any $g \in G$ and for any $z = (W, A, B) \in  \cal Z$:

$$\sigma_*(\rho_g.x, M_g.z) = \sigma_*(\rho_g.x, (\hat{\rho}_g.W.\eta_g^T, \rho_g.A.\eta_g^T, \eta_g.B)) =(\hat{\rho}_g.W.\eta_g^T). \sigma((\rho_g.A.\eta_g^T)^T.(\rho_g.x) + \eta_g.B)$$ 
$$ = \hat{\rho}_g.(W.\eta_g^T. \sigma(\eta_g. A^T.(\rho_g^T.\rho_g).x + \eta_g.B)) = \hat{\rho}_g.(W.\eta_g^T. \sigma(\eta_g. (A^T.x + B)))$$

Namely, it is enough that $\sigma: \mathbb{R}^b \to \mathbb{R}^b$ is $G$-equivariant (w.r.t. the action of $\eta$) in order to get $\sigma_*$ to be jointly equivariant. Indeed, this would yield:

$$\sigma_*(\rho_g.x, M_g.z) = \hat{\rho}_g.(W.\eta_g^T.\eta_g. \sigma(A^T.x + B)) = \hat{\rho}_g.(W.\sigma(A^T.x + B)), $$ 

as required. This justifies our claims about the pointwise activations and/or the use of a Norm-ReLU activation to ensure the joint-equivariance of $\sigma_*$ for different levels of complexity of $G$ (and its action on the intermediate layer). 

More specifically, if this action on the intermediate layer $\eta$ is taken to be trivial, any function $\sigma: \mathbb{R}^b \to \mathbb{R}^b$ is $G$-equivariant w.r.t. $\eta$  (indeed, $\sigma(\eta_g.x) = \sigma(x) = \eta_g.\sigma(x)$); i.e. no matter the inner activation function $\sigma$, under this action $M$, $\sigma_*$ is jointly-equivariant. 

In the end, a trivial $\eta$ assumes “no action of the group on the latent space”; and so we only act on the parameters on the dimensions defined by the data: e.g. $\rho$, that acts on $\cal X$, will act only on $A$ on the left; and $\hat{\rho}$, that acts on $\cal Y$, will act only on $W$ on the left as well. This readily generalizes to the case where $\sigma_*$ represents a whole multilayer unit, for which we take the action on parameters to be trivial on the hidden layers, but not on the input and output layers (where it acts according to the data). Indeed, no matter the inner activations, such $\sigma_*$ would be jointly equivariant with respect to such group action. 

As we also mention in our original rebuttal, a “trade-off” follows from considering such a “boring” group action $M$: the space ${\cal E}^G$ might not represent very “interesting” equivariant architectures. This doesn't seem too problematic since it wouldn’t be the use case if $\sigma$ didn’t respect any sort of equivariance to begin with. It does, however, allow to state that DA and FA are tightly related, even for unconstrained architectures. This also makes sense when we consider that DA and FA only interact with NN models through the input and output layers.

We hope to have clarified your question with this example, and we’re open to answer further questions if required.
I thank the authors for their thoughtful response and thorough work. I am satisfied with the changes promised in the global response and I will keep my positive evaluation.
Thank you for the comprehensive rebuttal. I am trying to technically understand the response on the joint-equivariance of single-hidden-layer NN when the group acts trivially on the hidden layer (regarding W1).

Following Line 65-67, let $\mathcal{X}=\mathbb{R}^d$, $\mathcal{X}=\mathbb{R}^c$, and $\mathcal{Z}=\mathbb{R}^{c\times b}\times \mathbb{R}^{d\times b}\times \mathbb{R}^b$, $z = (W, A, B)\in \mathcal{Z}$, $\sigma:\mathbb{R}^b\to\mathbb{R}^b$, and $\sigma_*(x,z)\coloneqq W\sigma(A^\top x+B)$. From Line 170-171, the joint-equivariance is written as $\sigma\_* (\rho\_g \cdot x, M\_g \cdot z) = \hat{\rho}\_g \sigma\_* (x,z)$ for all $g,x,z\in G\times \mathcal{X}\times \mathcal{Z}$.

Now let $M\_g$ be a trivial action. Then joint-equivariance is required as $\sigma\_* (\rho\_g \cdot x, z) = \hat{\rho}\_g \sigma\_* (x,z)$, or equivalently $W\sigma(A^\top \rho\_g \cdot x+B) = \hat{\rho}\_gW\sigma(A^\top x+B)$, for all $g,x,z\in G\times \mathcal{X}\times \mathcal{Z}$. As far as I am aware, this is not true in general; let $A=0$, $\sigma=\mathrm{id}$, $W=I$, and let $B\in\mathbb{R}^b$ have distinct entries, then we have $\sigma\_* (\rho\_g \cdot x, z) = B$ which is not equal to  $\hat{\rho}\_g \sigma\_* (x,z) = \hat{\rho}\_gB$ for faithful actions $\hat{\rho}_g$. Therefore I do not see how joint-equivariance holds for single-layer NNs if the group acts trivially on the hidden layer. Am I missing something?
Thank you for your response. My concerns have been addressed, and I'm happy to keep the current recommendation.

I have some more comments about Q4&5. For Q4, how many particles are there in the WI-initialized teacher? Originally I thought it was 5, as you mentioned in L1585. But the WI distribution $(\nu_{\theta^*}^{N_*})^G$ has support on 10 (5 pairs of) points. Do you randomly sample 5 particles from this distribution, or do you include all 10 points? If it is the latter case, then the teacher is indeed strictly equivariant.

For Q5, I think the difference between "discovering the symmetry of the data" and "optimizing model architectures compatible with data symmetries" is subtle. I feel these two objectives are often connected to each other. In [1], for example, one can tell from the optimized architecture what is the maximal symmetry in data. This is also done in [2] but in a quite different way, by solving the Lie derivative constraint wrt a non-equivariant network. Conversely, one can also start by discovering the data symmetry and design equivariant architectures accordingly [3]. These works may be relevant if you plan to further investigate the idea in Sec 4.2.

[1] Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution. ICML 2024.

[2] LieGG: Studying Learned Lie Group Generators. NeurIPS 2022.

[3] Generative Adversarial Symmetry Discovery. ICML 2023.