Propagation of Chaos for Mean-Field Langevin Dynamics and its Application to Model Ensemble
Authors: Atsushi Nitanda
Anzelle Lee
Damian Tan Xing Kai
Mizuki Sakaguchi
Taiji Suzuki
Venue: ICML 2025 poster (2025)
Keywords: mean-field neural network
mean-field Langevin dynamics
propagation of chaos
finite-particle approximation
defective log-Sobolev inequality
Primary area: theory->deep_learning
TL;DR: Propagation of chaos for mean-field Langevin dynamics

Abstract: Mean-field Langevin dynamics (MFLD) is an optimization method derived by taking the mean-field limit of noisy gradient descent for two-layer neural networks in the mean-field regime. Recently, the propagation of chaos (PoC) for MFLD has gained attention as it provides a quantitative characterization of the optimization complexity in terms of the number of particles and iterations. A remarkable progress by Chen et al. (2022) showed that the approximation error due to finite particles remains uniform in time and diminishes as the number of particles increases. In this paper, by refining the defective log-Sobolev inequality---a key result from that earlier work---under the neural network training setting, we establish an improved PoC result for MFLD, which removes the exponential dependence on the regularization coefficient from the particle approximation term of the optimization complexity. As an application, we propose a PoC-based model ensemble strategy with theoretical guarantees.

