Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approach

Mean-field Langevin dynamics (MLFD) is a class of interacting particle methods that tackle convex optimization over probability measures on a manifold, which are scalable, versatile, and enjoy computational guarantees. However, some important problems -- such as risk minimization for infinite width two-layer neural networks, or sparse deconvolution -- are originally defined over the set of signed, rather than probability, measures. In this paper, we investigate how to extend the MFLD framework to convex optimization problems over signed measures. Among two known reductions from signed to probability measures -- the lifting and the bilevel approaches -- we show that the bilevel reduction leads to stronger guarantees and faster rates (at the price of a higher per-iteration complexity). In particular, we investigate the convergence rate of MFLD applied to the bilevel reduction in the low-noise regime and obtain two results. First, this dynamics is amenable to an annealing schedule, adapted from Suzuki et al. (2023), that results in improved convergence rates to a fixed multiplicative accuracy. Second, we investigate the problem of learning a single neuron with the bilevel approach and obtain local exponential convergence rates that depend polynomially on the dimension and noise level (to compare with the exponential dependence that would result from prior analyses).

Paper

Similar papers

Peer review

Reviewer Hr7F7/10 · confidence 3/52024-06-25

Summary

Mean-Field Langevin Dynamics (MFLD) framework is used to solve optimization problem over manifold. The main contribution appears to be reducing a general optimization problem over signed measures to probability measures using lifting or bilevel approaches. Convergence rate of MFLD, when applied to both approaches, are investigated.

Strengths

It sounds interesting to use lifting or bilevel ideas to reduce MFLD to solve general optimization problems over signed measures.

Weaknesses

The scope of lifting or bilevel ideas should be made clear. For instance, lifting idea seems to be applicable when the signed measure can be represented as projections of probability measures; does the projection representation always exist? Also, reduction of the optimization problem over probability measure may change to a more difficult optimization problem, e.g., the objective functions (3.1) and (3.3) which might be more difficult to deal with. Some discussions on how to deal with these functions empirically will be helpful.

Questions

See *weakness*

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

None

Reviewer 3p9w7/10 · confidence 3/52024-07-05

Summary

The paper studies extension of mean-field Langevin dynamics (MFLD) to perform convex optimization over the space of signed measures. The paper considers the lifting and bilevel approaches and shows that the latter guarantees better convergence properties at a wider range of hyperparameters. MFLD-bilevel is shown to be amenable to a previously studied improved temperature annealing schedule over the standard $1/\log t$ rate. Moreover when learning a single index model, MFLD-bilevel is shown to achieve a local convergence rate which depends polynomially on the dimension and inverse temperature via an improved LSI constant.

Strengths

* The paper rigorously compares the lifting and bilevel approaches, and the result that the latter leads to stronger guarantees (at least in the low temperature regime) is quite surprising to me as many influential works on shallow NNs have built on the lifted formulation. * The paper is overall well written and detailed, and the analysis utilizes novel techniques (namely local LSI bounds) to go beyond the standard results in MFLD. * While the paper has some similarities with [Takakura, Suzuki, 2024], the setting is more general with a focus on more abstract optimization as explained in Appendix A.

Weaknesses

* The parameter space $\mathcal{W}$ is assumed to be a compact manifold. Does this mean compact without boundary, which precludes subspaces of Euclidean space? Can this assumption be removed by e.g. adding a confinement term and considering relative entropy regularization with a log-concave distribution? While the non-necessity of these elements are presented as advantages in Section 1.1, studying $\mathcal{W} =\mathbb{R}^d$ is simply necessary for certain optimization problems and it would be very nice to fill out the details.

Questions

* In Theorem 4.2, can the time complexity to achieve $(1+\Delta)$-accuracy be rewritten in terms of $\Delta$ and compared with the usual logarithmic annealing? The current rate expression in terms of $\delta$ makes this unclear. For example, it seems $\beta_t\approx O(1/t)$. * Besides the established positive results, are there any theoretical/empirical results showing the difference in optimization dynamics between using squared and un-squared TV regularizer? * What is the main issue when trying to extend the results of Section 5 to multi-index models or general target functions?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Discussed in corresponding sections.

Reviewer akPi7/10 · confidence 3/52024-07-09

Summary

This paper extends the well-known and recently extensively studied mean-field Langevin dynamics (MFLD) to optimization problems over signed measures (instead of probability measures). This has applications and relevance to the training of NNs or other problems in data science, such as sparse deconvolution, which are intrinsically of such form. As an example, the learning of one neuron is presented at the end of the paper. In order to fit the setting of signed measures into the classical probabilistic framework of the MFLD, the Authors leverage and explore two classical strategies. Namely, the lifting and bilevel reduction. They observe that the bilevel reduction behaves more favorably compared to the lifting reduction, allowing to obtain stronger and faster convergence.

