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

We thank the reviewer for their positive evaluation, and address their questions and comments below.

• Thank you for pointing out the potential confusion about the scope of these ideas. In the final version, we will further clarify that the lifting and bilevel ideas are both always applicable for optimization over signed measures. In particular, any signed measure $\nu \in \mathcal{M}(\mathcal{W})$ can be represented as a projection of a probability measure $\mu(\mathrm{d} r, \mathrm{d} w) = \delta\_{\\|\nu\\|\_{TV} \frac{\mathrm{d}\nu}{\mathrm{d}|\nu|}(w)}(\mathrm{d} r) \frac{|\nu|(\mathrm{d} w)}{\\|\nu\\|\_{TV}}$.
• The objective functions $F_{\lambda,b}$ (3.1) and $J_\lambda$ (3.3) are amenable to optimization using Wasserstein gradient flow (WGF) or MFLD, because they are defined over sets of probability measures, whereas the original optimization problem (1.1) is not. To solve the optimization problems (3.1) or (3.3), we rely on the WGF and MFLD literature which have studied the problem of optimizing convex functionals over probability measures in detail. We will further emphasize this fact in the final version.

We thank the reviewer for their thoughtful assessment, and address their questions and comments below.

• On compactness: $\mathcal{W}$ is assumed to be without boundaries. This assumption is missing on line 20. Thank you for pointing this out.

  The techniques presented in our work can be used without significant modifications for dynamics over $\mathbb{R}^d$ using additional confinement. We chose to focus on compact Riemannian manifold without boundaries because it is the natural setting for 2-layer neural networks with homogeneous activations (on the sphere), or for sparse deconvolution (on the torus). The absence of additional confinement is not an advantage in general; the phrasing in Section 1.1 can indeed be misleading and will be corrected.

• On the statement of Theorem 4.2: To state the time complexity of this theorem in terms of $\Delta$, we notice that $\delta = \Delta / \log(C/\Delta)$ for some constant $C$ that is polynomial in problem parameters (including $\lambda^{-1}$ and ${J^*_\lambda}^{-1}$). The time complexity would roughly read $T = \exp\big(O\big(\frac{\log(C/\Delta)}{\lambda\Delta}\big)\big)$. This can be compared with the classical logarithmic annealing procedure, which results in $T = \exp\big(O\big(\frac{1}{\lambda\Delta J^*_\lambda}\big)\big)$. Further, following the discussion on line 271, we can take $J^*_\lambda = O(\lambda)$. Thus, the annealing schedule of Theorem 4.2 leads to an improvement of order $O(\log(C/\Delta)/\lambda)$ in the exponent. Originally, we stated the time complexity of Theorem 4.2 using $\delta$ rather than $\Delta$ to simplify the presentation while having in mind that $\delta$ and $\Delta$ are equivalent up to logarithmic factors. We agree with the reviewer that the time complexity should be stated explicitly in terms of $\Delta$ for easier comparison with the baseline annealing procedure.

• On the difference between using squared and un-squared TV regularizer: Concerning the lifting formulation, our negative results can be generalized to the case of un-squared TV regularizer, and we will add a remark on this in the final version. For the bilevel formulation, using a similar trick as for the squared TV norm, the unsquared TV norm can be expressed as $\Vert \mu\Vert_{TV} = \int_{\eta \in \mathcal{M}_+(X)} \frac12 \int \frac{\vert \mu\vert^2}{\eta} + \frac12 \eta(X)$, but here $\eta$ is an un-normalized measure so the resulting bilevel problem is not amenable to MFLD.

• On the generalization of Section 5 to multi-index models: For Theorem 5.2, we rely on the rotational symmetry (orthogonal to the target direction) induced by the single-index model to simplify the Gibbs potential. It is an interesting open question to find examples of multi-index models with similar symmetries that lead to a polynomial LSI constant.

We thank the reviewer for their detailed evaluation, and address their questions and comments below.

• Thank you for the feedback on this part of the introduction, it will be adapted.

• "$J$" on line 44 should be "$F$".

• Thank you for the feedback on this part of Section 2, we will add more discussions on the reasonableness of the uniform log-Sobolev inequality (LSI) assumption to the final version. The cited reference [Chi22b] indeed contains sufficient conditions on $F$ ensuring uniform LSI, which include the training of two-layer neural networks under some technical assumptions such as bounded activation function. Part of our motivation for the present work was to remove this boundedness assumption.

