Propagation of Chaos for Mean-Field Langevin Dynamics and its Application to Model Ensemble

We thank the reviewer for reading our paper.

1.2 Mismatch between the ensemble methods in Sections 4 and 5.2

First, 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.

1.1 Choice of $N$

In 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 \\}$.

1.3 Comparison under Fixed Compute Budget

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

These results suggest that, under a fixed compute budget $MN=256$, the ensemble method achieves nontrivial improvements over joint training for a higher-rank model.

Furthermore, 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.

[1] S.Y. Liu et al. DoRA: Weight-Decomposed Low-Rank Adaptation. ICML, 2024

[2] D. Kalajdzievski. A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA. 2023

2. Benefit of the Ensemble Method

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

Intuition: 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.

Additionally, previous work [3] has shown that ensembles of independently trained networks can outperform joint training in specific scenarios, although their setting differs from ours.

[3] Z. Allen-Zhu and Y. Li. Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning. ICLR, 2023

Q. Technical Differences from Existing Studies

Our 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$.

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

[4] S. Chewi et al., 2024

[5] A. Nitanda, NeurIPS, 2024

[6] F. Chen et al., 2022

We thank the reviewer for reading our paper and for the positive feedback. We will revise the manuscript accordingly, following your suggestion.

We thank the reviewer for reading our paper.

Uniform directional LSI (UD-LSI)

We 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.) We now apply this result to the conditional distribution $\nu_{i|-i}$ in Example 3.5:

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

And 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) = \frac{2\lambda’}{\lambda}\exp\left( -\frac{R'^{2}}{2\lambda \lambda'} - \frac{4R'}{\sqrt{2\lambda\lambda'}}\right)$.

While this result is briefly stated in Example 3.5, we will include the above derivation in the revised version to enhance accessibility and transparency.

[1] S. Chewi et al., Uniform-in-n log Sobolev inequality for the mean-field Langevin dynamics with convex energy. 2024.

Experiment and comparisons with prior PoC error bounds

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

Furthermore, 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.

[2] F. Chen et al., S. Uniform-in-time propagation of chaos for mean field langevin dynamics. 2022.

[3] Suzuki, T., Wu, D., and Nitanda, A. Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.

We thank the reviewer for reading our paper.

Assumption 2.1, 3.2–3.4, Example 3.5, and Q3

Assumptions 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).

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

We will incorporate these concrete examples to improve accessibility and clarity.

[1] S. Mei et al., PNAS, 2018

[2] A. Nitanda et al., AISTATS, 2022

[3] L. Chizat, TMLR, 2022

[4] F. Chen et al., 2022

[5] T. Suzuki et al., NeurIPS, 2023

[6] A. Nitanda, NeurIPS, 2024

Q1/Q2

Assumption 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. Thus, the strength of nonlinearity $B$ controls the deviation from independence among particles.

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

[7] S. Chewi et al., 2024

Experiments (scalability w.r.t. $M,N,\lambda$) and Q4

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

Our 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$.

Global convergence and Q1/Q5

Our 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$.

Importantly, 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.

[8] L. Chizat et al. On Lazy Training in Differentiable Programming. NeurIPS, 2019

[9] G. Yang and E.J. Hu. Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks. ICML, 2021