Improved Particle Approximation Error for Mean Field Neural Networks

Summary

The manuscript investigates the mean-field Langevin dynamics (MFLD) and improves the particle approximation error by removing the dependency on the log-Sobolev inequality (LSI) constant. This is of relevance, as the LSI constant in general might deteriorate with the regularization coefficient, limiting the impact of prior bounds. The Authors illustrate the applicability of their result with three examples. 1. An improved convergence of the MFLD, 2. sampling guarantees for the final stationary distribution, 3. uniform-in-time propagation of chaos w.r.t. the Wasserstein distance for the MFLD.

Strengths

- Improving the estimates of the particle approximation of the MFLD is of interest given its appearance in the learning problem of mean-field neural networks, which has been an interesting research topic pursued by different research groups during the last years. The paper is of purely theoretical nature and improves upon the state-of-the-art by removing the dependence of the approximation error constant on the LSI constant. The technical tools to do so, seem novel. - The paper is well-written and -structured with the objectives and contribution clearly stated and pursued. - In my opinion, the content of the paper could be also of interest beyond the scope of the MFLD, as propagation of chaos results with favorable constants are of interest in a wide variety of fields. The Authors may want to consider commenting on this.

Weaknesses

Apart from some minor questions addressed below, there is one point of critique that I would like to raise. Namely, the Authors do not really motivate why one should expect that the particles approximation _does not_ depend on the LSI constant. This, in my opinion, would improve the reading experience of the paper.

Questions

- line 61: How does the LSI constant $\alpha$ deteriorate with the regularization parameter $\lambda$? Could you provide a sketch to give some more information? This could be also worth to be included and referenced in the manuscript. And some more minor comments: - In line 28, the Authors may want to add one further line of work, namely "Mean field analysis of neural networks: A law of large numbers" and "Mean field analysis of neural networks: A central limit theorem" by J Sirignano, K Spiliopoulos, to exhaustively cover the literature. - In regards of (2), the Authors might want to mention that $\nabla \frac{\delta \mathcal{F}}{\delta \mu}$ is also known as the Wasserstein gradient $\nabla_W \mathcal{F}$. - lines 35, 36: Could you add references here? - lines 99, 100: Write $\mathcal{P}(\mathbb{R}^d)$. - line 100: Not sure what you mean with "it follwos that" here - line 157: $N$ not $d$ - line 234: identical to - lines 266-268: This sentence sounds a bit complicated. - line 301: Discussion

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The Authors do point out some limitations in the Conclusions, which are reasonable and in my opinion justifiable. They could further emphasize some limitations of their result following from Assumption 3. In particular the uniform boundedness of $h$ seems to be a restriction, but since the class of covered neural networks is still fair, I would not consider it as a substantial limitation. Yet, it could be highlighted.

Summary

The paper presents an improved finite particle approximation error bound for mean-field Langevin dynamics. The main result establishes a $O(1/N)$ gap between the original and $N$-particle objective which is independent of the LSI constant and improves upon the existing $O(\lambda/\alpha N)$ bound. The is applied to obtain improved convergence error and Wasserstein propagation of chaos of MFLD for shallow neural networks.

Strengths

The independence of the particle discretization error on the LSI constant is a strong and useful result, especially in low temperature regimes where $\alpha$ can be exponentially large w.r.t. $\lambda^{-1}$. The proof technique of introducing the induced Bregman divergence to bound nonlinearity is both novel and simpler compared to previous analyses. The new error is easily incorporated into existing MFLD frameworks.

Weaknesses

* Besides Theorem 1, the remaining results seem to be a straightforward application to the existing analysis of time discretization of MFLD. In particular, the $\exp(-\alpha\lambda\eta k)$ convergence rate and $\eta/\alpha\lambda$ error term of [1] remain and hence the overall rate and error bound still suffer from the LSI constant, albeit decoupled from particle error. * The corresponding Wasserstein error bound also automatically suffers from the LSI constant due to using Talagrand's inequality. While the constant is improved from $\alpha^{-2}$ to $\alpha^{-1}$, the difference seems to be incremental as $\alpha$ is already adversarially exponentially dependent on $\lambda^{-1}$. * The standard assumptions of MFLD apply, for example the activation must be Lipschitz smooth as well as bounded to ensure that Assumption 1 holds. These implications should be explained in the text alongside Assumption 3 as currently Assumption 1 is only given for an abstract functional and not for the specific form (14). [1] Suzuki et al, 2023. Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction.