Propagation of Chaos for Mean-Field Langevin Dynamics and its Application to Model Ensemble
Atsushi Nitanda 123 Anzelle Lee 4
Damian Tan 312 Mizuki Sakaguchi 5
Taiji Suzuki 67
Abstract
training can now be formulated as the optimization over the space of probability distributions (Nitanda & Suzuki, 2017).
Gradient descent for MFNNS exhibits global convergence (Chizat & Bach, 2018; Mei et al., 2018) and adaptivity (Yang & Hu, 2020; Ba et al., 2022). To improve stability during training, one may consider noisy gradient training by adding Gaussian noise, giving rise mean-field Langevin dynamics (MFLD) (Mei et al., 2018; Hu et al., 2019). MFLD, with = 30, also achieves global convergence to the optimal solution (Hu et al., 2019; Jabir et al., 2019), with an ex- ponential convergence rate under the uniform log-Sobolev inequality (LSI) (Nitanda et al., 2022; Chizat, 2022) in the continuous-time setting.
Mcan-field Langevin dynamics (MFLD) is an optimiza- tion method derived by taking the mean-field limit of noisy gradient descent for two-layer neural networks in the mean-field regime. Recently, the propagation of chaos (PoC) for MFLD has gained attention as it provides a quantitative characterization of the optimiza- tion complexity in terms of the number of particles and iterations. A remarkable progress by Chen et al. (2022) showed that the approximation error due finite par- ticles remains uniform in time and diminishes as the number of particles increases. In this paper; by refining the defective log- Sog-ooolev inequality a key result from that earlier work- the neural network training setting, we establish an improved PoC result for MFLD, which removes the exponential dependence on the reg- ularization coefficient from the particle approximation term of the optimization complexity. As application, we propose PoC-based model ensemble strategy with theoretical guarantees.
However, the mean-field limit attained N = can- not be accurately replicated in real-life scenarios. When employing finite-particle system Px, the approximation error that arises has been studied in the literature on prop- agation of chaos (PoC) (Sznitman, 1991). In the con- text of MFLD, Chen et al. (2022); Suzuki et al. (2023a) proved the uniform-in in--immi PoC for the trajectory MFLD. In particular, in the long-time limit, they established the bounds C (14))) - L(u,) = O (~*). where a 2 exp (-0 ) is the LSI constant on proximal Gibbs distri- butions, À is the regularization coefficient, and (1("²) and L(v.) are the optimal values in finite- and infinite- particle systems. Subsequently, Nitanda (2024) improved upon this result by removing from the above bound, result- ing in O (+)) This refinement of the bound is significant as previously, the LSI constant could become exponentially small as A - O. While Nitanda (2024) also established PoC for the MFLD trajectory by incorporating the uniform--nn- LSI (Chewi et al., 2024): ((N) + (14) (N), ), this approach is indirect for showing convergence to the mean- field limit L(u.) and results in a slower convergence rate over time.
Introduction
A two layer mean-field neural network (MFNN) with N neurons is defined as an empirical average of N func- tions: Ex~ex [h (X = E M(a", where each Max', represents a single neuron with parameter x² and Px = N Ei 6R, is an empirical distribution. As the num- ber of neurons get infinitely large (N + oo). the field limit is attained: Px + 14, leading to MENN having an infinite number of particles: Ex [)),,)). Since distri- bution μ parameterizes the model in this mean-field limit,
Institute of High Performance Computing (IHPC), Agency for Science, Technology and Research (A*STAR). Singapore "Centre for Frontier Al Research (CFAR). Agency for Science, Technol- ogy and Research (A*STAR), Singapore 'College of Computing and Data Science, Nanyang Technological University, Singapore "School of Computing, National University of Singapore, Singa- pore 'Department of Artificial Intelligence, Kyushu Institute of Technology. Japan 'Department of Mathematical Informatics, The University of Tokyo, Japan " 'Center for Advanced Intelligence Project, RIKEN, Japan. Correspondence to: Atsushi Nitanda <at- sushi -nianda @ attaaa..attrruuus....
In this work, we further aim improve PoC for MFLD by demonstrating faster convergence rate in time, while maintaining the final approximation error O (?) attained ( = co. We then utilize our result to propose a PoC-based ensemble technique by demonstrating how finite particle systems can converge towards the fieldfield limit when merging MFNNs trained in parallel.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
1.1
Contributions
cle approximation error in Chen et al. (2022); Suzuki al. (2023a) and accelerates the optimization speed in Nitanda (2024); Chewi et al. (2024).
The PoC for MFLD (Chen et al., 2022; Suzuki et al., 2023a) consists of particle approximation error O (ov) due to finite-N finite-N-ppaticcles particles and optimization enror exp(- exp(-O(Aatt) This result basically builds upon the defective LSI: =5 > O,
We translate the PoC result regarding objective gap into point-wise and uniform model approximation er- rors (Theorems 4.3, 4.4, and 4.7). These results also remove the dependence on the LSI constant from the particle approximation terms in the previous model approximation errors (Suzuki et al., 2023a;b).
N (()) ((u))
8 + 2a.V
2a.
À
FT (
implicitly established by Chen et al. (2022) under the uni- form LSI condition (Nitanda et al., 2022: Chizat, 2022), where FI is Fisher information. The dependence on LSII- constant or in O ( of PoC is basically inherited from 8. In work, first remove the dependence on from 5 by introducing uniform directional LSI (Assumption 3.2) in training MFNNs setting. Based on this defective LSI, we then derive an improved PoC for MFL.D where the particle approximation error is O (+). Similar to Nitanda (2024), this improvement exponentially reduces the required num- ber of particles since the constant 2 exp (-O(+) can result
We propose an ensembling method for MFNNs trained parallel to reduce approximation error, providing the- oretical guarantees (Theorem 4.4, 4.7) and empirical verification on synthetic datasets. Moreover, going be- yond the theoretical framework, we apply our method to merge multiple LoRA parameters of language mod- els and observe improved prediction performance.
exponentially decrease as À - oc. Moreover, demonstrates faster optimization speed compared to Ni- tanda (2024); Chewi et al. (2024) due to different exponent in the optimization error terms: exp(-O(Xat)). In our analysis, a is a constant of the uniform directional LSI, which is larger than the LSI constant y1!")) appearing in the optimization error in Nitanda (2024); Chewi et al. (2024) (see the discussion following Theorem 3.8).
1.2 Notations
We use lowercase letters such as x for vectors and up- percase letters such as X for random variables Rd, re- spectively. The boldface is used for tuples of them like = (x² (x'.....,.NN) IN.., E R and X (X .x......X") X Given (z)) x denotes (x²,.... x x++1 .N). Il lla denotes the Euclidean norm. P2 (IR")) denotes the set of probability distributions with finite second moment on IRd, For probability distributions μa, v E P2(())) we define Kullback-Leible (KL) divergence (a.k.a. relative en- du tropy) by KL(p]p>) = d'(()) log de (z) and define Fisher information by Fi(Alle) = dd(=)]]V log de (x) Ent denotes the negative entropy: Ent(n)) = ((dz)) log da (z). We denote (). = f ((x)m((r) for (signed) mea- sure m and integrable function f on Rd. Given x = (x'....,,æ*) xN... E IR we write an empirical distribution supported on Px 15 Ezt.
Next, we translate the PoC result regarding objective gap into the point-wise and uniform model approxima- tion errors: [Ex=p, [N(X, =)] Ex=,, x)]] and Ex- [h(X, ]]l useful for obtain- ing generalization error bound on classification task (Suzuki et al., 2023b; Nitanda et al., 2024). Again, the bound con- sists of the sum of particle approximation and optimiza- tion error terms. Compared to the previous results (Suzuki al., 2023a;b), our bound tighter since the particle ap- proximation term is independent of the LSI-constant. This improvement directly eliminates the requirement for ponential number of neurons with respect to dimension d in their learning setup (e.g., k-parity problems (Suzuki et al., 2023b)). We also propose a PoC-based model ensemble method to further reduce the model approximation error and empirically verify its performance on synthetic datasets. To our knowledge, our study the first to provide theoretical guarantee for model ensembling of MFNNs using PoC re- sults. Going beyond the scope of the theory, examine the applicability of our method to merging LoRA parameters for language models.
|| [n(X,-)
du
2 Preliminaries
In this section, we explain the problem setting and give a brief literature review of MFLD and PoC. See Appendix for additional background information.
2.1 Problem setting
For a functional G (2( + R, we say G is differen- tiable when there exists a functional (referred to as a first variation): &G (2(IR x Rd e (H,x) + (G(x) 8p
(x) E R
8p
We demonstrate an improved PoC for MFLD (Theorem 3.8) under uniform directional LSI condition (Assump- tion 3.2). This improvement removes the dependence on LSI constant a 2 exp (-0 (1)) from the parti-
such that for any p. n¹ P2(4))
dG(p + elpl -1)
= - 88((u)
(z) (d)(d),
de
81
2

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
and say G is linearly convex when for any u. pL E (2(R)),
2Ad W on Rdy, where (W e220 is the standard Brow-
nian motion on RdN since NV F(x) = V ) (x²)
/ (x) (dxx).
G(r') 2 G(v) +
8y
(1)
(Chizat, 2022). Therefore, M4 the Gibbs distribution
(N) Law (X.) converges to
Given a differentiable and linearly convex functional Fo (2( + R and À > O, we consider the minimization problem of an antropy-regularizes convex functional:
dp 4pth) (x) x exp (- (-x FlOA))
dx
which minimizes the following entropy-regularizer linear functional defined on Pa for yM() E (2(IR
min {f(N)= = Fo(()) + Ex-x(((x)] + (En)( MEP2(R4)
(2)
C (u())) NE [F(xx)] + 1Enl( (5)
where r: Rdd is --sttrnglly convex function (e.g.. r(x) X|1=113 (N' 0)). We set F(u) (o( +
[x]]( A typical example of Fo is an empirical risk of the two-layer wrean-field neural network (see Example 3.5). Throughout the paper, we assume that the solution Ut E ( of the problem (2) exists and make the fol- lowing regularity assumption (Chizat, 2022; Nitanda et al., 2022; Chen et al., 2023) under which μlx unique and satis-
Next, we take the mean-field limit: N + oo under which Eq. (4) converges to the MFLD that solves the problem Eq. (2);
dX, = - V 8F 8p
(m) + V2ddW, 14 = Law Law(X,)
(6)
fies the optimality condition: μt, exp -- 6p
5F(u.) (see
where {We)o20 is the d-dimensional standard Brownian
Hu et al. (2019); Chizat (2022) for the details).
motion with Wo = 0. Under the log-Sobolev inequality
Assumption 2.1, There exists C1, C2 > 0 such that for
on the proximal Gibbs distribution á x exp (- 6
6y
any μ E (I))), x E TRdd V 6FF(u) 6x
(x)
S C1 and for any
Nitanda et al. (2022); Chizat (2022) showed the exponential convergence of the objective gap L(pe) - L(,)). where M. argmin (EP2 L(1))
H.p' pl E (IR I, x' E Rdd
V FF(())
(1) V =FFola)) (z')
(N) = Law (X+), is ex- Therefore, +ेr (1( (N)), where He pected to approximate L(ux) through the time and mean- field limit k - 0c, 0c,N + oc, leading to the natual question:
8p
S C2 (Wa(M, u²) )+ + 11 - a' 22112)
where W2 is the 2- Wasserstein distance.
What is the convergence rate of ((N)
(P()) to C(u.)?
2.2 Mean-field Langevin dynamics and uniform-in-time propagation of chaos
This approximation error has been studied in the literature of PoC. Recently, Suzuki et al. (2023a) proved the following mniform-in-time PoC for Eq. (3) by using the techniques in Chen et al. (2022):
First, consider the fmite-partiche setting Px N Σ i=1 dnt
N
for = (a=) E dR and the following noisy gradient de-
scent for (px). Given k-th iteration X₁ = (X), X
C ( (N)) C(ux) S exp (-Xank/2) 40 1 + 8, N
1
for each E (1,2,...., N}. we perform
(mk
5F(px.) (X() + v 2xme
X1+1 = X v/-
(3)
(Hb (1b))) -(p.) is the initial gap and 6,,NN (An+n?) XD2 (3D1, D: 0) is the dis- Aa
where A 1 (CM
where E ~ Ia) E {1, ..., NN) NY)
are i.i.d. stan-
aN
dard normal random variables and the gradient in the RHS taken for the function: oF(ex,) ) (.) Idd -> R. The 84
cretization error in time and space. The continuous-time counterpart (17 + 0) was proved by Chen et al. (2022).
continuous-time representation of Eq. (3) is given by the ( X )120:
The typical estimation of LSI-constant 2 cxp(-e(1//)) (e.g.. Theorem 1 in Suzuki et al. (2023a)) using Holley and Stroock argument (Holley & Stroock, 1987) or Miclo's trick (Bardet et al., 2018)) suggests the exponential blow-up of the particle approximation error XDz in Eq. (7) X ++ 0. aN
N-tuple of SDES (X.)
ax: = - - 5F(xx) (x))dt + v2XaW, 5u
(4)
where (Wi] E (1....,,N] N)) are independent stan- dard Brownian motions. Note that Eq. (4) is equivalent to the Langevin dynamies dx, -VV F(xx))() +
Afterward, this exponential dependence was removed by Nitanda (2024); Chewi et al. (2024) that evaluate the particle approximation error at the solution: Nv- C(N)((L) (1. - -(v.)
3

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
and optimization error: & (c)) ) (z. respectively. In the risk minimization problem setting, L(x+) < N (3C > 0) and Chewi et al. (2024) proved uniform-in- N LSI on (N) E RAN with the constant estimation à Z X X (-0 I + X + )x leading to the N- independent convergence rate of (1 (())) - ((r)) up to the particle approximation error C/N plus time- discretization error.
Here, we give connection between a conditional Gibbs distribution (+)=:(+(>>*) and proximal Gibbs distribution fxx using the Bregman divergence. Given x (x'...... x ) the following relationship holds as a probability distribution over x E IR:
Nitanda (2024) proved C
duuli (xlx=-) & exp dx
(- (~x F(enxx-)
N
N A
F(Px)+ + ( /SF(Px)
= =exp
-P PX
3
Main Result I: Improved Propagation of Chaos for Mean-field Neural Network
+ BF (Paux (,Px))
1FF(px)
N
x exp
(x)
Br BP(pxxx-1.PP
In this section, we present an improved propagation-of chaos for the mean-field Langevin dynamics under the uni- form directional LSI introduced below.
A
dîx x
(-NBr(Paxx
(x) exp
(8)
dx
Definition 3.1. For x = (x'....,, "...,,,'',,""" xx x E (1, 2, ..., N)), we define a conditional Gibbs distribu- tion (-1x((11--)) on d by
This comnection is useful for deriving the LSI on 1/1- from the uniform LSI on the proximal Gibbs distributions Px (see Nitanda et al. (2022)), in combination with Assumption 3.4 and Holley-Strook argument.
(x/xx - = I exp (- F dē²
d dæ
cxp (- N F(PzUx-:))
We give an example of training MFNNs that satisfies As- sumptions 3.2, 3.3, and 3.4.
where Prix-1 = Eiz; 8,s + 5zz.
Example 3.5 (Training MFNN). Let y C IR be a label
Assumption 3.2 (Uniform directional LSI). There exists constant > 0 such that for any E RdN and i E N}. //| --1- satisfies the LSI with the con- stant a; for all μ E P2(()) absolutely continuous w.r.t.
space, Z Rd be an input data space, h(x,)) -) Z +
R be function parameterized by x E R2, and ((.,)) R x R + R is a loss function. Given training examples (( C Z x V, we consider the empirical risk:
2- x it follows that
1 n n
e (Bx=n Po(u) = Σ
1 F(( 2a
(T.(plpp--<(11 ") S
and In-regularizzton r(z) = X||x=11- We assume that supaeRs.ee (M(x, 2)1 < R and that for any y E R, ((..)) is convex and T-smooth; there exists I 0 such that for any a,b E R. l(b, I E(a,y) + at(a.w) (b a) + 4 al2. Applying this L-smoothness with yj and taking n n... we get Bp, (Poux-t <
Remark. Wang (2024) also introduced the conditional Gibbs distribution and imposed Poincaré inequality on it.
We also make the following assumptions.
Assumption 3.3. A functional (o( is differentiable and linearly convex.
Ep 1n(x 2,)1, Ep (enx_,l((X, [N(X, 2,),,,,)1,4
Ep
average over j E (1. A(ri.a)) 2
The nonlincarity of F. is the key to the PoC analysis for mean-field models, thereby motivating the use of the Breg- man divergence associated with Fo (Nitanda, 2024); for distributions nt n1 (R2().).
Therefore, the Holley- Stroock argument (Holley & Stroock, 1987) with Eq. (8) implies an LSI constant a = on exp( 2LR² to Og as N + co, where oo = (exp(O(A-T] is the LSII- constant of the proximal Gibbs distribution Px.
4 ビラ 2n
<
LRA
2NT.
AN converging
21
(P, H() = Fo(()) Fo(H/) - 5Fo(m') 8p p (
The following defective entropy sandwich and defective LSI are key results in studying MFLD in the finite-particle setting. The proofs can be found in Appendix A. For X ~
Assumption 3.4. There exists a constant B > 0 such that for any X E TRNN x E IRd, and E (,, N}.
p( we denote by M (-Ix the conditional distribution
(N)
B
BF, (PaUX-t (Pp) S N2'
of X conditioned by X -
4

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Lemma 3.6 (Defective entropy sandwich). Suppose As- sumption 3.4 holds. Then, for any ((N) E (2T(RRN).
because fia (n()) n Br, = 0, and B in this case.
Therefore, Lemma 3.7 is instrumental in the computational complexity analysis of MFLD in the finite finite-particle set- ting as shown in the following theorem. We set A
KLL 11 ) + Ex [Brg (ox,
N
1 ((0) - -(we)
*े (C() (m000) ((())
N
Theorem 3.8 (Propagation chaos for MFLD). Suppose As- sumptions 2.1, 3.2, 3.3, and 3, hold and consider the Lz- regularization: r(z) = X|||= (X > 0). Then,
B
E Ex
(1x)))
1. MFLD in the continuous-time (4) satisfies
The lemma 3.6 can be viewed as the finite-particle counter- part of the entropy sandwich established in Nitanda et al. (2022); Chizat (2022):
B N 1 (( (2) -
+ exp((22.X) 2
N
(KT()))p.) S 2L()) L(x.) S (L(ulli)
MFLD in the discrete-time (3) with 1j2' < 1/2 satisfies
Eq. (8) says that the conditional distribution (-j=<(-(X--) approximates the proximal Gibbs distribution fix, where px is an empirical distribution consisting of X. Therefore, we expect (-y=(-(X--) to play a role analogous to the proximal distribution and lead to an entropy sandwich, In fact, we can confirm that Lemma 3.6 with (,(V) = n reproduces the above entropy sandwich in the infinite-particle system by taking N + oc under regular conditions.
1
( --((m))) B og at
-exp(-aankk))
N
where 8y = 8n((2 + (110 + Ad) +
+
Ellxxolla] + ( + xd) xd))
X²)
Proof. We here demonstrate the convergence in the continuous-time setting. The distribution ju (N) Law (X,) of Eq. (4) satisfies the following Fokker-Plance equation:
The defective LSI was originally established by Chen et al. (2022). We here derive an improved variant built upon the defective entropy sandwich. The proof can be found in Appendix A.I.
(N) 0pt
AV
(N) 14
log
du
(N)
Lemma 3.7 (Defective LSI). Suppose Assumptions 3.2, 3.3, and 3.4 hold. Then, it follows that for any (d()) PP((R
Öt
dp.
(NT
By the standard argument of Langevin dynamics (e.g., Vem- pala & Wibisono (2019)) and Lemma 3.7, we get
X N
KL KL(n 11 11 + Bx xx
[B [b)]
(1())) - C(u.)
B + 2aN - N
A
(1((())01)())))
di (C( (1)) - -NC(U) B) -XFFI( (2 (1p!")
d
N
I (a) (1. - NN(())- - B).
Lemma 3.7 gives an approximation error bound between y(N) and n which shrinks up to B/N error as !! + (N and shrinks to zero by additionally taking N co,
Then, the statement follows from direct application of the Grönwall's inequality. The convergence in the discrete- time is also proved by incorporating one-step iterpolation argument. See Appendix A.
meaning that each particle of the system (X'..... XN) ((N)) becomes independent to each other. Compared to the original result (Chen et al., 2022), the particle approxi- mation term B/ N is independent of a, similar to Nitanda (2024). Note that whereas Nitanda (2024) only consider the case of μu(N- = v.." our result allows for any distribution ((N) at the cost of the Fisher information FI(H( ())) Lemma 3.7 can be indeed regarded an extended LSI the finite-particle system and nonlinear mean-field objective, where Fisher information is lower bounded by the optimal- ity gap up to B/A B/N error. In particular, when Fo is the linear functional: (o()4) = [n (3f : Rd > R), the lemma reproduces the standard LSI on (N).
From this result, we see that MFLD indeed induces the PoC regarding KL-divergence. In fact, the following in- equality, which is a direct consequence of Lemma 3.7 with Theorem 3.8 in the continuous-time shows that the parti- cles (Xt) N + 0c:
14 become independent as L - oo and
A(N) A
KL 1I
< AN
B
+ exp(-22Ah
(9)
N
We note that the particle approximation term B/N in Theo- rem 3.8 is independent of LSI-constants, whereas the error
KL(M(N1114 (N)) <
1
(N)
O ((xx) obtained in Suzuki et al. (2023a) scales inversely
2a
5

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
with an LSI constant a' the proximal Gibbs distribu- tion which can be exponentially small as À - 0. Whereas the term B/N is comparable to Nitanda (2024); Chewi et al. (2024), our convergence rate exp(-2a.h) in optimiza- tion is faster since their results rely on the I.II constant o on p which is smaller than ag in Example 3.5'. For instance, Chewi et al. (2024) estimated the LSI constant ā 2 X' exp (-0 ( + XX + x
In particular, if y()) is exchangeable, we gel
N KITEEEEI s <KT(
Proof. The assertion holds by the direct application of Han's inequality. See Appendix A.2.
X
We here give the model approximation bound using (L(2(N) [ with the proof to show how the above lemma helps to control the correlation across particles.
4 Main Result I:: Model Approximation Error and PoC-based Model Ensemble
Proposition 4.2. Suppose y!h) is exchangeable. Then, it follows that for any Z,
In this section, we study how MFNNS trained with MFLD approximate the mean-field limit: Ex~μ. [h(X, =)]. More- over, we present a PoC-based model ensemble method that further reduces the error. Throughout this section, we focus on training MENNS (Example 3.5) and suppose sup [h(x, 2)1 < R.
Exapinn
(Ex IMXX 2)] Ex [h(X, =))*
4R
KL(ulN)luu0
N + 8R²
N
4.1 Point-wise model approximation error
Proof. We here decompose the error as follows.
We consider point-wise model approximation error and Ex-p. (M(X, on each point E 3, where X
E E (Bxxpx [h(X,, 2)] E Ex 2))²
between E Ex [X)] = N -
N
(X) ,...,, ,N) (N)), The error usually consists of the bias and variance terms where the bias means the difference between (²( and N and the variance is due to finite-N particles. In general, it is not straightforward to show the variance reduction as N + xo since X; (i = 1,2,...,NV are not independent, and hence can exhibit positive corre- lation. However, in our setting, PoC helps to reduce the correlation among particles, resulting in better approxima- tion error via the variance reduction.
= N2 Fx Ex
Σ (h(Xi,,) Ex [)))²
+ N2 Ex-p(M
Σ (n(X,, 2) - -x x))
((h(X);=) - Ex-m.IM(Xxx)
Using the boundedness of h, the first term can be upper bounded by 4R2 /N. The second term can be evaluated as follows. Set (X,) = h(X,, 2) - Ex [h(X,2)) 2)]. Then,
Since we are concerned with the correlation between each pair of particles, we reinterpret (L( ( as the gap between their marginal distributions. For each index subset
S C (l,...,NN, NJ. we denote by us the marginal distri- (N)
(N)
EX~MON (X.))
E -(xXX))-1 (N) (XX)))((X))
bution of (NN) on S and write 11: = [1)
"ig Mij - (i.j) for simplicity. We say the distribution (( is exchangeable if the laws of (Xot)):..,,Xat)
(N) (N)
(N)
-XxxX)) (X()))
(B (B (M)
(x. )x (XX,)
and (XI,... XN) are identical for all permutation a
[I, + (N.
SR 11122
4R² V 2KL (N) (1
Lemma 4.1. For any integers s, N E N such that s S N, it follows that
where TV is the TV-norm and we used Pinsker's inequality. Applying Lemma 4.1 with 2, finish the proof.
In the proof, we see that KL-divergence controls the cross term by absorbing the difference between marginal distribu- (N) and 22 By combining this result with the PoC for MFLD (Theorem 3.8), we arrive at the model approxi- mation error achieved by MFL.D
K < (L(
-
tions p Lej
However, we note PoC result obtained by Chewi al. (2024) applicable to non-hounded activation functions such as ReL.U.
6

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Theorem 4.3. Under the same conditions as in Theorem 3.8, we MFLD in the discrete-fime, with ni\ 1/2 and &N Then we get
total number of neurons, the bound suggests a non-trivial choice for the number of networks M. Rewriting the bound using M and K, we obtain O shows that M induces a trade-off, and the bound achieves its minimum value of O (AK))T)) when M = 1/3/1/33
Xo Mo
K + VAMK + .M. X This
E Xxupk
(N)
[Bxmpx 2)] - -x-xxx =11)
4R²
AN + bn
B
ad2 + +exp (ap(-a)nk) 4/4)
This phenomenon arises because training multiple networks enhances the independence among particles.
$ N
+
8R
V
Remark. Our result can extend randomly pruned net- works. That is, we consider randomly pruning (N - s)- neurons after training the MFNN of N N-neurons. Then, we get the counterpart of Proposition 4.2 as follows.
(N) Note that exchangeability of Mk is satisfied at all itera- tions because of the symmetric structure of problem and initialization with respect to particles.
(Ex=px [MXX =)]
Ex [M(X, =)]}²]
Model Ensemble We introduce the PoC-based model en- semble to further reduce the approximation error. We first train M MENNS of N-neurons in parallel with the same settings and obtain sets of optimized particles X, (j 1, 2, . .M) where each X; (X) .. ) represents each network and they are independent to each other. We then integrate them into a single network of M N-nurross as follows:
4R²
+ 8R
N
where Pi:s is the distribution of remaining neurons. More- over, Theorem 4.4 can also extend to the ensemble model of randomly pruned networks in the same way.
(N)
Ex Ex~px, IM(X,
1
Σ
M
N
n(X) 2).
4.2 Uniform model approximation error
M
MN
j=
We here consider uniform model approximation error over 2 E Z, which is more useful than point-wise evaluation in the machine learning scenario such as generalization analysis. The uniform bound essentially requires complexity evaluation of the model, and hence make the additional assumption to control the complexity. Assumption 4.5.
(10)
Because of the independence of networks (X, ) vari- ance reduction occurs, resulting in the improved approx- imation error. Indeed, by extending Proposition 4.2 into an ensemble setting (see Proposition A.I) and using PoC (Theorem 3.8), we get the following bound. The proof is deferred to Appendix A.2.
The data domain is bounded: Z C [-1,, C Rd
Theorem 4.4. Under the same conditions as in Theorem 3.8, we run M- -pralllll MFLD in the discrete time independently, with y 1/2 and X..0 Po
There exists 3 > 0 such that for any x E IR 2, E Z,
(j == 1,2,.... M). Then
|A(x,, 2) h( 2²)l S ||| 2'112-
[
27
For example, h(1, 2) 2 Lanhh(=))), (x* E Rd, (22 x²) E RR2,, used in Suzuki
R (tanh(x'T + x²) +
1
Bxx
E Epx [k(X,2)] - -E,, [h(X,
M
et al. (2023b) satisfies the above assumption with / = R due 1-Lipschitz continuity of tanh.
3
4R²
8R²
B
+ 8n
4 (N)
Given random variables (X,) (X, (X3 (X...... X²))), )),
+
+ eexp(-adnk)
MN
M
AN
ax²
we consider the empirical Rademacher complexity of the function class
B
8n
(0
F {r + h(x, 2) E 2):
+ 2R²
+ +exp(-aXft
XN
ax²
M
R (F) = E,
1 MN
ΣΣ
fxxx))
sup fEr
where X,,k is the particles at k-iteration for j-th network.
where the expectation is taken over the Rademacher random variables σ = (ot) which are i.i.d. with the probability Ploj = 1] P -1] = 1. Here, we utilize the uniform laws of large numbers to evaluate the approximation error of an ensemble model defined by μ note that the result for a single model is obtained as a special case M 1.
For simplicity, we consider the bound
2 MN
SR M
B V XN
2 obtained in the limit k + 20, + 0. This bound AN
indicates that increasing M offers better scalability than increasing N, as long as the second term dominates. In fact, under the constraint MN= = O(K), where K denotes the
7

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Lemma 4.6. Let X, ~ (j ..,,...., M) be indepen- dent random variables. For ó E (0. 1), follows that with high probability 1- ô,
networks. Finally, we further substantiate the applicability of our method LoRA training of language models.
5.1 Mean-field neural networks
Epx, [h(x,-)] By, [h(X,-)]
To demonstrate Theorem 4.7, we compute the log sup norm between the ouptuts of an (approximately) infinite width network and merged network, both trained using noisy gradient descent. A finite-width approximation of the mean- field neural network is employed with N = No while the merged network is obtained by merging M different networks N neurons each. This is repeated across various values of N and M. See Appendix B.1 for more details about our methodology.
[RN.s((F)] RV
21/((1/6)
21E 3
MN
Proof. The lemma follows directly by applying the uniform law of large numbers to the function class F (see, for in- stance, Mohri et al. (2012).
The complexity Ex, (F) can be then eval- uated by Dudley's entropy integral (Lemma A.2) under Assumption 4.5 and the boundedness [h(x,, 2)1 | $ R. By using the variational formulation of KL-divergence (e.g., Corollary 4.15 in Boucheron et al. (2013)), we translate the result of Lemma 4.6 into the approximation error of an ensemble model obtained by M independent parallel iterates X, M4 (N) (j 1,2...... ...,..... M)) of MFLDs. Com- bining these techniques with Theorem 3.8, we conclude the following theorem.
Classification setting We consider the binary classifica- tion of n data points generated along the perimeters two concentric circles with radius Timner and Touters where labels are assigned according to the circle data point belongs to.
Regression setting We consider regression on the multi- index problem. Each input sample 2; = (s....., = E Rd is generated uniformly within d-dimensional hypersphere of radius Let g Ted TR be link function, then y: g(2) Σド= tanh(z)) E I. where k S d, k I.
Theorem 4.7. Suppose Assumption 4.5 and the same con- ditions as in Theorem 3.8 hold. Run M -pparallel MFLD in the discrete time independently, with n < 1/2 and X50 4o
Double circle
Multi-index
&N (j 1,2. M). Then we get
-2.0
1
Lntsup norm)
E ((x
Ex~pu P I(X, Exx=x [h(X,
-3.0
M
10
10
2.0
-3.5
d
dB
dA
15
15
=
R
+
-4.0
MN
AN MN(A + MB)
-3.0
200
500
100
300
400
8n 1 A + MB ax² exp(( (xp(-o>nk)44
dAMN
+0
R
Figure I. Heat maps sup norm log Oog-sc between Noo and merged networks when varying M and N.
Here, the Ô-notation hides logarithmic factors. As for the concrete bound and proofs, see Appendix A.3. The term V d B + MN(A+ MBj represents the particle MN AN
From Figure I, we see that the merged networks converge towards the mean-field limit as M and N increase. This aligns with our theoretical findings which suggests that the sup of the approximation error decreases when more particles are added (ensembling in our experimental set- up) or increasing the ensemble size. See Appendix B.2 for supplementary experiments.
dA
approximation error due to finite N particles, and even when M = 1, they improve upon the bound in Suzuki et al. (2023a;b) by removing the LSI constant a from the corresponding term. And the upper bound shows the im- provement M increases.
5.2 LoRA for finetuning language models
Experiments
Beyond the scope of the theory, we empirically verify the applicability of our ensemble technique to finetuning lan- guage models using LoRA. Given a pre-trained parameter Wo E R of linear layer, LoRA introduces low-rank
We verify the validity of our theoretical results, by conduct- ing numerical experiments on synthetic data in both clas- sification and regression settings using mean-field neural
8

---

f with respect Propagation of Chaos for Mcan-Field Langevin Dynamics
Table Accuracy comparison of LoRA and PoC-based merging for finetuning Ilama models.
Model
Method
SIQA
PIQA
WinoGrande
OBQA
ARC-c
ARC-e
BoolQ
HellaSwag
Ave.
LoRA (32, best)
79.48
82.43
81.77
80.60
67.75
80.47
70.37
86.67
78.69
Llama2
LoRA (256)
69.95
69.69
69.61
61.40
47.44
61.15
63.73
47.27
61.28
7B
PoC merge
81.17
84.60
85.16
86.60
72.53
86.62
72.45
92.79
82.74
LoRA (32, best)
81.22
89.50
86.74
86.00
79.86
90.53
72.91
95.34
85.26
Llama3 8B
LoRA (256)
81.06
87.60
87.61
84.60
78.92
90.06
75.11
94.98
84.99
PoC merge
82.04
89.39
89.27
89.20
83.28
92.30
76.33
96.58
87.30
matrices A E Nxd and B E k and represents the fine-tuned parameter as Wo+>AW = Wo+yBA (7 > 0). Then, only and B are optimized, leaving Wo frozen. Us- ing the expression AT T (a!....,dd aa*) (a' E IR2) and B = =bb'.... b (bé E )R*), we can reformalize LoRA parameter 7BA with 7 1/N as the MFNN: Rd 9 + + h((aa, 6¹), 2) E & where h(((a' b)),)) = 2. Therefore, we can apply PoC-based model ensemble for LoRA parameters. Note that the ensemble model is reduced to single (k x d)-matrix ÁW due to the linearity of the ac- tivation function, and therefore it does not require additional memory and time for inference.
bounds, and derived PoC-based model ensemble method with an empirical verification. Moreover, we substantiated its applicability to fine-tuning language models using LoRA.
One limitation of our theory is that it cannot explain the asymptotic behavior as 4 This also the case in pre- vious work since the optimization speed essentially slows down exponentially, which is inevitable in general as dis- cussed in the literature. However, there might be room to tighten the particle approximation term XN to in the model approximation bounds. This term arises from the KL-divergence, which essentially controls the cor- relation among particles, as seen in the proof of Proposition 4.2. However, KI--iiergeece may be excessive for this purpose. Therefore, one interesting future direction is to utilize PoC with respect to smaller metric that alleviates the dependence on A.
We use commonsense reasoning datasets (Hu et al., 2023): SIQA, PIQA, WinoGrande, OBQA, ARC-c, ARC- BoolQ. and HellaSwag. and use language models: Llama2-7B (Touvron et al., 2023) and Llama3-8B (Dubey et al., 2024). We first optimize the multiple LoRA parameters ((A,,B))]
using noisy Adam where 22AmT (step-size Mik standard Gaussian noise {x) is added to each parameter update of AdamW (Loshchilov & Hutter, 2019). Then, we merge them into AW I
A Acknowledgements
This research is supported by the National Research Founda- tion, Singapore, Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the Ministry of Digital Development and Information under the AI Visiting Professorship Programme (award number A AIVP-2024-004), Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Sin- gapore, Infocomm Media Development Authority, and the Ministry of Digital Development and Information. TS was partially supported by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2115).
MN Ex, b TT E gexd, Hyperparameters are set to N = 32, M 8, 10-5 and the number of epochs is 3. In Table 1, we compare the accuracy of the merged parameter with the base parmeters of LoRA. For LoRA, the row "LoRA (32, best)" presents the best result achieved among cight different LoRA parameters based on the av- erage accuracy across all datasets. For both models, we observed that PoC-based model merging significantly im- proves the finetuning performance. For comparison under the same computational budget MN 256, we report the results obtained by LoRA with M = 1 and N = 256 in the row "LoRA (256)". We also verify the performance using one-epoch training to examine the effect Gaussian noise in parameter updates. See Appendix B.3 for more details.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Conclusion and Discussion
We established an improved PoC for MFLD that accelerates optimization speed in Nitanda (2024); Chewi et al. (2024) while achieving the same particle complexity O(1/N). We then translated this result into model approximation error
9

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
References
Dietterich, T. G. Ensemble methods in machine learning. In International workshop on multiple classifier systems, pp. 1-15, 2000.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., and Yang, G. High-dimensional asymptotics of feature leam- ing: How one gradient step improves the representation. In Advances in Neural Information Processing Systems 35, 2022.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., al. The llama 3 herd of models. arXiv preprint arXiv: 2024.
Bardet, J.-B., Gozlan, N., Malricu, F., and Zitt, P.-.. Func- tional inequalities for gaussian convolutions of compactly supported measures: Explicit bounds and dimension de- pendence. Bernoulli, 24(1):333 - -353, 2018.
Frankle, J., Dziugaite, K., Roy, D., and Carbin, M. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259- 3269, 2020.
Boucheron, S., Lugosi, G., and Massart, P. Concentration inequalities: a non asymptotic theory of independence, 2013.
Ganaie, M., Hu, M., Malik, A., Tanveer, M., and Suganthan, P. Ensemble deep learning: review. Engineering Appli- cations of Artificial Intelligence, 115:105-151, 2022.
Breiman, L. Bagging predictors. Machine Learning, 24: 123-140, 1996.
Gauthier-Carde T., Siriwardhana, S., Stein, E., Ehghaghi, M., Goddard, C., McQuade, M., Solawetz, J., and Labonne, M. Merging in a bottle: Differentiable adaptive merging (dam) and the path from averaging to automation. 2024.
Chen, F., Ren, Z., and Wang. S. Uniform-in-time propaga- tion of chaos for mean field langevin dynamics. arXiv preprint arXiv:22/2.03050 2022.
arXiv preprint
Chen, F., Ren, Z., and Wang, S. Entropic fictitious play for mean field optimization problem. Journal of Machine Learning Research, 14:111): 2023.
Goddard, C., Siriwardhana, S., Ehghaghi, M., Meyers, L., Karpukhin, V., Benedict, B., McQuade, M., and Solawetz, J. Arcee's ss mergekit: toolkit for merging large language models. arXiv preprint arXiv:2403. / /3257, 2024.
Cheng, J., Bibaut, A., and van der Laan, M. The relative performance of ensemble methods with deep convolu- tional neural networks for image classification. Journal of Applied Statistics, 45(15):2800 2818, 2018
Hansen, L. K. and Salamon, P. Neural network ensembles. Proceedings of the IEEE Transactions on Pattern Anal- ysis and Machine Intelligence, volume 12, pp. 993-1001, 1990.
Chewi, S., Nitanda, A., and Zhang, M. S. Uniform-in-n log- sobolev inequality for the mean-field langevin dynamics with convex energy. arXiv preprint arXiv:2409.1044/ 2024.
Hardt, M., Recht, B., and Singer, Y. Train faster, gen- eralize better: Stability of stochastic gradient descent. In Proceedings of International Conference on Machine Learning 33, pp. 1225-1234, 2016
Chizat, L. Mean-field langevin dynamics: Exponential convergence and annealing. Transactions on Machine Learning Research, 2022.
Holley, R. and Stroock, D. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics, 46(5-6): 1 159411944 1194, 1987.
Chizat, L. and Bach, F. On the global convergence of gradi- ent descent for wer-paarmeterizee models using optimal transport. In Advances in Neural Information Processing Systems 31, pp. 3040-3050, 2018.
Hu, K., Ren, Z., Siska, D., and Szpruch, L. Mean-field langevin dynamics and energy landscape of neural net- works. arXiv preprint arXiv:1905.0r 07769, 2019.
Chronopoulou, A., Peters, M. E., Fraser, A., and Dodge, Adaptersoup: Weight averaging to improve general- ization of pretrained language models. In Findings of the Association for Computational Linguistics 61, pp. 2054-2063. 2023.
Hu, Z., Wang, L., Lan, Y., Xu, W., Lim, E.PP., Bing, L., Xu, X., Poria, S., and Lee, R. LIm-adapters: An adapter family for parameter-efficient fine-tuning of large lan- guage models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp- 5254 5276, 2023.
Davari, M. R. and Belilovsky. Model breadcrumbs: Scaling multi-task model merging with sparse masks. In Euro- pean Conference on Computer Vision 18, 2024.
Ilharco, G., Wortsman, M., Gadre, S. Y., Song, S., Korn- blith, S., Farhadi, A., and Schmidt, L. Patching open- vocabulary models by interpolating weights. In Advances Neural Information Processing Systems 35, 2022.
Dembo, A., Cover, T. M., and Thomas, J. A. Information theoretic inequalities. IEEE Transactions on Information theory, 37(6))150) 1501-1110111 1991.
10

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A. G. Averaging weights lead wider optima and better generalization. In Conference Uncertainty in Artificial Intelligence 34, pp. 376-885,24 2018.
Nitanda, A., Wu, D., and Suzuki, T. Convex analysis of the mean field langevin dynamics. In Proceedings of International Conference on Artificial Intelligence and Statistics 25, pp. 9741-9757, 2022.
Jabir, J.-F., Siška, D., and Szpruch, t.. Mean-field neu- ral odes via relaxed optimal control. arXiv preprint arXiv: 191.:4911..05475, 2019.
Nitanda, A., Oko, K., Suzuki, T., and Wu, D. Improved statistical and computational complexity of the mean-field langevin dynamics under structured data. In The Twelfth International Conference on Learning Representations, 2024.
Jain, P. Kakade, S., Kidambi, R., Netrapalli, P., and Sid- ford, A. Parallelizing stochastic gradient descent for least squaers regression: mini-batching, averaging and model misspecification. Journal of Machine Learning Research, 223:1-42, 2018.
Ortiz-Jimenez, G., Favero, A., and Frossard, P. Task arith- metic in the tangent space: Improved editing of pre- trained models. In Advances in Neural Information Pro- cessing Systems 36, 2023.
Jin, X., Ren, X., Preotiuc-Pietro, D., and Cheng, P. Dataless knowledge fusion by merging weights of language mod- els. Proceedings of the 1tth International Conference on Learning Representations, 2023.
Ramé, A., Ahuja, K., Zhang, J., Cord, M., Bottou, L., and Lopez-Paz, D. Model ratatouille: Recyling diverse mod- els for out-of-distributton generalization. In Proceedings of International Conference on Machine Learning 40, pp. 28656-28679, 2023.
Kim, H.-C., Pang, S., Je, H.-M., Kim, D., Yang, S., and Kim, B. Constructing support vector machine ensemble. Pattern Recognition, 36:2757 22677777766, 2003.
Rotskoff, G. M., Jelassi, S., Bruna, J., and Vanden- E. Global convergence of neuron birth-death dynamics. In Proceedings of International Conference on Machine Learning 36, pp. 9689-9698, 2019.
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. In Proceedings of the 7th International Con- ference on Learning Representations, 2019.
Soares, C., Brazdil, P., and Kuba, P. A meta-learning method to select the kernel width in support vector regression. Machine Learning, 54:195-209, 2004.
Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Pro-
ceedings of the National Academy of Sciences, 115(33): E7665-E7677 2018
Suzuki. T. Wu, D., and Nitanda, A. Convergence of mean-field langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. In Advances in Neural Information Processing Systems 36, 2023a.
Mohammed, A. and Mohammed, R. K. A comprehensive review on ensemble deep learning: Opportunities and challenges. Journal of King Saud Computter and Information Sciences, 35(2):754-7744 2023.
Suzuki, T., Wu, D., Oko, K., and Nitanda, Feature learn- ing via mean-field langevin dynamics: classifying sparse parities and beyond. In Advances in Neural Information Processing Systems 37, 2023b.
Mohri, M., Rostamizadeh, A., and Talwalkar, A. Founda- tions of Machine Learning. 2012.
Sznitman, A.S. Topics in propagation of chaos. Ecole d'Eté de Probabilités de Caint-Flous XIX 1989, pp. 165-251, 1991.
Mousavi-Hosseint A., Wu, D., and Erdogdu, M. A. Learning multi-index models with neural networks via mean-field langevin dynamics. arXiv preprint S.Xi::2408.07254. 2024.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Neyshabur, B., Sedghi, H., and Zhang, C. What is being transferred in transfer learning? In Advances in Neural Information Processing Systems 33, pp. 512-523, 2020.
Nitanda, A. Improved particle approximation error for mean field neural networks. In Advances in Newral Information Processing Systems 37, 2024.
Utans, J. Weight averaging for neural networks and local resampling schemes. In Association for the Advancement of Artifical Intelligence 13, pp. 133-138, 1996.
Nitanda, A. and Suzuki, T. Stochastic particle gradi- ent descent for infinite ensembles. arXiv preprint 1712v 2017.
Razquez-Romero, A. and Gallardo-Antolin, A. Automatic detection of depression in speech using ensemble convo- lutional neural networks. Entropy, 22(6), 2020.

---

Propagation of Chaos for Mean-Field Langevin Dynamics
Vempala, S. and Wibisono, A. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems 32, pp. 8094-811888 8106, 2019.
Wang, P. Shen, L., Tao, Z., He, S., and Tao, D. Generaliza- tion analysis of stochastic weight averaging with general sampling. In Proceedings of the International Conference on Machine Learning 41, pp. 11142251466) 2024.
Wang, S. Uniform log-sobolev inequalities for mean field particles with flat-convex energy. arXiv preprint Stti::2008.0228888 2024.
Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., Gontjo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International conference on machine learning, pp. 23965- 23998, 2022.
Yang. E.,, Shen, L., Guo, G., Wang. X., Cao, X., Zhang, J., and Tao, D. Model merging in lims, milms,, and beyond: Methods, theories, applications and opportunities. arXiv preprint ar 2rXiv2240..0766 2024.
Yang, G. and Hu, E. J. Feature learning in infhhite-width neural networks. arXiv preprint arXiv:2011.14522 2020.
Yu, L., Yu, B., Yu, H., Huang, F., and Li, Y. Language mod- els are super mario: Absorbing abilities from homologous models free lunch. In Proceedings of International Conference Machine Learning 41, pp. 57755-57775, 2024.
12

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
A Proofs
A.1 Propagation of chaos for MFLD (Section 3)
Proof of Lemma 3.6. The first equality of the assertion proved by Nitanda (2024). We here prove the inequality by utilizing the argument of the conditional and marginal distribution of n(N) (Chen et al., 2022).
By the convexity, we have that for any E (2
[8P(xx)
11 px ) AEnt(u)) & L(v).
E (x)(())
F(xx) +
(11)
8j
To further evaluate the lower- boundd we begin with the following equation.
-Alog "(x min (EP:(R=)
(NEx~, (F(xxxx-]] + (Ent(())
(12)
where "(x- f cxp (- X FPrux :)) dæ is the normalizing constant of (-)_;(-|x-))
This equality is confirmed by reformulating the objective of (12) as follows:
NE (Fxxx(F((2xxx-)) + XEnt( = (( - Allg Z(X-))
Then, the lower-bound in Eq. (11) can be evaluated as follows:
F(xx) + (FF(ox)
px }
F X>
-
11
+ AFnn
N
(FF(Px) (X) 8F(xx)
= ExpplNs
F(px)+ +
Σ
Expa
(xi)
N
8p
1=1
(FF(px)
-PXuX-+ - -xx - -1) [F(xx)] AEn((H)
Σ Ex Bxrx F(ex)+ +
812
1=1
E EX) Bx Ex-n [F(exux-1) BF(XXUX- - (N -NN-1)EXxx (F(xx) + AEnt(R)
Ex (lg
N - (N- ()
(13)
where we used Assumption 3.4 for the last inequality.
We consider the following decomposition:
E Xxmu(N)
N Ex-n(N Entculf (-1x })}
N
| F(paux (
E
Expje(M)
Expuirt [log Z(x
(14)
N
The first term can be bounded by the well-known property of entropy.
N
Ent (1, (N) (-1xx-)))) 2 A (nt(((()))
E
(15)
N
N
13

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
The second term can be simply written follows.
(N) V ([Poux-))m(((((e))) -BBxxuin [F(x))]
E Ex-NN
(16)
Combining (13) (14), (16), and (15), we get
(F((x)
) (Ent(u))
Ex
F(px)+ +
, px)
Ex
Ex-R(NN
(N) (
)
B N + Ex [F(x)]
B
((()))
Ex
(17)
This concludes
(1 2 N + 1 ( L(u).
Ex
We finish the proof by setting μ Mla
Proof of Lemma 3.7. We denote by μ u the marginal distribution over X X- It holds that
log du(n) duc
E Xxx((t)
(X)
dy(N)
N Vas F(xx)
27
Ex
V log dx
(X) + A
i=1
Vas log du(N)
N
1n (3)
B Xx (m) (1x-1)
dx (X)
xz F(ex)
(18)
We write p_ the following equation:
du (N)
and pa-1(x|xx
du(h2) 61x 1s (x). Since
du dx (x) p_i(x (=)
we get
dx
(
Vo (p-i(x 1 P s (x (x
V VasPol- (x V log log
Φ log
(x)
dx
p_i(x-1 ((xxx))))-r(e (--
Pil-i (x)
Hence, Eq. (18) can be further bounded by the LSI on the conditional Gibbs distribution (Assumption 3.2) as follows:
N
N Fpxx)
B x
(M)
E Ximm
(
(N)
i-)
logm: 1x
2
(N) dp (X
Bx
(N)
E Ximm
(N) (-P 1)
V
Vlog
d
N
20
E x Sayu
(N)
(N)
(19)
14

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Combining this inequality and Lemma 3.6, get
du(a)) (
27
E Exmutto
Vlog log
(X)
20 A (-B + (())) Ne(w)))
(NN(()))))((())))
This concludes the proof.
Proof of Theorem 3.8. We here prove the convergence MFLD in the discrete-time by using the one-step interpolation
argument (Nitanda, 2024; Suzuki et al., 2023a).
We construct the one-step interpolation for k-th iteration: X R+1 follows: for E (1,})...,d};
X 1/V 6 (Xt) + V (i E 21,
dd))
dy = -V (P(r))
(2))d + 22ddW,
(20)
8y
where Y, (Yo....... ........ Y6a) (X k...., x and W₁ is the standard Brownian motion in Rd with Wo We denote by is (N)
the distributions of Y, Then, 1/ - -4 (= Law(Xx)). V, 0 - (= Law(X (i.e., Y, Xx+11.. this proof,
identify the probability distribution with its density function with respect the Lebesgure measure for notational simplicity. For instance, we denote by pe ( (y) the density of p (N)
By the proof of Theorem 2 in Nitanda (2024). we see for [0.
dL
(vo) $
X2
/ 1x()) Vlog
(y) Mr ((())
1/2
dy + N6,,
(21)
df
where 8, 8n(C3 + ( + Ad) 32 X2 (C3 X2)
( X 1
1/4 c
-xd))
Combining Lemma 3.7 with the above inequality, we get
dL(N)
dt (nd) s -ad ( (c())(()))) - NN((p. -B) No
dt (NN))-((((H))) -
d (v)) - NC((H))
N6p
S -aa (CM)(v)) NN((H)) (M)()))NN((L)) B
N8n
at
at
Noting Un = kk1 and vo = 14 (N)
(N)
the Grönwall's inequality leads to
No (N)
S exp((aan) (AMML)))) -NL(NN No
C (M NC(p)) B ad
ad
This inequality holds at every iteration of (20). Hence, we arrive at the desired result,
((4)) - C(H+) <
B
+
6
+
exp(-aank)
1 c(oi((bb))) -(mm)
s(.))
N
N
at
N aA
(N)
B
bn
1 ( C(H.))
+
+ expl--ani
N
aA
A.2 Point-wise model approximation error (Section 4.1)
Proof of Lemma 4.1. It follows that by Han's inequality (Dembo et al., 1991),
(N)
N / (d)) log du(N)
Ms (N)
(dxs)
log
døs
dxs (xs) 1 1
dx (x).
=
15

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Moreover, see
du 3k dxs
/ 14 (dx") log due (x*)
Σ 18 p
u (dxss log
(xx) = Σ Σ
dr
15| = iES
N N
14 (20) (dx")) (dxx)dlo)
log dli
(x²)
dx
1=1
N - 1) / (
d (x).
log (dbx)
dx
Noticing (2-1)-
N (2), we conclude the first statement which immediately implies the second statement.
Proposition A.I. Suppose (NN)) is exchangeable and X3 ~ ON (j = 1,2, 1,2,...,M) Then, it follows that for any E Z,
27
1 MMC
4R
SR?
KL(ulM 11
2L (N)
1 1
Ey, BXx - [(X,,]]
NM
M
N
N
j=1
Proof of Proposition A.I.
(x,)
E Exxpx, [(X,, 2]] Ex [h(X,
M
1 M²
M
Ex-px [xxx,)) - Ex 1x(X, 2
E (x 3mm
E (x,
1
Σ 3
Exxpx, [h(X, - Ex [M(x,2)) (Ex-px, [h(X, 2)] - Ex [x(X, [h(x,2]]) 21)
2
Using Proposition 4.2, we can upper bound the first term by Me
18a
+
SRa
The second term can be evaluated
M
as follows. Set H (X,) = Ex-ox, (M(X, =)] Ex- [1( (, 2)}. Then for j + k:
E [H(X))((Xx)) (Ex, [Xx)))1)
2
Ex,
Σ
h(x),
j: 2)
- Ex [h(X, 2)]
8'
im1
Exxpa (x(,,))) 2
(Rx=ns>) [I(X,2))
4R (M
(1 .14))
2R (111))))
2R KL(H
This concludes the proof.
A.3 Uniform model approximation error (Section 4.2)
We evaluate the empirical Rademacher complexity RN (F) by using Dudley's entropy integral. We define the metric
I V nn MM [f(x))=. We denote by (FF,s, (IN.M2) the e-covering number of F. with respect to
the || nnom
16

---

Propagation of Chaos for Mean-Field Langevin Dynamics
Lemma A.2 (Dudley's entropy integral). Given a function class F on R4,, we suppose sup II co. Then,
12
/,'
R Vlog 2N( e, llN.M.)dee
(F) S inf 830
48
+
VMN
Proposition A.3. Suppose Assumplion 4.5 holds and Xj nlN) (j = 1,2,.... M.) are independent. Then, gel
E [X (RN.M(F))
/ MN +
12R
(log2 + dlog (1 (1+2 28MR- EX= (N) 1|X112)
[x
MN
Proof. Since | 2 |l max,, [Xx))), it sufficient evaluate the e-covering number of F with respect to M We write r = max,y || |22 By Assumption 4.5, for any 2, 2' E 2,
max [h(x) h(x), 2 - max X||x; - 1l12 = Br|1 z/112,
see. V(F,e,ll (lN,M SN (2, || (12) (1 20r)
Therefore, by Lemma A.2 with 8 RVd(MN)) we get
(F) < /RV MN
+ 12R
log 2N RY d(MN)-I, II..
MN
d
MN (1og2 + dlog (1 (122rR--/MNN-)
V
+ 12R
MN
M
4R
+ 12R
log + dlog
1+28R-
Σ
|X 12
MN
MN
1=1
where we used r < Em [X 1|2. Finally, Jensen's 's inequality yields
E RX [RNM(F)) s 4R MN + 1 MN (1og2 + dlog (1 11+2MR--WMMN E xx-p(x)) (|/x)
Here, we give the complete version of the uniform model approximation bound.
Theorem A.4 (Complete version of Theorem 4.7). Suppose Assumption 4.5 and the same conditions in Theorem 3.8
hold. Run M- -prrallll MFLD in the discrete time independently, with ny 1/2 and X, ~ Ho soN (j 1,,,,..... ..,).. Then,
M
E ((x,,
E Ex~px [!(X,))] Exxp, [xx,))]
M
1=1
5CR
dB AN
dA
log C² (A+ "
+
+ +CR
4
MN
MN(A + MB)
dAMN
8n ax²
CCR
+ exp((ax)(b)4(1))
X MB
where 11 23MR- R .N [|x]
17

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Proof. For Xi,..., XM E Rd we set (x.,....,, .MM = sup the variational formulation of KL-divergence (e.g,, Corollary 4.15 in Boucheron al. (2013)), we get KL I| (
M Σ;=1 Ex-px, [h( - Ex~x. [x(XX22)]
By
1 log E Me www [exp(1g)] +
E Pinsos [g]
?
? MKL(HNN0110
1 log E NMM [exp(1g)] +
$
(22)
?
p.
For independent random variables X, N = .., M), by Lemma 4.6 and A.3. it follows that with high probability 5,
XM)
2 log(118))
(F) RV
< 2TE (2,
MN
S 8RV
d
+ +2RR
1
(1og2+ dlog (1+2MMR-VVNG
E
(x11 + R
2 2log(1/0 16)
MN
MN
MN
CRV
MN
where C is uniform constant and C² = 11 23MR E Ex
x~
SN
[|
This means
ddlo((C' 10)
IP p
GNN
g 1(......,, CRV
MN
MN12
P a
NNN
[9(X1 ......,,a)_) > 5 SOSsco C'exp
dd
g(X X 1 logt S C''exp'
(-
M t)²
P oNns
(ddc²R2
dC2R222
Using this tail bound,
E NNN [exp(19]
P 4r
ONN expp X 1..... XMr) t)d
100
I ONN
9(X1, X.)) > 7 log dt
MNo(lg)) 2)2
C' exp
RC2 dt
Rd
dC²
C.CRVV
MN
exp
4MN
Therefore, we get
RC² 1 log
Td C'CRTV
B [g) <
+
4MN
MN
7
Moreover, by applying Lemma 3.7 Theorem 3.8 to Eq. (22), we get
M
1
E
E Exxps [A(X,)))
Exxpx [b(
M
5=1
aC R2
1
log
(CCRRVV
nd
M
B
No
N
exp(-axnk))
4MN
7
MN
A
ax²
X
18

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Finally, by seting
CR
MN (1 +
MB
get
M
((x,,))
M Σ Exmpxy, [h(X,)] - Ex [X(
5CR d
dB AN
dA
log C' (A+ 2t
+
+ +CR
MN
MN(A + MB)
dAMN
8n
-CRV
+ )xx((=>2nk)41)))
À MB
ax² X
19

---

Propagation of Chaos for Mean-Field Langevin Dynamics
B Experiments
The code used in this work will be made publicly available later.
B.1 Pseudocode and training settings for mean-field experiments
For experiments concerning MFNNs, the output of a neuron in a two-layer MENN is modelled by: h(z;, 23) R tanh(x²) tanh(x) 24 + x;), where I; = (z) x, x E [R++1+1 is its parameter, 2; the given input and is a scaling constant. The tanh activation function placed on the second layer as boundedness of the model is crucial for our analysis. Noisy gradient descent then used train neural networks for T' epochs each. We omit the pseudocode for training MFNNs with MFLD since it is identical to the backpropagation with noisy gradient descent algorithm.
Algorithm 1 Generate the double circle data: D = (z 1= E R2, Mi R before splitting into Duain and Drest- We
set 200, Timer = 1, Foouter and use an 80-20 train-test split for the data.
Algorithm 2
Generate the k multi-index data: D = (2i, Vi) i=1 2j E Rd yi E R. key step normalizing and projecting
2, to the inside of d-dimensional hypersphere. We set n 500, d = 100, r = 5, k = 100 and R 100.
Algorithm Describes how we obtain and test the performance of merged MFNNs against (an approximation to) the mean-field limit by computing the sup-norm between both outputs. The relevant results are stored into a dictionary for plotting the heatmaps. The training procedure is identical for both the classification and regression problem. Let Mmax 20 and Nist {50, 100, .... 500} denote the maximum number of networks to merge and list of neuron settings to train in parallel respectively. We set the hyperparameters for training as follows:
Classification: R 10, Nx 10000, 7 0.1, X 0.1, 0.01, T 200 and loss function: logistic loss
Regression: = 10, No = 10000, = 0.01, = 0.1, A 0.01, = 100 and loss function: mean squared error
Algorithm 1 Generate data points along cocentric 2D) circles
Require: Timner: Touler
Ensure: Dataset D {(2, B:)
1: Initialize -
2: for 1 to n do
Uniform( (0, 2x) 4: Sample E &2 (o,mal(O, 0.1)
3: Sample
5: if n/2 then
6:
r Tinner
7:
--1
8:
else
9:
r + Touter
10:
Y/ f +1
end if
12: Compute Cartesian coordinates: 2; (r (o)) + E1, rin(r)))() Ex)
13: Add (2,, to D
14: end for
15: Randomly shuffle D
16: Split D into Durain (80%) and Dwet (20%)
17: return D Piiin Drest
20

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
Algorithm 2 Generate & multi-index data
Require: n, d, T, k. R
Ensure: Dataset - [(z, 1 Initialize
2: for 1 to do
3: Sample ( Normal(O,) 1)
4: < ((/d
{Get scaling constant
5: Sample ~ Normal (0, Ia)
6: 2 x 2/|²|
(Normalize (Project)
7:
zi < 2 X
8:
9:
9:
for
do
10:
< yi + tanh
end for
Vi 1 1/4 (R/k)
Add (zi, ys) D
14:
end for
15:
Split D into Dunin (80%) and Dies (20%)
16: return Prainn Dies
Algorithm Training and merging MENNS
Require: Dirain Diest (Zrest, )ees),, Noc, Nist Mmax
Ensure: Dictionary sur-norm_ maps N to the average sup_norm
h1 < Train MENN with No neurons on D
& Use h to predict on D rest
Initialize sup norm dic +- 0
for N E Niist do
[h .......h M -Train Maax MENNs with N neurons on Durann
Initialize sur_norm___. 1st <- []
for M E 21, } do
sup hup_norm_iue norm_otooalll for 50 iterations do
Randomly sample M networks from [n ni,,hhy..... ... Maxx
AMN Merge the M networks to form new neural network Use MMN to predict on Drest
sup sup_norm= + max (lg - 1)
sup norm total +- sup norm total sup norm
end for
Append sup sup_sorm norm_oooaaaa Jotall 50 to sup norm Ist end for sup norm dic[N] +- sup norm 1st
end for
return sup.norm_dic
21

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
B.2 Additional MENN experiments
Beyond examining the effect of both M and N on sup norm, we also compare the convergence rate of MFNNS using different À {10-1, 1002 10-3 1004) on the multi-index regression problem. Since the training dataset small and intend to investigate high X. have consider the low epoch setting to prevent deterioration of generalization capabilities. We train 20 networks in parallel and average the MSE (in log-scale) at each epoch, repeating this for N E {300, 400. 800). Figure 2 shows that higher À improves the convergence speed of particles and makes training more stable. Finally, we merge networks with the same hyperpurameters for comparison across different A. similar trend is observed Table 2, highlighting the efficacy of PoC-based ensembling when training for fewer epochs with a high X.
N = 300
400
N 500
3.0
3.0
3.0
= le-1
A= le-1
le-1
2.5
= 1e-2
2.5
A= 1e-2
2.5
A= le-2
= 1e-3
A= le-3
A= 1e-3
2.0
À= le.4
2.0
À le.4
2.0
A= le.4
In(MSE)
In(MSE)
In(MSE)
1.5
1.5
0.5
1
2
5
2
1
T
N = 600
= 700
N = 800
3.0
3.0
3.0
= le-1
A= le-1
le-1
= 1e-2
A= 1e-2
A= le-2
2.5
2.5
2.5
== 1e-3
d = le-3
A= 1e-3
2.0
= 1e-4
2.0
À 1e-4
2.0
A= 1e.4
In(MSE)
In(MSE)
In(MSE)
1.5
1.5
1.5
1.0
0.5
2
2
1
T
Figure 2, Averaged test In(MSE) of singular MFNNs, across different and for 5 epochs
Table 2. MSE comparison between merging M 20 networks across different N and À after 5 epochs.
A
300
400
500
600
700
800
10 1
0.9132253
0.9040508
0,9075238
0,9044338
0.9030165
0.9022377
1
1.2325489
1.2229528
1.2166352
1.1978958
1.1921849
1.1654898
10 1.5718020
1.5668763
1,5607907
1.5581368
1.5282313
1.5234329
10
1.6987042
1.6887244
199631799
1.6135653
1.5860944
1.5821924
22

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
B.3 LoRA for finetuning language models
To examine the effect of we perform LoRA and PoC-based merging by varying E {0, 10-5, 10-44 with one-epoch training. We optimize eight LoRA parameters of rank N 32 in parallel using noisy AdamW with the speficied A. Table 3 summarizes the results. For LoRA, the table lists the best result among the eight LoRA parameters based on the average accuracy across all datasets and also provides the average accuracies of the eight parameters for each dataset. We observed that for Llama2-7B with À = 0 and À = 10-s, the chances of the optimization converging are very low. Consequently, both the average accuracy of cight LoRAs and the accuracy of PoC-based merging are also low. This is because the regularization strength controls the optimization speed as seen in Theorem 3.8. On the other hand, by using a high constant = 10-4 the average performance was improved, and PoC-based merging achieved quite high accuracy even with only one-epoch of training. This result suggests using high reduce the training costs, provided it does not negatively affect generalization emror For Llama3-8B, one-epoch training is sufficient to converge, and while LoRA performed well and PoC-based merging further improved the accuracies.
Table Accuracy comparison of LoRA and PoC-hased merging for finetuning Llama models (1 epoch).
Model
Method
SIQA
PIQA
WinoGrande
OBQA
ARC-e
ARC-e
BoolQ
HellaSwag
Ave.
LoRA (best)
80.55
82.86
83.19
81.60
71.08
84.51
71.90
90.21
80.74
LoRA (ave.)
64.73
76.31
77.76
68.70
57.02
69.02
69.04
70.63
69.15
PoC merge
32.29
62.57
83.58
22.20
28.41
29.42
61.53
28.50
43.56
LoRA (best)
10-3
80.14
82.37
83.43
80.40
68.86
83.42
71.68
89.94
80.03
Llama2 7B
LoRA (ave.)
10
74.37
74.12
80.55
67.50
58.34
71.98
69.43
66.25
70.32
PoC merge
10-5
74.56
83.84
85.16
60.00
63.14
78.37
68.72
92.77
75.82
LoRA (best)
10
78.20
80.90
81.22
78.40
65.19
79.00
69.97
86.50
77.42
LoRA (ave.)
10
74.42
77.70
76.08
75.93
60.93
76.25
65.68
66.71
71.71
PoC merge
10
80.76
82.15
84.85
84.80
71.25
85.35
72.26
91.65
81.63
LoRA (best)
80.45
88.47
86.82
87.60
82.2
90.87
73.85
95.78
85.76
LoRA (ave.)
80.51
88.87
86.85
87.00
80.78
90.98
73.71
95.84
85.57
PoC merge
81.73
88.96
87.77
88.00
81.40
91.71
74.46
96.45
86.31
LoRA (best)
10-3
80.50
88.68
86.98
86.80
81.48
91.12
75.14
95.97
85.83
Llama3 8B
LoRA (ave.)
10-3
80.83
88.64
86.85
87.05
80.39
90.76
71.54
95.87
85.24
PoC merge
10-5
81.53
89.45
87.92
87.80
82.25
91.79
75.54
96.44
86.59
LoRA (best)
10-4
80.30
88.37
86.42
87.20
78.07
89.81
73.61
95.14
84.89
LoRA (ave.)
10
80.00
88.20
85.69
86.23
78.86
89.48
73.08
95.05
84.57
PoC merge
10
80.71
89.72
88.08
89.00
82.17
91.79
74.56
96.36
86.55
23

---

Propagation of Chaos for Mcan-Field Langevin Dynamics
C Additional Background Information
This section provides supplementary information about past works that are relevant to the paper. While not essential to the primary narrative, will provide readers with deeper understanding of previously established MFNN concepts and motivations behind our PoC-based model ensemble strategy.
C.1 Mean field optimization
Two layer mean-field neural networks provide tractable analytical framework for studying infinitely wide neural networks. As N + oo, the optimization dynamics is captured by a partial differential equation (PDE) of the parameter distribution, where convexity can be exploited to show convergence to the global optimal solution (Chizat & Bach, 2018; Mei et al., 2018; Rotskoff et al., 2019). If Gaussian noise is added to the gradient, we get MFLD which achieves global convergence to the optimal solution (Mei et al., 2018; Hu et al., 2019). Under the uniform LSI, Nitanda et al. (2022); Chizat (2022) show that MFLD converges at an exponential rate by using the proximal Gibbs distribution associated with the dynamics. MFLD has attracted significant attention because of its feature learning capabilities (Suzuki et al., 2023b; MousaviiHossseine et al., 2024).
As the assumption that N oo is not applicable to real-world scenarios, discrete-time finite particle system would align closer to an implementable MFLD iee noisy gradient descent. Nitanda et al. (2022) provides global convergence rate analysis for the discrete-time update by extending the one-step interpolation argument for Langevin dynamics (Vempala & Wibisono, 2019). Meanwhile, the approximation error of the finite particle setting is studied in propagation of chaos literature (Sznitman, 1991). For finite MFLD setting, Mei et al. (2018) first suggested that approximation error grows exponentially with time before Chen et al. (2022); Suzuki et al. (2023a) proved the uniform-in inntin-ii-ime propagation of chaos with error bound: O ( X ) suggesting that the difference between the finite N N-particle system and mean-field limit shrinks as V + o0. However, this also means that particle approximation error blows-up exponentially as À + 0 due to the exponential relationship between a and À (Nitanda al., 2022; Chizat, 2022; Suzuki et al., 2023a). Suzuki et al. (2023b) proposes an annealing procedure for classification tasks to remove this exponential dependence in LSII, wihch requires that A be gradually reduced over time and will not work for fixed regularization parameters. Nitanda (2024); Chewi et al. (2024) then proved a refined propagation of chaos independent of a at the solution as described in Section 1.1.
C.2 Ensembling and model merging
In recent years, efforts to improve predictive capabilities and computational efficiency in machine learning have revived interest in techniques such as ensembling (Ganaie et al., 2022: Mohammed & Mohammed, 2023) and model merging (Yang et al., 2024; Goddard et al., 2024). Ensemble methods improve predictive performance by combining the outputs of multiple models during inference (Hansen & Salamon, 1990; Dietterich, 2000). Although several fusion variants exist (Kim al., 2003; Soares et al., 2004), Cheng et al. (2018) shows that simple average voting (Breiman, 1996) does not perform significantly worse while still being highly efficient, with uses in several deep learning applications (Cheng et al., 2018; Yuzquez-Romero & Gallardo-Antolin 2020).
In contrast, model merging consolidates multiple models into a single one by combining individual parameters, showing success particularly in the optimization of LLMs (Ilharco et al., 2022; Jin et al., 2023; Davari & Belilovsky, 2024). An approach to merging models is to simply average the weights across multiple models (Utans, 1996). Taking the average of weights along a single optimization trajectory has been demonstrated to achieve better generalization (Izmailov et al., 2018; Frankle et al., 2020). Moreover, interpolating any two random weights from models that lie in the same loss basins could produce even more optimal solutions that closer to the centre of the basin (Neyshabur et al., 2020). These works then form the foundation of model soups in Wortsman et al. (2022) which refers to averaging the weights of independently fine-tuned models. Similarly, Gauthier-Carc et al. (2024) showed that basic weight averaging methods can perform competitively if constiment models are similar, despite the emergence of novel LLM merging strategies (Ramé et al., 2023; Chronopoulou et al., 2023; Yu et al., 2024).
Despite the widespread use of model merging in the current research landscape, theoretical results concerning the merging of fully trained neural networks are limited. For models trained with stochastic gradient descent, averaging the model weights during different iterations of single run improves stability bounds (Hardt et al., 2016) and variance (Jain et al., 2018) under convex assumptions. Stability bounds in the non-convex settings are then addressed by Wang et al. (2024). Ortiz-Jimenez et al. (2023) studied weight interpolation techniques for task arithmetic in aision-language models, demonstrating that
24

---

Propagation of Chaos for Mean-Field Langevin Dynamics
linearized models under the neural tangent kernel regime can outperform non-linear counterparts.
25


Reviews:
[{"review_id": "OC0AbdVsWa", "reviewer": "Reviewer_joZ7", "summary": "The paper establishes improved uniform-in-time propagation of chaos bounds for MFLD by removing the exponential dependence on entropy regularization, and applies this result to propose a model ensemble strategy.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading our paper.\n\n__Assumption 2.1, 3.2–3.4, Example 3.5, and Q3__\n\nAssumptions 2.1, 3.2–3.4 are all satisfied in several settings considered in the mean-field Langevin literature (e.g., [1–6]). Basically, these assumptions follow for two-layer NNs with smooth and bounded activation functions. For example, under the typical loss functions (e.g., logistic and squared losses) and L2 regularization, the two-layer neural networks with the following activation function $h(x,z)$ satisfy the assumptions [4,5]: (1) $\\sigma_2(r \\sigma_1(w^\\top z + b ))$, (2) $\\sigma_2(r) \\sigma_1(w^\\top z +b)$, (3) $\\sigma_1(w^\\top z +b)$, and (4) $\\sigma_1(w_1^\\top z + b_1 ) + \\sigma(b_2))$, where $\\sigma_i$ are bounded activation functions such as tanh and sigmoid, $x=(w_1,b_1,r)$ is the parameter of each neuron, and $z$ is an input data. The last form is also discussed in Section 4.2 of our paper. While our theory does not cover ReLU activations due to their unboundedness and non-smoothness, we note that such assumptions (bounded, smooth activations) are standard in the mean-field and Langevin literature (see [3], Limitation section).\n\nThe above models also satisfy the constraints introduced in Example 3.5, and thus meet Assumptions 2.1 and 3.2–3.4. Specifically, in Example we impose: (a) $\\sup_{x,z}|h(x,z)|\\leq R$, (b) $\\ell(a,y)$ is convex and L-Lipschitz smooth w.r.t. $a\\in \\mathbb{R}$, and (c) $\\sup_{|a|\\leq R, y \\in \\mathcal{Y}, x \\in \\mathbb{R}^d, z\\in \\mathcal{Z}} \\|\\partial_1 \\ell(a,y)\\partial_x h(x,z)\\| \\leq R’$. Typical losses such as logistic and squared loss satisfy (b). Given that $\\sigma_i$ are bounded and $\\ell$ is L-smooth (i.e, boundedness of partial derivative w.r.t. $a$), conditions (a) and (c) are also satisfied.\n\nWe will incorporate these concrete examples to improve accessibility and clarity.\n\n[1] S. Mei et al., *PNAS*, 2018\n\n[2] A. Nitanda et al., *AISTATS*, 2022\n\n[3] L. Chizat, *TMLR*, 2022\n\n[4] F. Chen et al., 2022\n\n[5] T. Suzuki et al., *NeurIPS*, 2023\n\n[6] A. Nitanda, *NeurIPS*, 2024\n \n\n__Q1/Q2__\n\nAssumption 3.4 quantifies the nonlinearity of $F_0$ with respect to the distribution. If $F_0$ is linear, MFLD reduces to a standard Langevin dynamics over $N$ independent particles. In this case, the joint distribution $\\mu_t^{(N)}$ is the product measure $\\mu_t^{\\otimes N}$ of each particle, implying $KL(\\mu_\\infty^{(N)}\\|\\mu_*^{\\otimes N})=0$ at the optimal joint distribution $\\mu_\\infty^{(N)}=\\mu_*^{(N)}$ attained at $t=\\infty$. However, in general case of nonlinear functional, there should be additional error as evaluated in Lemma 3.6; $\\frac{\\lambda}{N}KL(\\mu_\\infty^{(N)}\\|\\mu_*^{\\otimes N})\\leq \\frac{B}{N}$ at the optimal solution.\nThus, the strength of nonlinearity $B$ controls the deviation from independence among particles.\n\nAssumption 3.2 requires that the conditional distributions $\\nu_{i|-i}$ satisfy an LSI, ensuring concentration of distribution of each particle. This assumption is also satisfied under the setting in Example 3.5 (for the derivation see Lemma 6 in [7]). Here, the regularization $r(x)$ is essential to encourage such concentration, and the entropy term corresponds to the Gaussian perturbation in the method.\n\n[7] S. Chewi et al., 2024\n\n__Experiments (scalability w.r.t. $M,N,\\lambda$) and Q4__\n\nScalability with respect to $M$ and $N$ is empirically validated on two-layer NN with synthetic datasets (see Fig 1). The effect of $\\lambda$ is examined in Appendix Sections B.2 and B.3.\n\nOur merged method suggests a nontrivial choice of $M$ and $N$. Given a fixed computational budget $K = MN$, Theorem 4.4 yields the bound: $\\frac{1}{K}+\\frac{1}{\\sqrt{MK}}+\\frac{M}{K}$ at the solution, ignoring irrelevant constants for simplicity. This bound is decreasing on $M \\in [1, (K/4)^{1/3}]$ and hence the optimal choice is $M \\sim (K/4)^{1/3}$ that achieves the minimum approximation error $\\frac{1}{K} + \\frac{C}{K^{2/3}}$ for some constant $C$.\n\n__Global convergence and Q1/Q5__\n\nOur MFLD theory establishes global convergence of noisy gradient descent to the global minimizer of the un-regularized objective. Specifically, when $r(x) = \\lambda' ||x||^2$, the regularization (i.e., $\\mathbb{E}[r] + \\lambda \\mathrm{Ent}$) coincides with the KL divergence from a Gaussian distribution. Hence, minimizing $\\mathcal{L}(\\mu)$ leads to convergence toward minimizing $F_0$, up to a $\\lambda$-dependent error shrinking to $0$ as $\\lambda \\to 0$.\n\nImportantly, optimization in the mean-field regime is more challenging than in the NTK regime, as it involves solving a truly non-convex problem. In contrast, NTK theory effectively linearizes the model and neurons evolve near initialization. Actually, the mean-field regime is known to exhibit *feature learning* behavior [8,9], deviating from NTK-regime.\n\n[8] L. Chizat et al. On Lazy Training in Differentiable Programming. *NeurIPS*, 2019\n\n[9] G. Yang and E.J. Hu. Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks. *ICML*, 2021", "claims_and_evidence": "The central claim of the paper is that it establishes an improved PoC result for MFLD by eliminating the exponential dependence on the $\\lambda$ in the particle approximation error. This claim is clearly stated and mathematically proved in Theorem 3.7.\n\nThe derivation is supported by technical assumptions (Assumptions 3.2–3.4) and intermediate results such as Lemma 3.6. However, these assumptions are relatively strong and not well justified in practice. For example, the directional log-Sobolev inequality and the boundedness/Lipschitz conditions on model components may not hold in typical neural network architectures (e.g., ReLU activations or unbounded weights). Example 3.5 attempts to justify the assumptions but does not explicitly verify that Assumptions 3.2–3.4 are satisfied. Instead, it introduces additional constraints that further limit practical applicability.\n\nThe paper also claims a practical contribution via a model ensemble strategy derived from the theoretical insights. However, the empirical validation is limited:\n- The experiments do not show how the approximation error scales with $N$ or $\\lambda$, nor do they examine behavior in the small-$\\lambda$ regime where prior results are known to break down.\n- The ensemble setup uses $M$ independent networks of size $N$, which increases the total parameter budget and may lead to an unfair comparison under a fixed compute or model size constraint.", "methods_and_evaluation_criteria": "The theoretical methods used in the paper seem sound. Although I am not an expert in optimal transport and did not check the proofs in the appendices in full detail, the use of Wasserstein gradient flows and functional inequalities seems fine and consistent with prior literature.\n\nHowever, the evaluation criteria in the experimental section are limited and not well aligned with the core theoretical contributions. The paper does not empirically evaluate key aspects such as:\n- Whether the theoretical results still hold when Assumptions 2.1 and 3.2–3.4 are violated in practice.\n- How the particle approximation error scales with $N$ (number of neurons/particles),\n- How performance is affected by varying the entropy regularization parameter $\\lambda$,\n- Whether the improved $O(1/N)$ convergence rate in Theorem 3.7 matches empirical trends,\n\nMoreover, the ensemble strategy is only evaluated in the context of LoRA fine-tuning and lacks benchmarks on standard deep learning tasks (e.g., training deep neural networks from scratch on CIFAR-10 or ImageNet). The choice of merging $M$ independent networks of size $N$ each is not compared to more realistic alternatives under a fixed budget constraint (e.g., $M$ networks of size $N/M$), which weakens the practical relevance of the proposed approach.", "theoretical_claims": "No", "experimental_designs_or_analyses": "Please check in the section of Methods And Evaluation Criteria.", "relation_to_broader_scientific_literature": "This work removes that dependence by introducing a directional log-Sobolev inequality, but the lack of comprehensive experimental study weaken the connections to broader machine learning practice.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA.", "other_comments_or_suggestions": "NA.", "questions_for_authors": "1. What is the role of the regularization term $r(x)$ and the entropy regularization term in the theoretical analysis? Do the MFLD framework and the improved PoC result still hold in the absence of $L^2$ or entropy regularization?\n\n2. Could you provide more intuition behind Assumptions 3.2 and 3.4? Specifically, how do these assumptions contribute to the analysis, and in what types of neural network architectures or setups might they realistically hold?\n\n3. The assumption $\\sup_{x,z} |h(x,z)| \\le R$ seems nontrivial. In which practical scenarios or network parameterizations does this condition hold? Can you provide concrete examples?\n\n4. The ensemble strategy seems unclear in terms of fair comparison. You consider $M$ independent networks each with $N$ neurons, but under a fixed computational or model-size budget, this setup may be suboptimal. A more realistic comparison might involve $\\sqrt{N}$ networks each with $\\sqrt{N}$ neurons, totaling $N$ neurons overall. In that setting, the bound in Theorem 4.4 appears similar or potentially worse due to the additional error term. Could you clarify why the proposed ensemble strategy is justified and whether it offers a real advantage under fixed resource constraints?\n\n5. Does your theoretical framework for MFLD imply global convergence to the minimizers of the loss functionals $F_0(\\mu)$ and $F_0^{(N)}(\\mu^{(N)})$, similar to what is established in NTK theory? In the NTK setting, global convergence can be shown without requiring regularization. By contrast, your analysis incorporates both entropy and $L^2$ regularization. While $F_0(\\mu)$ is convex over the space of distributions, it is unclear whether this alone is sufficient to guarantee global convergence of the dynamics. Could you clarify what kind of convergence your results guarantee (e.g., global vs. local), and whether additional assumptions are necessary to establish global convergence in your setting?", "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": "GUc5XL9Cpn", "reviewer": "Reviewer_oQCo", "summary": "This paper improves the Propagation of Chaos (PoC) error bound for Mean-Field Langevin Dynamics (MFLD) by refining the defective Log-Sobolev Inequality (LSI) and introducing the Uniform Directional LSI (UD-LSI). Additionally, it proposes a PoC-based model ensemble method, which is supported by both theoretical analysis and empirical validation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading our paper.\n\n**Uniform directional LSI (UD-LSI)**\n\nWe can theoretically validate the UD-LSI in the setting of Example 3.5 by leveraging a known result (e.g., Lemma 6 in [1]): Let $\\nu \\propto \\exp( -H-V)$, where $V,H: \\mathbb{R}^d \\rightarrow \\mathbb{R}$, with $V$ being $\\alpha$-strongly convex and $H$ being $L$-Lipschitz smooth. Then, $\\nu$ satisfies the Log-Sobolev Inequality (LSI) with constant: $\\alpha \\exp\\left( - \\frac{L^2}{\\alpha} - \\frac{4L}{\\sqrt{\\alpha}}\\right)$. (Note that LSI constants in [1] is defined as a reciprocal number of our constant.) \nWe now apply this result to the conditional distribution $\\nu_{i|-i}$ in Example 3.5:\n\n$$\n\\frac{d \\nu\\_{i|-i}}{dx}(x|\\mathbf{x}^{-i}) \\propto \\exp\\left( -\\frac{N}{\\lambda n}\\sum_{j=1}^n\\ell(\\mathbb{E}\\_{X\\sim\\rho\\_{x\\cup \\mathbf{x}^{-i}}}[h(X,z\\_j)],y\\_j) - \\frac{\\lambda'}{\\lambda}||x^i ||\\_2^2 \\right).\n$$\n\nThe first term in the exponent is $\\frac{R’}{\\lambda}$-Lipschitz smooth since its partial derivative in $x$ is bounded as follows under the setting of Example 3.5:\n\n$$\n\\left|\\left| \\frac{N}{\\lambda n} \\sum\\_{j=1}^n \\partial\\_1\\ell(\\mathbb{E}\\_{X\\sim\\rho\\_{x\\cup \\mathbf{x}^{-i}}}[h(X,z\\_j)],y\\_j)\\frac{1}{N}\\partial\\_x h(x,z_j)\\right|\\right| \\leq \\frac{R'}{\\lambda}.\n$$\n\nAnd the second term in the exponent is $\\frac{2\\lambda’}{\\lambda}$-strongly convex. Therefore, we get LSI constant $\\frac{2\\lambda’}{\\lambda}\\exp\\left( -\\frac{R'^{2}}{\\lambda^2}  \\frac{\\lambda}{2\\lambda'} - \\frac{4R'}{\\lambda }\\sqrt{\\frac{\\lambda}{2\\lambda'}}\\right)\n= \\frac{2\\lambda’}{\\lambda}\\exp\\left( -\\frac{R'^{2}}{2\\lambda \\lambda'} - \\frac{4R'}{\\sqrt{2\\lambda\\lambda'}}\\right)$. \n\nWhile this result is briefly stated in Example 3.5, we will include the above derivation in the revised version to enhance accessibility and transparency.\n\n[1] S. Chewi et al., Uniform-in-n log Sobolev inequality for the mean-field Langevin dynamics with convex energy. 2024.\n\n\n**Experiment and comparisons with prior PoC error bounds**\n\nCompared to prior results [1], which yield a uniform-in-$N$ LSI constant: $\\exp\\left( - \\frac{1}{\\lambda’} - \\frac{1}{\\lambda\\lambda’} - \\frac{1}{\\lambda^2\\lambda’^3}\\right)$, our bound demonstrates significantly improved dependence on $\\lambda, \\lambda’ \\to 0$, leading to faster convergence in time.\n\nFurthermore, a major improvement over [2,3] is that our particle approximation error bound is independent of $\\lambda$. In contrast, earlier works suggested an exponential dependence on $\\lambda$, which was overly pessimistic. To support this, we empirically investigate the effect of $\\lambda$ under varying $N$ in Appendix B.2, and we did not observe such exponential blow-up, further reinforcing the practical relevance of our theoretical improvement. This highlights the importance of our contribution in tightening the gap between theoretical bounds and empirical observations. Although exactly verifying uniform bounds through experiments remains challenging, improving these theoretical bounds is an important fundamental research topic.\n\n[2] F. Chen et al., S. Uniform-in-time propagation of chaos for mean field langevin dynamics. 2022.\n\n[3] Suzuki, T., Wu, D., and Nitanda, A. Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.", "claims_and_evidence": "- The theoretical results heavily rely on Assumption 3.2 (Uniform directional LSI), but its justification remains insufficient. As I understand, the authors assume that the convergence rate of each conditional distribution of a Langevin particle is uniform, which indirectly ensures the network-wide uniform convergence and influences the effectiveness of the PoC-based ensembling method. However, the paper does not provide empirical evidence to support these assumptions. I encourage the authors to include numerical validation or theoretical discussion regarding the plausibility of UD-LSI in practical neural networks.\n- Additionally, I am uncertain whether the improved PoC bound is numerically validated. While the theoretical derivations are rigorous, the paper does not appear to provide direct numerical verification of the improved error bound. Empirical experiments demonstrating the practical impact of the improved bound, such as comparisons with prior PoC error bounds, would strengthen the paper’s claims.", "methods_and_evaluation_criteria": "- The proposed method is logically well-founded and builds on existing work in the PoC for MFLD.", "theoretical_claims": "- This paper presents rigorous theoretical proof for its main claims.", "experimental_designs_or_analyses": "- The paper provides empirical validation for the proposed PoC-based ensemble method, particularly in the context of LoRA-based fine-tuning.", "relation_to_broader_scientific_literature": "- While this paper makes an interesting theoretical contribution to PoC for MFLD, its relevance to the broader deep learning community remains uncertain—which is one of my main concerns.", "essential_references_not_discussed": "- This work has already discussed the related studies.", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "No.", "questions_for_authors": "- Please refer to the comments above.", "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": "bTXZd0eYZ0", "reviewer": "Reviewer_7CdW", "summary": "The paper proposes an improved bound on the convergence of neurons (of a network) under mean field Langevin dynamics to an iiid distribution. This argument is known as the propagation of chaos. The convergence of the empirical finite-N distribution to the limiting iid distribution is controlled by time and number of particles (neurons) $N$. The optimal approximation error is known to be order $1/N$. This paper further improves on this and proves exponential decay in time. The result is applied to the model ensemble where the variance due to finite approximation is further reduced.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading our paper and for the positive feedback. We will revise the manuscript accordingly, following your suggestion.", "claims_and_evidence": "Yes, the claims are well supported.", "methods_and_evaluation_criteria": "Yes. Toy models are good for theory papers.", "theoretical_claims": "no. I only read theorem statements in the main.", "experimental_designs_or_analyses": "Yes. The experiments in the main (with a multi-index model and two concentric circles) make sense.", "relation_to_broader_scientific_literature": "Mean-field limit dynamics of neural networks is an active area of research. Improvement in convergence speed is a very valuable technical contribution. Similarly, model ensembling shows better generalization capabilities and there is no complete theory that explains this (but I may be wrong). Though this paper does not study generalization, applying mean-field Langevin analysis to ensembling is valuable.", "essential_references_not_discussed": "The related literature is included and compared.", "other_strengths_and_weaknesses": "The results look solid but I am not an expert on this topic to catch a mistake if there was one. \n\nHowever, I did find the paper difficult to read. It would benefit from better organization along the lines:\n\n* Defective LSI comes on pg 2. This is not a trivial lemma. Please put this in a proper formatting (Lemma  ?) and cite exactly where it appeared in Chen et al 2022. \n\n* $\\Delta_0^{(N)}$ is introduced multiple times (pg3 and pg4). \n\n* Page 3. the paragraph starting with \"Afterward, this exponential dependence.....\" put this technical paragraph as a remark after the main result. \n\n* Assumption 3.2 looks like a Lemma.", "other_comments_or_suggestions": "see above.", "questions_for_authors": "NA.", "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": "FeD1xPwn5H", "reviewer": "Reviewer_i5RL", "summary": "The paper studies the propagation of chaos of two-layer neural network in the mean-field regime. The authors first obtain a uniform-in-time propagation of chaos (PoC) bounds that does not depend on the LSI constant, and maintain the \"original\" rate of convergence.  Then, the authors apply the PoC bounds to model ensemble problem and show that  in the mean-field regime, model ensemble help reducing the approximation error. Finally, the authors proposed a PoC based model ensemble methods, and did experiments to verify the usefulness of the methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reading our paper.\n\n**1.2 Mismatch between the ensemble methods in Sections 4 and 5.2**\n\nFirst, we would like to clarify that the ensemble method used in Section 5.2 is exactly the same as the one proposed in Section 4. Specifically, the ensemble in Section 4.2 is taken over model outputs of the form $x \\to  b^i\\_j a^{i \\top}\\_j x $, which reduces to parameter averaging due to the use of a linear activation function. That is: $\\frac{1}{MN}\\sum\\_{i,j} (b^i\\_j a^{i\\top}\\_j x) = (\\frac{1}{MN}\\sum_{i,j} b^i\\_j a^{i\\top}\\_j )x = \\Delta W x$ where the left-hand side represents the ensemble of model outputs, and the right-hand side corresponds to a single model with averaged parameters.\n\n**1.1 Choice of $N$**\n\nIn the LoRA setting, choosing $N=\\min\\\\{ d, k \\\\}$ corresponds to full fine-tuning, with no approximation error compared to the $N=\\infty$ case for fixed $d$ and $k$, thanks to the linearity of the activation function. Our goal is to close the performance gap from the full-rank case more efficiently by leveraging the ensemble technique under $N<\\min\\\\{ d, k \\\\}$.\n\n**1.3 Comparison under Fixed Compute Budget**\n\nFollowing your suggestion, we additionally evaluated LoRA with a higher rank (256) and found that performance is inferior compared to the ensemble of 8 lower-rank (32) models. Please refer to the table below for details:\n\n| Model | Method | SIQA | PIQA | WinoGrande | OBQA | ARC-c | ARC-e | BoolQ | HellaSwag | Ave. |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| Llama2 | LoRA (r=32, best)    | 79.48 | 82.43 | 81.77 | 80.60 | 67.75 | 80.47 | 70.37 | 86.67 | 78.69 |\n|               | LoRA (r=256) | 69.95 | 69.69 | 69.61 | 61.40 | 47.44 | 61.15 | 63.73 | 47.27 | 61.28 |\n|               | PoC merge     | 81.17 | 84.60 | 85.16 | 86.60 | 72.53 | 86.62 | 72.45 | 92.79 | 82.74 | \n| Llama3 | LoRA (r=32, best)    | 81.22 | 89.50 | 86.74 | 86.00 | 79.86 | 90.53 | 72.91 | 95.34 | 85.26 |\n|               | LoRA (r=256) | 81.06 | 87.60 | 87.61 | 84.60 | 78.92 |  90.06 | 75.11 |  94.98 | 84.99 |\n|               | PoC merge     | 82.04 | 89.39 | 89.27 | 89.20 | 83.28 | 92.30 | 76.33 | 96.58 | 87.30 |\n\nThese results suggest that, under a fixed compute budget $MN=256$, the ensemble method achieves nontrivial improvements over joint training for a higher-rank model.\n\nFurthermore, prior studies [1] have also observed that training higher-rank matrices does not always lead to better performance, (see Fig 4 in [1]) and [2] reported instability of LoRA with higher rank.\n\n[1] S.Y. Liu et al. DoRA: Weight-Decomposed Low-Rank Adaptation. ICML, 2024\n\n[2] D. Kalajdzievski. A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA. 2023\n\n**2. Benefit of the Ensemble Method**\n\nThe ensemble method helps reduce the approximation error from the mean-field limit compared to directly training a single network with $MN$ neurons. Our merge strategy provides nontrivial optimal choice of $M$. Given a fixed compute budget $K=MN$, Theorem 4.4 gives the error bound: $\\frac{1}{K}+\\frac{1}{\\sqrt{MK}}+\\frac{M}{K}$, ignoring constants for simplicity. This bound decreases with $M \\in [1, (K/4)^{1/3}]$, and is minimized when $M \\sim (K/4)^{1/3}$, achieving an error of: $\\frac{1}{K} + \\frac{C}{K^{2/3}}$ for some constant $C$.\n\nIntuition: The key in the mean-field approximation is the independence among neurons $\\\\{h(x_t^i,z)\\\\}_{i=1}^N$ since the variance of their empirical average (i.e., mean-field model) would decrease linearly with $N$ if neurons were independent. While the PoC ensures that neurons become approximately independent after convergence when $N$ is sufficiently large, the ensemble of independently trained networks introduces the independence across models, further reducing error.\n\nAdditionally, previous work [3] has shown that ensembles of independently trained networks can outperform joint training in specific scenarios, although their setting differs from ours.\n\n[3] Z. Allen-Zhu and Y. Li. Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning. ICLR, 2023\n\n**Q. Technical Differences from Existing Studies**\n\nOur proof strategy is significantly different from those in [4,5]. Specifically, [4] establishes a uniform in $N$ log-Sobolev inequality (LSI) by constructing a Lipschitz transport map from a Gaussian to the optimal distribution $\\mu_*^{(N)}$. [5] directly analyzes the variance of the mean-field model at the optimal distribution, leveraging the nonlinearity of $F_0$.\n\nIn contrast, our analysis is based on the argument of conditional and marginal distribution of $\\mu^{(N)}$ [6], which allows us to establish an improved bound that holds for any distribution, not just at the solution. Furthermore, our assumptions about LSI differ notably from [5]; please refer to our Assumption 3.2 compared to Assumption 2 in [5].\n\n[4] S. Chewi et al., 2024\n\n[5] A. Nitanda, NeurIPS, 2024\n\n[6] F. Chen et al., 2022", "claims_and_evidence": "I think the main claims clearly stated, and the proof are convincing.", "methods_and_evaluation_criteria": "This paper is majorly a theoretical paper. There are some experimental results which I think are not the major contribution of the paper: in section 5.1,  the authors consider training a two-layer neural network in mean-field regime on synthetic datasets. This experiments make sense to me as a sanity check for the theoretical results. In section 5.2, the authors consider incorporating LoRA finetuning with the proposed model ensembling method, I think the model used and the benchmarks are standard.", "theoretical_claims": "I check the proof of Lemma 3.6 and Proposition 4.6 in the appendix, as well as all the proof in the main context, and I think they are all correct. I didn't check the rest proofs in detail, but they are more or less easy to see given the other theorems mentioned in the main context.", "experimental_designs_or_analyses": "This paper is majorly a theoretical paper, and I think both experiments solid in experimental designs or analyses.", "relation_to_broader_scientific_literature": "This is paper studies the propagation of chaos for two-layer neural networks in mean-field regime. In general, propagation of chaos are widely studied in other fields such as stochastic analysis/ optimal transport, statistical physics and game thoery etc. While in other problems, theoretical formulation could be different from the one considered in this paper, similar ideas might be applicable.", "essential_references_not_discussed": "I don't find any important references that is not discussed.", "other_strengths_and_weaknesses": "**Strengths**\n I think in general the theoretical results are interesting:\n- Removing the LSI dependency in the PoC bound and reducing the order of $1/\\lambda$ in the convergence rate is a nice improvement. \n- Showing that averaging over $M$ independently trained nn reduces the approximation bound (improve the PoC bound) is interesting.\n\n\n**Weaknesses**\n\n1. The main weakness in this paper is the proposed methods in Section 5.2, there are following issues in my opinion: \n- 1.1  The authors take the LoRA rank to be $N$ and claim \"Therefore, we can apply PoC-based model ensemble for LoRA parameters.\" This is not very convincing, since in  general one need large $N$ to achieve good approximation in PoC bounds, and typically $N \\gg d,k.$ However, in practical application of LoRA, people let the LoRA rank $N \\ll d,k,$ thus I doubt that  the PoC results provide any insight on LoRA in general.\n- 1.2 There's a mismatch between this ensemble method and the one proposed section 4. In particular, the method proposed in section 4 averages the outputs, but the one in section 5.2 averages over the LoRA weights. \n- 1.3 I don't see the usefulness of  this weight averaging method from the experimental results, since the authors didn't control the variates when comparing the ensembled model and the individual model. In particular, the ensembled model increases the accuracy, but also requires $M=8$ times more computation resources comparing to only training one model, since it requires to train $M=8$ independent models.  Besides, the rank of the LoRA updates in each individual model is at most $N = 32,$ but the update of  the averaged model is $N \\times M = 256.$ I think it would be more interesting to consider the performance of the model ensemble method under fixed computational resources, for example,  compare with an individual model that trained $8$ times more epochs, or compare with another individual model whose LoRA rank is $N \\times M = 256.$\n \n2. While I believe the theoretical results are techinically interesting in the field of PoC,  I don't get much insight in model ensemble from the results. Theoretically, the authors show that by training $M$ independent models and then average over the output of the model, it improves the approximation error, however, (1) this results is very specific to the two-layer setting in the mean-field regime, (2) I don't see the benefit of this methods compare to directly training a large network with $M \\times N$ neurons.   Practically, I don't find the practical results very convincing as discussed in the previous point.", "other_comments_or_suggestions": "I do not have other comments or suggestions.", "questions_for_authors": "1. Could you address technical differences between the prove of the improved PoC results  and the proof techniques in [1,2] ?\n\n[1] Chewi, S., Nitanda, A., and Zhang, M. S. Uniform-in-n log sobolev inequality for the mean-field langevin dynamics\nwith convex energy. arXiv preprint arXiv:2409.10440,2024.\n\n[2] Nitanda, A. Improved particle approximation error for mean field neural networks. In Advances in Neural Information Processing Systems 37, 2024.", "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:

Meta-review: This paper improves the error bounds for the Propagation of Chaos in Mean-Field Langevin Dynamics by introducing a new condition called the Uniform Directional Log-Sobolev Inequality. Under this condition, the main result establishes uniform-in-time convergence, with a residual term that is independent of the LSI constant of the target distribution. Furthermore, the paper improves the convergence rates with respect to key parameters, compared to previous works that established similar uniform-in-time bounds. These theoretical findings are applied to model ensembling, showing that ensemble methods can further reduce finite-size approximation error. Although primarily theoretical, the paper includes experiments designed to validate the results and demonstrate practical relevance, including an application to LoRA fine-tuning.

Most reviewers weakly support the acceptance of the paper, and I align with the general consensus.

All reviewers agree that the paper provides an interesting improvement in uniform-in-time convergence bounds for the propagation of chaos in mean-field Langevin dynamics. However, several concerns were raised, particularly regarding the application of these results to model ensembling, a point with which I also agree. I strongly encourage the authors to carefully address the reviewers' feedback in their revision, clearly highlighting the main limitations of their theoretical results and their applicability.
Decision: Accept (poster)