• On why "bilevel" is better than "lifting" intuitively: For optimization over probability measures $\mathcal{P}_2(\mathcal{W})$, MFLD differs from Wasserstein gradient flow (WGF) by adding noise, which encourages exploration of $\mathcal{W}$. For optimization over signed measures $\mathcal{M}(\mathcal{W})$, which can be formulated as optimization over $\mathcal{P}_2(\mathbb{R} \times \mathcal{W})$ by lifting, there is no need to explore in the direction of the "weight" variables in $\mathbb{R}$, as the landscape is already convex in these directions. On the contrary, adding noise to both the weight ($\mathbb{R})$ and position variables ($\mathcal{W}$) turns out to be detrimental for convergence, by our negative result on "lifting" in Section 3.1.

  This suggests following an intermediary dynamics: adding noise to the WGF on the lifting objective, but only on the position variables. The bilevel approach is a two-timescale limit of this intermediary dynamics (see Eq. (3.4)), which is amenable to analysis as it is itself an instance of MFLD.

• Indeed, lines 197 and 198 are the central trick leading to the bilevel formulation. We will emphasize more on the role of this trick in the final version.

We thank the reviewer for their detailed and encouraging comments. We address each point in "Weaknesses" and "Questions" separately.

• On time-discretization over Riemannian manifolds: So far, the theory for time-discretization of MFLD is established when $\Omega = \mathbb{R}^d$ [SWN23], or when the objective functional $F+\beta^{-1} H$ is a relative entropy and $\Omega$ is a Hessian manifold [GV22] or a product of spheres [LE23]. In all cases, a crucial ingredient is the assumption that $F$ is smooth in the sense of (P1), hence our focus on this property. It would indeed be an interesting direction for future research to see how the analysis techniques of the aforementioned works could be combined, to show convergence guarantees for time-discretizations of MFLD on manifolds.

• On constraining $r$ to $[-R, R]$ in the lifting approach: Indeed, artificially limiting the range of $r$ would ensure (P1) and (P2). This is an interesting suggestion to make the lifting approach work. We did not find it natural to investigate this direction mainly for the following reasons:

  • Practically, this introduces a hyperparameter $R$, the tuning of which might be quite tricky: taking $R$ too small will exclude the true solution ($\mathrm{argmin}~ G_\lambda$) from ever being reached, and taking $R$ too large will significantly affect the smoothness and LSI constants.
  • Theoretically, the manifold $[-R, R] \times \mathcal{W}$ possesses boundaries, which adds significant technical difficulties to the analysis (in particular for time discretization).

  We note that a related modification of the lifting approach is discussed in our answer to Reviewer akPi (4th bullet point).

• On the technical difficulties compared with [TS24]: The dynamics they analyzed is (a variant of) MFLD-Bilevel applied to an objective $G$ of a specific form, as explained in Appendix A. The additional difficulties for us were $a)$ to establish (P1) and (P2) for the bilevel objective $J\_\lambda$ with no structural assumptions on $G$; $b)$ to analyze the convergence of MFLD-Bilevel beyond the crude upper bound provided by Theorem 2.1, corresponding to [TS24, Theorem 3.7]. We also note that the angle of that paper is quite different as it does not consider the lifting approach nor mentions signed measures. (We developed our results independently and concurrently with [TS24] which was announced on arXiv in March 2024.)

• On the time $t_0$ required until Proposition 5.1 applies: $t\_0$ can indeed be explicitly described by inspecting the proof. Inspection reveals that the localness requirement is $J\_{\lambda,\beta}(\eta_{t_0}) - \inf J\_{\lambda,\beta} \leq \delta$ for $\delta = \left( 2L_0 G(0) (1+\frac{L_0}{\lambda}) \frac{L_0}{\lambda} \beta \sqrt{\frac{2\beta}{\alpha^*}} \right)^{-2} \left( \log(1-\varepsilon/\alpha_\beta^*) \right)^2$, and that it is guaranteed for $t_0 = \frac{\beta}{2 \alpha\_\tau} \exp(\frac{L\_0}{\lambda} G(0) \beta) \log\frac{J\_{\lambda,\beta}(\eta_0) - \inf J\_{\lambda,\beta}}{\delta}$.

[GV22] Khashayar Gatmiry and Santosh S Vempala. “Convergence of the Riemannian Langevin algorithm” (2022).

[LE23] Mufan Li and Murat A Erdogdu. “Riemannian Langevin algorithm for solving semidefinite programs” (2023).

[SWN23] Taiji Suzuki, Denny Wu, and Atsushi Nitanda. “Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction” (2023).

[TS24] Shokichi Takakura and Taiji Suzuki. “Mean-field Analysis on Two-layer Neural Networks from a Kernel Perspective” (2024).