Questions

* Currently the dependence of the convergence rate on $\alpha$ in the MFLD framework may be unavoidable, however is there a possibility that the time discretization error or Wasserstein error can also be made independent of the LSI constant?

Rating

6

Confidence

4

Soundness

4

Presentation

2

Contribution

2

Limitations

The authors address limitations.

Summary

This paper provides an improved bound on the particle approximation error of MFLD under conditions of a log-Sobolev inequality (and a couple of extra regularity conditions). This bound improves on previous bounds by removing the dependence of the bound on the LSI constant. The authors then apply their bound to get improved convergence rates of MFLD in the finite-particle setting, sampling guarantees for $\mu_*$ and a uniform-in-time propagation of chaos result analagous to previous literature.

Strengths

- This work addresses an interesting problem - the convergence of MFLD - and provides a clear improvement over the previous best bounds in this domain. The proof method is also novel. - The work is extremely clearly presented, with all of the key contributions being well-motivated and explained. The relevant literature is also well summarized. - Overall, I consider this to be a very sound piece of theoretical work.

Weaknesses

I have no major concerns about this work. The bound consists only of a quantitative improvement on existing bounds, rather than a completely novel bound, and still requires the same log-Sobolev assumptions as previous works, perhaps limiting its general applicability. In addition, they also assume bounded activation functions, which is an additional restriciton. However, I consider both of these restrictions to be minor.

Questions

I have no clarifications to request.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors do provide a (quite brief) discussion of the limitations of their work in the final section. Beyond this, I foresee no other potential negative societal impacts of their work.

Summary

This work proves a novel particle approximation error for the continuous- and discrete-time mean-field Langevin dynamics (MFLD), which removes the dependency on the log-Sobolev inequality (LSI) constant in the number of particles, leading to a potential exponential in inverse-temperature improvement in the number of particles required. This new bound is used to improve prior optimization, sampling, and uniform-in-time propagation of chaos guarantees for the MFLD.

Strengths

The Bregman divergence based technique of analyzing the particle discretization error is novel and interesting, and can inspire further research in this line of work. The intuition behind the proof is also nicely explained in the main text.

Weaknesses

Currently my main concern is the verification of Assumption 4 in standard settings. The authors cite Lemma 22 of Kook et al., 2024 to obtain a uniform-in-N LSI constant for ${\mu_*}^{(N)}$. However, in the new version of Kook et al., 2024, Lemma 22 only proves an LSI constant for the proximal Gibbs measure. In the new version, Corollary 25 does prove a uniform-in-N LSI constant for $\mu^{(N)}_*$, but this requires $N$ to be exponentially large in the inverse temperature $1/\lambda$, which is what this paper tried to avoid. I would be happy to increase my score if this issue is resoled. Additionally, minor comments and questions are provided in the “Questions” section below.

Questions

* What is the main challenge in handling unbounded activations? Furthermore, does the boundedness of the first variation in Assumption 1 play an important role, or can it be removed as the bound does not enter any of the estimates? * In Line 162, the authors refer to Equation (11) as the continuous-time propagation of chaos error, while Equation (11) only considers the error at optimality, and the continuous-time propagation of chaos result is introduced in Theorem 2. This might need more clarification. * In Assumption 3, $\ell$ is said to be Lipschitz smooth, while in fact its gradient is assumed to be Lipschitz. It might be more clear to refer to $\ell$ as simply $L$-smooth. * The step size appearing on the RHS in Part 1 of Theorem 2 might be a typo, since it is a continuous-time statement.

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors have adequately discussed the limitations of their work. The main limitation is the fact that the number of iterations and therefore the computational complexity still depend on the LSI constant, which make the algorithm inefficient in low noise regimes in the worst case. Another limitation is assuming bounded activations, which could potentially be relaxed with a more detailed analysis.