Symmetries in Overparametrized Neural Networks: A Mean Field View
Authors: Javier Maass Martínez
Joaquin Fontbona
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Overparametrized Neural Networks
Mean Field Limit of Neural Networks
Symmetries in Neural Networks
Wasserstein Gradient Flow
Data Augmentation
Feature Averaging
Equivariant Architechtures.
Primary area: probabilistic_methods
TL;DR: Mean Field analysis of overparametrized shallow models under symmetric data and/or symmetry-leveraging techniques.

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

Symmetries in Overparametrized Neural Networks:
A Mean-Field View
Javier Maass Martínez
Joaquín Fontbona
Center for Mathematical Modeling University of Chile javier www.seawwwssssaaaaaaaoommmmma. com
Center for Mathematical Modeling University of Chile fontbona@dim. uchile cl
Abstract
We develop 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 class of generalized shallow NNs given by an ensemble of multi-layer units, jointly trained using stochastic gradient descent (SGD) and possibly symmetry-leveraging (SI.) tech- niques, 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 each single unit, corresponding, respectively, G-invarian distributions, and to distributions supported parameters fixed by the group action (which encode EA). This allows us to define symmetric models compatible with taking N namics of DA, FA and EA terms of Wasserstein Gradient Flows describing their MF limits. When activations respect the group action, 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 gets larger, in a teacher-student experimental setting, training a student NN learn from WI, SI or arbitrary teacher model through various SI. schemes. We lastly deduce data-driven heuristic to discover the largest subspace of parameters supporting SI distributions for problem, that could be used for designing EA with minimal generalization error.
and give interpretation of the asymptotic dy-
Introduction
Learning in complex tasks, employing ever larger datasets, has strongly benefited from the implemen- tation and training of Artificial Neural Networks (NN) with a huge number of parameters; as well as from training schemes architectures that can leverage underlying symmetries of the data in order to reduce the problem's 's complexity (see [33 34 for general reference). This raises questions, on one hand, of understanding the puzzling generalizability in overparametrized NN; and the other, of when and how symmetry-leveraging (SL) techniques (such as Data Augmentation, Feature Averaging or Equivariant Architectures), 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 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
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(WGF) of certain convex population risk on the space of distributions on parameters. Confluently, the incorporation of combined algebraic and probabilistic viewpoints have yielded more complete view of the benefits of SI. techniques under symmetry (see e.g. 13 27 59 and further references below); however, it 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 function symmetric w.r.t. a single orthogonal transformation. We consider our (independent] D work to largely broaden the scope and applicability of such initial contributions, as it provides a unified MF interpretation for both the use of SI. 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 class of generalized shallow models with multi-layer units on which we will focus, we recall WGFss 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, 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; also suggests a potential heuristic for discovering data-driven parameter-sharing schemes that lead to optimal equivariant architectures in MI. 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.
we study how SL
2 Preliminaries
2.1 Supervised learning with generalized shallow neural networks
Let X, )' and Z be separable Hilbert Spaces, termed as the feature, label and parameler spaces respectively. Typically, these are finite-dimenns dimensdimessional, c.g. X' = IRd and y = R (for c, d E No) 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 π E P(X x x)) denote the data distribution from which i..dd samples (X,Y) E X x y will be drawn, and f: y x y - R be a convex loss function. Consider also an activation function o. : X x 2 y. We introduce general class of shallow NN: Camples(X
Definition 1. A shallow neural network model of parameter 0 == (0.) E ZN is the function .N X + y given by T. (x) == N Ei= o. 0. 0z), Vx E X. Equivalently, if = NN > óo, the empirical measure associated with 0 E ZN, we can write Vx E Xx 2. (x) = (z: or, abusing notation, simply 4 = (o., v/ V ).
In the setting where X = Rd. y = Re and 2 = Redd x Rddd x Rd (for b E N²), if we consider, for (W. A, B) 2 and R + R², ox. 2) = Wo(ATx+D B); then &. (with N E N, 2* corresponds exactly to a however, these shallow NN models can represent settings that go far beyond this first example. In fact, a. can be taken to be an entire Multi-Layer NN model, in which case Φ."" will represent an ensemble of N such units trained simultaneously (see SuppMat 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 [66 62 69
neural network with N hidden units. Depending on a..
As we will also shortly see, for
This class thus allows for non-trivial internal units, while enabling the width N - Do consistently, and regardless of the possible underlying structure of the (fixed size) units represented by a.. Inspired
by this possibility, and by our writing of shallow NN models, we define more general notion:
Definition (Shallow Model). A shallow model any function of the form + (x) = (o. -), 11)
for some μ E P(2) (whenever the integral makes sense for all x E X). We write + == (o., u)) and denote the space of such models as Fo, (P)).
We became aware of the work 335 after first version of this paper was posted.
2

---

Classically, we want to find NN model that performs well with respect to and f. More precisely, having fixed an architecture (given here by N and o.), we consider the generalization error or population risk given by R(O)= = E [& d (X), Y)]. and look for vector of parameters 0 E 2 attaining min R(O). However, not only is this function highly non-convex and hard to optimize; but in practice we generally don't 't have access to π (and thus It) and we have to solve this problem only with a set of i..dd data samples {(X,, Yx) JKEN drawn from л. Thus, the usual approach to minimizing this population risk is to train a NN model 4g. through an SGD scheme (see e.g. [7):
First, initialize 09. Vi E (I,...,,,, NY. i.d. from a fixed distribution μo E P(Z).
. Iterate, for k E N, defining Vi E (1, (1. NN:
-f - (Xx, 0+) ViL(O
=
+ /230 5. (1)
Here, Sxx = ENE(KEN) is the step-size (or learning rate), parametrized in terms of < R+ + R+ a regular function and EN > O. Also, we have a penalization function r Z + R. regularizing Gaussian noise et itdd, [N00, Idz) independent from the initialization and data, and T. 2 When T, > O, the method is called stochastic gradient Langevin dynamics, noted SGLD ((74)), or simply noisy SGD, An infinite i...d. sample from π will be needed when letting later N - oo. When л is the empirical measure of finite dataset, we are performing empirical-risk minimization (which of course 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 models in F.. (P()), one gets the convex functional R : P(Z) + IR given by R(u)= == E, [ y (X), Y) The problem on ZN is thus lifted to the convex optimization problem on P(Z):
min R(p). (EP(Z)
(2)
Accordingly, this motivates looking at the evolution of empirical measures (NE))EN )kEN (N JAEN S
P(Z) instead of that of the specific parameters (0* keE) S Z The MF approach to NNs (see 16
5
62 67 aims at providing theoretical guarantees for problem ( justifying that 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.
2.2 Wasserstein gradient flow and mean-field limit of shallow models
We briefly recall some elements of Optimal Transport and Wasserstein Gradient Flows, referring to [ 63, 71 for further background. Let Z be a Hilbert space with norm || and. for E [1, oc), let pp Z := (u P(Z) Sz [[]])=(dd) < +o0 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: V), == [inf..eM(...)) [II - YIIi] Vp, v E (2) with II 1/) being the set of couplings between i and v (the infimum is always attained). The metric space (P,(Z), Wp) is Polish and called the p-th Wasserstein Space. In the remainder of this section we consider = 2 and = Rº.
We recall central objects for the sequel, including Lions' derivative
9 47 popularized in mean-field games (see e.g. 10 12 38) ) and shown (in (32) to coincide with the Wasserstein gradient (αI):
Definition 3 (Linear Functional Derivative and Intrinsic Derivative). Given F :P2(Z) + R, its linear functional derivative is the function (if it exists) OF Dom(F), ()), lim F((1 hho1-))uthh)-F(4
du Dom(() x Z + IR such that Vp, E - 2))(v-H)(z) and Jzz 8F (H, z)dp(z) = 0. The ou OF (H.-) is known as the first variation of F as p. Moreover, if OF Øu
fz aF (p,
function F', : :p E P2(Z) + exists and is differentiable in its second argument, we define the intrinsic derivative of F at μt to be: D 2) V, OF (p.z)). Abusing notation, we will write aF 8y
on
Sy : P2(Z)) x Z + R and
D.F P2(2) x 2 + 2, even if they are only partially defined.
This allows to define next Wasserstein Gradient Flow (following e.g. 116)
3

---

Definition 4 (Wasserstein Gradient Flow). Let R. - RR+ be regular scalar function and F; P:(2) + R be a convex functional for which the intrinsic derivative D is defined. We define a Wasserstein Gradient Flow (WGF) for F (shortened (GF(F)) as any absolutely continuous trajectory (1t)120 in (2 that satisfies, distributionally on [0, 20 x
dep = <(t) div F( -):)
Several authors (11
63 71 among others) have proven under various sets of assumptions that, given an initial condition μo E P2( the WGF(F) admits a unique (weak) solution, (He) In 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 [T6]
We interested in the where F is the following convex, entropy-regularizee population risk: RT (I)) == R(u) + rrdu + (H)) where T 2 0, À is the Lebesgue Measure on 2, r: 2 + R+ is penalization, and Hx defined as Hx( == f log( du otherwise, is the Boltzmann entropy of p. In this case, WGF(R reads as the PDE:
ax (=))du(=) if < À or +o0
0,p s()) [div ((D,, R(y,,) + - Vor) (4) + BBy4),
known as McKean-Vlasn equation in the probability and PDE communities (see the classic references 54 70]. and the recent review (I]) and popularized 'distributional dynamics' in NN literature (e.g. (3) ) When 3 > 0, a solution to (4) has a density w.r.t. À and is actually strong. Under rather simple technical assumptions (see Supp D.3] or [12 15] 38 57 69)), when T, 3 0 it is known that the WGF(R"*0) Wa conoerees to (unique) minimizer. When t, B = 0 a sort of converse holds (see 16) ): if WGF(R) R) converges in W2, then the limit minimizes R.
Proven by [6 53 62 67 and later refined e.g. by [14 22 23 51 66 69), the main result in the MF
Theory of overparametrizze shallow NNs states that SGD training for shallow NN, in the right scaling limit as N + co, approximates a WGF:
Theorem 1 (Mean-Fiee Field limit, sketch). For each T O, under relevant technical assumptions including regularity of a. and proper asymptotic behaviour o -> 0 as N xc, the rescaled empirical process given by n :=
(EE converges in law (in the Skorokhod space D2(2) (10, TT)) to μ =- (Ho)eefo.11 given by the unique WGF starting at po-
Despite the MF limit of NNs being 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 large, overparametrized NNs during training (as compared e.g. to the lazy-training regime described by the Neural Tangent Kernel approximation 17 Note that, for 3 >> 0, the entropy term H2 in WGF RT. (as well as the Laplace operator in equation is approximated, in practice, by the Gaussian noise term in the SGL.D (T. as N co.
2.3 Symmetry-leveraging techniques
We next discuss mathematical formulations of the main techniques to leverage posited distributional symmetries of the data at the training or architecture levels. We henceforth fix a compact group G of normalized Haar measure Àg, acting on X' and y, which we denote G Cp X, G Cp y. F A function
f: X2 + y is termed equivariani if Vy E G, f = f()) d We further say that the data (X
is equivariant, and write not enforced unless stated). The space of functions f X - y inttre-nttegrable (in Bochner sense) w.r.t xx = Law(X (X) is called L (X, Y); xx). Further relevant concepts
PC (X x)), if E G. (py.X, Y))
(this is
introduced as needed.
i.i.d Data Augmentation (DA): This training scheme considers {ou HeN ~ AG independent from the {(X,, }RE in ① and carries out SGD on samples ((pg -Xk; ke Pa.. (Y)) DA and the vanilla training scheme would thus be equivalent if PG Y)) One
show (see [3] 48) that, performing SGD with DA, results in an optimization scheme for the symmetrized population risk, (D) := E, [La e (2 (P,.X), Pg. Y) dg) Despite being effective in practice, DA gives no
guarantee that the resulting model will be equivariant. For deeper insights, see [1 21 39 46 48 52
.w.lo.gg by compactness, all G- actions are assumed to be via orthogonal representations

---

Feature Averaging (FA); Instead of focusing on the data, FA works with symmetrized versions of the vanilla NN models & at hand, averaging model copies over all possible translations through the group action. This amounts to constructing (or approximating) the symmetrization operator over T,,,X, 3/ xx) defined as (z) ==
(see (27]), and trying to minimize The resulting model will be equivariant, however, FA is inefficient, as as G times more evaluations are needed for training and inference.
R (0) == E, [e ((2c. x) (X) (see
46 48 ).
Equivariant Architectures (EA): Following [8], EA in multilayer NNs are configurations yielding models equivariant between each of the hidden layers (where G assumed to act). As stated in 75 and 2 18 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/g group convolutions between layers). In our context, assuming that G CM Z some group action, we require that ar X' Z + y is jointly equivariant, namely, V(g, x 2) E G' x X' x Z, (pg Mg.=) = Pg (z, z); to ensure G-actions over different spaces are properly related. Introducing the set of fixed points for G CM 2, E == (2 E 2 Vg E G, Mg. =}, a shallow NN model & thus has an EA if 0 E (EO)N. Under the right choices of 0. and M, the obtained EAs can encode interesting and widely applied architectures, such as CNNs [9 and DeepSets [77] (see SuppMat- C.1 for further discussion). We call 0 the subspace of invariant parameters, which a closed linear subspace of 2, with unique orthogonal projection Pgo : 2 - E explicilly given by Pga.z == SG My dAa (g) for E 2 (see (30)). We are thus led to solve: minge(e0) R()) or, equivalently, to find the best projected model b N.A.A. := (a., Pgo # ), by minimizing IEAA(0) (0) == E, ( ( (.4 (X), }. This can considerably reduce the parameter space dimension; however EA might generally yield decreased expressivity or approximation capacity.
29 30
3 Symmetries in overparametrized neural networks: main results
3.1 Two notions of symmetries for parameter distributions
The following notions regarding distributions from P(Z) are central our work:
Definition 5. Given μ E P(2), we respectively define its symmetrized and projected versions as
a = (M and _ tp. Moreover, we introduce two subspaces of P(Z): pc = {u E PZZ)= Vg E G, My#H= and P(EG) == {μ P(Z) p((E)- = 1}.
Example. For G (=1} acting multiplicatively on 2 RR one has E {0), hence P(EG)= {6o), while Po (Z) (=(+v(())) i E P( In particular, for 2 E Z, (62) (0, +0=z).
Definition 6 (Invariant Probability Measures). We say that μ E P(()) is:
Ieakly-Invariant (WI) if = G and Strongly-Invariant (SI) if L
We notice that: P(EG) C (2), G. E (2) and _ºeh E EPEE)). Thus, SI implies WI. Next result relates the symmetrization operation on P(2)) with the one on shallow models Fo. (P()):
Proposition Let t,, E Fo. (P(2)) with o. X 2 y jointly equivariant. Then:
(Qa o.
That is to say, the closest equivariant function (in ( (2 V.Tx)) to p is given by the shallow model associated to the symmetrized version of u..
Remark. In particular, b,, is equivariant as soon μ is WI only, Conversely, if -X-+0 y is an equivariand function, then $b = + yor i.e. it can be expressed in terms of WI distribution. This highlights a priority role of WI distributions on Z in representing invariant shallow models.
The alternative, 'projected model' & nee, in turn, is never the closest equivariant shallow model, to
b in ((x,),,)) unless equal to 4 o The latter rarely is the case (unlike commonly implied in the literature). In fact, the symmetrized version 2 of a shallow NN model Φ involves ( 1/ MN Pe., where Vz E Z, 12 is the orbit measure of the action 7 and has N GG|
detined Pp T: TAG where T' [g E G ++ Mg.0 2]
5

---

Z-valued parameters (possibly with [G] oo). This sharply contrasts (u E Opga. which has Σ N distinct parameters, all living in E So, in general, depending on a, and G, one might have (o.: (u * 2 notable case in which the equality holds is the class of linear models, which is discussed in SuppMar-E.J
Example. In the previous example. 6₂ and o, jointly equivariant, we have 4 = o. (-(,,)) 2), No D (a (0.,)) and + - =o, (-,(0) which are generally distinct if + O. Notice that $,, no is equivariant function without any of its 'parameters' living in EG.
3.2 Invariant functionals on P(Z) and their optima
In the same spirit as when defining the population risk R P(2) - in in 2 the risk functions associated with SL-techniques from Section 2.3 can be lifted functionals over P(2), namely to: (D)(M)) == E [Sa (op (p dAG (g)], R RF)(H)
Ex [e (Qa( (X)), Y)] and (m) = Es ( n x) respectively. This will allow us study these SL-techniques, in the overparametrized regime, under a common MF framework. We need the following assumption: Assumption 1. π E (2 x >)); L: y x y + R is convex, jointly invariant and differentiable with Vil linearly growing: and o, X'
+ y is bounded, jointly equivariant and differentiable.
The quadratic loss fly 9) = 21 011 is an example of such t. Having bounded and differentiable 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 commonly- used functions a. seems feasible, up to some additional technicalities (see SuppMat [A.2 for further discussion). Finally, having assumption: under the right choice of a. and M, any usual singlehhidden-lyee NN architecture can be made to satisfy it (see SuppMat C for deeper discussion). We also need:
be jointly equivariant (as defined in section 2.3 isn truly restrictive
Definition 7. A functional F: P(Z) + R is invariant if F(M = F(u) V(g,u) E E x P(2); equivalently, if it equals its symmetrized version F = FM (MMHH))ddo(o).
Proposition 2. Under Assumption 0 RDD,, RFA and REA invariant (and convex) and we have:
RDA(H)= = RROHH, (p((())R((()) and ((A(
In particular, R = RD if R is invariant. Moreover, Vu E (2)(). R(p) = RD^(P) RFA(H).
Last, if I E pp y) (the data distribution is equivariant), then R is invariant.
The proof relies on Proposition | and calculations as in [30], see SuppMat- E.2 Next result is general property of functionals over P(Z), which is key for the forthcoming analysis:
Proposition (Optimality for Invariant Functionals). Let F: P((Z) + W. be convex, C³ and invariant. Then: Vu P()), F(x))= S F(w); and so, unpa F() = infpep(z) F(u). In
particular, if F has unique minimizer over P(Z), it must be WI.
The proof relies on an ad-hoc version of Jensen's '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]
we have:
inf R. (A) inf Rd (DA(p)= (N) (EP(2)
R(u) inf R. (#) (EPG(2)
inf
inf R A²A(H).
(2PG(Z)
(EP
HEP(2)
Note that, on the other hand, REA only satisfies: inf 2) REA(H) inf NEP(Eo) R(y). In the case of the quadratic loss, Theorem 2 can be made more explicit:
Corollary 1. Under Assumption
I when the loss is quadratic and x is invariant, we have:
inf R((p)= = R. + MEPO(2) (2)
inf pEPO(2)
|0 - -f [=12 = RR,+ inf (2)
1 -6
where ʃ. = F. [ = and Ra, R., are constants only depending on n and fa. 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 7 hold and suppose क E (x y). Then, R is invariant and therefore: unf R(y)) = int (EPO(Z) R(u) = pef RDA inf pEP(Z) R Ay)) 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 simple example in which E is trivial:
Proposition 4. Even with finite group G acting orthogonally on X R. y = B and Z = (2)) with x being bounded and jointly equivariant; can have: infrep(z) R(p) < infuep(Eay R(v).
and equivariant; with being quadratic; and with a, being C
In fact, even if R is invariant, when EG is too restrictive, it might become impossible globally optimize R over SI measures (which amounts to using REA as proxy for R). This subtlety has to be considered when deciding to use EAs on problems where symmetries exist. Nevertheless, if &o has good universality properties, a true SI solution to the learning problem can be sought for:
Proposition 5. Let Assumption 7 hold, f. be quadratic and E +PC(2 x )) Ff. (P(EG) dense in (X, (2; mx) := 2a 1.2( (X, K y; xx)), then: inf (EP()) R(u)= = inf vEfEE) R(v)= R..
Remark. See e.g. 150
60 76 777 for conditions on EG and a- guaranteeing this 'restricted'
universality on Lig (x Tx). These allow for effectively solving the problem in fewer dimensions,
for which is key in successful EA like CNNs and DeepSets. See SuppMatt for a deeper discussion.
3.3 Symmetries and SL training dynamics in the overparametrized regime
We now study the MFL of the various training dynamics when Z = RR. We begin with the general:
Theorem 3. Let F: P(Z)- - R be an invariant functional such that (GF is well defined and has unique (weak) solution (He) If po E P (Z). then, for dt-a.e. t 2 0 we have pa P (Z).
The proof of Theorem relies on D F being equivariant (in a suitable sense) and (Mg14) )220 satisfying also, as a consequence, WGF(F) (See SuppMat- E.3 for the details). Note that μiu E P (2) is simply verified, e.g. by a standard Gaussian in Z. Specializing this result, we get:
Corollary 3. Let Assumption 7 lnn technical assumptions (as in ((6]) hold. Then, if R and r are invariant, WGF(R". starting from pa E P4 (2) satisfies: for dt-a.e. t 2 0, 14 E (Z). If moreover 3 > O, each μ4 has a density function invariant with respect to G Cw Z.
Remark. T 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 351. which addresses the case of wide 2-layer ReL.U networks with a target function that's symmetric w.r.t a single orthogonal transformation. The fact that strong solutions to WGF( can be sought among invariant functions to reduce the complexity when π is equivariant, was also first hinted in [3 The natural domain of invariant functions is in fact the quotient space of G CM Z (and not gG 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 po E (Z), WGF(RDA) and (GF(RFA) solutions are equal. 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 + oo?
We will now see that similar results hold for P(gG) instead of pG (Z). Notice that the entropy- regularized risk forces each μe to have a density w.r.. X in Z if 3 > 0. Therefore, if G Oar 2 is non-trivial (thus Ed is a strict subspace), we always have be & P(EC). It thus seems natural to restrain the noise in equation (I) to stay in EC; namely, to consider the projected noisy SGD dynamic:
0/-1 = 0. - (Veo=(XH.0)) 0/) VIE(ON (X.), + ~v((0)) + /285 Peo (5)
7

---

Note that projecting only the noise in 5) doesn't 't force 0: to be in E even if 6 was. Introducing the related projected-regularize Aga the Lebesgue Measure over E we get the following analogue of Corollary 3
functional: R/a (p) == R(n) + T rdu f + BHx (n with
Theorem 5.
Let Assumption and technical assumptions on R hold. Suppose that P and are invariant. Then, if po E P2(E (Pe)120 solution of (GGF(R2) satisfies Vt 2 0, Pe E P2(E
The result holds for 3 2 O. Its proof is based on pathwise properties of the McKean-Vlasse stochastic differential equation (studied e.g. in (22)) associated with the WGF( see SuppMat-DD2
Remark. Theorem is significantly stronger than Corollary (3 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 highly non-intuitive fact, and large N exclusive phenomenon, as our numerical experiments will show. See SuppMatD.22j1 a deeper discussion.
Remark. Notice that, despite the computation of E being generally hard (see (29)), Po E P2(EE) can be achieved by simply setting μ40 = óg- Moreover, since one can also take 3 = O, 'having access' to the noise projection Pga is never explicitly required, allowing for broader applicabiliity of the result. In particular, as we'l lll show in our experiments, usual shallow NNs with all parameters initialized at freely trained with 'noiseless' SGD, will satisfy Theorem S in the MFL.
Remark. Theorem hoolds too for the invariant functionals RDA, R and REA in the role of R. even if is not equivariant. Notably, DA, FA and EA procedures starting on 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 Shold. If po E P2(E the WGF(RFA), (DD( and WGF(REA) solutions coincide. If R is invariant, WGF(R) solution coincides with them too.
4 Numerical experiments and architecture-ddscoverr 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. [4 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 simple setting with: X' = y = R², Z = 2=x2 ~ RR, and G = C2 acting on X' and y by permuting the coordinates; and on Z via the natural intertwining action (for which E is explicit). We take the jointly equivariant activation a,, (z, 2) = (z z). V(z, 2) E X x Z with a : R - IR sigmoidal function applied pointwise: and consider normally distributed features, and labels produced from a teacher model f.. This teacher is given by shallow NN model, either f. 4 with
N. = 5 arbitrary particles 6* E N., or its symmetrized version f. = QG. N- (referred to as WI), with 10 particles. I Notice that the data distribution π will be equivariant only the teacher is. We try to mimic such teacher with student model, T with the same a., but different particles 8 E N that will be trained to minimize the regularized population risk RT (with quadratic loss and penalization). For this we employ the SGD dynamic given by Equation (I) (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. Fach experiment was repeated N, = 10 times to ensure consistency. Explicit values of the fixed training parameters are found in SuppMat
4.1 Study for varying N
We demonstrate how properties of WGF(R stated in Section 3.3 become apparent as N grows. We consider a teacher with - either arbitrary or WI, and different training schemes performed over N₂ epochs, all initialized with the same particles drawn from given (40 E PP222) that is either SI or WI. Figure 1 displays the behavior of the student's particle distribution, VNN. after training, in terms of certain normalized version' of the Wy-distance, which we call simply Relative Measure Distance, or RMD for short. [
We refer to SuppMat- for further insights and, additionally: a deeper
4 An additional variant, with f- = T. having SI particles, is given in SuppMat-
It is roughly equivalent W2 when far from the do measure, see SuppMat for details
8

---

(a) SI initialization and arbitrary teacher
(b) WT initialization and arbitrary teacher
(c) SI initialization and WI teacher
(d) WI initialization and WI teacher
Figure I: RMDs, at the end of training, for N = 5, 10, 50, 100, 500, 1000, 5000 and the vanilla (V),
DA, FA and EA schemes. The first plot of each position displays either RMD ( (NN)) or RMD² N ) depending on initialization (either SI WI); and the second shows the RMD between DA, FA and vanilla schemes.
analysis of the case of vg- being SI, comparisons between different techniques and EA,, and I comparisons between (.) and both f. and Qg.f+ (to illustrate Corollary
We first look at the ST-initialized training. Though from 30 we know that (exact) DA or FA during training will respect & without needing pass to the MFL. This is certainly true in general for the vanilla scheme, where the symmetry is never explicit for the model. We notice in Figure] however, that, as N grows big, the Sl-initialized vanilla training scheme, under only a WI teacher, does remain SI throughout training, as predicted in Theorem S This is absolutely remarkable, since there no intuitive reason why the vanilla scheme (were parameters can be updated freely) shouldn' 't escape 0 to better approximate f.. For instance, for an arbitrary teacher (with the same particles, but un- mn-symmetrizeed) vanilla training readily leaves E better approximate f.. Though this isn 't predicted behaviour from our theory. it motivates heuristic we present in the upcoming section. On the other hand, and as expected, both DA and FA consistently remain within E 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
Regarding the W1-initialized training, unlike the SI case, particles sampled i.i.d. from WI distribu- tion don't necessarily yield WI empirical distribution vi On the one hand, this means we require large N to see V being (approx.) WI; and on the other hand, it means we have no guarantee that DA and FA will be close unless we look larger N (where Theorem apppiees). The second column of Figure precisely shows these behaviours as N grows; both trend of V towards becoming WI, and clear coincidence between the DA, FA and vanilla schemes (the latter only for equivariant f=).
4.2 Heuristic algorithm for discovering EA parameter spaces
From these experiments, for non-equivariant f., the Sl-initialized WGF is seen eventually escape gG In turn, for equivariant f., Figure 2 shows that a training scheme initialized at F C g (i.e. 18, E P(E)), eventually leaves E, but stays within go (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 Z supporting SI measures. We hence consider some (large) N, a shallow NN model with e.g. o. (z) (2.2), and numerical thresholds J)) C R4 We define E. {0) < & initial subspace and initialize viy. vg P(Eo). Then, iteratively proceed as follows:

---

Student Particies' Postioons
Student Particles Positions
Student Positions
Teadeer Purlois Fanidess
Step.0
Toop ve Panoes Stadert Fartilss
Teadter Paaiok
Stup
Soude Pantdas
40
41
Figure 2: Heuristic method applied on teacher (squares) and student (dots) particles. Row I: acrial view of hyperplane ε Row 2: parallel view, to verify student particles always remain in E (red line). Column 1: step j = 0 after training; particles leave F. {0). Column 2: initialization of step j = 1 on F1 = (vE.)- Column 3: step j = 1 after training: particles leave F1 (Row 1), but not 3
For j = 0.l.... initialize a model at Vi E P(Ej), train for Ne epochs, and check whether RMD² (vN PE, #vN.) S 0j. If that is the case, the training didn't escape Ej, and one could suppose go = Ez. Otherwise, it left Ej (so EG * Eş) and one can set e.g. Eg+1 == E, @ va,, with eg == N E (0.-- PE. .)..). Allegedly, this scheme would eventually leave all strict subspaces to reach the 'right' g Figure 2 indeed illustrates this behaviour in our simple teacher- student example (see Supp.Mtt-F2 for further details). Notice that we start knowing close to nothing about data symmetries (Ei- {0}), and end up 'discovering' data-based parameter-sharing scheme (E. = 3 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 722 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 sunmetry--beaking in ML problems. A deeper comparison between both approaches shall be found in SuppMat- F22
5 Conclusion
In the light of theoretical guarantees given by the MF theory of overparametrized shallow NN, 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 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 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.
heuristic to
10

---

Acknowledgements
JM thanks partial financial support from ANID Magister Fellowship 22231325 and CMM-Basal Grant FB210005, ANID-Chile, JF thanks partial financial support from Fondecyt Regular Grants 1201948/1242001 and CMM-Basal Grant FB210005, ANID-Chile. Both authors thank Roberto Cortez, Daniel Remenik and Felipe Tobar for comments and discussions which motivated part of the numerical experiments developed in the paper, as well as the refinement of some ideas. The authors also thank Daniela Bellott for her assistance in running part of the numerical experiments.
References
[1] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows in Metrie Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser, 2. ed edition, 2008. OCLC: 254181287.
[2] Jimmy Aronsson. Homogeneous vector bundles and g-equivariant convolutional neural networks. Sampling Theory, Signal Processing, and Data Analysis, 20(2), jul 2022.
[3] Dominique Bakry, Ivan Gentil, Michel Ledoux, et al. Analysis and geometry of Markov diffisiion operators, volume 103. Springer, 2014.
[4] Raphaël Barboni, Gabriel Peyré, and François- Xavier Vialard. Understanding the train- ing of infinitely deep and wide resnets with conditional optimal transport. arXiv preprint arXiv:2403. /2877, 2024.
[5] Andrew Barron. Universal approximation bounds for superpositions of a sigmoidal function. ieee trans. on information theory 39, 930-945. Information Theory, IEEE Transactions on, 39:930 - 945, 06 1993.
[6] Benjamin Bloem-Reddy, Yee Whye, et al. Probabilistic symmetries and invariant neural networks. Journal of Machine Learning Research, 21(90):1-61, 2020.
[7] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223 311, 2018.
[8] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges, arXiv preprint arXiv:2 2021.
[9] Pierre Cardaliaguet. Notes on mean-field games (from P.-L. Lions lectures at Collège de France). 2013. Available at: https://wwwn
httpa://sww.cerremade ceremade...eeeeeeee dauphine fr/ htth//aaraaalgguee
MFG20130420 pdf
[10] R. Carmona and F. Delarue. Probabilistic Theory Mean Field Games with Applications Mean Field FBSDEs, Control, and Games. Probability Theory and Stochastic Modelling. Springer International Publishing, 2018
[11] Louis-Pierre Chaintron and Antoine Diez. Propagation of chaos: A review of models, methods and applications. I. models and methods. Kinetic and Related Models, 1565) 2022.
[12] Fan Chen, Yiqing Lin, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for kinetic mean field langevin dynamics. Electronic Journal of Probability, 29:1-43, 2024.
[13] Shuxiao Chen, Edgar Dobriban, and Jane H Lee. A thoueeth framework for data augmentation. Journal of Machine Learning Research, 1-71 2020.
[14] Zhengdao Chen, Grant Rotskoff, Joan Bruna, and Eric Vanden -ijdde-. A dynamical central limit theorem for shallow neural networks. Advances in Neural Information Processing Systems, 33, 2020.
[15] Lénaïc Chizat. Mean-field langevin dynamics: : Exponential convergence and annealing. Transactions Machine Learning Research, 2022.

---

[16] 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.
[17] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in Neural Information Processing Systems, 32, 2019.
[18] Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant CNNs on homogeneous spaces. Advances in Neural Information Processing Systems, 32, 2019.
[19] Y. Le Cun, B. Boser, J. S. Denker, R. E. Howard, W. Habbard, L. D. Jackel, and D. Henderson. Handwritten digit recognition with a back propppropnnttion network, page 396-406 404. Morgan Kaufmann Publishers Inc., San Francisco, CA. USA, 1990.
[20] George V. Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of Control, Signals and Systems, 2:303-314, 1989.
[21] Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. kernel theory of modern data augmentation. In International conference on machine learning. pages 1528-153.. PMLR, 2019.
[22] Valentin De Bortoli, Alain Durmus, Xavier Fontaine, and Umut Simsekli. Quantitative propaga- tion of chaos for SGD in wide neural networks. Advances in Neural Information Processing Systems, 33, 2020.
[23] Arnaud Descours, Arnaud Guillin, Manon Michel, and Boris Nectoux. Law of large numbers and central limit theorem for wide two-layer neural networks: the mini-batch and noisy case. arXiv preprini arXiv:2207. 202.22277.2724 2022.
[24] J. Diestel and J.J. Uhl. Vector Measures. Mathematical surveys and monographs. American Mathematical Society, 1977.
[25] C. Drutu and M. Kapovich. Geometric Group Theory. Colloquium Publications. American Mathematical Society, 2018.
[26] Paul Dupuis and Richard Ellis. weak convergence approach to the theory of large deviations. John Wiley & Sons, 2011.
[27] Bryn Elesedy and Sheheryar Zaidi. Provably strict generalisation benefit for equivariant models.
In International conference on machine learning, pages 2959-2969. PMLR, 2021.
(28) Stewart N Ethier and Thomas G Kurtz. Markov processes: characterization and convergence. John Wiley & Sons, 2009.
[29] Mare Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing cquivariant multilayer perceptrons for arbitrary matrix groups. In International conference on machine learning, pages 3318-3328. PMLR, 2021.
[30] Axel Flinth and Fredrik Ohlsson. Optimization dynamics of equivariant and augmented neural networks. arXiv preprint arXiv:23033. /3458, 2023.
[31] Qiang Fu and Ashia Wilson. Mean-field underdamped langevin dynamics and its space-time discretization. arXiv preprint arXiv:2312.16360, 2023.
[32] Wilfrid Gangbo and Adrian Tudorascu. On differentiability in the wasserstein space and well- posedness for hamilton-jacobi equations. Journal de Mathématiques Pures et Appliquées, 125:119-174, 2019.
[33] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. org
Deep Learning. MIT Press, 2016.
[34] Philipp Grohs and Gitta Kutyniok, Mathematical aspects of deep learning. Cambridge Univer- sity Press, 2022.
12

---

[35] Karl Hajjar and Lénaïc Chizat. On the symmetries in the dynamics of wide two-layer neural networks. Electron. Res. Arch., 31(4):2175-2212 2023.
[36] B.C. Hall. Lie Groups, Lie Algebras, and Representations: An Elementary Introduction. Graduate Texts in Mathematics. Springer, 2003.
[37] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks, 2(5):359-366, 1989.
[38] Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. In Annales de I' I'Institut Henri Poincare (B) Probabilites el statistiques, volume 57, pages 2043-2065. Institut Henri Poincaré, 2021.
[39] Kevin H Huang, Peter Orbanz, and Morgane Austern. Quantifying the effects of data augmenta- tion. arXiv preprint arXiv:2202.091 2022.
[40] Ningyuan Huang, Ron Levie, and Soledad Villar. Approximately equivariant graph networks. Advances in Neurul Information Processing Systems, 36, 2024.
[41] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances Neural Information Processing Systems, 31, 2018.
[42] Olav Kallenberg. Probabilistic Symmetries and Invariance Principles. Springer, Dordrecht, 2005.
[43] Olav Kallenberg, Random Measures, Theory and Applications, volume 77. Springer, 01 2017.
[44] Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International conference on machine learning, pages 2747-2755. PMLR, 2018.
[45] Leon Lang and Maurice Weiler. A wigner-eckart theorem for group equivariant convolution kernels. In International Conference on Learning Representations, 2020.
[46] Zhiyuan Li, Ruosong Wang, Dingli Yu, Simon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora. Enhanced convolutional neural tangent kernels. arXiv preprint arXiv:191I.008 2019.
[47] Pierre Luis Lions. Cours au College de France. 2008
[48] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem- Reddmy On the benefits of invariance in neural networks. arXiv preprint #wXiv:2005.00/78 2020.
[49] Haggai Maron, Heli
Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint 1rXiv:1812.09900 2018.
[50] Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant networks. In International conference on machine learning, pages 4363-4371. PMIR 2019.
[51] Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pages
PMLR, 2019.
[52] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pages 3351-3415 PMLR. 2021.
[53] Song Mei, Andrea Montanari, and Phan-Minhh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115633):E7665- E7671, 2018.
[54] Sylvie Méléard. Asymptotic behaviour of some interacting particle systems; McKean- McKean-Vlasov Vlasov and Boltzmann models, pages 42-95. Springer Berlin Heidelberg, Berlin, Heidelberg, 1996.
13

---

[55] Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. Mathematical Statistics and Learning, 6(3):201-357, 2023.
[56] Tomohiro Nishiyama. Convex optimization on functionals of probability densities, arXiv preprint arXiv:2002.064480 064888 2020.
[57] Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pages 9741-9757. PMLR, 2022.
[58] Otto and C. Villani. Generalization of inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis, 173(2):361-400) 2000.
[59] Mircea Petrache and Shubhendu Trivedi. approximatioon-eeeraaiiiznn trade-offs under (ap- proximate) group equivariance. Advances in Neural Information Processing Systems, 36, 2024.
[60] Siamak Ravanbakhsh. Universal equivariant multilayer perceptrons. Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7996-8006. PMLR, 13-18 Jul 2020.
[61] Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parameter- sharing. In International conference machine learning, pages 2892-2901. PMLR, 2017.
[62] Grant Rotskoff and Eric Vanden-Eijnden Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75(9):1889-1933 jul 2022.
[63] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[64] John Shawe- Taylor. Threshold network learning in the presence of equivalences. In J. Moody, S. Hanson, and R.P. Lippmann, editors, Advances in Neural Information Processing Systems, volume 4. Kaufann, 1991.
[65] John Shawe-Taa Taylor. Sample sizes for threshold networks with equivalences. Information and Computation, 1 118(1):65-77) 1995.
[66] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3);1800-185,, 2020.
[67] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):725-752. 2020.
[68] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. Mathematics of Operations Reseurch, 47(1):120-152, 2022.
[69] Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynam- ies: time-space discretization, stochastic gradient, and variance reduction. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[70] Alain-Sol Sznitman. Topics in propagation of chaos. In Paul-Louis Hennequin, editor, Ecole d' E'éé de Probabilités de Saint-Floun XIX 1989, pages 165-251, Berlin, Heidelberg, 1991. Springer Berlin Heidelberg.
[71] Cédric Villani. Optimal transport. Old and new, volume 338 of Grundlehren der mathematis- chen Wissenschaften. Springer- Verlag, Berlin, 2009.
[72] Rui Wang. Elyssa Hofgard, Han Gao. Robin Walters, and Tess E. Smidt. Discovering symmetry breaking in physical systems with relaxed group convolution. In Proceedings of the 41st international conference on machine learning (ICML2024). 2024.
[73] Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable CNNs. Advances in neural information processing systems, 32, 2019.
14

---

[74] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML2011) I), 2011.
[75] Jeffrey Wood and John Shawe-Taylor Representation theory and invariant neural networks. Discrete Applied Mathematics, 69(1):3360, 1996.
[76] Dmitry Yarotsky. Universal approximations of invariant maps by neural networks. Constructive Approximation, 55(1):007-474, 2022.
[77] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in Neural Information Processing Systems, 30, 2017.
15

---

Table 1: Summary the main abbreviations employed throughout the paper
Abbreviation
Meaning
NN NNN
Neural Network / Neural Networks
SL
Symmetry-Leveragin
MF
Mean Field
MFL
Mean Field Limit
WGF
Wasserstein Gradient Flow
SGD
Stochastic Gradient Descent
DA
Data Augmentation
FA
Feature Averaging
EA
Equivariant Architectures
WI
SI
Atrongly-Invarian
A General Considerations for the Reader
This section presents summary of recurrent notation, abbreviations and key concepts used in work, as well as discussion on its limitations, scope and possible extensions. We thank anonymous Reviewers for suggesting us to add this section to the original manuscript.
A.1 Summary of recurrent notation, abbreviations and key concepts
In this section we summarize the main abbreviations and notation employed throughout the body of the paper, as well as simple definitions of fundamental concepts from probability theory and algebra that might be useful for understanding our work
Table serves as a glossary for the most used abbreviations in the body the paper. Fable@contains 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 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 xx is some measure X, then say f is square-integrable (in Bochner sense) if it satisfies: Ix (x) (Z)dxx(()) < 00.
This integral also respects closed linear operators, shown by Hille's theorem (see Theorem I1.2.6 (24). Namely, if y > y closed linear operator over } we have: Sx I.fdTx = L. Lx fdrx. 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 due clarification may be required. We assume G be 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 2, which we denote G C Z, whenever there exists map:
T:GxZ + Z
2g.2) + T
that satisfies 2'(GG,))= and T(g T (g2. =)) = T'911922, =), Vg1.92 E G, Vz E Z. We always assume the actions be continuous; namely, T' continuous with respect the product topology. Alternatively, by denoting T'2 = -'(g, -), T. G T, Sym is group homomorphism and, if the action is continuous, Tg an homeomorphism Vg G. If we further assume that Vg E G, T, is linear, we call T. a group representation. Further, if (T, ( BEG are orthogonal transformations, call I. an orthogonal group representation
16

---

Table 2: Summary the notation employed throughout the paper Represented Object
Notation
X, K.
Respectively, feature, label, and parameter spaces (separable and Hilbert). Usually just X' = Fd y = IR and Z = IP' for c, d, D E N".
(*))
Space of Borel probability measures on the metric space (*)
T, NX
Data distribution P(X' x )) and its marginal on X.
l
Convex loss function L x
a.
Activation function or unit, a, 2' x + V. Not necessarily corresponds to the usual activation function from traditional NN implementations. Shallow neural network model with N units and parameter 0 E g
Dirac measure on point.. :. Recall that, for a function yo integrating gives (4p, 6.) = (p.). Empirical measure of 0 in P(Z). i.e. vg = + > Ci-1 da,
Shallow model with associated measure μ E (P(Z). Also denoted (ax,14
R(w)
Generalization error or Population Risk. Abusing notation, it either represents
E. [((@)) (X), Y)] for 0 E ZN or E, (X), for μ E P(Z)
(P(Z)
Space of shallow models with unit 0+ and measures P(Z)
P2(2)
Probability measures μ E P(2) with finite p-th moment: 1z (1(0))(dd) < +20.
p-th Wasserstein distance, defined W (,(,, == [inf E,[ YIP]] 3 V4, v E (2) where [I(,, v) is the set of couplings between μ and v.
oF
Linear Functional Derivative of F (2) - R as in Definition
dy
D
Intrinsic Derivative of F (2) I, in Definition
WGF(F)
Wasserstein Gradient Flow for F P22 [)) - IR, as in Definition
Some penalization function r Z + E+
A
Lebesgue Measure Z
p <
μ. E P(z) is absolutely continuous w.r.t X.
HA
Boltzmann entropy μı; Hx( := log( du (z))dp(=) if < X or too otherwise. dd
R
Penalized population risk, given by R P (M) = R(n) T rrlp + BH)()), with
8 2 0 the regularization parameters.
Space of functions f X' > y square-integrable (in Bochner sense) w.r.t xx Step-size (or learning rate) used during training. Parametrized as s = EN (kEN)) with a regular function IR+ - R2, and some eN > O; with k; N E N. Compact group and its normalized Haar measure, which the unique right and left invariant measure w.r.t. the group multiplication, Orthogonal representations of the action of G over X. y and respectively. We also denote these actions as G Op X. G Cp y and G C. Z respectively.
G, Ag
(43).
P, P.
Qa
Symmetrization operator over (xx, (Qc.f)(z) == Pg. x(d))) zd)aa(g) for any f E (X V Tx).
gG
Space of fixed points for OM Z (i.e. 2 2 s.t. Vg AgeGr G Mg.
Pga
Orthogonal projection from onto EG; Pgo.2 .2 =
Pushforward measure μ E P(Z) via ʃ 2 -- Mg. d) for all E 2. given by f#u(.) (-)) Symmetrized and projected versions of p. E P(Z): n = (M, uDdAG and
_ = Peo
(4())
Space of G-invariant measures over space e.g. for G Cm Z, the space of all E P(Z) s.t. Mo## for all E G. Space of measures concentrated on 8 i.e. μ E P(Z) s.t. ((EG)= =. Population risks that are optimized when applying different SL techniques; DA, FA and EA respectively.
P(EG)
RDA, RPA
REA'
Space of functions f: + y square-integrable (in Bochner sense) w.r.t 7x that are also G-equivariant. In other words, shorthand for Qo (L2 ((,,)))xx) Almost everywhere w.r.t. the Lebesgue measure dt on R4,
dt-a.e.
Aga Rid
Shorthand for the Lebesgue Measure over EG
Projected-regullaizzee functional given by R5d (14) == R(n)+ +r rddu+331 (10²)
RMD
Sllative-Maasur--Distaaneee RMD2 (H, v) M2 4
W where M3 = 2 2 du)z).
Ar
PP(2)((0,2))
Skorokhod space of càdlàg functions from [0, to P(Z).
17

---

and denote the group action by G Cr 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. 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 equivariant linear maps also as intertwining maps or infertwiners. This is also why we refer to the layer-by-layer action in EAs for shallow NNs as intertwining action (see e.g. SuppMar.C.T)
[ 48 Beyond NNs, whole field of mathematics is of course
Finally, it is well known that a compact group G admits a unique normalized Haar measure Ào E P(G) (see (25). which is left and right invariant, finite 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 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 (Ma) MeE P(2), we say it weakly converges to some point fr E (P( if, for any continuous and bounded function f Z - R, P). Notice that this type of convergence is weaker than W, we have (J, Hm. n1000
convergence for p 2 1 (which additionally also requires the convergence of p-th order moments of the involved measures).
Notice how we write (f, (1) to denote J2 fdu. This notation is heavily used (and abused) throughout the core of our work.
A.2 Limitations, scope and possible extensions of our work
For convenience of the reader we here present a discussion of some limitations and of the scope of application of the present work. This subsection does not provide any mathematical results required for the sequel.
Some of these limitations are of technical nature, and regard specific assumptions made in order for the specific proofs of our theoretical results to hold. In absence of these conditions, or under less restrictive ones, some of our results (or weaker forms of them) might still hold true, but further research is needed in order to properly establish their validity. Other limitations are the object of more general research questions in this area.
On the infinite i..d. data sample: The MF theory of NN makes the assumption that it is possible to take an infinite i.dd. 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 [ which holds for ēN = o(/NN, ensures that a sample size of the order of IN is required in order to approximate the WGF up to ("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 data sample will be required in order to attain predetermined generalization error level. On the other hand, even for π with finite support, the MF approximation will end up minimizing the empirical- risk. This, as pointed out in Section 2.1 is not the same as minimizing generalization error w.r.t. some underlying data distribution, but is widely interpreted (in the application of most NN-based machine learning algorithms) as 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: requires that the width N goes to +oo), 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
Despite the MFL being a theoretical tool (which
18

---

MF analysis of NN are, in fact, unaccessible from fixed N standpoint (as is the case for Theorem [5. Further asymmtottiiic and quantitative answers to practical questions (e.g. how many neurons are needed in order to attain given level of generalization/popullaa error at given computational cost?) could also be obtained from MF theory, via quantitative propagation of chaos results (see e.g. [2).
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 f, as well as the boundedness o., 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 c.g. or 38 for some alternative conditions).
Regarding the assumption that is G-invariant, it known that traditional loss functions naturally satisfy this condition. The joint-cquivariance of o., on the other hand, is much less constraining than it may initially seem. In practice, depending on the choice of a and M, it might even end up being a trivial constraint. A deeper discussion on this very assumption shall be found in Ca.
On the generality of shallow models + (P():) These allow for modeling a wide range of situations (including some variants of multi-layer models). In fact, d, can by itself encode a complex deep architecture (see SuppMat C.I) and the resulting shallow model can represent way more interesting structures than single-hideenllayer 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, 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 Jand 4 can be extended to some of those settings, which is a question we will leave for future work.
Universality guarantees for shallow models $,, E Fr. (P(2)): In this work, we have only explored the simple setting of tensor-of-order- / 'NNs as modelled through a., see SuppMat- C.I 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, limitation the applicability of Proposition 5). Exploring other interesting situations that could be modeled with our current framework, or plainly reformulating it 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 related discussion.
The assumption of being G-equivariant implies that x is G-invariant as well (6. This implication can seem 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 17 30 48 ), assuming गं to be G-invariant 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 't explicitly rely on having क E P x )), but rather on the risk functional R: P(Z) + IR 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 G-equivariant data distributions:
On the numerical experiments: The suite of experiments that were presented in Section 4 and Sup 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 he 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 currently limited computational resources,
19

---

leave these inquiries for future work. Similarly, taking our experiments to larger scale (e.g. N 100.000) is also challenge left for future inquiries (our current results, with N S 5000, might still be considered as small N). A significant problem for scaling our experimental verifications the fact that we relied on calculating Wasserstein distances (which is usually really computationall--exeesssii to provide rigorous numerical evidence of our theoretical findings.
Furthermore, we 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. ReLU, tanh, and many others), as well as variants of the optimization algorithm (e.g. Adam, RMSProp. etc.).
Finally, our heuristic needs be tested on 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 dynamie: In this work we focus mostly 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 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 and 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 work straightforwardly in the setting of [22]. where the MFL is established even with 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 and 6 which involve stronger notions of symmetry, don't 't seem immediate generalize many other kinds of asymptotic dynamics. Studying the applicability of our results under different variants of the training process, is surely an interesting question be tackled as part of our future work.
B Symmetries in functions, measures, data and shallow models
In this section we state some useful, basic results on the effect of symmetries acting on measures, functions and data, that will be used in the sequel. We also explain how symmetries of interest can be incorporated into the generalized shallow NN setting from Definition] | 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 a compact group with Haar measure Àg, such that G Gp X, G Cj y and G CM Z. Also, let E C Z be the linear subspace of parameters that are fixed points of the action G over Z, and Pgo the orthogonal projection onto it.
B.1 Differentials and integrals of equivariant functions
The following lemma characterizes the differential of jointly equivariant functions with respect the action of some group G. Here we can assume G be a lesH group w.l.o.g. and consider representations that aren't 't necessarily orthogonal (we denote them differently to avoid confusion).
Proposition 6. Let G Cx X, G Cg Z, G Ce y via some representations X, X and x respectively (not necessarily orthogonal). Let J: X x 2 + y be jointly G-equivariant with respect these actions (ie. Vg E G, Vx E X, Vz E 2, ((,,) = 1(,, Xy-=)) and Fréchet-differeniie on its first argument. Then:
Vg e G. Vz X, Vz Z, Xx Dzf(,22)))
20

---

Proof of Proposition 6 Indeed, we know that A2 E 2 (+,(())) is the unique function that satisfies, Vž E X:
[I + 2) -(3)-
Da.
lim h-0
=0
|
Since we want to prove that Ví X. V2 2, Ag E G:: Dzf Xg. KfDaf( x)xg', will be enough to check that:
11 + h xg²) J(xg Xg.2) Dr (2x, 2)
lim h 0
||/||
which by uniqueness implies the result. Thanks to the joint equivariance of f, we have vh + 0:
(X h h, [(xg Xg.2) X Dz 2 x)1
|//|
1f((g-(i + Xg .1), x2 - f(xg Xg.2) - X Dz (i, 2)x hl
[1/]
1ig + Xg 1xx,,,,h =) X f
Dz x)x
|//
1xg. [f(ē Xg h, 2)
f(i, 2) Du Duf(ā, 2 2x h1]
IlxgXXo .h1l
Now, recall that for every g G. the operator X9 is bounded, i.c. it has finite operator norm 0 < I1xell < oo (non-zero as Xg is invertible). By defining = Xg .h, we have:
1xx- [[( + h, [(,, x) Dz DZ,(
[Ixgllll(( + h, f(8,2) D 2)
h1l
Ilx
Multiplying by 1|x xobl the last term seen be bounded by [M]l]
Xx II5( h, f(Ø, =) D D,,
|A
Since Xg and Xg
bounded operators, have that: h + = h X₂ + O. Thus
llf(x. + h, xg²) - (Xg. R. 2) - Dx f20 2)x; '/ll
lim k-0
|
- lim 1x 12-1|| 11ʃ(E + h, - =(,
h|ll
This, in particular, allows to characterize the differential of equivariant functions.
Corollary 4. G Cx X, G Cg V, and f : X + y is a G-equivariant and Fréchet-differentiable function, then:
Vg E G, Vo E EX = D=)x2(x)>
Proof of Corollary
Direct.
We can also get some interesting integral properties of jointly equivariant functions.
Proposition 7. Let X,y y and Z be (separable) Hilbert Spaces and G a lesH group. Let Cy X, and X respectively (not necessarily orthogonal).
G GO2 Z, G Ox y via some representations
Let f: X Z + y be a jointly G-equivariant function (with respect to these actions). Consider measure μ E P(Z) and let f be Bochner integrable on its second argument with respect to μ. Then:
VI X, Vg G, (f )) (f (x)z;) 0,xotp)
21

---

Proof of Proposition 7 Let P(2) and be as stated. Notice that, V2 E N, Ag E G, we have:
Ka (, Xa ()du(()))()))))))) = SfaXa Xaf(2, 0)d,((),
where we' ve used the linearity of the Bochner integral under continuous linear operators (as is Xg). It follows, from the joint Geequivariance of f and the definition of the pushforward measure, that:
J (2
0)dp(() Xx((x,0)du(0)
/ f(gg X, - f(xz2, ()d(()))(()(()))
We conclude the desired result.
B.2 Properties of symmetric measures
Consider again compact G acting orthogonally over the different spaces. Recall that, given E P(Z), we defined μb fa (M pJddGG as its symmetrized version and Fº := Pec## 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(EG) (2). Also, for any μ E P(2) the following equalities hold Vy E G: o = =Mg##u) = (u = -1 and (Mz# = n.
Proof of Lemma 7 Let E P(EG) (i.e. 1(EG)= =1, ge E and consider any positive measurable f: Z + R. We can see:
| [(M [(Mgz)u(dz)) Isc s(>))((=)= = (oniii)s_((ewwaa [()))((a)).
So, that Vg E G. μ = My and thus μ. E (2). Regarding the equalities, consider μ E P(2) and A E Bz some borel set. Since the AG is right- right-invariant, we have Vgg E G, V2 E Z
Peo = 1 MM(MM))dDa(()) / (Mn = [( M1z)dda(hh) PRaz.
Then, for g E G. P (A) = (2P = (Pec) and so:
(Mg#1) (A) = M(M, Ped (A)) = 1 ((PEd(A)) = (A),
(A),
and:
(A) (A)= PTd (A))dAa(g) / (lPdd (A))dAo(o) = plP (A)) 1000 (A).
On the other hand, since _ E pp (2) C (2), (.) leaves it unchanged: (108)c- = 10000
For the last equality, let g E G and f 2 R4 be measurable. We have:
(S. (.(M = / (S. M, (2 1 (...M) -1 once again by the right- instriinvariance of AG. Namely, (M, cc
Proposition 8. For P(Z), we have: fhth E (Z) and Sd P(EC).
Proof of Proposition Indeed, let h E and E Bz (borel set of 2), using the properties of M and the left-invariance of AG, we get that:
a ((M5+(B) - [
( (0)= = (B))
, So, Vg E G, _ = Mg implying that प E (Z). On the other hand, by definition we have
(as the projection is surjective) いº (EG) u(Pdd (EG)) μ(Z) 1, that n E P(EG). D
22

---

Remark. It's not hard to notice that, on 2 = RD and with X being the lebesgue measure. if a E P(Z) is such that μ < X and has density u : Z + R, then: μ (2) has density u S u (g) w.r.. À (whereas u² doesn't 't have a density w.t. À unless the action is trivial). This follows from the O(D) invariance of X and some standard calculations.
Since we will be working on P2(Z) on Section 3.3 it's useful also notice that:
Remark. If Pp(2), then n² 11 E Pp(2). Indeed, it follows from the fact that M| I 1 Ag E G (since the representation is orthogonal) and || S (since Pgo is a projection).
Also, we have that:
Proposition 9. :(E)))=P(EE P and PC (Z) = pG (2) are closed and convex subspaces of (2) (under the lopology induced by w,).
Proof of Proposition Convexity is direct by definition of the involved spaces. For closedness under the Wasserstein topology, recall (see e.g. (I] that Wp Wp((nnHH) (b) = > is equivalent to: M4n n-00 f n-
Sz d(()) Since for f E Co(Z), f Mg (for g E G) is continuous and bounded, (Un)neN E pC (z) implies Vg E G. Mg## = μ (namely, J E pG (Z)). Similarly, f Pgo is continuous and bounded, and so if (un)nEN S P(E)), then Pea μ (i.e. μ E P(EG)).
(weak convergence) and f& (()))))((0))
B.3 Properties of equivariant data
We are representing the idea of 'data being symmetric" with respect to the action of G, by assuming the data distribution be equivariant. We will next see that this is natural generalization of more intuitive, though restrictive, notions of data being symmetric, for instance when X is r.v. on X' with invariant law म and Y ft (X) + & for some equivariant function fr : X + Y and some centered independent noise Ç. Indeed the following results tell us that, assuming π E pa (X' x V), implies such 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 π E P(X x y) be an equivariant data distribution such that E, ||||| co. Then f = F. ( is an equivariant function.
Proof of Proposition 10 Indeed, as [[ < co, we know the conditional expectation [[ well defined and there exists measurable f X + y s.t. f'X) []] Now, by properties of the conditional expectation): Given any h X' + R square integrable, we will show that: W.,r, (Y h(X)] Ex 2,
Indeed, notice that by Fubini's theorem (as l E (22 ); mx) and h square integrable), linearity of the integral and G-invariance of m:
E, ((,x)M(())X = E, V. Dg' (r Xa0a)))))
Pr FE (X - -x)]dA6(()
G 1 ..hpyghhp................ -yhlpgg-1xx](c -x)]jAAc(9)
E. (Y(((()))))))))) E, (V)(X)]
By uniqueness of the conditional expectation. we know therefore that fxX)
P ( ()) In particular, a.e. f = Qo(f*), making f* G-equivariant.
As a particular example, if we assumed that the data distribution was given by some function f - y, i.e. y= f(X), taking to be equivariant would be equivalent to assuming that क. is invariant and f an equivariant function (which is the setting of the data simulated our numerical experiments; see Section 4).
23

---

We notice that Proposition 10 can also be used to recover a celebrated result from 227 (later generalized by [40], in the general setting where data symmetry encoded by the condition that E P2(X x J). Define the symmetrization gap of learning problem with quadratic loss as:
A(f, Raf) == Ex - f(x)133) - E (Raf)
The following extension of mentioned statements from 27
40 is not needed for our results, but we
provide proof of in SuppMat-Gi in view of its potential, independent interest:
Lemma 2 (Symmetrization Gap Characterization). Consider the quadratic loss and P(X x y) such that En || < oc. Also, assume that σ|x G-invariant, but π is only H H-invariant with respect to some H S G (closed). Then, the generalization gap satisfies:
(s, 2cf) 22f" ( (P(X))x))))) + llfc ([
(eo)=Rer where - [V|X = x²] is the conditional expectation functionn, and fā == = Qas. In particular, if x is G-invariant as well, we get A(f, Qcf) lIfé ll
C Concrete realizations of shallow models
The class of models we have introduced in Definition T allows for taking an arbitrary, common parameter space 2 for all hidden units, as well as an arbitrary function Ø. : X x 2 + y. As noted in the description EAs in Section 2.3] by requiring only that o,, is joinly-equivariant, we moreover ensure that G CM Z is. in some sense, properly related to the actions G Cpp X' and G C3 y. This abstract property of a. allows us, in fact, to encode a wide range of situations and interesting results, without delving into the specifics of particular choice of architecture. We next analyze this notion the concrete example of the setting of a traditional single- hidden-layer shallow NN.
C.1 Traditional single layer neural networks and large ensembles of multi-layer units
Recall the sitite-dmensionall setting of geggle-hidde--aee neural networks. We considered X' = Rd = R and Z Rexb x Rdxb x R², and defined, for = (W, A, B) E Z and Rº + R (=, == W Wo(ATo + B). This allows us to express shallow NN with N hidden units, of parameters 0 (0 E 2 with 02 (W,, A,, B:), as the function T. X - y given by:
Va E X, 1N (x) = N Σ Woodd N (A x a+B),
(6)
1=1
where we write the expression by blocks, considering
B
W = (WI..... WN) E RoxbN A= -A AN) E RdxbN and =
E R
BN
This corresponds exactly the usual single-hiddeenllyeer setting (see e.g. 23 53
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 (99) or DeepSets (see (77)) into the shallow NN framework that we propose (see Appendix C22 for fully developed examples).
62
67
We now also consider G-action the intermediate layer, G Cn R² (so that G Cnaidy N This allows to define the natural intertwinning action G on 2, which given by:
My.2 Mg-(W, A, B) == (Py.Wn) 17 Pg.An B).
for any g E and Z = AW, B) E Z. This is exactly the action under which the fixed points (i.c. {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:
'The name is inspired from the traditional definition of 'intertwinning linear maps' in representation theory, SuppMat A.1 for an explanation.
24

---

Proposition 11 (Joint equivariance of a,, for single-hidden-layrr NNs). In the setting of single- hidden-laye NNs described above, if o F² + R² is G-equivariant (with respect to the action given by n)) then is jointly equivariant.
Proof. This fact follows directly from the specific definition of o., the linearity and the orthogonality of all the relevant group representations. Indeed, for any equivariant o R + R², any E G and any 2 (W, A, B) E 2, we have:
T+ (pg , My. =x (Pg.t W.g. Pg.A..g 7g B))
= -ipg W. ) (Pg... A.ng )T (pg x) + B) Dg. - (Mg-A A (Pg -pg 7g. B))
(W. .0 (n ( (A
B)))
= PW. My O x + B))
Po. x B))
= (z,
Thus, in this particular case, in which a. represents the unit of a single hidden layer neural network, we require only for a RD - TR' to he G-equivariant (with respect to the action given by v) on R2) in order for the joint equivariance of o, hold (and, in consequence, most subsequent results from our work).
In particular, if 17 is chosen to be trivial (i.e. 1 = idge), any single-hidden--ayer NN is jointly- equivariant regardless of o : R - R Namely, all of the results contained in the core of the paper can be applied to arbitrary single-hidden-Laye -ayer NN; most importantly, those relating vanilla, DA and FA training. The disadvantage of having a trivial 17 is that the space Eo might not encode very interesting EAs. This might make Theorem 5 lose part of its impressiveness, but it takes credit off the rest of our results (such as Corollary 3 and Theorem 4
Analogously, if acts via permutations of the coordinates in RR', it is enough to consider that is the pointwise application of a scalar function. In practice, this usually isn't It a restrictive condition, since most commonly used NN architectures are naturally built following this pattern. Therefore, almost any single-hidden-llyerr architecture can yield a jointly-equivariant a, for some of the most common and interesting finite groups (e.g. Sn,. Cn. among many others; see [3 for further discussion).
For a more complex (possibly infinite) compact group G acting orthogonally on the data and parameters (with non-trivial 17). for a, to be contty--quivariant we have to start properly restraining
: 2 + Rbb. For instance, choosing an O(b) -uiivequuaar a (e.g. a Norm-Rel.U) would grant Proposition TT to always hold; but such restraining choice could potentially harm the model's 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 IT) readily generalizes to the multilayer case. Namely, if a. encodes Multi-Layer Perceptron (MLP) with multiple hidden layers whose parameters live in some linear space Z, we can define OM Z as the intertwining action between each successive layer (similar to the previous example). a. can be made jointly-equivariant by making every activation function on each hidden layer equivariant (see 30). Then, eo corresponds exactly to the parameters that make the entire MIP an equivariant architecture (in the sense that every layer is an equivariant function). With this, & 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 a. representing a multilayer architecture can be -quiva--quivariiat.
In the upcoming sections we will further develop these ideas, show that some of the most relevant and widely applied equivariant architectures can be realized as part of our setting.
Notice that this fact is explicitly used in the proof of Proposition
25

---

C.2 Shallow DeepSet models, shallow GNNs and shallow CNNs
C.2.1 DeepSets
An emblematic example of neural networks with equivariant architecture are the Deep Sets, introduced by 177;; which correspond, in practice, to NN architectures designed to be invari- ant/equivariant to the action of the group of permutations G = Sn
Consider that our NN processes sets of size n, which contain real-valued vectors of dimension E N. We can represent this input space simply as X' == Igadd Say we wanted to build wide single-hidden-layee network that is invarian//qqiivariatt to the action of Sn, and returns a new set of n vectors, but now of dimension ē E N (i.e. == R=*x), For this, let's consider the same architecture as in the previous section (replace d = nd,, = no), where we'll let our intermediate layer be simply another set of vectors, now of dimension 5 with b E N (such that b = nb above), and repeated N times (as in the multiple hidden units that we want to achieve). With this, our network flows as follows: 4 IRadd + [(xx))x - Raxe Notice how S, acts naturally on each hidden space by simply permuting the vectors of the set (i.e. p. P and 1 are defined in this way).
Following the same structure as in the previous section, we will have a parameter space given by: 2 == R (1xxx)xxxxx) x ((axd)x(nxxb x R=xb and a unit that acts on 2 = (W, A, B) E Z as a. 2x) == (o x + B), with _axb + gnxb some usual activation function applied pointwise (which, as mentioned above, will be immediately equivariant to the defined action of G via 7)) 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. p. P and 1 acting on the sets by permuting their elements, and M defined as in SuppMat-C.l) Ed corresponds exactly to the set of parameters that yield -qquivarri shallow neural network (which can be interpreted interested reader, we will make this connection explicit in the rest of this section.
DeepSet). For the
As shown in (77). the only way to have an -qqqiiaari affine transformation between two spaces Raxd; and Raxd2 is if the parameters Å E ((nxdi)((xxdb) and B E Rnxds (from the definition
AT x + B) are of the form:
A=00l+384
B= (1,....,.
(8)
Where a, 3 E Rd xd2 y E Rd are the truly trainable parameters of the layer; 1 Idax and J= I7 are two n x n matrices; and & is the usual tensor product. More explicitly, writing the matrices by blocks, this is:
a
0
0
R (xxx)))
@81=
0
0
a
B
B
8
3
3
8
=
E R ((nxxd))(xxd))
B
8 B
7 8 (1,...,,1) (7,..., E (mxd))
In particular, Equation { gives us an explicit expression for our space EG Indeed, an element
z = AW,A,B B) E Z will live in go if and only if there exists wn, W2 E Rēx aj 0g E R and b E - such that:
W = wi 1 w2 8 J, A= = a @ I a 8 J, and B bj X (I,...,
In sense, we can think of e simply as being equivalent to ((xb)) x 1 (dxx)) x R² In particular, we went from having D = dim Z) = n² b. (ā + d) + free parameters on each unit, to simply D dim 2 bb 5.6 (ē + d) + b, which should be casier to manage in general.
26

---

Now, recall our construction from SuppMat- C.1 and consider the matrices W, A and B from Equation @ We notice that, in this example, they are of the form:
W E WERI ( A E ((nxd)x(xxxx) and B E IE (nxbN).
Namely, sensibly replacing di = d, d2 = AN for A and B: and dj = č. d2 iN for W in Equation (8); we also get an explicit condition under which W, and B yield a globally S,- equivariant architecture. By properly writing these matrices by blocks (as in SuppMat-C.I and separating the action of each of the N units, one shall notice that the condition for W. A and B to be -qquiivaria corresponds exactly every 02 (We, A., B,) E Z. ...... .N, being of the form given in Equation (9). That is to say, T has an S, -quequvvariant architecture (a DeepSet from (77) if and only if Vi E [1....,NN,66 0g E E 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. A Raxb + R& is the usual linear GAP operation, we might want a NN that flows as DN BRad + RR(xbjxNN A RON + R This is no trouble within our frumework, since we can simply consider the 'adequate' unit: a. (r. 2) == AAo((AT> B)) for 2 (W, A, B) E 2 Raxh x x(xx))x(mxxh) x RR and all of our results would be applicable.
The main disadvantage of doing this is that are forced encode some priori knowledge of the symmetries of the problem into our choice of architecture. While this isn' I useful for 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 a, properly. Namely, set o. to be a whole multi-layer structure, with parameters in 2. and E being the subspace of those that make the architecture equivariant (see e.g. Equation (
In such case, the shallow model D. with 0 E ZN, will represent an ensemble of N multi-layer
units, each one given, for every E (1,..., N), by o.(,0:)) + y. The oulput of the ensemble simply the average of the outpuls of each one of the units.
C.2.2 GNNS
Generalizing the ideas from DeepSets to GNNs is fairly straightforward. Namely, consider the input of a layer to be a graph, represented by a set of features, coupled with an adjacency matrix that also contains relevant edge-features; and the output to be analogous. Namely, let's say the input space is X = (mxdd)) x (*x and the output space is X2 ((xdda) x ((xx))xd, Consider also the natural S,, 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 find an analogous characterization to Equation (8). but for affine layers that are S, -qquiiiarr between graph spaces (see e.g. 22 49 for more explicit construction). From there, it is not hard construct, following the same steps as for DeepSets, an explicit characterization of EG for single-hidden-layer GNNs, analogous to Equation ( Also as before, more complex GNN structures (with multiple layers, pooling operations, etc.) can be encoded in this setting through the unit o. (with the possibly same drawbacks as in DeepSets). For brevity, we don't delve into GNNs in their full complexity and leave that to the curious reader.
C.2.3 CNNS
Finally, following the same logic as before, we can model another one of the most emblematic traditional equivariant models: CNNs. We here consider only ID-CNNs for simplicity, but all arguments can be readily generalized 2D 3D CNNs.
In this setting, take an array vectors of dimension d N as input (X := x and, analogous to DeepSets, say that the output and hidden layers have the same structure, so that ) = =xē; and n for i, ē E N. We consider a, simply before.
Z R ((x)) (inddxx
The single main difference with the study of DeepSets is that, this case, we consider the action of Cn (i.c. the cyclic group of order n. also denoted B,,) instead of Sn, In particular, the natural action
27

---

of Cn on each space (through á and 7)) consists simply of shifting the array's coordinates in a given direction (modulo n). As before, this makes the set g correspond exactly the set of parameters that yield each layer -equivarrin (as in a ID-CNN).
characterize (analogous Equation (8) how single Cn-equivariant layer looks like between two spaces 2 = ((xxd)) and X2 (nx under the natural action. It is well known (see (8) that C, -ququivariant affine layer, + A . B, between Ai and X2 of parameters A E [g(xx)))((xxaa)) and B [2xda, must be of the form:
More specifically,
Ā = C(a and B 3 s (1,
with 3 E IR,, 0z E Rdaxxd, Vi (o,..., 1). and the associated circulant matrix being defined (by blocks, and considering the indices modulo n)
ag
01
a_
00
01
C(ao....,
E R ((xxdx)((xx)))
Okg
Ot-1
As consequence, we get analog of Equation
to explicitly describe Ec: have that
(W, A, B) E and only if ((w)) [-] C R (a:) C Bdxx and B E R such that:
W C(wo. 1n
Clao....... an-1). 1), and
3 @ (1....,
Thus, the parameter space goes from having d2m(2) = n² b (2 (2+ d) each unit, to simply dim(EG) n- (?+ d) b. One might also notice that the obtained parameter space g contains 8S" (from DeepSets), which is expected from the fact that C,2 S Sn In global terms, the shallow CNNs 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.
free parameters on
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 unil d.
D Further elements from the MF theory of shallow neural networks
In this section we study several theoretical notions required in the MF approach to overparametrized shallow NN. For the purpose of our results, Subsections D. and [.22D are the most relevant ones, as we establish therein some useful properties
formula that will be explicitly required. Subsections [D.3 re4iee some results and recent literature regarding well-posedness and long-time convergence of WGFs, which are relevant to the MF interpretation of the training dynamics of shallow models.
D.1 Linear functional derivatives and intrinsic derivatives
Let X and Z be separable Hilbert spaces. Recall the definitions of the linear functional derivative and intrinsic derivatives:
Definition (Linear Functional Derivative (First Variation)). For a functional F :P(()) + R, its 0F linear functional derivative (Ifd) is function: op :P2(2) 2 IR such that Vju, P2(2))
lim F((1 -hpp hv) F(R)
[
OF (M, 2)d(v - ()),, and [ [ 0p
OF (H, (H,
Dy
The function F" P2(Z) OF (H, is also known as the first variation of R at p. By
Definition 9 (Intrinsic Derivative). Whenever oF P2(2) x 2 - R exists and is differentiable on
its second argument, the intrinsic derivative of defined as: D2, = V Op (u,
Example. To better illustrate the notion of the linear functional derivative and the intrinsic derivative, consider the following usual examples:
28

---

1. In the important case of the Boltzmann entropy F Hy, defined for & X by Hx [ log( ( (z)) (and too otherwise), we have that (modulo an additive constant that doesn't 't depend on 2, see (56]):
OF 0u (1, 2) = log
dg (x +1 and (,, dx
dp Yp
(z).
2. Whenever F(u): [2 ((2)dp(()) for some bounded continuously differentiable function
Φ: 2 + R, it is well known that
OF du (14. 2) =0(z)=
Φdp and ((,, =v=((3)
The most relevant example, in our case, regards the linear functional and intrinsic derivatives of the population risk functional, R(p) = R [e( (al(((i((() ()), We can consider the general setting presented in [66,, where some Hilbert Space
considered, and is assumed that F: P()) RR can be written as F(n) = I( p)), where + Z + 7 is parametrization of elements in H; H + R some loss functional, and the integral (Ø, μ1) is a Bochner integral on H. This generalizes the shallow NN setting, one might consider H L² ".r,V,mxii πx), L: H IR given by L(F) En [X( Y) and : + H defined as Vz Z, ((2) = (.(:)); so that R(u) L((O,))) In this setting, we can prove the following result
Proposition 12. Let H be a separable Hilbert Space and F(u) := L( 11) for some function that's Gateaux-differentin
L H I every direction and of continuous differential; and Φ: Z + H such that Vu E P(2), ((D,,)lla, < cc.
Then V2 E Z, Vp E P(Z):
OF
0p (11, 2) D1,L((P, -
Cr.
D,,F(u,)) (DAL((0, p))(D (.
Here, CF. == (VnL((D,,)) 1))) (Ø, plrr is exactly the constant needed to avoid ambiguity in the definition; (.)- denotes the adjoint operator and, in particular, Vz = (D,)(z))* H - Z.
When Z = R this corresponds to the usual definition of the gradient.
Proof of Proposition
12
We know that, Vu, P(Z), h [0, 1]:
F((1- h) + hw) F(u) L((Ø, (1 h) + hv)) L((
h
L( 1)) + 1/ n)) - L((0,
h
Let's denote by qu := (Ø, 1) (analogously qu-p = (4) - pi)) and Sp, = hqu-uu so we can write:
F(( h)μ + hv) F(u) L(gu + (pmn L(qu)
h
h
As L is Gateaux differentiable, we have the following first order Taylor expansion Vx, H. Vt e R:
L(x + +ts) L(x) + Dx {[t]]ss]))
which allows us to write:
F((1 h)p hv) F(n)
aa-p L(ga + hgpppp) Llgn). h.Dp -u-p o(hhlll@-=ll))
h
As |1q--pl| by hypothesis, we can say that: o(hhlllu-=ul) o(h). Therefore, taking the limit with h + 0, we get that:
F((1 h)p + hv) F(P) D2 DnL(qu)
lim o(h) Dn.L(qu)qqup) h-50
lim
h
h
Where the gradient V. Vaf(x) is the unique vector in 7. representing the action of Dz D= R
29

---

Now, developping this last term (using, for instance. the linearity of the Bochner integral, as we know DaL(a, --()) to linear and bounded), we get that:
Dj, Lgg Qup D2, Dr,T(en). n.) (D -14,
and so by definition of the gradient of
F F((1 - h hw) F(u)
lim
(VnL((0, (0)) ((v (z)
From here we deduce that:
0F
0μ (1 (V) ()), Cr
where Cr,u is a fixed constant, given by:
CF { OF 0u (n, z)d(()) / (V.L((0P))) D())) d(z)(2) (V,I((#, p)), (#, (),
On the other hand, for the intrinsic derivative, notice that Dz OEE Jy (H, =)) : + R2 a bounded linear
functional over 2, so (by Riesz Representation) F(,, 2) = V= Op (1, (1, E 2 such that:
aF
(
OF 0p (1,z)
OF 0u (1, z) (z)
Vz E
Vz
:2
D2
However, develop the RHS, and the constant CF, doesn't depend z, get that:
OF Dp (1,=))
D2
(=) Da ((V ()))) (z)
Now, by the chain rule and the definition of the adjoint operator of (z):
OF
D2
8u (H, (1,2)
= (.)))
(VLL)))))))))))((()(()())(()))) (4)(()))))
So, as they coincide for every 2 E Z, we conclude that:
D₁ F(u, = (D) (V.T((,))))
Proposition considering
12 applies directly to our population risk functional R(μ) = E, ()) by
The Hilbert space: = 1XX,
L 7 + IR given by L(J) = E, (X((()) Y)). which is Gateaux-difffrentiaal on every direction in H if assume É ) y R to be continuously differentiable its first argument, with v, linearly growing. The differential (which is continuous) can be explicitly computed be:
DaL(f(() = E. ((T((((()))) x(X))
+ Z P(Z), [1 p) < oc under the assumption of a. being bounded and continuous.
defined as Vz E Z, (=2) as (332), which satisfies Vy
Corollary 5. We can explicitly compute the linear functional derivative and the intrinsic derivative for the learning problem's population risk:
0R
0u (n, 2) = Ex R(,, E,r
+ (constant not depending on 2)
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:
30

---

Proposition 13. Let F P(Z) - R be invariant and of class C\. Then: Vz 2, Aa E P(2), Vg G OF
OF
dp (Mg Mg.2) 0u (1,2) 2) and D,F(M,HH.M = M,
i.e. OF By
jointly invariant and D jointly equivariant.
Proof of Proposition 13 To prove this, recall that the linear functional derivative of is the only Bμ (2) Z + IR satisfying V,, P(Z):
function OF
(p,2)d(v (2) and / aF Ōμ In particular, as is G-invariant (and M, linear), we can write Vn, E P( wh + 0 and y E G: F((1 h)p + hv) F(u) F(( h)( + h(M,#w) - F (M,
F((1 h)p hv) F(P) / OF
lim
Ōy (p. 2))2((z) 0.
h0
h
Taking the limit as h on both sides, we get:
F((1 f) + hv) F())
OF
lim h++
(M z)d(M (z#1)(2)
h
2 0p
S 0p
OF
(Mg Hu, Mg.2)d(v (2),
and also:
OF
OF
MM u, My.z)dp(z)
MM 2 dd 2) O.
0y
0p.
So, by uniqueness, get V/g E G:
OF
OF (M H1 Mg.z);
Ōn (H, 2)
Oy and, from proposition 6 (since OF is jointly invariant), we get 7/g E G:
OF
0F
DyF(M Mg.z)= v.
(Mg-ts, Mg.2) Mg. V.
(1, D,
0p
0p
Finally, for convenience, let us denote by Hº the function defined on P(Z) by Hº (p) Hxgo (1E)) Hxgg Pgo # (11) (presented in Section 3.3). We can straightforwardly compute its linear functional derivative and its intrinsic derivative in that space:
Example. (LFD and intrinsic derivative for H²²) By definition of the linear derivative Hxga on P(EG) and the form we know it takes (see the examples from Appendix D.I we see that,
SHago on of
whenever u u 46 Aeo, we have:
lim Hxga ((1 c2 + -hrt h
Hxyo (1E9)
HH
:x)d(UEo
u00 (x)
h
ga
On
dyec
log
diec (x)
+C
diveo
20 50 (x) (x)
du adga (PE0.t) + d(v -)))),
log
z
a which yields that Up (1,2) = log expression for the intrinsic derivative follows, which is given by
c Trga (PEa.2) +C' (for an appropriate constant), A formal
duca
C2 (1,2)= =
1
Pf
(Pga.z)
_1 (z)
ddga
31

---

D.2 Expression for the WGF of the regularized population risk
In the case of the regularized population risk functional RT P(2) + R. we can explicitly write its intrinsic derivative. Consider a slightly more general functional, denoted by R7 where the entropy calculated against Gibbs measure y < such that v(dz) = (=((()) X(dz) for some function U : 2 + I (as in (38)). We have Vp E P(2) S.L. & v, E
V Vau(z)) [(z)
Du (1, 2) + (V. +
so that WGF(R as in definition reads:
O.pe <(t) div (D (pe:) p4) s(t) [div
_V) 14) -V +
3Apl.
We recover the expression for WGF(RT*P) in Equation (4) by considering U = 0:
0ep = s(t) (div (DD, ) + v.r) p4) + 3Bt4)
(10)
We can see that Equation (4) corresponds to a Fokker-Planck equation, which can be interpreted in terms of non-linear SDE system, representing the behaviour of the type parameter: the McKean-
Vlasov SDE [T 54 (70) (also known as the Mean Field Langevin Dynamics (MFLD) in the NN
literature). In the case of RT. it reads:
dZ, s(t) [D 2,) + + + 223dB, with 14 = Law(Zi),
(1I)
where (Bt) is D-dimensional standard Brownian Motion. indeed standard check (by applying Itô' 's formula to (Z,, t) for y smooth function, and taking expectation) that μ4, = Law( 2,) a weak solution to (TOp. Under mild regularity conditions, both formulations are equivalent. See
11 14 70] for details. The previous correspondence also holds true when / 0 (in which case an ODE). The process (TT) or variants of will prove useful establish some the relevant
results of the paper.
D.3
Global convergence in the regularized case
For the example we just presented of the reguaie-eeullrrized population risk, multiple authors (see 57 69) among many others) have studied the properties of WGF(R*.9), particularly, , the global convergence results that can be obtained. For instance, consider the following results from (where we look RT. for generality). First, define:
12
15
38
38
By Definition 10. We say that a functional R P2(2) + R is of class C if aft (p,) )) is well defined and
bounded for every μ (Z) and the function (p, 2) E Pp
+ Bu (p, z) is continuous.
Now, from 15 38
we get the following key result. We include the proof for completeness:
Lemma 3 (as
38 [5). Assume that R. Pp(Z) - is convex and of class Ct. Then, for any
u. E (2), we have;
{ OR
R(1') R(n))
(n, 21,2)d(k' (2)())
2z Op
Proof of Lemma (from 38 Define μz = (1 - t) u + eji'. Since R is convex, we have
<(R(y') R(n)) 2 R(n R(H) = 1- [ ÖR Op
(u(,))d('' ())d)) ds.
Since the map E [0, 1] + μrº is continuous, it is of compact image (denoted [e, v(l). In particular, 8R is continuous and bounded on its second argument, we get that, it is bounded on [j,, n²] Z. The as Dy dominated convergence and Lebesgue differentiation theorems (as s - 0) allow us to conclude. O
Consider now the following assumption:
32

---

Assumption 2 (As in [38 ) U: 2 2- ARRRE is assumed to be C with VU Lipschitz continuous, and
such that 3Cv > O, 30y we will also assume that r: Z - IR. satisfies these conditions.
such that vx 3 VU(z) .2 2 Cr|aa|| C. When required.
Notice that these conditions imply that 30 S C' S C s.t. Vo 2, C/| C s U(x) s C(1 + ||=]]²) (i.e. U has quadratic growth) and (AU
Since R includes an entropy term, it guarantees strict convexity, weak lower semicontinuity
and compact sublevel sets (see c.g. [8 or [26]). On the other hand, assumption 2 iiplles that U (or r) will have quadratic growth. Namely, we get (see (8 for a detailed proof):
Proposition 14 (Existence and Uniqueness of the minimizer (regularized case)). Let R be convex. of class C² and bounded from below. Let v be the Gibbs measure with potential U. Then, R7 has a unique minimizer, pr,",,," & When either U satisfies assumption 2 this minimizer also belongs to (2(Z))
P()), absolutely continuous with respect to Lebesgue measure
For establishing global convergence results further assumptions are requred
Assumption (Assumptions for well definedness (from 12 and 38 )) Assume that the intrinsic
derivative R, P(Z) x Z + 2 of the functional ft: P(Z) + R exists and satisfies either one of
the following:
1. (From 381
Assume:
Dr R is bounded and Lipschitz continuous, i.e. 30n > S.s. V2,2 E 2, Yp, E P2(
(D,R(u,2) - -D,R(, / S CRllZ - W2(H, R(,,)) Cx) )
VD,R: P(2) + 2 x Z is jointly continuous.
2. (From fT2y, who relax some differentiability conditions at the cost of boundedness assump- tions; this allows them avoid altogether the coercivity condition from assumption 2 is used in (381):
which
Vx E Z, m' E P,(
D2, ,2) S M W W1(m, m') for mm 2 0 (i.e. it is lipschitz on the measure argument).
some constant M R
Suppose that
sup sup VD, x) & M (EPa(Z) DEZ
for some constant M 2 i.e. RD,,R(H, x) is uniformly bounded.
This allows to establish a traditional global convergence result from the MF Theory of NNs:
Theorem (from [1] and (38)). Let μo E P2(Z), and let assumption 2 and hold; then:
Vt O, dt (R. -(t) D + +Vr(2) (V +
Vu (z) + BVU(=) dpe(2)
ue
where un denotes the density of pt := LXw(X the solution to equation (4). i.e. following the WGF makes the regularized risk decrease 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:
[(p]le) = /
du V log
(z)) du(z)
dv
Then, almost by definition, we get:
=
dt
This allows us to characterize the stationary points of the dynamic explicitly, as done in 669.
38
33

---

Theorem Zimplies that the WGF converges to the unique global optimizer of the regularized problem:
Theorem 8 (from (8]). Let R be convex, bounded from below and C'; also assume that assumption 2 and 3 hold. Consider po Up>2Pp(2) and let (μe))20 be the WGF(R starting from po- Then, the equation has a stationary distribution, (loos that satisfies:
Moo := arg min 12,,((4) and 1-000
lim W2(Ue, =0
(EP(2)
Remark. Global Convergence Results such as Theorem 7or Theorem S Shave been established as early in (53) (for the quadratic loss). However, settings such those of [ 38
57 69 75V
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 μ E P(Z) satisfies the Log-Sobole Inequality with constant 1 > 0 (in short, LSK()), if for any v E ((2) such that v <x 14, we have:
du D(vlla)= = / log(
du (=)))u(z) S
1
/
du log( Vlog
V
(3))
du(z)
1 (vllm))
20
dp
20
where D(vlla) is the KL divergence and (plly) is the Fisher divergence.
(see 3 188 for background on functional inequalities and [15 12 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 PPZ)) - R is assumed to hold:
57 669 for applications of it to the
Assumption 4 (Uniform LSI from [1 12 57 69].. There exists 0 > 0 such that Vi E P2((Z), satisfies LSI()). Here f is the probability measure with density w.t.. X given by (slightly abusing notation, and considering U the potential of Gibbs measure used in the entropy, which is 0 if == A):
A(z) x exp (-1
(u. (-33
1 YR [1RR 0p
- U(2)
Remark. This LSI is recurrent element in the literature of WGF and Optimal Transport in general. In particular, it implies (see (3)) Poincaré Inquulltty: Vo E (Z). Vara(Ø) S 1 V and
(58)) the Talagrand's s Tæ-trunsport inequality as well: Vu E P2(2), (W (.) < D(vll()
Beyond the characterization of the decay (from 388 have the following guarantee: Theorem 9 (from 12, 151). Let R be convex, C³ and bounded from below, and let assumptions and hhold Then, for some to 2 0, Mıo has finite entropy and finite second moment; then Vt 2 to.
3
D(pillps)) < R (11) RT (HOO) S (R, (Hrr) - R,; R. (poo))e -250 Jig s(s)ds
where Poo = p 8, = arg min RT, (H)). That is, the value function following the WGF thus converges exponentially fast the optimum value of the problem, and this implies exponential convergence in relative entropy.
One thus recovers, under the right technical assumptions, version of Theorem 4 from 153 and actually a quantitative improvement of it. By Talagrand's 's inequality this also implies exponential W2 convergence of (4 to Mos We note that the result in 112 is established in the setting with 7 0, 3 1 and 1; however, one can show that the result holds as stated by using standard arguments.
D.4
Conditions for well-posedness of WGF
Most of the technical conditions that will be here presented are directly taken from both [16 and
22 We only adapt them slightly to fit into our notation.
Regarding the existence of weak solutions to the WGF presented in equation 4 guarantee it under the following assumptions (more general assumptions might sought but these are relatively standard in the MF context::
16 are able to
63
34

---

Assumption 5 (Assumptions for existence and uniqueness of the WGF solutions (taken from (T6)). Consider a setting as described in proposition 72 with R(μ) := L(( + V(u), with V(u) .z rdp.
1. Let Z to be the closure of a convex open set within some finite-dimensional euclidean space.
2. Let : - RR be differentiable, with a differential dL that is Lipschitz on bounded sets and bounded on sublevel sets.
3. Let 2 2 H be differentiable and V Z + IR+ be semiconver (i.e. YE E R: V+11-12
is convex).
4. There exists a family (Qr)rso of nested nonempty closed convex subsets of Z such that:
(a) {u E S2; dist(u, Qr) S r) C Qrtr' for all r,r' > 0,
(b) T and V are bounded, and d is Lipschitz on each Qr
(c) 3C1, C2 0 such that sup (u))l + [[0V (u)]) S C1 + C2r for all r >
where | (u)ll stands for the maximal norm of an element in (u).
On the other hand, [22 are able to prove (based on Theorem 1.1. from (70) the existence of strong solutions with pathwise uniqueness for McKean-Vlasov SDE given by
dz, = btt, Ze. dt) + Et, Hi)
where b and Σ satisfies the conditions of B2 (presented below) and for all 2 0, μr = Law(Z)) E (RD)), (B,)22 is r-dimensional Brownian motion (with E N potentially different from D N*), and Zo has the (fixed) law n P2(RD). For this, consider the following technical assumptions (B1 and B2) which have been taken directly from (2]]
Assumption (Assumptions for the existence and uniqueness of solutions in 12
Consider:
BI. There exist a measurable function g RR.. x W - R. M1 2 and no E P2(RDD such that for any N E N. the following hold.
(a) For any w1, u2 E RDD and = E W we have
[1g(w1,)) 2) - g( 2)1l S ((>)[lu) - w2ll. and (lg(w,, 2)1l S ((2)
with Sw (2) dTw < to0
(b) by E C(R+ x RD x P2(R D) and En E C(R+ x RD x P2(RD), RD
(c) For any W1, W2 E R and μl, #2 E ()
| (t, w1, p1) bn (t, w2, p2)1l + (t, Ex W2, m2)ll 120
sup
$ Mi [wu] (1w1-ww - w2ll + ( Jo. 20((,)),11) - 2). ( 1/2 }).
+ (,)((, S MI.
120
B2. There exist M2 2 0, k <> 0, bE C(R_ x R x P2(RD), and 2 E C(R+ x RD x
P2 D), TR. Dxr) such that
sup W20 (RD) (ED)))))))))))
(1b w. blt, w. (EE w. - E(t, w. (M)1)) S M
Proposition 15 (Proposition 11 in 12 Assuming BI and B2 Given dº' E P2(RD) as fixed
initial condition; then, there exists an -a0aaaddd process (Zz)120 that is the unique (pathwise) strong solution of the McKean- Vlasov SDE:
dz, = btt Ze, me) dt + Et,, Ze, Ms d.B
Additionally, it satisfies for each T > 0: supre
|1
35

---

E
Proofs and discussions of main results
E.1 Proofs of Section 3.1
Proof of Proposition
By definition of the symmetrization operator, we know that Vo E N:
(Qa (.))),)===== / P $ 0 ((((g))))AA((())
For E G. since o, is equivariant and Mg is invertible, we can write:
$ 21 = (a- (p 1) = (o- (p, -), (MM #u)) = Po (x. -), M, #1)
)
where we' ve used Proposition iin the last equality. In turn, write (via the inversion-invariance of Ag):
(2) Ic (2g HM) = =/6 (x, # (r. _20
/
As mentioned in Section 3.1 simple case where we will have @b no uco is when o, linear: Proposition 16. If X x Z + y is jointly equivariant and ax Vx E X' 12+ ox 2)) is a bounded linear operator, then, for any μ P(Z): & Dgo = (o,,,1 (o., (0.,1e") $ nea.
Proof of Proposition 16
A straightforward computation yields (using Fubini's theorem and the linearity of integrals and o.). A4 P(Z), Ve E X' (Tx):.
(a. (1.1) = / o. (o:) = / z (e)))
f o. ( Mgzdda(())dlll)
1z o, Pea. (o, (1,
Example. Any usual linear model written in terms of some feature function 2 X Z (where Z is possibly an Reproducing Kernel Hilbert Space) enters this framework, by defining: a (x, 2) (2, (x)). This won't satisfy Assumption 7 since not bounded; but it still serves as an illustration.
E.2 Proofs of results in Section 3.2
E.2.1 Proof Proposition 2
Now consider, as a shorthand notation, Vo E X. Vy E y the functional Lu. P(Z) - given by Aμ E P(Z): Lz.z(u) = ( (.)))) The following lemma that shall be useful for later stages. Lemma 4. Let a be jointly equivariant and be invariant. Then, Vg E G, Vo E X, Vy E V, Au E P(Z),
L L p. (M (L.
Equivalently, the map P(Z) + (2 x V, given by L(u) + ((x,y) + Le is equivariant (under the appropiate G-anpiattt---aiooon G-actions).
Using the joint equivariance of a, (via proposition 7 and the invariance of f, straightforward computation yields, for all x E X, y E V, and g G:
Proof of Lemma 4
L (M
= ((a.(agg .), Mg#p), Py-v)
= E (a. (I, p) Pg.))
l ((o,(2, :) 1.) Lzz
'In this case, Vg E G. let 9.u = My#u and g.f = fo given by Ve X, Vy E 3, (*(,,) [Pg 2x,,, b ..)
36

---

Notice that we will basically utilize the equivariance properties of or and in Assumption ([ (the convexity of the functions comes directly from the convexity of R when & is convex, together with the linearity of (-)s, (.)c and (.)da.....
With this we can prove Proposition? 2
We can readily see that: (B) Ex [C uusa (x),)) R(4e²)
Proof of Proposition
On the other hand, o, is jointly equivariant, from Proposition we have:
RFA = Ex [
Ex =E, (e (X) R(N
Next, using Lemma
Fubini's theorem and the inversion-invariacce of AG, we get:
(w) = Err Jo e ( (0 X). dAclo))
= FT [/ T.p, X. x
/ Fx [pp.. (n)] dga(g) /.- TFx [..p.....p. (M #n]] d)a(g)
Ex [xxY(M, #1)] d)a(g)
En [Lxx (M dNc(())
= R(Mg))))Dol0)) = R
From these expressions we can quickly verify that RDA, RF'S and REA are invariant. Namely, by Lemma for y E G and μ E P(2) we have: RF (M R((M@H1) = (((C)= RFA (n))
and: R (M, invariance of ÅG implies:
R((M+#1)
R(n²²) (w). On the other hand, the right-
RDA [ (M(MM = / (((M)#))ddbo (R)))
We can also see that whenever invariant, we have that, for μ P(2) and g E.G.G G, C(M R(p), so that: RDA (P) (H) = Jo R(Mg
Sc R(u)ddo(g) = R(p).
Also, if F (Z), we have, for all g E G, p Sa R(Mg#H)ddAd = So (lu)ddo(g) R(u)= R(x = (F)
G Mg#H so that: RDA (p)
Finally, we verify that our population risk R: P(Z) R is invariant whenever E PC ))- Indeed, Vg E G and V/ E P(Z), by the invariance of f and x, together with Proposition from the cquivariance of o., we get:
RM E. [E((o, (X::), Y)] E, (o. (x:))
= Er ((((oo(p,)x:)),111
k):),,),000) Y)] [)] = Er (X((,, Y)) = R(u)
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 P(Z) + IR and r : Z + R are invariant (in their respective sense), then RT.8 invariant.
The result can be proven for R7.8 with some G-invariant measure (such À for orthogonal representations). Notice that r(0) = ||0 [2 is an example of invariant function for orthogonal
representations.
Proof of Corollary 6 is enough to notice that (()) [ r(0)du(())) and H,) = J2 log( ( )du
(with/ μ << v)) invariant when : Z E and E P(Z) invariant (in their respective sense):
37

---

1. For V, notice that for 9 E G:
V(M##1) = ((0))(()+))(0))))) = r(M)))ddu(0) r(0))u((0) - V(()
thanks to the invariance of r. i.e. V is invariant.
2. For Hu. notice that, for E G, as y is invariant, we know that d(M,#1) Therefore:
(x) 4 (Mg
dv
ddMg#H)) (0) d d(Mg#t)(0))
Hy Hu(M
log
du
/ log
dqu du
(M)10) d( d(M)#t)
d4c
(M1 (22) du(0) Hu(w) du
log
Which proves that H, is invariant.
We can readily conclude, since: (p) R(U)) + T rdp + (H))))), for μ E P(2).
E.2.2 Proof of Proposition 3
we require a version of Jensen's inequality that's suited for our context. Such result might exist in the literature, but since we couldn't 't find complete proof under our assumptions, we provide our own.
In order to prove Proposition 3
Proposition 17 (Jensen's Inequality). Let F :P(Z) - IR be such that Lemma 3 holds. Let be some measurable space, À P(S) and E S Me P(Z) measurable function. Define J P(Z) as the intensity measure: j = SsHl dX(s) E P(Z). Then, Jensen's 's inequality holds:
F(R) < / F(ug)dX(s)
JS
Proof of Proposition 17 Since Lemma holds, we have that 44z, P2 E P(Z):
F(m1) 2 F(p2)+ /
OF
0μ (u2, 2))(1) (2)(2)
Let & E S be arbitrary and consider f12 = (s,) and po = -4z. Then:
F( ((saX(s)) S F(us) - - OF
(ʃ (H.d((ss2)) d Hs (Ms / (odd(s)) (2)
0p
Integrating the inequality with respect to (on 5):
I F (ʃ (Iod((s) da(s)
1/ F(ug)d((3)
S ( OF (Luaa)))))) d a ( ax(3)
SA:=
pu
We notice that the LHS doesn't depend on 8, that Is F ( (s)d)(ss) dd(s) = F (J pdd)(ss)) On the other hand, the right-most term in A can be developed as:
1 [1 OF
( (ndd(s).))
d (11s
/ madd(ss) d2(3)
145
( OF
OF
(iaNNs)))) d ( ((aa)s)) d2(3)
(,dd(())))
dpz
0p
Op
( 0F
( [m.dx(o):
d(ua))
d2(3)
0p
( OF
0p ( (uddiss..) a ((uaasss) ( ax(3)
38

---

Notice that the linear functional derivative is chosen in such way so that it satisfies ty E P(2), Iz aF (v, z))u(z) = In particular, the second term of the previous expression vanishes. öu We get that
1 122200 OF Op But, by definition: vf z + R integrable,
(IoHadN(s) dus(z)) )
dd(33
v (.dx()))))) - (IIN))a)))) -I ((e)das))))) dX(s)
So this is, by definition, and applying the same convention on the definition of the linear functional derivative
= [
OF /
(J [1s12) d (f (z) =
8n
With this, we conclude that A -s F(p,)dd(s)) and so, we get
that:
( (add(ss) F(ua)d(ss)
F F
which corresponds to Jensen's inequality.
Remark. We believe that the C hypothesis can be lifted. Understanding what happens when equality holds should be of interest both in context and in general scenarios.
Thanks to this Jensen inequality, readily get the following result:
Corollary 7. If F :(()) - R convex, C³ and invariant, then Au E P(2): F(P F(u)
Proof. Direct from the definition of (-)6 and Proposition 17
With these results in place, we are ready to prove Proposition 3
Proof of Proposition Evidently, since (2 C P(Z), we have:
inf F(x) 2 inf F(n) (2)
nEP(2)
For the other inequality, take (JN))EN S P(Z) to be an infimizing sequence for F; i.e. such inf u2)())) F(u)). Such sequence always exists. By Corollary 7 have Vn E N, F(ug) S F(u)); thus, Vn E N:
that F(Hn) 2 F(Hn+1) and F(Pa)
inf F(n) S (( S F(un)
(EPC(2)
Which allows us to infer, by taking co, that: inf (c)) F(u) inf wer() F()). In turn, can conclude that:
inf F(J) inf (2)
F(n) μEP(2)
Notice that there was some minimizer μl E arg min F(p), then by Corollary 7 we would also have 1 E arg (PP()) R(μ). Namely, if such a minimizer was unique, then would satisfy: μ. = - E (Z). That is, the unique solution would be WI.
10 Notice that the convention on the linear functional derivative's definition isn't truly important for the proof. since in the end simply corresponds the term Sz OF (J us, 2) d [J padA(3)) (z) being substracted to itself.
39

---

E.2.3 Proof of Theorem 2 and Corollary 1
Notice that, under Assumption I from Corollary we know R is of class C¹ (as well convex). This properties actually transfers to the functionals RDA RPA and REA, as shown by the following result:
Proposition 18. If RR: P(Z) + R is a convex and C¹ functional, then RDA, R and REA are
convex and Ct as well, with linear functional derivatives given by:
ORPA
[
0R 0p
OR
[c
2R ( (Mg.2)dAo(s)
0p (p,2) 2)
(M #1,
(p, 2) = Ōu
Ōp
OREA Ōy (1,, 2)
2R (nto .Peo.z)
and
And intrinsic derivatives given by (when well defined):
(p, I My. DD,.(MMMMHH #M Mg.2))(a(9))
MT JM DN, (g.2)d)o(g) and D Peo.D .Poo.))
D A,,
In particular, from Proposition 17 have that: Vy P(Z), R RFA(N) S (P)
Proof of Proposition 18 We can calculate the linear functional derivatives (I.d.f. for short) as follows.
Let 1 P(2), and consider:
(PA((11 - h)p + hv) - (1)
lim
h-40
h
Jc R(M ((1 h) h R(,,
lim
A-50
R((1 h) + hM - R(MgtHu) (
lim
h-10
h
( R((1 M + NM - R( dda(g)
Ja
lim
h-10
OR
(My 2. =)) M2
0p
[ 0R
(M # =) -))==))((())))))
JG J 0p
0R
#))2).
2 JG 0p
We have used Fubini's theorem, which is applicable thanks to the fact that R is of class C², and we' ve used the definition of the linear functional derivative for R. Also, we see that (using Fubini's theorem once again, well the definition of the linear functional derivative of R):
[a OR Dp. (MaH MG Mg.z)dda(()ddff)
/o
dR (M, (M 0. 0p
We can then identify:
ORDA
[
0R Ōμ
8p (1, 2) (M,
(M H Mg.2)d
and, by taking the gradient:
D DA (p,z) / MI R(, Mg.
II particular, as for any fixed E P(2) the function g E + Mo# continuous (thus, of compact image), then the function (g, z) G Z (M, p. M9 22) is bounded
40

---

We analogously calculate the expression for the lffdd. of R let U. -
P(2):
R h)p + hv) (M)
R((1 hu +
lim
lim
4-0
h
h-00
1
ÖR
Op
0R
(n .Mg.)ddd(()ddte ())))) n)(2),
z G dq.
and also:
12
0R 0p
dR | 2))d(M
| /
(n M,zz)dAa(s)du(z))
2 Op
ÖR (n .2)dp = O.
8u
So, by the definition of the lineal functional derivative, identify:
DR
0p (1,2) = -o
0R 0p
(re MM.2)dDa(l),
and taking the gradient we get:
DaRFA(p, / My DupR 2Mgg2)
Lastly, the l.f.d. of RFA calculated similarly, noticing that:
REA ((1 - h)4 + hw) - (B)
R(Pge ((1 - (11-)ja + hv)) - R(PsoH)
lim 6-
lim
h
h
h
OR
(u :z)d(Peo#v Pęo#Hu(())
Ōp
ôR (1Fe -Pea))d(u
0p
and that:
[ (NE0 :Pec.2)du(2)= = / 0R
ÖR 0u
2 0p (2)d(pFe()(2) = O.
We can thus identify:
DREA
0R E4
0p (n, 2)
3p , Pgg.2) and D,,RAA (H, = =Pe DD R(He ,Pea.2)
The last remark direct from proposition
With all of these different elements in place, we are ready prove Theorem 2
Proof of Theorem Under Assumption R is convex and of class C² from Corollary 5 and so are R RFA and R.A from Proposition 18 Since Proposition 2 ensures that the latter are always invariant, Proposition implies that R RFAA and RAA 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, RDA and RFA coincide over (2).
In the case of the quadratic loss, one can employ the properties of QG from 27 to show Corollary
Proof of Corollary 7 Notice that, for μ pG (Z). 4 is G-invariant function (i.e. &
(X, xx)). Also, a simple calculation (see e.g. 51) allows us to write: R(u) E, ||| 1b (X) Y = R. + Exx (xx fz (x) R. + 1 f-llEA(x.) with
41

---

- - (X) being independent of We can thus write (simplifying subscripts for simplicity):
R.
R(n)= R. + | Qc.f= + Qa-f- - f- -x
= R. + 1 Qc.falli +1 - 21 - QG.ff, (f))c)1
Qa-f.. We notice that, since Φ, and Q6-f. G-equivariant functions, we have that (($) 1 2a-f.), (f.)) 0. That is, for any l (2), we have R(p) = R. 11 F
where (f.)
Qa-f-1122. where R. == R. (I( ( is independent of y (and doesn't 't
intervene in the optimization). Finally, get:
inf R(p) R. + inf (Z)
1 -c.ff-lla(x NEPO(Z)
E.2.4 Proof of Corollary 2
When is assumed to be equivariant, we can summon our previous results to prove Corollary
Proof of Corollary 2 From Proposition we know that equivariant data implies R: PP(2) is invariant, and also that this makes R RDA.
We conclude using Theorem
We can readily extend this to the regularized case by recalling Corollary 6
Corollary When R : PP(2) - and r + TR. are G-invariant, minimum for the regularized
population risk RT.9 can be found within pG (2). When 3 > 0 such WI minimum is unique.
Proof of Corollary 8 Direct from Corollary 6 together with Proposition 3 (as in Corollary 2 The
uniqueness comes from the strict convexity of the entropy term (see proposition T4).
E.2.5 Proof of Proposition 4 and Proposition 5
Proof of Proposition 4 Consider the group G C4 acting R22 via 909 rotations. Let K
B(0, 1) C IR2 be compact set. Consider random variable ~ NO, Id2) (i.e. given by X Z1 zek for Z ~ (d, and set Y ||||2 Notice that defined this way is compactly supported.
Clearly G is finite (thus compact) and it can be seen its ortogonal representation:
{ ( 9) ( - (o 9) (1 1) c 022). A [Id,} S 0(1) (trivial repr.)
PG
By the definition of our r.v.s, it is clear that:
(d)
P Vg E G because X ~ N Id2)
(d)
(pg Pg Y) (since ý the trivial representation, is enough notice that XII X 2 for all g G.
Ag E G. (X, Y)
Therefore, π = Law(X,Y) is G-invariant (and compactly supported). Consider a shallow NN given by: T. RR2 - R Nxb _ R (with b E N and some action Cipy TR. as: $ (x) =
Di=i Wio(Af x B,), Va E Rdd where 0: = (W,, A,, B.) E 2 Rixb x R2xb x R We let Cm Z as described in appendix C.T
M2 = (p, 7/ Po A, nig Ml B,) = (W, 1/ Pg A₁ nr ng B,)
We can assume, for instance, that b = 1 and ry is the trivial representation (so that no condition is required for to be jointly G-equivariant) and recall that: 0, E EG
+ Vg E G, MgO. 0,
42

---

A 13
However, if we assume that: Vg E G, PgA, A, then, in particular: 0 -1
This in turn implies, as A! A and A - that A A 0. i.c. A, = Thus, any 8.
W/ EG has A 0. Therefore, if we choose any activation (e.g the sigmoid activation
A,
or tanh, both C and bounded) and we choose N No and E E Vi ,,,,, N; then:
1 N Σ Wo,(B,),
Vx ER
D N.EC
(x)
Woo x+ x+ N
1
E 1=1
which constant independent of x. i.e. any equivariant architecture in this context is a constant function (whereas Y= ||||| is not). Notice, in particular, that any shallow model 4; with v E P(EG) will also be a constant function. In particular, notice that we will never do 'better' than minimizing over all posible constants:
inf R(0g
inf R(v) 2 inf Bijy - CP = inf |||| CI VEP(EG)
0,E80
CER
CER
The problem on the right has a known answer, which is C" = R [ > O. On the other hand, consider fully conected neuronal network. By the universal approximation theorem (which applies for the chosen a, as in [37 5), as is compactly supported (in particular, n (K) 1); we consider the parameters that approximate the function f(x) ||=|||² 2 in K B(0, 1) to precision
8 > i.e. For & E (0, V,"), we know: = NN N, 3a, ..., aN E TR?, Fun WN E TR such that:
|1 fll,,, sup 1 (x) f(x) <e<vo- zeK
Then:
|E||Y 2/ (x) < ([Ssup 1 (x) f(x))) =00 C
mek
But, in particular, 3ug E P(Z) such that:
B[JY 2 (x)1 < C
and so:
inf R(x) inf BIJy ( < inf R(v) (EP(Z)
8EZ
vEP(SC)
In particular, we can't expect an optimum of the learning problem to be achieved within P(E)).
To overcome situations in Proposition 4 the usual setting is assume some universality condition. This leads to Proposition 5] which we will now prove:
Proof of Proposition A standard calculation from the quadratic loss case (see the proof of Corol- lary 1 for details) yields that, for any E P(Z):
R(u)= = B. ||| 1 n (x)13) = R. -,,EE111 [1ll - (X)133
where f((x) = Ex Y = x| and R. is the Bayes risk of the problem. From Proposition 10 we know that 1' (X,
(X, V x. and so by universality of Fo, (P(EG)) onto that space (as well as that Fa. (P(Z))) we conclude directly that: infver(eo) R(v)) = R. pef(2))) R(H).
Remark. Works such as 150 60 76 77 precisely provide conditions under which universality of Fo. (P(EG)) can be guaranteed (modulo some adaptations from their setting to ours).
Particularly, our single-hidden-layer NNs, with the width N + oC, 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 (501). Despite this, 'first order universality' has already been established for some of the most important examples of equivariant architectures, such as Deep Sets (176) 771) and CNNs (1761).
Adapting our setting, order to eventually allow for arbitrary order tensors the MF formulation, is part of the future challenges to make our work more broadly applicable.
43

---

E.3 Proofs of results in Section 3.3
Having laid out all the different relevant elements for our work, we can now procede with the proofs of some of our main results.
E.3.1 Proof of Theorem 3 Corollary and Theorem
We start by proving Theorem 3 on the general case.
Proof of Theorem
We know that a family (Me)220 S P2(Z) satisfies WGF(F) in the weak sense Vip E (Z (O,
S. ((((),,)))((()),,_.))) =), V=4((2,)) d)),)))))))))))))))))))))((()()) dt
Now, profiting from the uniqueness of the solutions of this equation, it will be enough to show that, given a solution (pr) C (2) of WGF(F), then (ui
(uf))20 )=20 P (2) is also a solution. Indeed,
consider, for E G. fr = Mg MMg#H4 He, and notice that for p E CC (Z x (0,T)):
[." (0up(2,1) (s(1)D, 2), Vas((,
= =[o | (Oig(M (<( (M Mg. V. (p(() =:
Now, we can define 8 E CZ x (0, T)) given by V(z,t) E 2 x (0,T) (2,t) = ((M which satisfies:
(z, t) = O:p((Mg. t) and (z,t(z,t)= = My v (Mg
So that, by also using proposition 13 and the orthogonality of the group action, get:
[ [ (M,s)))(M)))),,F( Mg (2, (2, (1)(2) dt
[ (0:00 -
(s(1)D (100(()))((())))))))))))))) =). V V (z, b))) (x)(=) dl
Where the last equality comes from the fact that (#t) is solution to the WGF.
In particular, as we also have that jo = #M = po (because E (2)) by uniqueness we can conclude that this means that Vg G, Vt E (0, X-a.e be = =M
This may seem weaker that what we want to prove. Nevertheless, as our group is compact and has unique normalized Haar measure, positive and measurable function. Given g E G. take Ω S [0. 7] a full measure set where holds that μ4 My#Hu 14. In particular, f== f(t,)) 2 R is positive and measurable, that:
proceed follows: let f [0, x Z + R4 be any
Vt Qgg (f. (fr, Mg = (fr Mgs
and we can integrate this equality to get: 'n" (ft, Mt) Soo (fr Mg. Pc)dt Now, by integrating both sides with respect to the Haar measure, and applying Fubini's theorem (because everything is positive) we get:
[ (o 1 So. (fe, (Fe.e))dtdaa(o)) [c. (Se Ma H.)dddda(() J. (f
-
Implying (by standard argument) that Vt E [0, a.c. = pe and therefore: Vt E [0,T] a.c. 14 (pC (2).
Proof of Corollary 3 From Corollary
Corollary and Corollary & 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- Pakker-Planck equation exists. Furthermore, this solution is known to be strong if / > 0.
44

---

In particular, theorem 3 applies and allows us to conclude that if pr E (2),), then Vt 2 0 (a.e.) H4 P22 (2).
When 8 >> 0, since solutions are strong, we conclude that the densities (u)) are all G-invariant functions (A-a.e.). This follows from the remark about densities of invariant measures provided in Supp.Mat-B. B.2
Remark. When B > O, we have a unique weakly-invariant minimizer (from proposition 3 anddor corollary 3) and also, under mild assunptions, a global convergence result. That is, independently of the network's 's initialization, we will converge to the G-invariant solution, An interesting question in this setting then: At which point does the WGF enter the space P (2)?
We also prove Theorem
Proof of Theorem This proof follows from the fact that Vz E Z, Vu E PZ)
Dr, =) = D2 R (1,
R.
Indeed, notice that, from proposition 18
(p. / My D R(y Mg.2)dAc(9) | My. R(,
while also:
[ My Mr.
J My. M(g
D Dp.RDA
Now, let (uf (e=0))>0 and (uP )220 be the WGF solutions starting from po for A and RDA respectively. As RFA is G-invariant, by corollary 3 (a.c.)tt 2 0, urAa E pG((). Now, let's see that this process actually also satisfies WGF(R )), forcing both processes coincide by uniqueness.
DA),
Indeed, we know that (nEA)120 satisfies: V4 E C (Ex(xxx)) (0,7)):
S T / (Oup(z,t) (0)p(2,t) (s(t)),, AnE FA 2)), Vep(z,t)) t))) (2) =0
Now, as (a.e. V 2 0, uf E (2), we have V2 E 2: D (M = DD 2). In particular, (EE (AA) Je20 satisfies Vip E C.0 x (0,T)):
J'" (((((),)))((()),,. RD (uf FA ,2,) Ve t)) duf (x) dt =0
Implying that (nL (1.FA))20 solves WGF(RPA) starting from no: thus by uniqueness: (nF (1FAA)220 (n.")120
The last part of the theorem comes from Proposition 2 since if R is invariant, its WGF will exactly
coincide with that of RDA (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 applications.
Remark. Since RPA, RFA and R all coincide on PG one could expect Theorem 4 to hold for R even without assuming to be equivariant. However, the invariance of R is crucial for such result: if isn't I invariant, nothing guarantees that its WGF process will stay within PG (2), whereas WGF WGF(RPA) and WGF(R" A) always do so.
45

---

E.3.2 Proof of Theoremmaan and Theorem 6
We can now provide the proof for the stronger result, Theorem 5 stating that the WGF of an invariant functional will respect E all along training. This proof uses the McKean-Vlason non-linear SDE (Equation (TT) presented in Appendix (see [22] for reference). Namely, we will consider the following projected McKean- Vlasov SDE (with (B:)220 a BM on 2), given by:
dz, = (D,, + T (Vrr(z)) dt + P2 with Law (2) Pt.
(12)
which corresponds to the MF limit dynamics arising from performing the projected noisy SGD scheme from Equation (5. This can be shown by adapting relatively standard arguments from the MF literature on NN, see e.g. 12 23 53
Proof of Theorem 5 The proof has two steps. The first will consist in showing that the process [2
satisfies vt 2 O, Mr E P(EG). In the second step we will check that (44)220 (Law (z) ))220 is solution to the WGF (Rgge) presented in Section 3.3
Step 1: The (pathwise unique) solution of the projected McKean-Vlason SDE satisfies a.s. for all & 2 0:
12 Z = (Ze) (Z))
Z Zo Zo 1/ D(s) (Me: Z.)ds + 23 / c(s) and Zo = Eo (initial condition)
(13)
Here, is such that Law(E)) u and R := R+ (r, is being used shorthand notation. We first let g E G be arbitrary group element, and we study how the process 2 = (Ze)120 (My satisfies this same equation (13).
Denote Vs- == My Ut, as the law of Z., we want to show that for all 2
2, a- 0. % S D(s) (Us, Z,)ds + 223 L/ P(SdB,
(14)
Indeed, first notice that:
1. Let Ω be the full measure set where to E 0 (which we can do since no E P(E or, equivalently: PC E E = 1). Then, Vo E Q Za(w) = fo(w) E EG In particular, Yu E EΩ Ag E G, Zo(w) = My Zo(w = M&Ca(w) = Eu(w) = Zo(w), That is, Zo Zo-
2. Now, the equation is satisfied by (Zi)z2 and therefore, for 2 0, we have:
2, = M2Ze = MyZo M, (f, s(s) (us,
(us, Z,)ds) + V23M, 1. ((s)) 2/0 / s(s) D (Ms) ds)) + 28 L s(s)Mg.Peddl))
((s)D,, (M, MoHtu Mg 2o) + /23 ( s(s))
-20
1C
- L/ R. (v., Z.)ds + /23 f ((s)Pg
Here, we used the linearity of the integral (and the stochastic integral), the fact that Vg e G, M, Pea Peo, and Proposition 13 which holds for vo E Z.V4 E P(2)) (in particular for Zw(w), AA E Ω and μ = Law Zz)). Thus, Ag E G. 14 holds.
By the pathwise uniqueness of the solution (2t) we have (following, for instance,
28
sup 12 2:11 =
In particular, as g E G was arbitrary, we have that:
Vgg E G, sup 14 My
8.9.
(15)
t00
46

---

We now want to be able to interchange the Vgg E G with the probability measure. Fortunately, we are dealing with a compact group with a normalized Haar measure Ag. Indeed, from equation (15) we deduce that Vgg E G. vt. 2 0, P( Z - Zel 0) = 1.
Now, notice that, for any t 2 0 and Ω:
[[Z](w) - (ww)l = Zi(w) - [c Mg. (o
- $ |2 - Mg. d)c(g)
We can integrate both sides by P to get (using Fubini as functions are positive and measurable):
- / [Z((w) M, Z[, d'c
0 < .2
[i((w) Pgo Zi(W)
$ Ja 9 [Z)(w) M2 M,
(g)
where in the last step we have used the fact that Vg E G, Vt 2 0, P(] M2 0) = 1, so
that Vt 2 0. Vi E G. Jo (w) Zi dT((w) =0
This implies that Vt 2 O P-a.s. Z,- = P&e Zi i.. P(Z, E EQ) p.(E = 1, or, in other words, Vt 2 0, Ae E P)E as required. Note that all arguments work as well in the case that 3 = 0
Step 2: We now prove that (pe)12a, studied in the previous step, is weak solution to equation 3 that is, to the WGF(F), with (using the previously introduced notation R ):
F(u)= = R&d (1) = R' (1) + (H (1²²)
Also recall the notation He (4) := Hxga (n = Hxga (p) presented in the end of
Appendix D.I. as well as the calculations for its intrinsic derivative.
It is standard to check, applying Itô's formula and taking expectation, that the family (pz)120 (Law( (22) satisfies, Vip E C (2 x (0,T)):
S." J. (Bis(e. - (s(t) (<(()), R (1e,
=), V V2,(2,)) f) + 3 Pr D2 f)) dpe(2) dt =0
with denoting the trace of a square matrix and D the Hessian matrix acting on the z variable. Notice that the process P'g B is classically a Brownian motion in E Together with the fact that pr supported on 0 this implies that, for 0, br has density w.r.t. Age.
is clear from the case 3 = 0 that the terms involving first order spatial derivatives of oo exactly give rise to the terms associated with functional R' in the definition 3 of the WGF (R Therefore, we just need to check that for all t 2 O, the distribution defined for every ø E (z) by 6 +
Iz tr(Pge (=)dx))) and the distribution div (D HT (4) are equal. In fact, for all such
we have:
D,.H(), (HL, dpe(z)
= =- f (V=o(2), P.. v
duf
(
duf
dAge (Peo.s) du((2)
(Pra.z))
dAga
Since μ4 is concentrated in E9, we have Peo#14 M4. This and the equality Pfa Pea imply the previous expression is equal to:
dufaa ddeG
duf dAga
(V ( V.
(Pea.2))
(Peo.z)
due(2)
duf ddga
duff diga (0))
(V (VeΦ(2)) PTC V
(xi)
dp((x)
ge
duf
=-/ (Pto V -((()), PTc V=
(x)))deea
ddga
47

---

Noticing that O V. is the gradient calculated on E", integrating by parts with respect to Aga, and using again the fact that dt is concentrated in E straightforward calculation yields the desired expression Ja gg D (=)]ddu This concludes the proof.
Remark. Notice that, by a.s. continuity of the McKean- Vlasov diffusion closed, Step of the previous proof actually shows that P(Zz E EG Vt
12 and the fact that E is
Notice also that Theorem 5
bears some resemblance to Corollary I in 130 which states that E is stable under the traditional gradient flow of the augmented risk (0 E zZ + RDA((0) E RR, as in Section 2.3 Our result shares similar flavor, but for the MF dynamics of freely-trained NNs with equivariant data.
Remark. Unlike with WI distributions, initializing a shallow NN with μo E P(EE) isn't as straight- forward as using a normal distribution. Effectively (and efficiently) computing the space E is actually quite challenging (as noted in (291).
A natural way to ensure that pa E )(EG), independently of the form of E is to initialize all parameters to be O. The question of whether under such initialization the parameters will eventually exit {0} (or some larger, strict subspace E, S EG) and find values over the entire space EG 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 EG which would be interesting to analyze, in particular in view of potential theoretical guarantees for bebhitectureddscovveee heuristics as suggested in Section 4.2
Note that there is no need for seeing g as subspace of an ambient space Z. When training with EAs, we simply force our parameters to live on E since we fix the architechture beforehand. Namely, our "whole space' is 2 = EG (regarded directly as a vector space g RD)) rather than 2. Thus, the relevant population risk is the restricted version of the original: i :2= -R)p(z) PP2) + R; and we can apply the usual results from the MF Theory when the relevant hypothesis are satisfied by 2 and R.. Notably, we can have global convergence of îe to inf LEP(2) ie (n) in unfeep(e)) Rig (1)
Remark. As shown in [3] (see Proposition 19 in Appendix Gioo the details) the regularized versions of the involved functionals (i.e. RT. and Rgoo I-converge to the original R T, - 0; meaning that, for small values of the regularization parameters, we should expect the achieved optima to ressemble inf (2) R(u) and enggg R(v) respectively (or, under Proposition 5 both to R., ).
We will also leave the exploration of how this approximation behaves as future work.
Finally, we provide a proof for Theorem 6
Proof of Theorem 6 The proof structure is very similar to that of Theorem 4 Namely, it comes from noticing that for μ E P(EG) and z E gG we have:
Dr R,, =) = D D CARAAAU, (1,,) D2, R R.,.EA RE,(,, 2).
We already know the first equality, as seen in Theorem 4 (since μ E pD( from lemma 1). We only need to show the last equality. Indeed, notice that:
Dr (w, = [ Mr. D R(n, 2)da(()) Pgo. Peo.D,,RI.. 2(,,)),,
while also, from Proposition
18
D DnREA (p, = Pfa Dp.RD Pec.2) = Pge D,,
Knowing this, the rest of the proof is analogous to that of Theorem 4 Let (uff)) )220 (uPA)r20 and (ufAA) be the WGF solutions starting from pa for RFA RDA and REA respectively. Since No E P(EC) C (Z). from Theorem 4 we know that (nt
(A) )e20 and ( coincide. Let's see
that, w.l.o.g., (ff ) coincides with (1 EA )=20
As RF A is G-invariant, by theorem 5 VI 2 0, uf E P(EG). Now, let's see that this process also satisfies WGF(REAA, forcing both processes to coincide by uniqueness.
48

---

As before, we know that (e20 satisfies: Vp C 2 x (0,T)):
J | (0up(2,t)
T
2), V Vep(z, t))) duf (2) dt =0
Now, as Vt 2 O, of E P(EE), we can restrict our integral to E Also, we have Vo E E Dj C..RAA(M FA, 22 DEREA(EE 2). With these properties, (nf (r2o satisfies Vip C2 (0,T)):
(2 x
P J. [e -
((()),, 2). Vep(z,t))) duf duiA(22)
dt =
Making the integral over 2 once again (we can since A P(EG) for all 0), get that (pE"))2o solves WGF(R starting from Po; thus by uniqueness: (Hl FA ()120 = (mt FA J:20-
The last part of the theorem comes, once again, from Proposition 2 exactly coincides with that of RDA.A
since if R is invariant, its WGF
F Experimental setting and further experiments
All the different experiments run Python 3.10. Google Colab session consisting (by default) 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 R2 and 2 R2x z TR.. We let G C2 acting on X' and y via the coordinate transposition
(i
and Z via the natural 1
action (i.e the group generated by the orthogonal matrix
0
intertwining action (i.e. Mg.2 2 = 2pg We also consider the jointly equivariant activation given by (2, 2) o(2 x) Vz R², Vz R2x2 with - E a sigmoidal activation function (which is C and bounded) applied pointwise. Under this setting, E can be explicitly computed
as EG
which 2-dimensional subspace of the ambient
dimensional space. It's projection operator PEa is also explicitly known.
We consider a teacher model f. D with N. fixed particles, such that vN is either arbitrary, WI or SI. Let 0.5 be scale parameter. The arbitrary particles were chosen be, for N,
07 0. 0.(-1,0,0,0.55)
8 1,00.5,,,0,,1"
0 d. 0.(-0.5,033,1,0)7
0 = 0. (0,
05 = (07.,,0.7,0.500077)
This was fixed in order to make the task non-trivial and interesting. The WI teacher distribution was simply chosen to be (vg²-) with 0* as just described, so that the corresponding teacher function resulted be f- N. In other words, the WI distribution has 10 particles, corresponding to each of those of 0", 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 EG (i.e. only providing parameters). Particularly, they were fixed be N_ and, denoting them by a" (at) avoid confusion, explicitly described as:
a = 0.. 0n (..5 T)TH az 0. of 0.(0. -1), og 0..77....7.
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) ~ n will satisfy X ~ A'0 o Id2) (with o, = 4), and Y fx).. 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 Oa as a projection (which require मx to be G-invariant, as shown in (7)).
49

---

We will try to mimic the teacher network by using student networks, which will be given by Φp'; namely, with the same σu. but varying values of N and 0 E ZN, We will train them to minimize the regularized population risk RT- given by a quadratic loss, Ely. 9) = lly 01 and quadratic penalization, r(2) = || ||= For this purpose, we employ minibatch variant of the SGD training scheme provided in Equation ①T) (possibly projected, as in Equation ( 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 be (unless explicitly stated otherwise):
Step Size: ç = a > 0 (with = 50 in most experiments), EN N. so that sxx = g. This was convenient, since corresponds the usual implementation of SGD on most common NN frameworks in Python (namely, pytorch and jax).
Regularization parameters: 7 -= 10-4 and ,3 = 10-6.
Batch Size: was chosen to be 20.
Number of Training Epochs: In line with the statement of Theorem 1 to observe phe- nomena at a MF scale, we need an amount of iterations (commonly known as epochs in the ML literature) that 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 No N ' 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 | steps. Notice that Ne 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, E 2N are initialized i.i.d. from some 140 that is chosen be either WI or SL. When WI-initialized, they are sampled from a random gaussian Z ~ (0, ) When Sl-initialized, particles are taken to be P&c.2 ZZ with Z as before.
Number of Repetitions: Each experiment was repeated total of N,- 10 times to ensure consistency. Each repetition, a different random seed was employed to: initialize the student's '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 Ne 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 I even huve dedicated GPU). This is the reason why 5000 a single training (with the above hyperparameters)
don't scale our experiments beyond the
N = 5000 case. As reference, for of single model feither of vanilla, DA. FA or EA) took ze 15 minutes (which quickly amounts to large amounts of running time for the Nr = 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 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 EG), 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 fundamentul.
Remark. Notice that the N, 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, = 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' Ill comment below).
To facilitate the implementation of the ideas behind our EA models, we use the group and represen- tations tools from the emlp repository provided as part of [29]. This code is openly available and
50

---

has an MIT License for unrestricted access. We employ it to numerically (and efficiently) determine the space E (namely, its basis), as well as Pge. We do remark that these calculations were correct only up to 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 W 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 vig" (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 RRelative-Measure Distureee This is (H.H))
defined as: RMD D) = M2+M, where M2 2B| Z [2] for Z ~ μ (so that 0 S RMD S 1).
The RMD provided a good metric for the experiments here presented. Notice that, as N + oo, 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 quamtity approaching (finite, non null, in our case) constant, we 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.
F1 Study for varying N
Beyond the analyisis already provided in Section 4 we here provide some meaningful insights. We want to observe to what extent the properties proved in Section 3.3 for the WGF of R can be observed practice. For this purpose, we observe, for N E (5, 10, 50, 100, 500, 1000, 50003. different relevant quantities to evaluate the different combinations of teachers and students.
For this set of experiments, WI-initiallzed students were trained with the usual SGD scheme from Equation (T); while, Sl-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.
SegureBdispaae some comparisons between particle distributions in terms of RMD the end of training, knowing that they were all initialized with the same particles drawn from Po E P(EG). 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 (since RMD? 2 ( )00) is small) increasingly more as N becomes large. This fact is absolutely remarkable, since, for WI teacher there should be no reason why the vanilla training (that's completely free, principle) shouldn's I excupe E to achieve a better approximation of fz. On the other hand, in every single teacher setting, almost independently of N, both DA and FA consistently remain within EG (as expected, even if fo isn't t equivariant). For an arbitrary teacher, we see that the vanilla training distribution readily leaves &G to better approximate f., which isn't I a predicted behaviour from our theory (we have no guarantees of 'leaving EG 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 aad-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 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 visualization of the final particle distribution, after an SI-initialized training under a WI teacher fa, 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 rigorous comparison at all, and providing better quantitative comparisons between the methods
51

---

Cistrce fequcted vasion
Oliancc Projctad
Uistanca Podectedd
FDM
flum
134
20
508
Co prrintion betwee DA seceious
Consefbre scvere Ekend tretteee
Concersen beweer D4 tecmriuss
4-2-
3000
(a) arbitrary teacher
(b) teacher
(c) ST teacher
Figure 3: RMD comparisons between training regimes, for different values of N, the end of an SI-initialized training for Ne epochs. Each column corresponds to a teacher with, respectively, an arbitrary, WI and SI distribution. Row / displays RMD? (uN, (NN )ee for the different regimes, in order to evaluate to what extent the training remained within E Row 2 displays the RMD between DA, FA and vanilla training regimes; and Row 3 does the same for each them against EA.
is to be considered for future work. As a counterfactual, we provide in Figure 5 the results of an SI-initialized training that is performed under annequivariant f.. We can see that the vanilla model readily leaves EG to achieve better approximation of S., while the DA, FA and EA schemes 'stay inside' (roughly coinciding between them as well).
T also displays RMD comparisons between particle distributions at the end of training, but for (identic) initialization with particles drawn from No E p Now, unlike the SI case, with particles sampled i.i.d. from WI distribution, nothing ensures that the resulting vo will be WI as well. Namely, in this case the limit as N - oo becomes significantly more important visualize the theoretical results. Indeed, since vi" isn't necessarily WI, we no longer have guarantee that the finite-N networks trained with DA, FA vanilla methods will be close to each other any sense. We do however notice on Figure 6 that, for increasing N, the tidd-oo-raiinng distributions of DA, FA and vanilla schemes (the latter only when f. is equivariant) become increasingly closer their symmetrized versions (namely, RMD (N ( ) becomes smaller, though never small in the SI-initialized experiments). Also, as guaranteed by Theorem 4 for large N we see that DA and FA become indistinguishably close, no matter the teacher's 's properties; also, when f- is equivariant, they both "coincide" with the vanilla scheme. A comparison between the EA scheme and the result from projecting the FA scheme on the last step is also presented. is used simply to illustrate that, in principle, directly training on EG isn't necessarily comparable to performing 'free-training' and projecting the resulting particle distribution only on the last step (even when an SI. technique such as EA is used).
52

---

Student Particles' Positions (vanilia)
Student Particles' Positions (DA)
Teacher Partidess
Teacher Particies
Student Partcies
Student Particies
24
44
0.
0.
03
0.2
23
-05
-08
06
66
-22
04
04
0.3
0.
Q4
02
04
63
24
Student Particles' Positions (FA)
Student Particles' Positions (EA)
Teacher Particies
Teacher Particles
Student Partiee
Student Partidies
0.4
0.A
23
23
-02
002
-0
000
88
02
02
92
12
x2
04
42
-01
Figure 4: Visualization of the NN particles after training under the vanilla, DA, FA and EA schemes, for a single realization of the experiment. Squares represent the teacher particles (which are WT), dots represent the student particles, and the hyperplane is gG The bigger plots show aerial view of the global partiele distribution after training: and the minor plots below them show a viewpoint at the level of (and parallel to) g The student particles were all initialized be SI (and coincide at initialization between the different schemes), and trained with equation (5) correspondingly applying the proper SL technique.
53

---

Student Particles' Positions (vanilla)
Student Particles' Positions (DA)
Partides
Teacher Partidies
Student Parüides
Student Partides
0.2
23
-02
-02
404
0.
4
-04
-0.
-2
10
03
Student Particles' Positions (FA)
Student Particles' Positions (EA)
Teacher Partidiss
Teacher Particies
24
34
0.6
0.5
0.4
0.A
02
0.2
-38
-08
54
04
0.
0.1
02
02
9
84
02
22
64
02
22
25
24
24
Figure 5: Visualization of the NN particles after training under the vanilla, DA. FA and EA schemes for single realization of the experiment. Squares represent the teacher particles (which are arbi- trary), dots represent the student particles, and the hyperplane is EG The bigger plots show an aerial view of the global particle distribution after training; and the minor plots below them show a viewpoint at the level (and parallel to) EG. The student particles were all initialized to be SI (and to coincide at initialization between the different schemes), and trained with equation (5) correspondingly applying the proper SL technique. Notice how the particles for the vanilla scheme readily leave & (despite the noise being projected onto it) and seem to approach the teacher particles.
54

---

Oblonck Symmitiiv Vevien
Distrcee Symmerteed vasion
Obbarr 2o Symmattree Virseos
200
Diitance technraan inskiing
Derance betweeen ELveriiques traiiing
Datance brween tectrittu after trantin
227
(a) arbitrary teacher
(b) WI teacher
(c) S1 teacher
Figure 6: RMD comparisons between training regimes, for different values N, at the end W1-initialized training for Ne epochs. Each column corresponds to teacher with, respectively, an arbitrary, WI and SI distribution. Row displays RMD² (M (V )²) for the different regimes, to evaluate to what extent the training remained WI. Row 2 displays the RMD between DA, FA and vanilla training regimes; well as comparison between EA and the projected particles of FA.
364
1000
Detave Spaa
20.8
(a) arbitrary teacher
(b) teacher
(c) SI teacher
Figure 7: Approximation of T,2 dis-taneeee between each model and its symmetrized version for increasing values of N. Each column corresponds to a different teacher as before. Row corresponds the ST-initialized experiment and Row 2 to the WI-initialized one.
55

---

Techer
Oalanca
646
12
150
24
50
180 500
500
arbitrary teacher
(b) WI teacher
SI teacher
Figure 8: Approximation of disddtttece between each model and the corresponding teacher network fe. for increasing values of N. Each column corresponds to different teacher as before; Row corresponds to the Sl-initialized experiment and Row 2 to the Wl-initialized one.
659
1200
1860
to Symmettrrd
Syemerttes
12
1900
sis
1860
arbitrary teacher
(b) WI teacher
SI teacher
Figure 9: Approximation of Li-ddistace between each model and the symmetrized teacher network f.. for increasing values N. Each column corresponds to a different teacher as before; Row corresponds to the Sl-initialized experiment and Row 2 to the WI Wl-initialized one.
Now, beyond the analysis of the underlying particle distributions after training, we turn our focus to comparisons of the resulting models. We measure some distances in L Y) πx), by approximating ||. (l.. with a Monte-Carlo sample of 100 random data points drawn from m.
Figure shows that the observed behaviour for the underlying particles, vi of each model, is consistent with the behaviour of the obtained model Φ That is, as particles become close to being symmetric some sense, the resulting shallow model also becomes increasingly equivariant well.
Finally, Figure iilureeie quite well the observations from Corollary [ When our teacher isn't 'tn'ten equivariant, models trained using any kind of SL technique end up suffering from the inductive
56

---

bias introduced by the symmetric assumption (something that's hinted by the symmetrization gap characterization from Lemma presentee in SuppMat-B.3) On the other hand, the vanilla model thrives in approximating S. as it is capable of breaking the assumed symmetry, unlike DA, FA and FA. The training regimes that use SL. techniques are effectively approximating La-f., as shown in Figure | (and proven in Corollary D. We also notice that, for the WI-initialized experiments, EAs end up suffering from their constraint of staying within E as they can't approximate f. (or f.. well as DA or FA (even when the teacher is SI). This isn't the case for the Sl-initialized experiments, where the performance of DA, FA and EA (and vanilla only for equivariant f-) is quite closely comparable (once again, hinting Theorem 6 We also notice general trend showing that, for bigger N, the approximations of fx (or, eventually, Qc.f+) become increasingly better (specially for the WI-initialization..
F.2 Heuristic algorithm for discovering EA parameter spaces
The proposed heuristic that we infer from the results on the previous experimental setting is quite
thoroughly described in Section 4.2 We only notice that, for this particular setting, the learning rate was chosen to be 20 (to better approximate the MFL conditions). Beyond the description of the proposed heuristic and the simple example visualized in Figure 2
we also provide Figure 10 here,
illustrating a possible threshold choice in that setting.
Considering E for j 0, I.... as the spaces that are discovered on each step of the heuristic, Figure 10 displays the values of: RMD2 (vg". # ) and (MD?(v and after training could be fixed to detect whenever the obtained particle distribution after training stayed on E,. In the case of this example, on steps 0 and we would decide that the training left the original space Ej. but we wouldn't 't do so on step 2, allowing us to fix E == Ez. As shown by the values of RMD"(UN,, Pgo # ), we wouldn't 't be too far off with our prediction.
); both before given heuristic step j. The red line simbolizes possible value of 8 that
Despite this proposed heuristic being potentially interesting for real-world applications, acknowl- edge that the setting where is applied here might be too simple, synthetic and idealized. On one hand, this provided 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
method presented in [72]. 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-ali mn-aligmment 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 'break' overtime. Their method seemingly works in a single training iteration, while ours iteratively constructs the invariant linear subspace &G by adding one new 'symmetry-breaking dimension' at a time. Our Theorem 5 guarantees that, the MF scale, our method won't leave EG; but we have yet to establish brrnakry-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 compat- ible with data symmetries: either finding the "right" subspace go or the "right" weights for each group convolution filter. Identifying the true underlying structure of data symmetries a much harder problem that is yet to be tackled in both cases.
G Further theoretical insights
The following result provides consistency guarantees when the regularization parameters T and are small, and is a slight extension of result in (38
57

---

Distance to Original Subspace when applying the Heuristic
RMD Comparison
vs. P_ P_D(V)
Relatie Keastine Distance (IMD)
10
10-3
10.3
fased e
fase
tasel
Heuniticc Step
Figure 10: RMD comparison between the empirical student particle distribution, rg" to both PE, #ug and (vg) eet (where is the heuristic step). These performed the beginning and the end of training every fixed heuristic step. The red line is placed the value 10-2 and represents possible threshold ö, to be used in the heuristic to determine whether training left Ej or not.
If w conttimuuus, is Gibbs measure of potential U, and both U and satisfy assumption (or alternatively, equal to 0), then Ry I T-converges to R when T, B & O. Particularly, given R, &.TT.B the minimizer of R B we have
Proposition 19 (T (evvovvereecce, in Z= of....fRiis.
Tim R( T,8-0
inf R()).
NEPa(Z)
In particular, every cluster point of (,"",,,, "),, is 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 T 0, and so we do include for completeness.
Let (Tr)EEN and (Sn)neN be two positive sequences decreasing to 0. On the one hand, since R is continuous (weakly if
in W, for other p 1) and (H) D(ullp) 0, for all Hn, + p (in the appropiate sense), we have
lim inf RT (mn) 2 lim R(p,) = R(x). +10+000
On the other hand, given μ (Z).)) consider p to be the heat kernel in Z RD and Pm 3.-Dv<x v/xx////// B.). In particular, from M (as the heat kernel has finite p-th moments) we know that a in W, (or weakly if it is the case).
M4 = * Pn
Now, since the function h(x)= == x log(x) is convex, from Jensen's inequality we get that
J Muspan )dr S [ h (pa(x u)) u(dy)do / h(pp(x)))xx===== = | h(p(x))de D logl loglV2221 23,,,,
Besides, we have (denoting here g(x) = --U(w));
| (n* Pn) log(g)dr =-
1aslg)dr=-/ - u(dy) / Pn On()U(x v)dr -c [+ / lui luPPuct)))
The last inequality is due to the quadratic growth of U: and by the same argument on m:
/ (ux p.)ddrdd = p(dy) / pa(x))(x- y)dz s C 1 / (vl
Notice that whenever = 0 or O, despite them not satisfying assumption we still get the same inequalities (since the leftmost term would be 0).
58

---

Now, as R is W Wp-cootinuous, R(Hr) n-+00
R(H), and:
lim sup R. (n Un) +-++00
R(u) lim sup Tn (f a-#oo
( (u+ pu)rde
[ (1* pri log(g)dz)
+ lim sup Bn +o0
And, as lim Bn (logoo/8/ng and the rest of the terms are bounded, we conclude that:
lim sup
R R (μ Pn) R(u)
+00
In particular, denoting by μ. B., the unique minimizer of R then from the previous expressions get Vn N and Vp P, (2):
R(p R(p...a. S Rininn (u. S R Rn pm)
So that,
lim sup R(p < lim sup a+0
He (u pr) S R(H), for all E P2(Z).
20000
Finally, we provide, for completeness, a proof of Lemma -
Proof of Lemmal 2 (based on (27] and (40)) As H < G is a compact group and H-invariant, from proposition 10 we know that f = E, Y X -] lives in f E L (X' V; xlx). Consider any f E (Xi x))x by Lemma from 27 (which applies since т x is G-invariant), we can decompose f as f = fG + fa. where Ja Qa is its symmetric part and fgJ f Qa its antisymmetric part. A standard calculation of the population risk under the quadratic loss setting (see the proof of Corollary for further insight) gives: R(f) R. Il f' and so:
(SS, Qa-f) R()) - R(Qc.)) B (I))()))))))))))---))) -(X))
-
which can be written as:
AJJ = F. (X) - -2010))-222(0))77700020000000030 (1 - Exx
-
Where we used that Fa. fa) (xxx))x) 0. The first term on the right hand side, -2f (xx 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 G-equivariant, by proposition giving us the desired result:
10
2G = f", and so: fd
Alf Qaf) = llftl 1 (xxx))
Lemma essentially says that we try to symmettrize a model with respect to group that has 'more symmetriess than what are actually observable in our data (i.e. in itself is only H-invariant, but we symmetrize with respect to > 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 's strict generalization benefit from choosing symmetric model to tackle our learning problem (which (27). In particular, whenever f is non-zero (on strictly positive mmeasuur set) there's 's strict gain in generalization power from using the symmetrized version of the model.
gives the name to the paper
59

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Within the limited space we could consecrate to those sections, both the abstract and the introduction give quite accurate descriptions of the context of the work (Mean-Field approach to overparametrized NN), the questions and problems addressed (role of symmetries and impact of symmetry-leveraging techniques, Field limit), our setting (the type of neural network models considered), the specific results obtained (as precisely it was possible in words, namely the precise description of the Wasserstein gradient flows describing the limiting training dynamics in different settings), the novelty of the paper (what can be learned from this viewpoint), and the contents of our numerical experiments. The main results are clearly understandable from the abstract, and the introduction gives a rather general description of the paper's goals, since it also leaves room for the context description and the outline of the paper. See Section for details.
from the Mean-
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We acknowledge and discuss many of the key assumptions required for both our theoretical and experimental results. The need for these assumption limits the applicability and practical impact of our work. We explicitly provide relevant insights in this regard, in dedicated section, SuppMat-(A22 A.2 Also, and without including every single instance, also discuss our limitations in: assumption T, after presenting equation (T. section and in many remarks in the SuppMat (e.g. in E.2.5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
60

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All mathematical statements presented are given complete, rigorous mathemat- ical proofs in the Supplementary Material (mainly in SuppMat-FF. We also provide complete sets of assumptions for each result, either the statement itself or (for some results where the precise assumptions are very technical and lengthy to state) previous to the corresponding proof in the Supplementary Material. Even proofs of some elementary facts or statements are provided for completeness. Furthermore, nearly every mathematical result stated in the main paper is given a comment or discussion regarding its context, its interpretation (beyond the pure mathematical value) and/or idea of what tools arguments its proof will rely upon.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Together with our submission, we include ZIP file with the necessary code (and instructions, following the guidelines provided by the conference) for replicating all of our experimental results. We also provide the necessary details understand and reproduce our work both in Section and (more extensively) in SuppMat-P
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same datasct, provide access to the model. In general. releasing code and data is often
61

---

one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: As mentioned in our previous answer, we include ZIP file with the necessary code and instructions faithfully reproduce the main experimental results. We closely follow the NeurIPS guidelines for our code submission, providing the exact commands and environments required 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 NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is 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 NeurlPS code and data submission guidelines (https
7/
gewss///gueeeeeeee CCodeSoodiiieeee ssi mollicy 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. only 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
62

---

Answer: [Yes]
Justification: In Section 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-P
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 appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The numerical experiments considered in this work, involved the application of noisy SGD dynamic (as in equation (T), 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 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-[n Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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 the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 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 figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In SuppMat-F we indicate both the computer resources employed to run different experiments (see the first paragraph) well 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:
63

---

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 compute required for each of the individual experimental runs well 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 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 httppe//peourpp.cc//ooool////cia
Answer: [Yes]
Justification: This paper rigorously conforms to the NeurlPS Code of Ethics. By its mathematical//hcorrtici 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 toirt-to-mid term. Since our numerical results come from simulations, there are no major concerns to be had related the nature of the employed dataset (in terms of privacy, consent, fair use 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 NeurlPS 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 special consid- eration 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 studies abstract mathematical models of artificial neural networks, and their properties and behaviors in specific, idealized situations and contexts. Therefore, no direct indirect social impact is expected in the short-to-mid term. It is, however, expected that the theoretical results presented could have positive impact in the long term, 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 might be useful eventually, our work 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 the paper.
Guidelines:
The answer NA means that there no socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploymeet 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate
64

---

to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: As previously mentioned, this work is mainly theoretical, and thus poses no major threats that could require the implementation of such safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We employ mostly public ML librarics and packages (e.g. objax, pyot, among others). We do, however, also involve the use of the emlp package provided by 22) These assets are all properly credited in our work (see SuppMat- and the corresponding use licenses (particularly for the use of emlp) are properly mentioned and respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the batasets..paperswtt
package should be provided. For popular datasets, paperswi thcode. cow/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
65

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Beyond the provided code reproduce the experiments presented in the paper (which is accompanied with the details and instructions needed to execute it), no assets are introduced in this publication, hence, specifie documentation needs to be provided in that regard.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codc/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
66


Reviews:
[{"review_id": "sYiZKLeL9t", "reviewer": "Reviewer_j4Mu", "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:\n\n- In considered neural networks, a unit computes $\\sigma_*:(\\mathbf{x},\\theta_i)\\mapsto\\hat{\\mathbf{y}}$ and the network computes $\\Phi_{\\theta}^N:(\\mathbf{x},\\theta)\\mapsto \\frac{1}{N}\\sum_{i=1}^N\\sigma_*(\\mathbf{x},\\theta_i)$.\n- Overparameterization means the width $N\\to\\infty$, yielding a limiting ensemble $\\Phi_\\mu$ where $\\mu$ is understood a probability measure on the space of parameters $\\theta_i$.\n- The parameters $\\theta$ are optimized using noisy SGD.\n\nThe mean-field theory of the neural networks considered above [1] aims to analyze the training dynamics of the limiting ensemble $\\Phi_\\mu$, given that its optimization using a convex loss function becomes a convex optimization problem (unlike $\\Phi_{\\theta}^N$ which induces non-convex optimization) and that its global optimum may be approximated by training $\\Phi_{\\theta}^N$ with a large width $N$. Specifically, it is known that SGD on $\\Phi_{\\theta}^N$ approximates Wasserstein gradient flow on $\\Phi_\\mu$ in the scaling limit $N\\to\\infty$ under certain assumptions, which then, under certain assumptions, converges to the global optima of the convex optimization problem.\n\nGiven the background, the paper presents a mean-field analysis given that the data generating distribution is possibly under a symmetry described by a compact group $G$, and symmetry-leveraging techniques are possibly used. The authors consider data augmentation, feature averaging (i.e. symmetrization; also called the Reynolds operator), and equivariant architectures, in accordance to literature [2]. The idea for the analysis is to distinguish between two types of symmetries for the measure $\\mu$ on parameters: weakly-invariant and strongly-invariant, where weakly-invariant refers to an invariant measure on the subspace containing possibly non-invariant elements, and strongly-invariant refers to a measure on invariant subspace. The authors first show in Section 3.1 that, given that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant, $\\Phi_\\mu$ is equivariant if and only if $\\mu$ is weakly-invariant. Then, feature averaging (and data augmentation) can be understood as being related to weakly-invariant distributions obtained through symmetrization, and equivariant architecture is related to strongly-invariant distributions.\n\nThe authors then, in Section 3.2, begin mean-field analysis on optimality under a set of assumptions including that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant (Assumption 1), and show that:\n\n- (Proposition 3) Optimizing with respect to an invariant loss function yields a weakly-invariant optimum if it has a unique minimizer,\n- (Corollary 1) If input data distribution is invariant, optimizing under data augmentation and feature averaging are equivalent and corresponds to approximating the symmetrized version of mean of the labeling function,\n- (Corollary 2) If data distribution is equivariant, optimizing under data augmentation and feature averaging provide no advantage,\n- (Proposition 4 and 5) But even when data distribution is equivariant, equivariant architectures (i.e. strongly invariant measures $\\mu$) may not achieve the optimum with respect to a given invariant loss function, unless universality properties that are \"good enough\" for the given data distribution are assumed.\n\nThen, in Section 3.3, the authors provide mean-field analysis on training dynamics, and show that:\n\n- (Theorem 3, Corollary 3, and Theorem 4) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\\mu_{0:t}$ following the Wasserstein gradient flows yields a unique weakly-invariant solution $\\mu_t$. If the data distribution is equivariant, from Corollary 2 this result applies to freely-trained neural networks without symmetry leveraging techniques, and data augmentation and feature augmentation again provide no advantages.\n- (Theorem 5 and 6) Under an invariant loss function and Assumption 1 (among others), a training dynamics $\\mu_{0:t}$ starting at a strongly-invariant initial condition $\\mu_0$ and following the Wasserstein gradient flows yields strongly-invariant solutions $\\mu_t$. Importantly, this means that if the data distribution is equivariant, $\\mu_t$ stays strongly-invariant throughout training even though it is not explicitly forced or constrained to be invariant. Furthermore, equivariance of data distribution can be dropped if one of the symmetry leveraging techniques are used, which all yield coinciding solutions.\n\nIn Section 4, the authors provide synthetic experiments focusing on validating Theorem 3-6, and then propose and demonstrate an empirical method for finding good equivariant architecture for the task at hand (i.e. the most expressive among invariant subspace) by training according to Theorem 5 and 6 until convergence.\n\n[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.", "questions": "Please see the Weaknesses section.", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- S1. The paper targets an important and challenging problem of understanding training dynamics of overparameterized neural networks under symmetries and their leveraging. The technical approach that applies mean-field analysis of generalized shallow neural networks [1] to learning under symmetries is original as far as I am aware.\n- S2. The paper considers three types of symmetry-leveraging techniques: data augmentation, feature averaging, and equivariant architectures. These three are sufficiently general to represent currently used techniques [2]. The descriptions as well as strengths and weaknesses of each technique discussed in Section 2.3 are correct.\n- S3. The authors' approach of imposing distributional symmetries on the measure $\\mu$ to describe equivariant architectures as strongly-invariant (and feature averaging as weakly-invariant given that the unit $\\sigma_*(\\cdot, \\cdot)$ is jointly equivariant) is interesting and could be potentially useful in future work.\n- S4. The results presented in Section 3 and 4 are interesting, in particular Theorems 5 and 6 which shows that overparameterization and equivariant units, combined with equivariant data distribution or symmetry leveraging, leads to the property that strongly-invariant parameter initialization stays strongly-invariant throughout the training even without explicit constraints to do so. (But also see W1.)\n\n[1] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.", "weaknesses": "- W1. I have concerns regarding Assumption 1, which is a key assumption underlying Section 3 and assumes that the unit $\\sigma_*:(\\mathbf{x},\\theta_i)\\mapsto\\hat{\\mathbf{y}}$ is jointly equivariant with respect to the input $\\mathbf{x}$, the parameters $\\theta_i$, and the output $\\hat{\\mathbf{y}}$. This assumption states that we have already (partially) baked in the symmetry constraint to our model class. This may be natural for analyzing equivariant architectures, but does not coincide with the practical uses of data augmentation and feature averaging (e.g. please see [2]), as they are generally applied to an arbitrary, non-equivariant function, while in the setup of this work they are applied on already equivariant units. This limits the usefulness of the results regarding data augmentation and feature averaging. Perhaps the setup of this work could be relevant to symmetry breaking in already equivariant architectures [3], but I am not entirely sure.\n- W2. In Line 92-93, the authors claim that \"Note that requiring an infinite and i.i.d. sample from $\\pi$ is key when letting later $N\\to\\infty$. This is not truly a restriction, since one can always sample from the empirical distribution of finite data points\". This seems problematic since (1) if we can infinitely sample from $\\pi$, then we have access to it, which violates basic assumptions in machine learning (Line 82), and (2) if we infinitely sample from an empirical distribution supported on finite set of data points, the resulting distribution would not be $\\pi$, and consequently the SGD scheme in Equation (1) would not be optimizing the objective in Equation (2). (Please correct me if I am wrong.)\n- W3. The writing of the paper could be in general improved. For example, the paper uses an excessive number of abbreviations (MF, MFL, NN, SL, DA, FA, EA, WI, SI, WGF, RMD, V), which is hurting readability. It would be also helpful to provide an example of a measure and its (non-equivalent) invariant symmetrization and projection in Definition 5; these are important concepts that are used in all later sections, but I was not able to directly grasp them from reading the text. For example, considering a uniform distribution on $[0, 1]\\subset\\mathbb{R}$, and a multiplicative group {1, -1} acting on $\\mathbb{R}$ by multiplication, may be sufficiently informative.\n- W4. The architecture discovery algorithm proposed in Section 4.2 was a bit questionable to me, since the architecture of the ensemble is already specified (at least up to a significant portion) by the jointly equivariant unit(s) $\\sigma_*(\\cdot,\\cdot)$, and the algorithm searches for their parameters within a strongly-invariant regime. I am not sure why this algorithm can be understood as doing architecture discovery, instead of doing the usual parameter optimization.\n- W5. Currently, it is not immediate how the study of the ensemble considered in the paper would benefit the development of practical neural networks that leverage symmetry. Showing how some of the currently used symmetry-leveraging deep neural networks can be viewed as such ensembles could be informative in this directions.\n\n[2] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00178, 2020.\n\n[3] Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, and Tess E. Smidt, Discovering symmetry breaking in physical systems with\nrelaxed group convolution, arXiv preprint arXiv:2310.02299, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful reading and relevant feedback on the paper. \nWe will incorporate an explicit Limitations section in the Appendix of the revised version, as well as a simple yet clarifying example of the different notions of symmetric distributions (e.g. if ${\\cal Z}=\\mathbb{R}$ with the multiplicative action of $G=$ { -1, 1 }, setting $\\mu = \\delta_x$ for $x \\neq 0$, we get $\\mu^G = \\frac{1}{2} (\\delta_x + \\delta_{-x})$, and $\\mu^{{\\cal E}^G} = \\delta_0$). The beginning of our Global Response mentions other clarifying sections we plan to add.\n\nWe next address the perceived weaknesses:\n\n**W1**\n\nYour concerns regarding the “a priori knowledge” of symmetries that’s encoded in the fact that $\\sigma_*$ is assumed jointly-equivariant, are generally justified, but miss a key point in the setting of the paper. We refer the reviewer to our Global Response, where we provide thorough revision of the assumption of $\\sigma_*$ being jointly-equivariant (also partially addressed in Appendix A.4; we will be sure to make it more explicit in the core of our work). In short, $\\sigma_*$ can describe many configurations beyond single-hidden-layer NNs, and the property of it being jointly-equivariant doesn’t necessarily imply that “everything inside it” needs to be equivariant as well. For instance, if we take $\\sigma_*$ to be the “traditional” single-hidden-layer NN and we consider the group action on $\\cal Z$ to be trivial on the hidden layer; then all our results apply with no constraints on the architecture whatsoever (see our Global Response and the proof of Proposition 4). In particular, DA, FA and “vanilla” training with distributionally equivariant data are still tightly related in terms of their optimization and WGFs (regardless of the properties of the inner activation). More interestingly, when $M$ is taken to be a more intricate action (and the unit $\\sigma_*$ respects it), the EAs are also more attractive, and our results on DA and FA work just as well.\n\n**W2**\n\nWe agree with you on the fact that performing SGD with an infinite i.i.d. sequence sampled from $\\pi$ is not equivalent to sampling infinitely many times with replacement from a (finite) given empirical measure $\\pi_k$, “previously” obtained by sampling $k$ points from $\\pi$. What we wanted to stress is that, from the theoretical point of view, our treatment of the first and second cases (often respectively called “population SGD” and “empirical SGD” in the literature) is indistinct and can be applied to both situations. We will rewrite said statement, to make our intent clearer. We also agree that having access to an infinite i.i.d. sequence of a given law can be considered a way of “accessing” the whole distribution but, of course, we do not expect this to be the actual application setting. Our results, as well as the whole MF approach to NNs, should be understood as an approximation that is valid for a large enough i.i.d. sample.\n\n\n**W4**\n\nWe believe that our discovery heuristic can be nicely justified through our numerical experiments. As mentioned in our Global Response, our experiments are run with a traditional single-hidden-layer NN without any specific regard for the architecture (by default, code frameworks often apply activations pointwise). In particular:\n- We assume only that the symmetry group acts via permutations, and little to no care is put into truly encoding anything into $\\sigma_*$ (see our Global Response for further discussion). In particular, we never assume to “know” the group $G \\leq S_n$ encoding data symmetries.\n- We simply initialize our architecture with all parameters being 0, and let it train unconstrained for a set number of epochs. Theorem 5 guarantees that such training will never leave the subspace ${\\cal E}^G \\leq \\cal Z$, which is possibly of much lower dimension than $\\cal Z$ (e.g. imagine convolutional layers as a subspace of all possible linear layers between images).\n- Our heuristic iteratively discovers ${\\cal E}^G$, by capturing the directions on which the trained parameters “break” the symmetry (under the assumption that this indeed happens). Despite not explicitly “discovering” new ways of building NNs, it does provide the “most expressive weight-sharing scheme” that respects data symmetries. It also yields a “basis” for such space, allowing to later build NNs with $dim({\\cal E}^G) << dim(\\cal Z)$ parameters.\n- Indeed our heuristic is only doing parameter optimization, but our theory (cf. Theorem 5) guarantees that it will never leave a strict, possibly much smaller subspace of $\\cal Z$.\n\nIn short, the “discovery” part is that we start knowing little to nothing about our data symmetries (e.g. only that they are given by a finite group acting via permutations), and we end up with a way of building strongly-equivariant NN (i.e. more “parameter-efficient” than what we started with) in the smallest possible parameter subspace signaled by the data itself.   \n\n**W5**\n\nArchitectures such as DeepSets, GNNs and CNNs (which are massively used in practice) were key pieces of inspiration for the development of our work. Namely, we know that our “unit” $\\sigma_*$ allows for modeling such equivariant architectures both as wide single-hidden-layer networks, and as large ensembles of multilayer NNs (both in such a way that $\\cal Z$ is a large linear space and ${\\cal E}^G$ represents exactly the desired EA). We will be sure to include a section in the Supplementary Material where some of these practical examples shall be described in detail. Knowing this, and despite the theoretical nature of our work, we believe that our results in Section 3 could provide key guarantees and ideas for users of practical equivariant architectures.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "655kkZCOUN", "reviewer": "Reviewer_2QUz", "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.\n\n- DA and FA are equivalent in terms of optimal values and also equivalent to standard training (free training) for WI measures (Theorem 2).\n- Furthermore, this equivalence is shown over the space of general measures when the data distribution is equivalent (Corollary 2).\n- WGF for an invariant functional preserves weak invariance of the measure (Theorem 3 and Corollary 3). Moreover, there is an equivalence of WGF for DA, FA, and free training under suitable conditions.", "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?\n- Is WGF essential for the results in Section 3.3? Do similar results hold for other dynamics (e.g., underdamped version)?", "limitations": "Limitations are addressed in the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes the notion of WI and SI and proves several equivalences between DA, FA, and free training. These results indicate when these techniques (DA and FA) are meaningful.", "weaknesses": "- This work mainly focuses on two-layer neural networks; hence, it is unclear if the results can be applied to deep neural networks.\n- The results, especially for WGF, seem to fully rely on the mean-field limit $N\\to\\infty$. Hence, it’s unclear whether the theory has some implications for finite-neuron settings.\n- (Minor) The paper is somewhat hard to follow because of the numerous notations. I suggest including a table that summarizes these notations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their relevant feedback about the paper. We will make sure to address the detected weaknesses in our final version, specifically:\n\n**W1: Regarding our work mainly focusing on two-layer neural networks.**\n\nAs mentioned in our Global Response, the full MF theory for deep NNs is an open, active research endeavor, in which sustained progress is being done. We think that a significant part of our results can be extended to settings in which MF results on deep models are (or will be) available; this has been left for future work, however, since the involved technicalities are usually quite intricate. Nevertheless, as noted in our Global Response, our current setting already allows to go far beyond two-layer NNs (and we will be sure to further stress this in the core of our work). Indeed, our \"activations\" $\\sigma_*$ are general functions on parameter spaces that can, by themselves, encode complex deep architectures (including e.g. kernels in CNNs, heads in Graph Attention Networks, etc). Thus, the resulting \"shallow models\" can represent large stacks of such (jointly-trained) multilayer “units”; i.e. deep models which, although not fully general, are way more interesting than single-hidden-layer NNs. Finally, as put forward in most of the growing literature on the MFL of overparametrized NNs (see references cited in the paper), the MF behavior can already become apparent for reasonable/realistic finite values of $N$; which speaks to a possibly broad applicability of such framework. Along this line, see our comments below regarding W2.\n\n**W2: On our results fully relying on the mean-field limit as $N\\to \\infty$ and the theory having no clear implications for finite $N$**\n\nAlthough not fully conclusive, the MFL and its associated WGF is one of the most promising approaches to mathematically explain the generalizability capacity of the SGD training dynamics of overparametrized NNs. Furthermore, it is the asymptotic approximation of NNs known to better describe their feature learning capabilities [36, 37, 38, 46]. In practice, as shown both by experiments in those papers and by our numerical results, the MF viewpoint provides truly useful insights (and new results) on the training dynamics, that manifest themselves even for relatively small numbers of neurons (1000 was enough in our case). \n\nBut not only does the MF theory provide new insights for finitely many (though large numbers of) neurons, that would be unaccessible from a fixed $N$ standpoint (as is the case for Theorem 5). Ongoing research in this field (based on theoretical tools such as propagation of chaos, optimal transport, WGFs and functional inequalities, see e.g. [9,18]) also has the potential to provide non-asymptotic estimates and quantitative answers to concrete questions with practical implications (e.g. How many neurons and training steps should be enough in order to attain a given level of generalization/population error? Which architecture/training method takes better advantage of given, limited resources?). We thus expect the MF theory to increasingly bring useful insights and concrete implications for finite-neuron settings, including for the issues addressed in this work. We will give further details on non-trivial implications of our theoretical results, for the finite-N setting, in the final version of our work.\n\n**W3: On the paper being somewhat hard to follow because of the numerous notations**\n\nAs stated in our Global Response, we will include a table explaining our relevant notation. \n\nWe now address the proposed questions:\n\n**Q1: Regarding the projected noisy SGD dynamics.**\n\n- Indeed, as we also mention in the core of our work, computing the projection onto ${\\cal E}^G$ is a highly complex problem (see [23, 24]) that can’t always be efficiently solved in practice. However, our alluded results hold even taking $\\beta = 0$; that is, without including any projected noise (and thus, without ever needing to compute it).\n- As mentioned in our remarks and our experimental results, any shallow NN with all parameters initialized at {0}, being trained freely with traditional SGD and without noise, will satisfy Theorems 5 and 6 in the MF limit. We believe that the previous conditions (modulo the shallowness of the NN and taking a big N) are all fairly easy to achieve in practice. In fact, this is exactly the case of our numerical experiments (where we never assume to know G and never explicitly compute such projection; see our Global Response for details).\n- The projection of the noise is only included as a tool to allow for easy “global convergence” guarantees for the dynamic within the space of SI-distributions (following the traditional MF literature on the topic [9, 28, 38], among others): a “general” noise cannot be used, as it would always remove the dynamic from ${\\cal E}^G$. In short, the projection of the noise in (4) is far from being fundamental in the proofs of Theorems 5 and 6, and their further application.\n\n**Q2: Regarded dynamics beyond the WGF**\n- In general, we remained within the domain of the classical WGF dynamics following the cited (standard) MF literature on the topic; without looking much into other variants of the dynamics, but also without discarding the applicability of our results therein. \n- Indeed, we believe Theorems 3 and 4 to be “natural” results, and expect them to hold for many potentially different training dynamics (e.g. Wasserstein sub-Gradient Flows [13], underdamped dynamics [FW23], annealed dynamics [12], among many others).\n- From [24], we know that Theorem 5 holds in a finite-N setting when exact DA is applied during training. The interesting MF phenomenon is that it also holds under free training with randomly sampled data (Theorems 5 and 6). This result doesn’t seem immediate to generalize for other asymptotic dynamics (e.g. annealed or underdamped), and we thus consider it an interesting question to be tackled in future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gAS7dlCo99", "reviewer": "Reviewer_4cbm", "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.", "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?\n* Proposition 5 states that equivariant architectures can lead to optimal solutions provided some universality conditions. How restrictive are these conditions?\n* Is the property of data equivariance (defined in L150) equivalent to the marginal distribution on $\\mathcal X$ being invariant and the function $f: \\mathcal X \\to \\mathcal Y$ being equivariant, given that the mapping from $\\mathcal X$ to $\\mathcal Y$ is deterministic? If so, would it be too restrictive to assume the invariance of $\\pi_\\mathcal X$?\n* In Section 4.1, the authors show that SI-initialized but freely-updated models remain in $\\mathcal E^G$ for large $N$s. This is indeed an interesting phenomenon. I wonder if this is related to the exact WI-initialized teachers. Can you provide, in addition to the arbitrary particles in Appendix D, the WI-initialized particles? Most likely, sampling from the pushed-forward empirical measure would result in a non-symmetric teacher, but there's a chance that the WI teacher might contain two pairs of symmetric parameters (e.g. $\\theta_1^*, g\\theta_1^*, \\theta_2^*, g\\theta_2^*$) and another singled out, making it mostly equivariant. After all, if the teacher is not equivariant, then the data is not equivariant, and the assumptions in Theorem 5 do not hold.\n* Section 4.2 reminds me of another work on discovering symmetry [1], which initially specifies a large candidate group and discovers the subgroup symmetry by inspecting the weights of relaxed group convolution filters. Can you comment on the connection between your method and theirs?\n\n## References\n[1] Wang, Rui, et al. \"Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution.\" Forty-first International Conference on Machine Learning.", "limitations": "The authors have discussed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The math in the paper is rigorous and detailed. Complete proofs and abundant examples and remarks are provided. As more of a practitioner myself, I can understand the main ideas of the paper without much difficulty, though I did not check the proofs carefully.\n* The experiments are compact and well-designed, clearly supporting the theoretical results.", "weaknesses": "* This is not essentially a weakness, but the theoretical results don't bring many interesting insights to me. They mostly match my existing intuitions. E.g. for equivariant data, freely-trained models tend to have the same training dynamics as symmetry-aware models. I will leave it to other reviewers to comment on the theoretical significance of these results.\n\n* I suggest the authors provide more detailed explanations about some mathematical terminology, e.g. the pushforward of a probability measure; dt-a.e. standing for almost everywhere wrt dt. This could make the paper more accessible to the general audience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer’s insightful feedback. We will address the signaled weaknesses in our final version, including in the Appendix a table with relevant notation and further reference for the main concepts. We will also further stress those aspects that, we believe, go beyond the usual intuition on symmetrically-trained NNs. These include the non-trivial, quite interesting fact that models freely-trained on equivariant data, stay for big N within the space of equivariant parameters; and the non-necessarily intuitive difference between two types of parameter distributions that naturally produce equivariant models: WI measures (which define them in general), and SI measures (which encode traditional “Equivariant Architectures” and have possibly fewer approximation guarantees).\nWe next answer the questions posed:\n\n**Q1: About the connection between jointly equivariant and equivariant**\n\nIn the setting of traditional NNs, the group action $M$ is determined by the action on data and on the model architecture (as displayed in Sect. A.4); and indeed a NN model parameterized by $z$ s.t. $M_g z = z$ for all $g\\in G$ corresponds to an equivariant architecture (thus, an equivariant model). A slight subtlety is that a shallow NN model can combine different $z_i$, e.g. $(1/N).\\sum_{i=1}^N \\sigma_*(x,z_i)$, thus being an equivariant function without all the $z_i$ satisfying $M_g.z_i = z_i$ (i.e. without explicitly having an \"equivariant architecture\" in the traditional sense). This is the core principle behind WI distributions. Further details on the assumption of $\\sigma_*$ being jointly equivariant (i.e. equivariant in both arguments) are given in our Global Response.\n\n**Q2: Regarding Proposition 5**\n\nDepending on the group considered and the architecture itself, these conditions may be generally satisfied. As mentioned in the Sect. A.4, important architectures such as CNNs and GNNs (see [35, 58, 59] for further discussion) are known to be “first-order-universal”, which implies that, taking the infinite-width limit in our setting, universality in the target class of models is guaranteed. We will elaborate more on this in Section A.4  in the final version.\n\n**Q3: On the property of data equivariance in L150**\n\nThe reviewer is right. In general, $\\pi$ G-equivariant implies $\\pi_{\\cal X}$ G-invariant [BT20]. Further, if $\\pi$ is equivariant and such that $Y = f(X)$ for some function $f: \\cal X \\to \\cal Y$, Prop.10 in Sect. A.3 implies $f$ is a G-equivariant function $\\pi_{\\cal X}$-a.e. Conversely $\\pi_{\\cal X}$ G-invariant plus $Y = f(X)$ for $f$ as before implies $\\pi$ G-equivariant. As discussed in the literature (e.g. [10, 24, 34]), assuming $\\pi_{\\cal X}$ G-invariant might seem restrictive in some settings (e.g. in image classification, one shouldn’t assume that “images can arrive with any possible orientation” when training), but such assumption is reasonable when “exploiting symmetries” does make sense. As mentioned in remarks after Theorems 3 and 5, one could also just neglect the assumption of $\\pi$ being equivariant, and introduce the inductive bias by applying data augmentation (which forces the marginal on $\\cal X$ to be G-invariant). Notice also that $\\pi$ “equivariant” is usually more general than the situation discussed above, including models like $Y = f(X) + \\xi$ with $f$ as before, and suitable random noises $\\xi$ possibly correlated to $X$.\n\n**Q4: On SI-initialized but freely-updated models remaining in ${\\cal E}^G$ for large N**\n\nThe exact WI-initialized teacher particles are a key part of our empirical result, as they make the teacher equivariant and thus the data distribution equivariant too. We mentioned the WI-initialized particles only implicitly in Appendix D (in terms of the symmetrization of the empirical measure), but they are exactly as stated by the reviewer: for every particle $\\theta_i$, the symmetrized version has both $\\theta_i$ and $g.\\theta_i$, with $g$ the non-trivial element of $S_2$. Explicitly (omitting the scale parameter and the transposition for clarity), they are given by:\n\n$\\theta_1 = (-1, 0, 0, 0.5)$, $g.\\theta_1 = (0.5, 0, 0, -1)$\n\n$\\theta_2 = (0.5, 1, 0, 1)$, $g.\\theta_2 = (1, 0, 1, 0.5)$\n\n$\\theta_3 = (-0.5, 0.3, 1, 0)$, $g.\\theta_3 = (0, 1, 0.3, -0.5)$\n\n$\\theta_4 = (0, -1, -0.5, 1)$, $g.\\theta_4 = (1, -0.5, -1, 0)$\n\n$\\theta_5 = (0.7, -0.7, 0.5, 0.7)$, $g.\\theta_5 = (0.7, 0.5, -0.7, 0.7)$\n\nNote that, besides them being symmetric pairs, none of these particles lie within ${\\cal E}^G$ (given by vectors of the form $(a, b, b, a)$ for $a,b \\in \\mathbb{R}$), making the empirical demonstration of Theorem 5 quite remarkable from our point of view.\n\n**Q5 on Section 4.2 and work on symmetry-discovery**\n\nWe were unaware of [WHG+23] and will properly reference it. As in our work, their method starts with the most constrained architecture (in our case, the null space), respecting symmetry w.r.t. the largest possible group; and symmetries are then \"broken\" as models are trained on data. Their method seemingly works “one-shot”, while using a weighted combination of group convolutional filters to find “symmetry-breaks”; ours iteratively constructs an invariant linear subspace of the parameter space by adding dimensions as symmetries are broken. Our Theorem 5 guarantees that our method won’t leave ${\\cal E}^G$, and part of our future work is establishing symmetry-breaking guarantees for our heuristic (comparable to their Proposition 3.1). The latter is discussed around L.1535 in Section C.3.2. in light of the MF interpretation. Finally, both methods aren’t really discovering the “underlying symmetry” of the data, but rather “optimizing architectures” compatible with such symmetries (be it the “right” subspace of equivariant parameters, or the “right” weights for each convolutional filter). Identifying the true underlying structure of symmetries is a much harder problem that is yet to be tackled in both cases.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NY24H299SZ", "reviewer": "Reviewer_aSNH", "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.", "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?\n- Can GNNs or Transformers also be encoded in your class of equivariant architectures? If so, these example would significantly broaden the target audience.\n- In future work, in which ways could Assumption 1 be relaxed?", "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:\n\n- The provided results are only asymptotic and only consider gradient flow on shallow models as opposed to edge of stability dynamics in finite deep models.\n- The practical relevance of the findings and proposed algorithm at finite width for practical architectures, datasets and optimization procedures could be empirically verified. In particular, does vanilla training remain SI throughout training beyond the simple student-teacher setting?", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper provides several deep and quite general theoretical results on learning with symmetries in the mean-field limit. Although the results are quite technical, the authors manage to present them in good clarity and embedded in the related literature of mean field limits of shallow models. Considering the examples of Data Augmentation, Feature Averaging and Equivariant Architectures seems natural and provides several interesting insights, for example the equivalence of DA, FA and even free training under symmetric data.", "weaknesses": "The paper already seems quite complete and well written and I could not identify significant weaknesses. The following weaknesses are only minor points:\n\n- The description of Figure 1 is hard to parse. Markers like (a)-(d) could reduce ambiguity of what is meant by a column.\n- The markers in Figure 2 should be larger and it is unclear to me which additional insight to draw from the second row.\n\nAs I am not an expert in the mean-field and Wasserstein gradient flow literature, I did not verify the correctness of all proofs, and did not double-check the claimed novelty.\n\n**Typos:** 22: lastly, 65 and 66: b=N?, 75: Definition name in bold for consistency?, 120: […], 275: will, 329: a heuristic", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough reading and the relevant comments and suggestions.\nRegarding the detected weaknesses of the paper, we will make sure to address them by providing clearer descriptions of the relevant figures, and correcting the detected typos. For instance, we shall complement the description of Figure 2, emphasizing the role of the second row as a “parallel” view of the plane ${\\cal E}^G$, that allows us to visually check that particles indeed remain inside it (see Figure 5 for an example where this doesn’t occur).\nWe next address the questions posed by the reviewer:\n\n**Q1: Regarding practical NNs, the entropy term in practice, and extensions to deep networks and edge-of-stability (EoS) dynamics.**\n- In our Global Response, we provide a deeper discussion regarding the applicability of our setting to practical NNs. Namely, we stress that the MF limit is closer to real applications than other asymptotic regimes, and that our numerical experiments correspond to quite “practical” settings (a “usual” single-hidden-layer NN architecture trained with the “usual” SGD). That is, despite our results being “asymptotic”, we believe they are applicable to reasonably practical NN implementations (with relatively normal width $N$).\n- Regarding the entropy-regularization, despite it being mostly a theoretical tool to ensure uniqueness and global convergence, it also corresponds in practice to considering a “noisy/Langevin” SGD dynamics. Indeed, when taking the limit $N\\to \\infty$, the presence of properly scaled noise in the SGD dynamic exactly translates to the appearance of the entropy term in the WGF. Therefore, in a sense, the entropy term is truly realized in practice through the introduction of noise.\n- We most surely hope to extend our results to general deep neural networks, for which we believe that some of the symmetry-related theorems (Section 3.3) might still hold. As discussed in our Global Response, however, this requires a more mature (and unified) mean-field theory of deep neural networks, which is yet to be established.\n- We haven’t specifically looked into SGD with large learning rates and EoS dynamics, but it indeed seems to be an interesting direction for future work. Now, we know that some of our proofs (e.g. for Theorems 3 and 4) can be applied in the setting of works like [18], where the MFL is established under a fixed learning rate (which may be large, and doesn’t necessarily decrease with N). This can be seen as a first preliminary result in the suggested direction.\n\n**Q2: Regarding applicability on GNNs and Transformers.**\n\nAs mentioned to other reviewers (see Global Response), very commonly applied equivariant architectures (such as CNNs and GNNs) can be exactly modeled in our setting: either as “shallow” models, or as large ensembles of multilayer models; and such that ${\\cal E}^G$ exactly encodes the desired EA. To clarify this point, we will include a section in the Supplementary Material where some of these relevant examples will be explicitly modeled in our framework (e.g. a “shallow” GNN, under the right action of $G = S_n$ on $\\cal Z$).\n\nAs for Transformers, it is well known that they can be seen as a particular type of GNN; and so, in particular, they can also be modeled using our framework (again, as \"shallow\" Transformers or also large ensembles of Transformers). \n\n**Q3: Regarding Assumption 1.**\n\nWe believe that only a couple of our assumptions are truly fundamental for achieving our results; namely, the convexity and joint invariance of $\\ell$, and the joint-equivariance of $\\sigma_*$. On the other hand, some of the more “technical” assumptions we believe can be relaxed. For instance, it is known (see e.g. [9, 51]) that the boundedness of $\\sigma_*$ can be replaced by assumptions on the data distribution $\\pi$ (e.g. bounded “finite-moments”, or compact support). Also, passing through the Wasserstein sub-Gradient Flow (as in [13]), and considering weak solutions, we believe that some of the differentiability constraints on $\\ell$ and $\\sigma_*$ could be potentially lifted.\n\n**Regarding the Limitations**\n\nGenerally speaking, we plan on tackling many of the detected limitations of our results in upcoming works. Namely, we will be sure to continue experimenting with our results (particularly, Theorems 5 and 6) on more realistic and practical datasets (i.e. beyond the simple student-teacher setting). We also comment on the applicability of our work in our General Response.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their insights. We here present information to help transversally clarify some of their questions. References from the manuscript are mentioned as presented therein; new references follow the format at the end of this response. 

First, for better understandability of the paper, in the Appendix we will add: a table explaining relevant notation/abbreviations; a summary for a wide audience on fundamental notions used in the paper; a summary of Limitations clarifying its applicability; and an explanation on how some well-known equivariant architectures (CNNs, GNNs, DeepSets) enter our setting.

We next give general context on how close our setting is to practical applications of NNs:
- Despite the MF limit of NNs being a theoretical tool, it is the asymptotic regime that most closely describes the actual behavior of large NNs during training (as compared e.g. to the Neural Tangent Kernel), and we believe truly useful insight can be obtained from it: our experiments show that the predicted MF behavior emerges in practice already for finite, not too large $N$ (1000 was generally enough), in  reasonable practical settings of shallow NNs (standard pointwise sigmoid activation and objax’ default SGD training were used).
- A fully unified, satisfactory MF theory for deep NNs is still an open, actively tackled question (for advancements on it see e.g. [SS21],[NP23]). We believe some of our key results (e.g. Theorem 3) can be extended to such settings, undoubtedly in a future research line. Notice though that $\sigma_*$ can by itself encode a complex deep architecture (see below), 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).

We also clarify further the “jointly-equivariant” assumption on $\sigma_*$:
- We defined our shallow models through an abstract, general $\sigma_*$, termed “activation” in the MF literature, but more complex in general than the “usual” activations. Since no constraints (beyond being orthogonal) are put on the action of $G$ on $\cal Z$, this “joint-equivariance” plays a key role in connecting the G-action on the data with the one on parameters. It doesn’t, however, imply that all “inner components” of $\sigma_*$ are equivariant.
- The joint-equivariance of $\sigma_*$ can encode a wide range of situations and  allows us to find interesting results beyond particular choices of architecture. It also encodes possible trade-offs between “freedom” in the architecture and “intricacy” of the group action $M$.
- In “traditional” single-hidden-layer NNs (cf. Appendix A.4) $\sigma_*$ is defined in terms of a “usual” activation function $\sigma:\mathbb{R}^b \to \mathbb{R}^b$ on the hidden layer and $M$ in terms of G-actions on the data (via $\rho$ and $\hat{\rho}$) and on the hidden layer (via $\eta$). These choices determine how constraining the “jointly-equivariant $\sigma_*$” assumption is. For instance: If $\eta$ is trivial (i.e. $\eta \equiv id_{\mathbb{R}^b}$), any single-hidden-layer NN is jointly-equivariant, regardless of $\sigma:\mathbb{R}^b \to \mathbb{R}^b$ (see e.g. the proof of Prop. 4). Similarly, any multilayer architecture under a G-action acting trivially on the hidden layers (but not on input/output), also results in a jointly-equivariant $\sigma_*$. Thus, our results apply to general architectures if the G-action on $\cal Z$ is well chosen. The catch is that, under such G-actions, the equivariant architectures encoded in ${\cal E}^G$ might end up being uninteresting (and so, Theorem 5 loses part of its punch); however, results relating DA, FA and vanilla training still apply.
- If G acts via permutation of input coordinates (as commonly for finite groups), $\sigma_*$ will be jointly-equivariant for any activation function $\sigma:\mathbb{R}^b \to \mathbb{R}^b$ applied pointwise (a common practice). The same holds in many other interesting cases involving finite groups (e.g. $Z_n \times Z_n$ acting on square images, $S_n$ acting on graphs/sets, etc). This discussion is also tackled in [24] regarding a similar assumption. For a more complex (possibly infinite) group acting orthogonally on the data and parameters (with non-trivial $\eta$), for $\sigma_*$ to be jointly-equivariant we have to restrain the chosen $\sigma:\mathbb{R}^b \to \mathbb{R}^b$. For instance, an O(b)-equivariant $\sigma$ (e.g. a Norm-ReLU) would grant our results to apply; but this particular choice could potentially harm the model’s expressiveness and applicability.

Last, we thank the reviewers for pointing out relevant references. We shall properly cite [WHG+23] in Section 4.2. We have also become aware of the work [HC23] describing a result comparable to Corollary 3, but in the particular case of ReLU activations and under a single symmetry transformation. We think our work branches far beyond both mentioned papers, and proper references for them will be included.

**References**
- [BT20] Benjamin Bloem-Reddy and Yee Whye Teh. 2020. Probabilistic symmetries and invariant neural networks. J.Mach.Learn.Res. 21, 1, Article 90 (January 2020), 61 pages.
- [FW23] Qiang Fu and Ashia Wilson. Mean-field Underdamped Langevin Dynamics and its
Space-Time Discretization. In: arXiv preprint arXiv:2312.16360 (2023).
- [HC23] Karl Hajjar and Lénaïc Chizat. On the symmetries in the dynamics of wide two-layer neural networks. Electron.Res.Arch., 31(4):2175–2212, 2023.
- [NP23] Phan-Minh Nguyen and Huy Tuan Pham, A rigorous framework for the mean field limit of multilayer neural networks. Math.Stat.Learn. 6 (2023), no. 3/4, pp. 201–357
- [SS21] Justin Sirignano, Konstantinos Spiliopoulos (2021) Mean Field Analysis of Deep Neural Networks. Maths.Oper.Res.47(1):120-152.
- [WHG+23] 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.

Meta-review: 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.
Decision: Accept (spotlight)