Strengths

- Extending the MFLD via a bilevel approach to optimization problems over signed measures seems to be an interesting and relevant direction given that this captures several problems of interest in machine learning and data science. - After discussing two potential directions to approach the problem (lifting and bilevel), the paper adapts and employs the convergence analysis of [Chi22b] to the new setting of the bilevel MFLD, which transpired to be the way to go due to the assumptions holding in more realistic scenarios. - Despite the paper being of purely theoretical nature and in parts quite technical, it is overall well-written and good to follow. The organization and structure of a general introduction, revisiting the classical MFLD as well as the two reduction strategies in Section 2 and 3, respectively, helps in this regard. - Overall, the technical contributions seem to be rigorous and convincing.

Weaknesses

The experimental exploration of the proposed approach is a bit limited. The experiment in Figure 1 seems very academic. I was wondering, if the Authors could comment on whether there are experiments already conducted in the literature that could prove the practicability of the approach. A remark on that instead of the experiments in Figure 1, which could be presented in the Appendix, would be appreciated. Yet, as the paper's contribution is predominantly of theoretical nature, I _do not_ see the limited experimental investigation as a reason for a lower score.

Questions

- The description of the lifting and bilevel reduction approaches in lines 50--55 feels a bit loose and vague and not much insightful. I would recommend to the Authors to either skip this part here or be more elaborate by giving more insights into how these approaches allow to transfer from the signed measure to a probability measure. At the same place, some citations of the papers proposing these ideas would be welcome (i.p. for the lifting). - What is $J$ in line 44? - uniform log Sobolev inequality: At first sight, it seems unclear, whether this assumption is reasonable or quite strong, as it has to hold for $\mu_t$ for all times $t$. Could the Authors comment on that? Are there sufficient assumptions on $\mu_0$ and $G$ that ensure such assumption at all times? Moreover, does it hold in interesting settings such as the training of one-hidden layer NNs? As far as I know, the cited reference [Chi22b] addresses parts of this. It would be nice to include a comment. - Is there any intuitive insight that could explain, why the bilevel approach works much better than the lifting strategy? - And I have a clarifying question concerning the bilevel reduction: The central observation allowing this reduction is the fact that the TV-norm can be expressed as in line 197? In addition to the change of infima in line 198, which allows to recast the optimization problem over the signed measures into an optimization over probability measures? (If so, maybe it would be nice to put this as an equation (in addition to the Proposition) to make it better visible as the central "trick".)

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

I do not see any unadressed limitation. Moreover, I like how the lifting approach is discussed, despite turning out not to be the way to go.

Reviewer hSwQ6/10 · confidence 4/52024-07-12

Summary

Mean-field Langevin dynamics (MFLD) has been developed for optimizing convex functionals over the space of probability measures. This work extends MFLD to convex problems defined over the space of signed measures. The authors consider two approaches: lifting and bilevel approaches, and prove the superiority of the bilevel approach. The lifting approach cannot satisfy the two required conditions for MFLD convergence under weak regularization, whereas the bilevel approach satisfies both conditions simultaneously. Additionally, an improved convergence rate is given for the enhanced annealing schedule of the temperature.

Strengths

- This work successfully extends MFLD to convex optimization problems over the space of signed measures. Although recent work [TS24] also considered a similar approach, their method is a special case of the MFLD-bilevel proposed in the paper. - This work proves an improved convergence rate for MFLD-bilevel with annealing temperatures. Although the rate still exponentially depends on the regularization strength, it may be difficult to avoid such dependence.

Weaknesses

- Since this work considers (compact) Riemannian manifolds as the particle space, time discretization is essentially challenging. I’m curious about how to discretize the dynamics in time and guarantee the convergence rate. - I guess the lifting approach satisfies both (P1) and (P2) once we limit the range of $r$ to be bounded, like $r∈[−R,R]$. If this is true, it is worth considering whether the bilevel approach is truly superior to the lifting approach. - As the authors commented in the paper, [TS24] studies a quite similar method. I would like to see any additional technical challenges or difficulties compared to [TS24] if possible.

Questions

- In Proposition 5.1, the iteration complexity for entering the sublevel where MFLD converges faster is not specified. Is it possible to explicitly describe the time required to enter this level?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are addressed in the paper.

Reviewer 3p9w2024-08-07

Thank you for the reply. I will maintain my high rating of the work.

Reviewer akPi2024-08-07

Thanks for your reply and for addressing my questions, in particular also for providing some intuition about the two approaches, which does sound very reasonable. After reading also the other rebuttals and the Authors' replies, I remain very positive about the paper and maintain my positive score.

Reviewer hSwQ2024-08-09

Thank you for the reply. The authors have addressed my concerns well. I will increase the score. I would encourage the authors to add discussion on the lifting case with the constraint $r \in [-r,r]$.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC