Mean-Field Langevin Dynamics for Signed Measures via a Bilevel Approach
Authors: Guillaume Wang
Alireza Mousavi-Hosseini
Lénaïc Chizat
Venue: NeurIPS 2024 spotlight (2024)
Keywords: two-layer neural network
Wasserstein gradient flow
noisy particle method
mean-field Langevin dynamics
logarithmic Sobolev inequality
bilevel optimization
Primary area: optimization

Abstract: 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 polynomial 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).

Mean-Field Langevin Dynamies for Signed Measures via a Bilevel Approach
Guillaume Wang"!
Alireza Mousavi-Hossein"
Lénaïc Chizat!
École polytechnique fédérale de Lausanne
2 University of Toronto and Vector Institute
guillaume. wangtepfl, ch.
mousaviscs. toronto. edu, Senaic.chizat@@pee
Abstract
Mean-field Langevin dynamics (MLFD) is a class of interacting particle methods that tackle convex optimization over probability measures on 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, investigate how 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 higher per-iteration complexity). In particular, 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 annealing schedule, adapted from [SWON23J, that results in improved convergence rates fixed multiplicative accuracy. Second, we investigate the problem of learning single neuron with the bilevel approach and obtain local exponential convergence rates that depend polynomially the dimension and noise level (to compare with the exponential dependence that would result from prior analyses).
1 Introduction
Let MMW) be the set of finite signed measures on a compact Riemannian manifold without bound- aries W and let G M(W) + R be convex function, assumed smooth in the sense of Assumption 1 below. In this paper, we investigate optimization methods to solve
min G2(v). vEM(W)
Gx(1) = G(v)
+ - 111..v.
(1.1)
2
where || |/ is the total variation norm and X >> 0 the regularization level. This covers for instance risk minimization for infinite-widh width 2-layer neural networks (2NN) [BRVDM05; Bac17] by taking
W = gd the unit sphere in Id+11 or W [d++1 and
G(v) = (aa))~ [(h((,,).
where
h(v,)) = / (((r, w))dv(w). (1.2)
Here 90 R R is the activation function, h(,,,) -) is the predictor parameterized by v, G is the (population empirical) risk under the data distribution p P(Rd++ and & is smooth (uniformly
"Equal contributions, authors ordered randomly.
'The square exponent on ||- |lzy might appear unusual, but is convenient our subsequent developments. We show App. A that the regularization path the same with or without the square.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

in y) and convex in its first argument. These 2NNs will be our guiding examples throughout, but note that the class of problems covered by Eq. (1.1) is more general and includes for instance sparse deconvolution via the Beurling-LASSO estimator [DG12] or optimal design [MZ04].
To tackle such problems, interacting particle methods use the parameterization apply gradient methods in well-chosen geometry [Chi22c; YWR23; GCM23J. They have recently gained traction thanks to their scalability and flexibility, and in the context of 2NNs, the usual gradient descent algorithm is an instance of such a method. On the downside, global convergence guarantees remain difficult to obtain due to the nonconvex nature of the reparameterized problem and existing positive results require either very specific settings [LMZ20], or modifications of the dynamics which often limit their scalability².
rrówe I and
In a related, but slightly different context, mean-field I Langevin dynamics (MFLD) solve entropy- regularized problems of the form
min MEP(W)) Fa(u),
FØ(n)= = F(u) + +(
(1.3)
where P(W") is the space of probability measures on a manifold W (typically R²), F: :(W)) R is a (sufficiently regular) convex functional, (()) = f log(du/d vol) is the negative differential entropy and / > O. These dynamics are obtained as the mean-field limit of noisy interacting particles dynamics [MMN18 1M HR$S21] and converge globally at an exponential rate [NWS22: Chi22b], under two key conditions on F: (i) notion of regularity, which we refer to as displacement smoothness (see P1 below) and (ii) a uniform bog-Soooeev inequality (LSI) condition (see P2 below). These mean mean-field, continuous-ting time guarantees have been further refined into computational guarantees for fully discrete algorithms [CRW22: SWN23J. The favorable properties of MFLD naturally lead to the following question:
Can we efficiently solve problems of the form Eq. (1.1) using MFLD?
At first, it is not obvious that MFLD can be applied at all since it is originally defined only for problems over probability measures. However, we can find in the literature two general recipes to reduce a problem over M(W) to a problem over P(W"), ). thus amenable to MFLD. The first one is a lifting reduction, that takes W = IR W where the extra dimension serves to encode the signed mass of particles [CB18, 18, Section A.2] [Chi22c]. The second one, that takes W = W, is a bilevel reduction [Bac21; TS24] that uses variational representation of the regularizer || ll common in the multiple kernel learning literature [LCBGJ04]. A first task is thus to compare the behavior of MFLD on these two approaches. Furthermore, MFLD involves an entropic regularization which is absent from Eq. (1.1). A second task is thus to analyze the behavior of MFLD in the large / regime, when the regularization vanishes.
In this work, we tackle these two tasks and make the following contributions:
In Sec. 3, we introduce the lifting and bilevel reductions and compare the "displacement smoothness" (P1) and "uniform LSI (P2) properties of the resulting problems. These properties play a central role in the global convergence analysis of MFLD. Specifically, we consider a large class of lifting reductions and show that none satisfies simultaneously (P1) and (P2) unless À is large. In contrast, the bilevel reduction satisfies both under mild assumptions. So in the sequel we focus on MFLD applied the bilevel reduction.
In Sec. 4, we investigate what convergence rates can be obtained for the problem (1.1) by using MFLD on the bilevel formulation. While a classical simulated annealing technique yields convergence in Olog log // log t), we show that the structure of the bilevel objective is in fact amenable to more efficient annealing schedule, adapted from [SWON23], that reaches a fixed multiplicative accuracy, say 1.01 inf Gx. in time e e(A log > instead of (((-²) for the classical schedule.
In Sec. 5, to obtain more complete picture, we investigate the problem of learning single neuron. Here, using a Lyapunov type argument, we show that the local convergence rate of MFLD applied to the bilevel formulation scales polynomially in 8 and d, odds with all previous MFLD analyses which had exponential dependencies.
All proofs are deferred to the Appendix.
"Such as forcing the particles to remain close to their initial position [Chi22c], or adding new particles using potentially hard linear minimization oracle [DDPS19].

---

1.1 Related work
Particle methods and mean-field limits. Interacting particle systems have been studied for decades in various fields, see e.g. [Szn91; CD13; Lac18]. Their more recent connection with the standard training of 2NNs INS17; SS20; RV22; MMN18] has suggested new settings of analysis, where convexity of the functional plays key role, and has led to many developments. In particular, the case of MFLD (under study here) quickly progressed from nonquantitative guarantees [MMN18; HRSS21J, to mean-fiele reean-field convergence rates [NWS22; Chi22b] and fully discrete computational guarantees [CRW22; SWN23; KZCE+24] in the span of a few years. Recent progress also address its accelerated (underdamped) version [CLRW24; FW23], which could also be of interest our setting.
Multiple kernel learning and bilevel training of NNs. The lifting reductions consider inspired by the unbalanced optimal transport literature [LMS18], while the bilevel reduction comes from the Multiple Kernel Learning (MKL) literature [CVBM02; LCBGJ04; RBCGO8] (see [Bac 19] for an account). While the latter is usually studied with a discrete domain W (see also [PP21; PP23] for recent computational considerations), it was suggested for the training of large width 2NN in [Bac21] and used in conjonction with MFLD in [TS24] (more details below). Relatedly, recent line of work studies the (noiseless) training of 2NN in a two-timescale regime, where the outer layer is trained at much faster rate than the inner layer [BMZ23; MB23; BBP23]. This implicitly corresponds to optimizing the bilevel objective and leads to improved convergence guarantees.
The work that is closest to ours is [TS24], which considers the MFLD on a 2NN with weight decay where the outer layer is optimized each step. They interpret the resulting dynamics learning dynamics and study properties of the learnt kernel and its associated RKHS. While they do not formulate explicitly the problem Eq. (1.1), it can be shown that our approaches are equivalent when considering W = Rd+1 in Eq. (1.2) (and adding an extra regularization). The details are given in Sec. A.2. Key advantages of our formulation with W gd are that we cover the case of unbounded homogeneous activation functions (such as ReLU), and can obtain improved LSI.
kernel
2 Background on guarantees for mean-field Langevin dynamics
The MFLD is defined as the Wasserstein gradient flow (HE))ER, in P(2)) of an objective of the form Eq. (1.3). It characterized as the solution to the partial differential equation (PDE)
dep div VF + +
Ho E P(L).
(2.1)
where F'[n] : f2 - R. is the first variation of F at μ [San15, Sec. 7.21, defined by lim 1 (F(( + e(Hí - -+)) - F(P)) F'luld(p'' μ1) for any nt E P()). This PDE corresponds to the mean-field limit (N + 00) of the noisy particle gradient flow wi E SN::
vis N, dut = -V me F(N) (wi WN) dt + 2B3
lid. wi
po
where P(N) "o......" = F N Emsr 6.) and the B; are N independent Brownian motions on Ω. The convergence guarantees for MFLD rely three key properties:
(PO) (Convexity) F is convex and is such that F3 admits a minimizer 11g-
(P1) (Displacement smoothness) is L-displacement smooth, in the sense that
Vy E PP(O),, Ve E Ω max SET,1
22 F')u[(s, s) S L,
| S1
and Vu. E P2()), Vo E Ω [VF V L Wa(,,)))
where 2 denotes the Riemannian Hessian.
(P2) (Uniform LSD) There exists a > such that Vt 2 0, F3 satisfies local a-LSI at H4g, as in Def. 2.1.
"Strictyy speaking. (PI) is only a sufficient condition for displacement smoothness (see details in App. B). We refer to (P1) as displacement smoothness in this paper for conciseness only.
3

---

Definition 2.1 (Local LSD, We say that a functional F3. = + 8-1 H satisfies local a-LSI μ P(t) Z == Sn exp (-3F"(p)) Z (xp(-3FF<p] E P(2)) satisfies a-LSI, that is
oo and the proximal Gibbs measure p ==
Vyl E P(2), H (n'lA) -1(((
2a
where the relative entropy and relative Fisher Information respectively defined as du' (w) dp'(w)
dp
I(plA)) S log du
H (v'1A)= log
dy'.
dp
and || ||w denotes the Riemannian metric.
We review some useful criteria for LSI in App. B. In particular, the uniform LSI property (P2) holds for example when training two-layer neural networks with frozen second layer, under some technical assumptions such as bounded activation function. In fact that case, the proximal Gibbs measures fi even satisfy LSI uniformly for all μ E P(2) [Chi22b; NWS22].
Note that the Riemannian gradient V and the Laplace-Beltrami Beltrami operator appearing in (2.1), as well as the definition of Brownian motion, depend on the Riemannian metric of Ω. This dependency is reflected (PI) and (P2).
The global convergence of MFLD is guaranteed by the following theorem, with a rate.
Theorem 2.1 ([Chi22b, Thm. 3.2]INWS22, Thm. 1J). Consider F P((2) R. and (14) as in (2.1).
If (PO), (P1) and (P2) are satisfied then for 2 0 it holds
8- ( S Fo(ux) - Fa(u3) s exp(-23 (Fa(Ho) Fz(uz))
Note that although the L-smoothness constant does not appear in Thm. 2.1, , does appear in the discrete-time guarantees of [SWN23], and is thus important quantity in practice. In this paper, limit our analysis to the mean-field dynamics (2.1) because its dineeiiscretization has not yet been studied Riemannian manifolds. In continuous time, the proof of Thm. 2.1 translates directly to Riemannian manifolds thanks to our definition of (P1), see App. B.
Reductions from signed measures to probability measures
In order to apply the MFLD framework solve initial problem over signed measures (1.1), must first recast it as an optimization problem over probability measures. In this section we build two such reductions, and discuss the properties (PO, Pl and P2) of the resulting problems.
3.1 Reduction by lifting
Reductions by lifting consist in representing signed measures as projections of probability measures in the higher dimensional space Ω R W. This construction involves the 1-homogeneous projection operator* (9)) + M(W) characterized by
Vip E C(W, TR), /,, ((u)(())(du)) = ros(ww)uudrr dw))
where (22) is the subset of P(22) for which ʃ [ dw) < +oo. For instance, it acts on discrete measures as h 2 = 1 Em ryou, We also define, for E [1,2] and E (I)))), Vo(u) == (Jn |r| drj,, w))
2/b The objective functional of the lifted problem is then
defined, for μ E Po(
Fxb(p)= = G(hp) > Vo(u).
(3.1)
is equivalent to minimize GA or FX, b. as shown in the following statement.
*We could consider more general p-mooggeeouus projections as in [LMS18], but we show in See, C.2 that always bring ourselves back to the p = up to change of metrie.

---

Proposition 3.1. Let M(W). For any E Po(W) such that hu = u, it holds Fxo(H) 2 G2(v), and equality holds for p(dr, dw) = 6yw (dr) (e((w)) this p. when b > 1). In particular, if G admits a minimizer then Fx,6 does too, and it holds
||l| where fww = ||l|l du dlul (w) (and only for
min FXb(H)= min Gx(r). MEP
VEM(W)
is not difficult to see that FA.5 satisfies (PO) as long as GA admits a minimizer. In order to study (P1) and (P2), we need to define Riemannian metric on Ω. Following [Chi22c], we consider general class of Riemannian metrics on J2" = - x W, parameterized by grs Tw E IR and I > 0, defined by
<
òri ow1
drz 6w2
åriōr2 2
r-1 1r/0r
|r|00 (ōw1, ow2)
(3.2)
(m,w)
This indeed defines an inner product on T Tir,w)S2* R x Tw W that varies smoothly, and so equips S2" with (disconnected) Riemannian manifold structure [Lee 18.. Intuitively, the parameter I will govern the relative speed of the weight or position variables along gradient flows; larger I' means faster weight updates.
Two particular cases of this construction appear (sometimes implicitly) in the literature on 2NN:
(i) when q- 2 and Que O, the metric (3.2) extends to the product metric on Ω W. With W = Rd+11 this corresponds to the usual parameterization of 2NNs and is the setting of most previous works applying MFLD 2NN (with a weight decay regularization on the second layer for b = 2 and À > 0).
when 9r = qu = 1, S2" is isometric to the union of two copies of the (tipless) metric cone over W [BBI01] (via the mapping (r,w) + (sign(r), [|r],w)). This is the natural setting for optimization over signed measures; and with W $d, is equivalent to the parameterization of 2NNs with ReLU activation and balanced initialization [CB20, App. H].
Issues caused by the disconnectedness of Ω On the level of the equivalence of variational problems, one can check that the statement of Prop. 3,1 also holds if Ω = IR x W is replaced by Q* IR. W. However, when the manifold Ω is truly disconnected, then P(22) not connected in the sense of absolutely continuous curves in Wasserstein space. More precisely, Ω is the disjoint union of S2
I. x W and Ω² R x W, and one can show that (for certain choices of (w.) if (1 is Wasserstein gradient flow (or any other absolutely continuous curve), then (x()) = no (f2) for all
Moreover, supposing for simplicity that GA has a unique minimizer v and that b > 1, then Fx,6 has unique minimizer μ², and μ² (92) 24 (W) ||||y where y 1/4 V_ is the Jordan decomposition of v. Therefore, Wasserstein gradient flow for Fx,6 can only converge to nº if it was initialized such that po( n² (Qt). In terms of particle methods, this means that the fraction of the particles (r:, wx) initialized with > O must be precisely μ (S2.). ). A similar problem arises if we apply MFLD to Fx,o, since it is nothing else than Wasserstein gradient flow for FX,6 + 8- H; but it is more tedious to discuss formally, as FX,6 + 3-1H H does not have a minimizer in general.
In order to bypass this limitation, one may focus on settings where the ratio 24 (W)/ || for the optimal v is known advance, e.g., the problem (1.1) constrained to non-negan non-negative measures, or on choices of gr. 9w for which Ω can be extended into connected manifold, such as the product metric q. = 2, 40 = O. However, even in those cases, MFLD on Fx, presents other limitations.
Incompatibility with MFLD. We now show that, in spite of the degrees of freedom given by the parameters gr: 9w and b. satisfying both (PI) and (P2) requires restrictive assumptions. This suggests that the lifting approach is fundamentally incompatible with MFLD.
'This issue also occurs in the case q. = qu. 1. even though 92" can be completed into topologically connected set by adding an element 0 "bridging" the two cones 24 and 92.. Indeed, any particle reaching 0 remains 0 for all subsequent times. Besides, this completion is not itself a manifold, 0 singularity.
5

---

Proposition 3.2. Consider & Eq. (3.1) and Qx equipped with the metric (3.2). Suppose G' [v] is continuous for all v and that there exists v such that G²) is not constant equal to O. Then
qf * or que + or brrb5 + 1, then (P1) does not hold.
If gr - -w = b = 1, then for any μ E PI( there exists X > O such that Fx,b + B-1 H does not satisfy local LSI at μs for any X < No (in particular (P2) does not hold unless À is large enough).
When g. = qu b = 1 and is large enough, then it can indeed be shown that Thm. 2.1 applies under natural conditions, see for instance [Chi22b, Sec. 5.1].
Remark 3.1. For functionals of the form GA,* G(v)+ + A |I instead of (1.1) which corresponds to s 2, one can formulate similar reduction by posing (p)=(m) = (So Ir|* dp(r, w))*/b and (p) = G(hu) + A D. (1)) The statements of Prop. 3.1 and Prop. 3.2 hold true with Gx replaced by Gs,s- and Fx,6 by Fx,6,s. for any S b S 8, can be shown by very simple adaptations of the proofs (only the second inequality in the proof of Lem. C.I, and the definition of X in (C.2), need be adapted). Note that the problem considered in [Chi22c] is of the form G(v) + A >||//TTV and they analyzed Wasserstein gradient flow on 1A,1,1 with 9r = Qu = 1 (in particular the issues caused by the disconnectedness of l are bypassed thanks to the choice b = 1). The above discussion shows that applying MFLD to that problem would only yield convergence guarantees for À large enough.
3.2 Reduction by bilevel optimization
We define the bilevel objective functional Jx for 17 E P(W) as
(n) inf G(v) A |01 vEM(W)
(3.3)
Jw "
can be derived using the variational representation of the squared TV-norm [LCBGJ04; Bac for any v E M(2),, one has | minnep(w) Sw |v|2 By exchanging infima, it thus holds inf inw) (x(r) inf inf G(v) + 2/2 f
n nnww Jx(1). Moreover, the objective minimized in (3.3) is jointly convex in (1, 2) and partial minimization preserves convexity, so JA is convex. Let gather these crucial remarks in a formal statement.
|o| 7
Proposition 3.3. The bilevel objective Jx is convex and inf JA inf Mow) Gx. Moreover, if Gy admits minimizer v E .M(W), then arg min Jx = {{ TTllzy
1 |u] |ا|,, E arg Trgminn. G.).
Link between the lifted and bilevel reductions. The equality case in the statement of Prop. 3.1 shows that we can restrict the lifted reduction to measures lr E (22) of the form p(dr, dur) = ((r))) (dr)))(dw)) for some f W - R and 7 E P(W). Since they satisfy hp(dw) f()))ddw). the lifted reduction with b = 2 thus rewrites
A min min G(fn) + f f(w))d)(w). NEP(W) SELP(g)
After the change of variable (v, 1) = (fy, 11), the outer objective is precisely (x)... Thus, Wasser- stein gradient flow JA can be seen two-timescale optimization dynamics: it is the Wasserstein gradient flow on xx in the limit where I - oc. In the context of 2NN training with the parametriza- tion (i), this amounts in [BMZ23; MB23; BBP23; TS24]. This remark allows to implement the bilevel MFLD numerically by discretizing in time the system of SDEs, for fixed large N and r,
training the output layer infinitely faster than the input layer. done
vi < N, drž- - V dwf - V. (re w/)dt + / - -vVC(l/uuud + 22
w()d
- (G'(v)((u))) Art) di
(3.4)
'We use fw |u|2 as shorthand for Jwv ( (w)) dd((w)

---

where μ4 = 1. E 6 (ri,a:)) and v/2 = E rio and taking N absence of noise term on the weight variables r: it reflects the fact that MFLD for the bilevel objective not a limit case of MFLD for the lifted objective, as the noise would prevent reach optimality in the inner problem.
Ei= 8mp. Notice the
Compability with MFLD. We now show that, in contrast to the lifting reduction, the bilevel reduction amenable to MFLD. The main assumption on (1.1) is as follows.
Assumption 1. G M(W) + R is non-negative and admits second variations, and for each i E {0, 2, there exist L,, B; < oo such that V''G"v](w, w')l I L, and [ S I || B, for all E EM(W) and w, w' E W. Moreover there exists T < oo such that I v.r G" w'll S L2 for all N, w, w', Furthermore, W is compact and the uniform probabil- ity measure 7 on W satisfies LSI with constant a,,.
Concrete settings that satisfy Assumption 1 are discussed in Sec. 5. The following proposition confirms the compatibility with MFLD and gives quantitative bounds on the LSI constant.
Proposition 3.4. Under Assumption I, Jx satisfies (PO), (P1) and (P2). More precisely, for any
P(W), Jx + 3----BBB H satisfies local LSI at n with the constant Osş = a, exp (- Lo)J((n)) Further, Jx B H satisfies c-LSSI uniformly along the MFLD trajectory (nt)). with the constant a= a, exp
' } LoB min (G) (No)) + 3- (NolT)})
In view of the negative result of Prop. 3.2 for the lifting reduction, and the positive result of Prop. 3.4 for the bilevel reduction, the sequel we focus on MFL.D applied on Jx, which we will refer as MFL.D-Bilevel.
4 Global convergence and annealing for MFLD-Bilevel
While the bounds from Prop. 3.4 along with Thm. 2.1 allow to establish global convergence to minimizers of JA + + H, our aim to minimize the unregularized bilevel objective JA. This can be achieved by annealing the temperature parameter 8-1 along the dynamics. Namely, Theorem 4.1 of [Chi22b] guarantees that by choosing Bg = clog(t) )) for an appropriate constant c, the annealed MFLD trajectory
O.77 div (n V (m))) + 8 An
satisfies J (nt) - inf J=E=AO0 -0
log loget log!
This is very slow rate however.
In this section, we show that the structure of J/ originating from the bilevel reduction can be exploited go beyond the generic guarantees from [Chi22b, Thm. 4.1]. Namely, we study in detail an alternative temperature annealing strategy, and we show that it improves upon the classical one 3, log(f) terms of convergence a fixed multiplicative accuracy.
4.1 Faster convergence to a fixed multiplicative accuracy
Definition 4.1. Suppose 0 & arg min G. so that J == inf JA > O. We will say that MFLD-Bilevel with given temperature annealing schedule (3:)20 converges to (1 A) A)-muliplicaativ accuracy in complexnyy Ts, for a fixed positive constant A (say A = 0.01), if (T(T)) I (1 + A)J;
Note that in machine learning settings where the problem (1.1) corresponds to learning with overpa- rameterized models, it is realistic to assume J to be small (as long as the regularization À is small), and TA is the time it takes for the annealed MFLD to achieve a suboptimality of at most AJ
For ease of comparison, let us report the time-complexity Ta that can be achieved by simply running MFL.D-illevll with constant but well-chosen B, based on the bounds from Prop. 3.4 and Thm. 2.1. Proposition 4.1 (Baseline "annealing" schedule: constant 3t). Under Assumption 1, let A > 0 and assume that A < LoL, Then, MFLD-Bilewe with the temperature schedule Vt, B
Ad AJ
log
CB converges to (1 mumuliippieeative accuracy in time AJ
C" log
CB
(C'LoG(O)
CB ) .log
(2G(0)
+ C'H (mol-) AJ
Ta S
AJ
AJ
exp
X AAJ
log
AJ.)
where B = poly(Lo. L1, B1, G(0), and C, O' are constants dependent W (and d and ac
7

---

For the annealing schedule Bg ~ log(t), the compleomyyy Ta that can be guaranteed from inspecting the proof of [Chi22b, Thm. 4.1] has the same dependency on d, x and J as for the baseline 3. = est.
Improved annealing schedule. Recall the result of Prop. 3.4: for any 3 > O, JA + 8 H satisfies local Lsg-LSI vi with Okş = ar exp(- fa (.xx())) (17)). Informally, if we manage control Jx(n) along the annealed MFLD trajectory and show that it decreases, then we can increase 3r at the same rate, while retaining the same local LSI constant. This observation and the resulting annealing procedure were introduced in [SWON23], in a 2NN classification setting with the logistic loss. There the optimal value of the loss functional, corresponding to our Js is O, and the annealing procedure yields favorable rates for global convergence. Here we show that this procedure is also applicable for MFLD- D-Blleellll, as soon as satisfies the mild Assumption 1, yielding favorable rates for convergence to fixed multiplicative accuracy."
Theorem 4.2. Under Assumption 1, there exist constants B = pLly B, G(O), 1 and C; dependent only on GO)), H (710), W (and d and a+)) such such that the following holds. For any A < B LLFLD-Bilevel with the temperature schedule (3:)12 defined by Vk S K, Vt E [tss tx+1 Be = 2 where to 0 and K [2 (B/s2(BB(A))) (B/))) and
J
1-1 1 1- C,2 * k exp (Iud
(Cs
log (
(C3 A
B
+ C23 C₂))
AJ
achieves (1 + A) multipiiiie accuracy, with time-complexity
C4
B
2
Lood
C3
B AJ
+ cc
TA S Stt+1
log
exp
log
AJA
AJ
A
Note that assuming that G admits minimizer Po and that min G = 0, as is typically the case in over-
parametrized machine learning settings, then by the envelope theorem J = inf (G + 2 I1-12v)= =
|ollir X (A). So in the regime of small X, ignoring the subexponential factors, the time complexity 2
bound achieved by the annealing schedule of Thm. 4.2 scales as exp (cA- -log >) for a constant c. This improves upon the time complexity bound of the classical annealing procedure 3, ~ log(!) (the same as Prop. 4.1), which scales exp(c') c'x-22) -2).
5 Local LSI constant at optimality for learning a single neuron
Devising temperature annealing schemes for global convergence, as illustrated in the previous section, relies on bounds on the local I.SI constant at every iterate 1/t of the (annealed) MFL.D. Such bounds are readily provided by the widely applicable Holley-Stroock perturbation argument, on which for example our Prop. 3.4 is based, but may be overly pessimistic. Indeed in this section, we demonstrate that for MFLD-Bilevel, the LSI constant at convergence can be independent of B, X and d, instead of exponential as global analysis would suggest.
More precisely, we are interested in a", the best local LSI constant of Jx,8 = Jx + B- (:IT), at 17,,8 == arg min JA,8. 8. In fact the proximal Gibbs measure of the optimum is the optimum itself: Tx.s a, so a" is precisely the LSI constant of 7x,88. A bound on a is of interest, especially in the regime of large 3 (low entropic regularization), for two reasons. Firstly, it directly implies a local convergence bound on MIFL.D-Bilevel, as shown in the proposition below. Secondly, characterizing the dependency of a" on 3 may open the way to more efficient temperature annealing strategies; but this is out of the scope of this paper.
Proposition 5.1. Under Assumption 1, suppose nx, satisfies LSI with some constant a'g. For any E > 0, there exists a sublevel set of Jx,8 such that, for any initialization no in this sublevel set, (x,)( inf Jx,s I (Jx,s(m0) - inf Jx,3) e (agg- tc))
"In fact, the annealing procedure of Thm. 4.2 would also yield rate of convergence for any J PPW) + R with J (w, w') uniformly bounded and imf J > O, instead of Jx; but the resulting bound on Ta would have additional factor of (inf J2//2 inside the exponential. See Sec. E.2 for detailed discussion.
8

---

3.0
0.30
Noisy GF on Blieve
MFLD-Biles
GF Biievel
2.-
MFLD- L-iiiinnnn (Conic)
0.08
Noise stopped
MFLD D-Littnng (Canonical)
2.0
1.5
0.01
1.0
0.02
0.5
0.00
0.0
100
1000
2000
3000
4000
5000
10000
20000
30000
40000
50000
iterations
Iterations
(a) A comparison of Wasserstein GF on the bilevel objective with (i.e. MFLD) and without noise.
(b) A comparison of MFLD applied to the Bilevel vs. Lifted formulations.
Figure 1: The regularized training loss (A()) (1.1) of a 2NN with the ReLU activation, learning teacher 2NN with the 4th degree Hermite polynomial as its activation. In both plots, d' 10 and A == 3- = 10-3 The implementation details are provided in Sec. F.4. Plots are averaged over 5 experiments. G is the best value achieved each experiment. In Fig. (1b), "Conic" refers to using the metric (3.2) with gr = 1, = 1, while "Canonical" refers to the choice of q. 2, 9w 0.
For the local LSI analysis, focus a specific setting of (1.1), namely, least equstesquares regression using a 2NN with a normalization constraint on the first-layer weights, and a single-neuron teacher network. See Fig. for an illustrative numerical experiment. Note that Assumption 2, with additional bounded-moment assumptions on şo and p. is a special case of Assumption 1, as shown in Prop. F.4.
Assumption 2. W- 2d is the Euclidean sphere in R++1 and there exist p a covariate distribution over R d+| y E (R a fixed target function, and i R + R C² activation function such that G(v) = B 10 - r(x)){² where De = Sw (( (((,,=))ddu ())).(w)).
Under the above assumption, we show in Prop. F.1 a simplified expression for the bilevel objective and its first variation,
(n) - (v. (K,, + (id) 1)... (b)(w)) - ()(((,))), (K,, xid) 1)) 2)
where K, is the integral operator in L3 of the kernel xx, f / ()tr,,)))(( () )d)(w))))))))))))))))))))))))) and id is the identity operator I.2. Additionally, make the following assumption on the data distribution P and on the response y.
Assumption 3. P is rotationally invariant and the labels come from a single-index model: y ((((,z)) x)) for some fixed E W.
With the above assumptions, we can state the main theorem of this section.
Theorem 5.2. Under Assumptions 2 and 3, there exists a function g: [-1,+1] - TR+ such that J (w)) = ->g((w, v)) for any w E St. Suppose that X I 1 and that there exist constants c, C. > 0 such that for all +11,
c1 S (r) S Ci, g²(r) C2, (0"()(11 S Co.
( s Ca.
Then there exist constants av. Do (dependent only on the c, C.) such that for any 3 > Dodx-1
5, x e e-JJ116ll satisfies -LSL Furthermore, if additionally te Bz-p || || lleA() ao for i E {0, 1, where [1 == J -[lps( invariant), then there exists constant a² dependent only on those constants and on the c,2 C, such that, provided that 3 2 (oly X T),3 satisfies a -LSI.
(independent of w p is rotationally
The proof is based on the observation that 7/,3 RS arg J2 = 6, the Dirac measure at v, for certain regimes of 3 and X, the Wasserstein metrie. Thus we show that [[ is amenable to a Lyapunov type argument inspired from [MS14; LE23], and then transfer its properties to J [T)...]

---

We now verify the assumptions of Thm. 5.2 for a class of smooth, non-negative, and monotone activations which includes some popular practical choices such the Softplus (p(2) = In + +*) and sigmoid ()²) = 1/(1 + e e)) While we only consider smooth activations here for simplicity, certain non-smooth activations such as a leaky version of ReLU can also satisfy the conditions of Thm. 5.2.
Proposition 5.3. Suppose Assumptions 2 and 3 hold, and bi(d + 1) S [=] S B||l=|| < bu(d 1) for constants b1. b2
Let m 2b /bb. Suppose ip and u' are non-negative, insm (p(2) A (z)) > 0 and || < 30 for i S Then, ip satisfies the assumptions of Thm. 5.2 with constants that only depend on by, b2. and sp
Conclusion
In this paper, we investigated how mean-field Langevin dynamics (MFLD), an optimization dynamics over probability measures with global convergence guarantees, can be leveraged to solve convex optimization problems over signed measures of the form (1.1). For large class of objectives G, we highlighted that MFLD with a lifting approach necessarily runs into some issues, whereas the bilevel approach always inherits the guarantees of MFLD, leading to convergence guarantees for GA via annealing. Finally, turning to 2-layer NN learning task which can be stated as instance of (1.1), we showed that the local LSI constant MFLD-Bilevel can scale much more favorably with and than generic analysis would suggest.
Another approach to tackle (1.1) could be to build noisy particle dynamics directly in the space of signed measures, complementing the MFLD updates with, for instance, a birth-death process. A challenge then is build such dynamics that can efficiently discretized. It also an interesting question for future works to find other settings to which MFLD can be extended, beyond signed measures.
References
[AH12]
Kendall Atkinson and Weimin Han. Spherical harmonics and approximations on the unit sphere: an introduction. Vol. 2044. Springer Science & Business Media, 2012. Francis Bach. "Breaking the curse of dimensionality with convex neural networks". In: The Journal of Machine Learning Research 18.1 (2017), pp. 629-681. Francis Bach. "The "n-trick" reloaded: multiple kernel learning". In: (2019). Francis Bach. "The quest for adaptivity". In: (2021).
[Bac17]
[Bac19]
[Bac21]
[BGL14]
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators. Vol. 103. Springer, 2014.
[BRVDM05]
Yoshua Bengio, Nicolas Roux, Pascal Vincent, Olivier Delalleau, and Patrice Marcotte.
"Convex neural networks". In: Advances in neural information processing systems 18 (2005).
[BMZ23]
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. "Learning time-scales in two-layers neural networks". In: arXiv preprint arXiv:2303. 00055 (2023).
[BBP23]
Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. "On learning gaussian multi- index models with gradient flow". In: arXiv preprint arXiv:2310.1 197993 (2023). Nicolas Boumal, An introduction to optimization on smooth manifolds. Cambridge University Press, 2023,
[Bou23]
[BB101]
Dmitri Burago, Yuri Burago, and Sergei Ivanov. A course in metric geometry. Vol. 33. American Mathematical Society Providence, 2001.
[CD13]
René Carmona and François Delarue. "Probabilistic analysis of mean-field games". In: SIAM Journal on Control and Optimization 51.4 (2013), pp. 2705-2734. Olivier Chapelle, Vladimir Vapnik, Olivier Bousquet, and Sayan Mukherjee. "Choos- ing multiple parameters for support vector machines". In: Machine learning 46 (2002), pp. 131-159.
(CVBM02]
[CLRW24]
Fan Chen, Yiqing Lin, Zhenjie Ren, and Songbo Wang. "Uniform-ine time propaga- tion of chaos for kinetic mean field Langevin dynamics". In: Electronic Journal of Probability 29 (2024), pp. 1-43.
10

---

[CRW22]
Fan Chen, Zhenjie Ren, and Songbo Wang. "Uniform-in certtorm-i--time propagation of chaos for mean field langevin dynamics".. In: arXiv preprint arXiv:2212.03050 (2022). Lénaïc Chizat. "Unbalanced optimal transport: Models, numerical methods, applica- tions". PhD thesis. Université Paris sciences et lettres, 2017.
[Chil7]
[Chi22a]
Lénaic Chizat. "Convergence rates of gradient methods for convex optimization in the space of measures". In: Open Journal of Mathematical Optimization (2022),
pp. 1-19.
[Chi22b]
Lénaïc Chizat. "Mean-Field Langevin Dynamics: Exponential Convergence and Annealing". In: Transactions Machine Learning Research (2022).
[Chi22c]
Lénaic Chizat. "Sparse optimization on measures with wver-parameterized gradient descent". In: Mathematical Programming 194.1-2 (2022), pp. 487-532.
[CB18]
Lénaïc Chizat and Francis Bach. "On the global convergence of gradient descent for over over-parament ser-parameterized models using optimal transport", In: Advances neural information processing systems 31 (2018).
[CB20]
Lénaic Chizat and Francis Bach. "Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss". In: Conference on Learning Theory. PMLR. 2020, pp. 1305-1338.
[DG12]
Yohann De Castro and Fabrice Gamboa, "Exact reconstruction using Beurling min- imal extrapolation". In: Journal of Mathematical Analysis and applications 395.1 (2012), pp. 336-354 354.
[DDPS19]
Quentin Denoyelle, Vincent Duval, Gabriel Peyré, and Emmanuel Soubies. "The sliding Frank-Wolf algorithm and its application to super-resolution microscopy". In: Inverse Problems 36.1 (2019), p. 014001.
[FE12]
Christopher Frye and Costas J Efthimiou. "Spherical harmonics in dimensions". In: arXiv preprint arXiv: met.: (2012).
[FW23]
Qiang Fu and Ashia Wilson. "Mean-fiele field Underdamped Langevin Dynamics and its Space-Time Time-Timmeaa Discretization". In: arXiv preprint arXiv:23/2.16366 (2023).
[GCM23]
Sébastien Gadat, Yohann de Castro, and Clément Marteau. "FastPart: Over- Parameterized Stochastic Gradient Descent for Sparse optimisation on Measures". In: (2023).
[GV79]
A. Gray and L. Vanhecke. "Riemannian geometry determined by the volumes of small geodesic balls".. In: Acta Mathematica 142 (1979), pp. 157-198.
[GGGM21]
Charles Guille-Escured Manuela Girotti, Baptiste Goujaud, and Ioannis Mitliagkas. "A study of condition numbers for first-order optimization". In: International Confer- ence Artificial Intelligence and Statistics. PMLR. 2021, pp. 1261-1269.
(HS86)
Richard Holley and Daniel W Stroock. "Logarithmic Sobolev inequalities and stochas- tic Ising models". In: (1986).
[HR$S21]
Kaitong Hu, Zhenjie Ren, David Šiška, and Lukasz Szpruch. 'Mean-field Langevin dynamics and energy landscape of neural networks". In: Annales de I'Inttitt Henri Poincaré (B) Probabilités statistiques. Vol. 57. Institut Henri Poincaré. 2021, pp. 2043-2065.
[KZCE+24]
Yunbum Kook, Matthew S Zhang, Sinho Chewi, Murat A Erdogdu, et al. "Sampling from the Mean-Field Stationary Distribution". In: arXiv preprint arXiv:2402.07355 (2024).
[Lac18]
Daniel Lacker. "Mean field games and interacting particle systems". In: preprint (2018).
[LCBGJ04]
Gert RG Lanckriet, Nello Cristianini, Peter Bartlett, Laurent El Ghaoui, and Michael I Jordan. "Learning the kernel matrix with semidefinite programming". In: Journal of Machine Learning Research 5.Jan (2004), pp. 27-72.
[Lee18]
John M Lee. Introduction Riemannian manifolds. Vol. Springer, 2018.
[LE23]
Mufan Li and Murat A Erdogdu. "Riemannian langevin algorithm for solving semidef- inite programs". In: Bernoulli 29, 4 (2023), pp. 3093-3113.
[LMZ20]
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. "Learning perametrattrie two- layer neural networks beyond NTK". In: Conference on learning theory. PMLR. 2020. pp. 2613-2682.

---

[LMSI8]
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. "Optimal cetropy-trasppor problems and new HellingrrKKattrovich distance between positive measures", In: Inventiones mathematicae 211.3 (2018), pp. 969-1117.
[MB23]
Pierre Marion and Raphaël Berthier. "Leveraging the two- no-timmsccal regime demon- strate convergence of neural networks". In: Advances in Neural Information Process- ing Systems 36 (2023).
[MMN18]
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. "A mean field view of the landscape of two-layer neural networks". In: Proceedings of the National Academy of Sciences 115.33 (2018), E7665 E7671.
[MS14]
Georg Menz and André Schlichting. "Poincaré and logarithmic Sobolev inequalities by decomposition of the energy landscape", In: The Annals of Probability 42.5 (2014), pp. 1809-1884.
[MZ04]
Ilya Molchanov and Sergei Zuyev. "Optimisation in space of measures and optimal design". In: ESAIM: Probability and Statistics 8 (2004), pp. 12-24.
[NS17]
Atsushi Nitanda and Taiji Suzuki. "Stochastic particle gradient descent for infinite ensembles". In: arXiv preprint arXiv: 1ttp:1/22.05555 (2017).
[NWS22]
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. "Convex analysis of the mean field Langevin dynamics". In: International Conference on Artificial Intelligence and Statistics. PMLR. 2022, pp. 9741-9757.
[OV00]
Felix Otto and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality". In: Journal of Functional Analysis 173.2 (2000), pp. 361-400.
[PP21]
Clarice Poon and Gabriel Peyré. "Smooth bilevel programming for sparse regulariza- tion". In: Advances in Neural Information Processing Systems 34 (2021), pp. 1543- 1555.
[PP23]
Clarice Poon and Gabriel Peyré. "Smooth over-param parameterized solvers for non-smooth structured optimization". In: Mathematical Programming (2023), pp. 1-56.
[RBCG08]
Alain Rakotomamoniy, Francis Bach, Stéphane Canu, and Yves Grandvalet. "Sim- pleMKL". In: Journal of Machine Leurning Research 9 (2008), pp. 2491-2521.
[RV22]
Grant Rotskoff and Eric Ganden-Fijndent "Trainability and accuracy of artificial neural networks: An interacting particle system approach". In: Communccations on Pure and Applied Mathematics 75.9 (2022), pp. 1889-1935.
[San15]
Filippo Santambrogio. "Optimal transport for applied mathematicians" In: Birkäuser. NY 55.58-63 (2015), p. 94.
[SS20]
Justin Sirignano and Konstantinos Spiliopoulos. "Mean field analysis of neural net- works: central limit theorem". In: Stochastic Processes and their Applications 130.3 (2020), pp. 185221852.
[SWN23]
Taiji Suzuki, Denny Wu, and Atsushi Nitanda. "Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction". In: Advances in Neural Information Processing Systems 36 (2023).
[SWON23]
Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. "Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond"'. In: Thirty- seventh Conference on Neural Information Processing Systems. 2023.
[Szn91]
Alain-Sol Sznitman. "Topics in propagation of chaos". In: Ecole d' été de probabilités de Saint-Flour XIX-1989 1464 (1991), pp. 165-251.
[TS24]
Shokichi Takakura and Taiji Suzuki. Mean-field Analysis on Two-layer Neural Net- works from Kernel Perspective. 2024.
[Ver18]
Roman Vershynin. High- Diigh--Dimniinal Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2018.
[Vi109]
Cédric Villani. Optimal transport: old and new. Vol. 338. Springer, 2009.
[YWR23]
Yuling Yan, Kaizheng Wang, and Philippe Rigollet. "Learning gaussian mixtures using the Wasserstein- Fish-r--iseaaa gradient flow". In: arXiv preprint orXiv:2001.077 (2023).
12

---

A Details for Sec. (introduction)
A.1
Using Il-1ivv vs. lill as the regularization term
The optimization problems we consider in this paper are of the form (1.1), that is, for ease of reference,
A Gx(z)= G(v) + ||l
min vEM(WW Gx(v);
Note the regularization term à tion problems
2 ||1|2v
This is to be contrasted with the more usual form of optimiza-
min (2(1), PEM(W)
(x(1) = G(v) x Apllzz
which uses ||r| as the regularization. On the level of variational problems, these two classes of problems equivalent, in the sense that
{0} U U
arg min G {0) U U
arg min G3
A20
X20
where "0" refers to the zero measure on W. Indeed, note that by convexity, the argmins are determined by the respective first-order optimality conditions, so that
v(dw) [r(dw)]
U
arg min GA [MW); Yw, G'[v((w) + A A|l/|TTV
= 0, A E R
X
arg arrgmin min G3 { WE M(W); Vw,, (w))
v(dw) [v(dw]] O, N R }
U 100
To see that the set on the first line is contained in the second, let v E arg min Gx, then v satisfies the first-order optimality condition for Gj with = A ||ll.r. Conversely, if v E arg min G3 then either
arg min GA with X Tllyy
In terms of optimization convergence guarantees, when using the reduction by lifting, the problems with I|-lizy vs. with 11-1 regularization give rise to similar analyses, as discussed in Rem. 3.1. However when using the reduction by bilevel optimization, it seems that only the problem with ||-1|ry regularization amenable to a precise analysis. This is perhaps most apparent our derivation of the simplified expression for the bilevel objective, Prop. D.2.
A.2 Detailed comparison with Takakura and Suzuki [TS24]
In this subsection, we show that the learning dynamics considered by [TS24, Sec. 2, 3] an instance of a variant of MFLD applied to the bilevel reduction of (1.1). We do this by recalling their setting (in the case of single-task learning for simplicity) in notations that are compatible with ours.
For set of first-layer weights w; E W = Rd and second-layer weights R (for < < N). and an activation function &
R, the associated 2NN is defined as Note that in our notation of Sec. 3.1, this also writes x fw spw x))(h))
x + EN i aup(w: x).
For E P(R W), the associated infinite-widthh 2NN is v. + SEXw ap(w r)dp(a w).
Consider data distribution p(dz, dy) P(RT Ryy). We may define the Hilbert space of predictors 7 1.2 (IR)), and the "single first-layer neuron predictor" mapping W H by (x)) sp(w" Tr). The predictor associated to an infinite-width 2NN parametrized by u is then Jax ad(w)du(a,uu
13

---

Consider a loss function ((U.8) : Rv x R. - R, inducing risk functional over predictors given by R(h) E ((r)))>pp ((((2))))))) ( y)). We may define the unregularized risk functional
over (infinite-widle winitte-widrh) 2NN weights by
(Ve = = R
ab(w)du(a,w) - (l,, ((ww)d(hul(w))
Accordingly, let the operator /b :M(W) - H. such that tu Sw ((w)dd(()) and
G(v) = -((bv) R (S,, )(ww)du(w))
Then the unregularized risk is C(x) G(hue).
The regularized risk functional considered in (TS24, Sec. 2.1] is 1
F(n) R (Jew ad( + f
/
addu(a, w) + 20 JaW
|w|| du(a, w)
3xW
(A.1)
{ d'du((a, w) + 20 JRWW
G((pp) + 2 JRXW
| da(a,
(More precisely, "F(f,")) (1) in their notation corresponds to F (da) (dd))d()) their "Xa" " corresponds to our À and their "Xw"" corresponds to our 1/o².) Note that, in our notation of Sec. 3.1,
{ [|w]] du(a, w).
F(P) Fx,(())
202 JRXW
The bilevel limiting functional, which is the main object of study of [TS24, Sec. 2.1]. is then defined the mapping P(W)) R such that
g(n)
inf (())())(d))((du)) foW-iR
corresponding precisely
((m) = Jx(1) + 20 / ||m||2
1
[|]l]/22((uu dn)(w)
in our notation of Sec. 3.2 (see the paragraph "Link between the lifted and bilevel reduc- tions"). Interestingly, the convexity of G is almost immediate with our presentation, is expressed as a partial minimization of convex function, whereas the proof the convexity of in [TS24] is quite involved.
They also introduce functional "U" which corresponds precisely to our (x)). and which is an important auxiliary object in their analysis.
The learning dynamics studied from Section 2.3 onwards [TS24] (except for the label noise procedure in Section 5), is precisely MFLD for G(n):
0.7: 3- An + div (n (g'(n.)
1 VJ[md T
= An + div
(A.2)
(and their constant "X" corresponds to our 3-1).
"MFL+ + confining" dynamics. The PDE (A.2) can be interpreted as a variant of MFLD for J/ two ultimately equivalent ways: one is as the MFLD PDE (2.1) with an added "confining" term 2 w, which intuitively encourages the noisy particles to remain close to the origin. Another is as Wasserstein gradient flow for the regularized functional
1 {
Jx,3, = Jx + B H
||l|| dn(ww)
20 J
= J - (1/-112007
where B-1/20g -(0,8-30aa),,
14

---

whereas MFLD for JA is the Wasserstein gradient flow for the functional regularized by entropy only. Jx,8 Jx+33 = J*+3 H est. Unsurprisingly view of this second interpretation, the distribution 0-1/20y plays similar role the analysis convergence of (A.2) [TS24, Lemma 3.5], as played by the uniform measure in our paper: the local LSI property of JA,8,a (resp. Jx, is obtained by applying the Holley-Stroock argument using 3 3/1/2ay (resp. as a reference measure.
Note that the additional confining term 1 ur in (A.2) cannot captured straightforwardly by any additional penalty term on the objective G from (1.1). Indeed, informally, the three terms in (A.1) each have different homogeneity in the variable a. Rather, the confining term should be viewed as corresponding another regularization term added to (I.3), besides the entropy one in B-1
In short, while our work considers MFLD i.e. Wasserstein gradient flow for + B- H as the main "algorithmic primitive", the work of [TS24] considers a MFL+confining dynamics, i.e. Wasserstein gradient flow for F B- H (-)022).
Summary of differences. On a technical level, the learning dynamics considered by [TS24] corresponds a special case of variant of the MFLD-bilevel we consider from Sec. 3.2 onwards. Namely, they focus on instances of the problem (1.1) where G has particular form, corresponding learning with 2NN; and they consider W R and use an additional confining term - w in the MFLD dynamics, while we consider settings where W is compact Riemannian manifold, and no additional confining term is needed.
We also emphasize that, while our work and that of [TS24] cover some similar settings, our focus is quite different. In that work, the key object of interest is the kernel that learned by MFLD in 2NN
setting ((x,2" + <xx" w)p(x" w' )dn(ww in the notation of our second bullet point above). By contrast, our main motivation a general optimization question: how to use MFLD as algorithmic primitive for problems of the form (1.1). In particular we do not assume particular form for G except in Sec. 5, and we pay special attention the bounds on the local LSI constants of J/ along the MFLD trajectory. instead of using the global uniform LSI bound (compare Prop. 3,4 and [TS24, Lemma 3.5)).
B Details for Sec. (background about MFLD)
B.1 The displacement smoothness property
For MFLD (Eq. (2.1)) to be well-posed, we require that F is L-smooth along Wasserstein geodesics for some L < +00. More precisely, for any constant-speed Wasserstein geodesic (H4)eEjo.1) C (2(
with W2(Hg, (1) = 1, / + ((p)) should be L-smooth in the usual sense of continuous optimization. This property ensures that the PDE defining MFLD has a unique solution [Chi22b, App. A], and also helpful to ensure convergence of explicit fime-dicrretization schemes [SWN23]. The following proposition gives a practical sufficient condition.
Proposition B.1. Suppose F (2((9,, g)) + R is twice differentiable in the Wasserstein sense. Let 0 < L < oc. Suppose that F satisfies (P1), iee,
Vu E P2()), Vo E Ω
max SETLR
2 F (H]((s)) L
and Vu, HL P2(Ω), Yo E Ω [V
V'' S L W2(p.(e
where 2 denotes the Riemannian Hessian. Then F is 2T-smooth along Wasserstein geodesics.
The first condition can be stated F'[n] Ω IR having Lipschitz-continuond gradients in the Riemannian sense [Bou23, Coroll. 10.47], whereas the second condition can be interpreted as displacement Lipschitz-continuity of μ + ["(n[[w)) for each uniformly.
Proof. Let a constant-speed Wasserstein geodesic (H4)ee(o.1 C C2() with W2 p1) = 1, and pose [(() F(r)). We want to show that f is 2T-smooth the usual sense of continuous optimization, for which it suffices to show that Vt, If" < 2L.
15

---

sd.t = div( div(Vooilu Vorl4)
By (Vi109, Eq. (13.6)) there exist functions @r : Ω - such that
and
03 - VV,
d V = W3 (20, (00,A111 p1) = for t. So we can compute explicitly:
d F(m))= = =f dye (VF"(u)), Vát))
(()) = dt
(1) a(Oa)u)) (VF')u),, Vol) + f dum d (VFlun), Vod
d
di
di
dμt (v((VF"(m)). Von Vos) + / ((FFliu), Vor
d
d VF'jul,, Vode)
dt
dt
/ duu V2 F' Vo.)
+ +f v² de (VF')ml) voi) + / du (VF'lu), vo.o.)
-fam dt VF' Vor)
Vol [w)[l Now the first line can be bounded using the first condition of (P1): writing (t(w) and w,
for all
If du 22 Fu Vo.)
- - f due ||V 2 (o si) s L. / duc V L.
Moreover, one can show by direct computation that the second line is zero, using that 0.6r / 2 For the third line, we have
|aam ( [ Vo.
(
) dde V|Vlll
v F* (u
sup sup xE xe wEll
di
dup | 1. Finally, let us show that the second condition of (P1) implies bound the last quantity: for all S2 by applying the assumption μ 44 and u' Ps.
since (f de, \V < f
[VF'(,.)(w) - [F'(
L (2(Hss
<
8-t
since (Ne), is constant-speed geodesic with W2(Ho, /1) = 1. So by letting - obtain that v F [m]][]] < L. for all [0, 1). E Ω. Thus have shown [f)([)] s 2T, and F is
2 2--smooth along Wasserstein geodesics.
B.2 Classical sufficient conditions for LSI
For of reference we reproduce here classical sufficient condition for probability measure E P(2) to satisfy LSI.
Lemma B.2 (Holley-Stre Stroookk bounded perturbation argument [HS86]). Let Ho E P(2)) such that is absolutely continuous w.r.t. Poo Suppose that μo satisfies LSI with constant and that -M < log dpg (w) + S M for all w E supp(po). for some E It and M > 0. Then μ satisfies LSI with constant ae M
du
C Details for Sec. 3.1 (reduction by lifting)
C.1 Proof of Prop. 3.1
Here we present slightly stronger version of Prop. 3.1 that uses the p-homogeneous projection operator for arbitrary p > 0, preparation for the next subsection, where we show that one can restrict attention the case as done in the main text.
16

---

Recall that we let s = R x W. For any O. we denote by h" P(2)) + MMW) the signed
p-homogeneous projection operator [LMS18] defined by
Vip E COW,R) / ((())( ()((w)
S/1 sign(r) |r| ((w)(((d an
More concretely, for atomic measures, h" Eド= ((r,,))) En sign(r,) lryl Sws-
Lemma C.I. For b E [1,2] and p > 0, let Vop P())) + IR U {+00} defined by Vo.((1) =
2/6 if (S), and too otherwise. Then (Jn 101% ( (u(r,w))
(a || min s.t. h"
Moreover, i = 1 then the set of minimizers is {P E P(W); hPp = v and Yw, supp(u(-lw)) C R.+ or supp(u(-lw)) C cR-)
and if > 1 there is unique minimizer which is (dr) (ul(dw) where f(w) I 1/2 dip (w). lilllv
Proof. For any E P(2) such that h" = V.
max J sign(r) ((w))drr S / InP du(r, @-W-1-
|| |1ppal2 S (1II dutr; w))
(C/
b 2/b
(f (eddddd(r)) (o)
2/6
(b,(())
where the first inequality follows from the triangle inequality. and the second inequality follows from Jensen's 's inequality since t +++ t" is convex on RT. Note that the first inequality above holds with equality if and only if there exists e W - [-1,,] such that sign(r)((w)) 2 0 for all (r; w) E supp(u), i.e., if the conditional distribution (ddr(w) is either supported on TR or supported on R_ for each w. Conversely, the value ||r/ is attained by letting p(dr, dw) = (dr) [k((dw)
MMlly
where f(w)= = | ITV dioj
du (w). This proves that min (p) = ||ri
For b = 1, t + t = t is linear, so equality always holds in Jensen's 's inequality. So the set of minimizers is all of {μ P(W); h²4 y/ and Vw, (uu,sup((((-ww) E+ or supp(p(-[w)) c. R_
For b > 1, t t is strictly convex, the second inequality above holds with equality if and only there exists a constant e such that |r|² = c for all (r, w) E supp(u). So for μ to be a minimizer, the conditional distribution s(dr(w) must be concentrated on (ct/p, -c/pp for each w. Moreover, for the first inequality above to hold, the conditional distribution at each w must be either supported on R+ or suported on R_, so there exists a function f W + {cl/p, -cl// such that u(dr, dw) 5(( (dr) (dwc where 1 E P(W) denotes the marginal distribution. Since hºp μ = v, then for all fixed w, J2, sign(r) |r| p(dr, dw) = sign( (()([[(())))) (dw) (dw). du du (w)) = de (w) and (dww 1 |p| (dw) since n² a probability du measure so non-negative, and integrating on both sides over S2 shows that e = ||llay. Hence the [r((w) where f(w) = c/p de dlrl (w).
So sign( (w)) sign(
divj
only minimizeris is p(dr, dw) (w) Tilly
Prop. 3.1 follows directly as a special case of the following proposition with p p= 1. Proposition C.2. Let any p > O and b E [1, 2) and let Fo,p P(S2) + R U {+20} as in the lemma above. Consider the optimization problem over probability measures, with X > 0.
(1) where Exb.(u) G(h"m) + 2 MEP(O)
min
(C.I)
Then minp(a) Fx..p- min Gx-
{
Moreover, if > 1 then arg min F
v E arg min G and otherwise 8 1ill//e ato (dr) |pllzy arg F = {y b24 E arg min G and Vw, supp(n) C TR, or supp(n) C R. Furthermore, F is convex.
17

---

Proof. The fact that minp(a) Fx, = miw)(w) GA can be seen directly as follows:
G(h" p1) + A (bp(u)
min PEP(R)
F(u)
min MEP(()
G(h"u) > (orp(H)
min (EM(()
min (εEP()):b>=
min MEM(S)
min G(v) + [P((S))hP= VoL((1)
G(v) + X |11| vEM(()
min
min Gx(v) vEM(S)
where we used the lemma above at the fourth equality. The characterization of arg min F in terms of arg min follows from the characterization of the minimizers of the inner minimization (min [x) in the third line, which is given by the lemma above.
Furthermore, Pxp is convex since G and Va.p are.
C.2 Equivalence of using (cp, ogr, Cq. T/c²) for any > 0 by reparametrizing
Equivalence of Riemannian structures on S2" for (cqr, Cqw 11/c²) for c > 0. Recall that we consider equipping f2" I=" W with Riemannian metric of the form (3.2), reproduced here for ease of reference:
((
8ri 6w]
8r2
or 5r18r2 -r_1 |r/90
|r 0
+|r|"M (fw1, 5wz) w i.e., D(r,w)
ow2
x2
|r Iw
(n,w)
The following proposition shows that, in fact, different choices of gr: 9u and lead to the same geometry, up to reparametrization of the form (a, w) (r", w) (for >> 0). Namely it is equivalent to use the metric with exponents (4r, tw) or with (9- "=), up to adjusting r.
|r|
Proposition C.3. For any gr, Qws denote by qu the metric M(rww)
0
|r| ym
on Ω* g' x W. Then for any gr: qw R and r, a > 0, the map Ta Q 9 uc u .ar)) defined by Ta w) (sign(r) |r|" w) is an isometry.
(9²
Proof. Since Ω² is disjoint manifold: Ω Rt x XWU x W, and since To (R. W) R xW, it suffices to check that the restricted map T (R x W. W.gaa.....) R W. g[ aM ..ar)
is isometry (as well the analogous statement for the restricted map T,. but it will follow analogously).
Indeed, denote by ÿ the metric on Ri W induced by T It given by, for (a, w) T w)) (r.* w), so da dr
5r 8w
8r
8a 8wn
(aaatrr) )
aaa Taa-or2
I(
ów2
twz
=
dw;
d'2
0
so [(a,w
Da
0
(r.w)
2090/a-2
0
al gao
So ğ is precisely 90 a aar) on RC x W. which proves the claim.
Equivalence of the Wasserstein gradient flow of Fx,b,p for (cp, eqrs Cq., L/c²) for any C > O. Proposition C.4. Let T (Li, mu)) (12, 912) isometry between Riemannian manifolds. Let F P())) - R (sufficiently regular) and (pa) Wasserstein gradient flow for F, i.e., 0.44 = (div().VF"(p.)) (where denotes Riemannian gradient (21, g(1)). Then, (ï): = (Tgre) Wasserstein gradient flow for F. P(li)) - R defined by F(ū) F(T," p).
18

---

Proof. First note that 9/2) is given by, for all T(x) E 92 so dy = (T(d)dr where D denotes the differential,
oy J2ly dy = òx 9/10 dz' (yT Mpl (DT(z))-
so T = (DT (pो ((DT()))-))
Also note that F'( = FI j (T (y)). as one can check directly by computing lim F(u + Eū) - F(u) lime 1 F(T "(TT' " b) F(T:" [)]
In particular
DF [z]( = DF; ( (T-
Then for any 4p S2 - R,
/ ædji dt / ((T(x))du((z)
dt an
/ (T(T(=))DT(x)) कां DF (F()()))))))
Sa De(t) TØे DF (F)
That is, 0,ā = div DF (il), i.c. (it)s is a Wasserstein gradient flow for F.
Proposition C.5. Consider the functionals Fx.b.p over P(2)) from Prop. C.2 and the Riemannian metrics aw over Q from Prop. C.3, where Q = x W and Ω = R x W.
Fix gr, Tw R. r, P. > and [1, 2). Let (He)) the Wasserstein gradient flow for Fx.b.p over 9a.....:))) r)), starting from some μo P(S2))
Let a > O and Ta Q2" - Q2 defined by Ta(r, w) (sign(r) |r|0, ,w). Then (īe)) == ((Ta)alt)e coincides with the Wasserstein gradient flow for Fig over (St", 99 starting from poo (Ta)sHg, where
p = P
9r g.
qu
9w a
F=
==
b=b. b.
a
Proof. The proposition follows from an application Prop. C.4 with = To, 2 (S 91g..
Q2 (22' Dat. -a) and F Fx,.p. Indeed the fact that Ta is an isometry from G2 522 was shown in Prop. C.3. It only remains to show that F o T, Fx, And indeed for any Jr E P(2"),
Fx (Ta)) (p) () (AT 1)a)) -G Gh 1)=ū)) + Mopp ((Ta- 1))ūa,
and he (Ta- )Ef = hp/aa iu. since for any yp: ::W- + R,
Inn sdd 1)=ā
/ p(w) sign(r) |r [[T, 1) (dr, dw)
[ / ((w) (fg /// jidF, ((dF,dw) = Sr sdd [heiap]
Jw
and
(T. 1)=a)) = (1r1² ((Ta 2/6 = (1772/// dit(, w)
2/6
This confirms that F o T = F and concludes the proof.
Thus, it is equivalent to consider the lifting reduction with the hyperparameters (p, gr, 9w I') with (cp, Cgr, ohe, r/e²) for any 0.
Remark C.I. The choice p = 9r = you plays a special role, as Wasserstein gradient flows (Ft) on
P(RT x W) for functionals of the form f G(h²j) (1) then correspond to gradient flows (Mx), on M+ (W) for G the Wasserstein-Fisher-ee correspondence is lost however for functionals of the form of Fx in Prop. C.2 with X + 0.
geometry [Chi22c, Prop. 2.1 via in hhp.. uc. This
19

---

Equivalence of MFLD of Fx.b.p for (cp, (cp,cqr, Cgr, Cqw l/c²) for any > 0. Since MFLD for Fx,b.p is the Wasserstein gradient flow of Fx.bp + + then by Prop. C.4, by proceeding similarly as in the equal // itself, up an additive
proof of Prop. C.5, suffices to check that µı + H ((Ta-))=()) constant. And indeed, since T. is invertible, by data processing inequality for differential entropy. we have H ((T, )) H (ū) for all P(S2").
C.3 Proof of Prop. 3.2
Lemma C.6. Let F defined in (C.1) and Q W x R. For any P E P(Q), Fxplll(r, w) = sign(r) G" + X 1mpe
(C.2)
where X' = - V.. (1))
Proof. For any n² P(S2),
1 lim - [G o h²) en²) - oP) = lim c-0
1
G(hh' + G(hPma)
-
= J (") (h' (w)
sign(r) ("""""""""""
BXW
and so (Go heY [n](r, w) sign(r) ( ("))) Moreover
(p)) = (S, spdddplr w)
Vo,plul(r, (r, w) = b (// |r/rbb dutr', .w)) 3-1
2
|r|
(o)p(1))= 111600
Summing the results of these two calculations gives the first variation of Fx,b.P = Gohd+ + d Vo..r
Lemma C.7. Let f R W + R defined by f(, re r²o(w) X for some p, X b E (1, and 5 W -> TR. Assume that 2 Φ is not constant equal to O.
Consider R2 W equipped with the Riemannian metric (3.2). If f has Lipschitz-continuuus Riemannian gradients, then necessarily b and p 9. Jw. or b and p = q./2 = qu. and
V²ō(w) I'p² (ō(w) + +X) Iw for all w.
The proof of Lem. C.7 is technical, so it is deferred to the next section.
Proof of Prop. 3.2. Let prove the first item in the proposition. Suppose by contraposition that Fx, does satisfy (PI). Let any v E M(W) such that V² is not constant equal to O, and consider some P())) to be chosen such that hµr v. Then by the first condition of (P1), f == Fx,o(1] RE
the restriction of FX,o to E7 x W must have sopschitz-coninnuon Riemannian gradients. More explicitly, by (C.2), [(r,, w) = rG'(v](w) + Xa where X - ( So by Lem. C.7. necessarily b = 1, and so X₁ = (x) If = G'[v) satisfies (²Z(w) ryp (δ(w) + +x.) gm
for all w, pick any other μí such that hy' = V and V (1()) * (1()) the existence of such μ' follows from the first step the proof of Lem. C.I. Then by applying the above reasoning to
instead of f, since X * X/ also have by Lem. C.7 that p qr qw- This
R xW
shows that if Fx, satisfies (PI) then (gr, qw, b) (1,,, which was the announced necessary condition.
We now turn to the second item of the proposition. Suppose that gr aw = b = 1. For any 4 E P1)), denote
G' [1))((w)
Aop = sup wEW V (/())//2
20

---

Let us show that if À < Aou then Fx,1 does not satisfy local LSI at Suppose that À < Aays i.c. there exists wp W such that
V (I(1) < (G)))))(woo)
Let us distinguish cases between (wo) 2 or (")) <
First suppose G' (mo) 2 O, that 4 (1)) (uo)))(wo)- By continuity of G' [hp)p let N C W an open neighborhood of wo such that Vw E N, vi ( /2 G'[hp](wx). Then, since Fx, w) = |rl (sign(r)G" (w) + ( by (C.2),
Vr E TR_,, Vw E N, F1 w)=
w) |r (arl(-G (-"(hu]((w) (v) 00
J - drdw
JR_ J 3F (F922,lml0:: drdw
and so
f J 1 drdw +oc.
This contradicts the exponential integrability condition in the definition of local LSI, and FX, does not satisfy local LSI at
Likewise, now suppose that ("[))](wo) < 0, so that v (2) < (G"(hul(un) By continuity of G" [hp],, let C W an open neighborhood of wp such that E N, (1()1/233 (h)) Then
E E R4, AW E N, Ex,lal w) Irl (G'(hu)(w)) (1)(1))112 0
and so Saln --9F%.lllmm)) drdw 2
/ / e-BFx,lbllorddddd
JR JN
/ 1 drdw = +00.
2
R+
As in the previous case, we conclude that Fx,1 does not satisfy local LSI at μ.
C.4 Proof of Lem. C.7 via computing the Hessians under the lifted Riemannian geometry
We start by general lemma. We use D to denote differentials, and for function f : W - R,
we will write D,f 8f(r.w) and D = the
of(r.w))
Or
Lemma C.8. Let (W,g) Riemannian manifold. Let Lt
Rt W and consider
[a(r)
9(r,w)
0
3(r)-1
for smooth positive functions a. 3. R. + R This defines a smooth Riemannian metric g on Dt
Denote by (r,ww V, I, - the Riemannian metric, gradient, Christoffel symbols, resp. Hessian on Mi and by 9w: V.I. 2 the corresponding objects the original space W.
Let f + smooth scalar field. Write for convenience fr(w) fr, w), so that for example (w), Fix a local coordinate chart on W. This induces local coordinate chart on Q2, by adding the index for the variable r. Then the Riemannian Hessian f at (r, w) given in coordinates by
V (w:) = gT. frr w), and note that D, f.
1 V2 000 = a(r) (r)))
1 2 piu V2 pk = at)) 2 (r))
V - 2 (w) -1 (r)))(r)) D.F ()
Proof. We will use uppercase letters for indes ranging over [0, d) and lowercase for [1, d). with the index 0 corresponding to the variable for example ff w) D,r)), We will use
21

---

Einstein summation notation freely. With slight abuse of notation we denote (g"))s - = inverse matrix of the metric (gr))ig 9. and likewise for g'd IIJ. so that for example nth = (r)) We start by using that [Lee18, Example 4.22, Eq. (5.10)
1 for the
M of Ouo d.L
o
V fe; w))l = - g
KI. JM
AM MK
Ōうgı Ow
ØūKJ 0g
and
2"
Ow
Ow.
where (r, w²), and that the analogous formulas hold for fr : W Christoffel symbols of W.
TR. for all and for T the
By direct computations, we find that for all i,j, m E [1,d]. (r)
1o 0 0
rh 1
2 2a(r)
(r)
00 2 a()
B(r) ghj
1 m = 1 (r) 6.Th
F = 0
F = m
2 (r))
So by direct computations, we find that
00 a(r)
(
2 (
V p 7 مoر = a(()) + (+VV,))
f" = B(r)² v2 f.(w) _a(r) D,f 9",
as announced.
Corollary C.9. Let f . R2 x W R defined by [(r, ()) r" "D)(w X for some p b [1, 21, X 2 and Φ : W
Consider Ω equipped with the Riemannian metric (3.2). Then the Riemannian Hessian of f is given in coordinates by
72 000 = 122 - 2r) ((w) + pbx qrr/2)r2 2-))2222--022 2q.-pb pb
10
V for = [pp qre 2)r tVp(W()
V pاd = p-22 29 2 ((w) + Tgw 40 (pro)(w) + pbX' (g (gm))
Proof. Continuing with the same notations as in the proof of the lemma above, have
D,f pr ((w)) + pbX pb-1
D& p(p 1)r 2 ((w) + pb(pb 1) V [r()" rh v2 H(w)" VDrf(ww)' =pr-" =pr"-''('w)
= r
and so
72 000 = a(r)² - 11) (5(w) + pb(pb 1)x76-1)1'-0-- 1)X7700-2) + 1 ar())) (pro- + pbd'r obX'roo-1)
= =a(r)p ((r)(p - 1) + ~ ) p. (5(w) a((r)pbX ((()) 1) - rra((r) yob-2
V p0 v2 yor a())8(r) pr ( Volw)* + (r()))))))))) (Vōlw) 1
=a(r) ((())p
() - Vō''"(()"
2 3 -()) K(w) (()))'(r)) (pr (w()) pb2'r gb.
1
By substituting - and ((r) = rww, i.e. a(r) rv2- and ((r) = -w we
obtain the announced formulas.
0
22

---

Proof of Lem. C.7. Continuing with the same notations as in the proofs of the lemma and of the corollary above, note that f 2 R. W having sipschitz-contimmnn gradients in the Riemannian sense is equivalent to [Bou23, Coroll. 10.47]
v f(w)" 95K8 < oc.
sup sup wERR SET.O ||
Rewriting everything in coordinates, this means that the matrix H(w) R(++1)x(++1) must be bounded, uniformly in w E 07 where (Vgrz) = v denotes the square root the positive-definite matrix 7 (pointwise for each w). Concretely, for all i,j [L.d]
(Vāik f(w) vū.L) KL
Von = -
Vo =0, Vo = -=(()r-1///) V0.5 = - Vii 0-/221,
and
H(w)on = 900 7 for
-Lp(p- q, 12) 2 ( Ipb 9r /r/2)r tb
[ = VJw V 9zi 72 pd0
= (p qw 12) 4-- 2- 12t8 /g Vo(w)*
[(W) VoL v ig V2 أD
PP- 2m
9i V )(w)
+ 2 (prog(u) + pbX'r00) 8k1.
(Note that here the indes do not respect the covariant/contravriiatt convention, i.c., "Vōrk" and " "T kT." " do stand for covariant tensors: really manipulate everything in coordinates explicitly.)
Now, note that the desired condition means that (W)) KL should remain bounded both for r + to0 and - 0. That is, the exponents of in the non-zer terms must all be 0. Thus, since we assume that X' * 0, and that 2 is not constant equal to and so in particular Ø and Vø not constant,
Uniform boundedness of the second term in H (w)xt implies that 1. Indeed X * and the first term (in V2 cannot cancel out both the term in p(w)r ş. and the term in yb 4r if they scale differently with This also implies that either qu gr that tw gr and ( [TquP (Ō(w) + x) gDd for all w.
Uniform boundedness of H (w) implies that gr or p = 9r/2.
Uniform boundedness H (w)j0 implics that = t= or = qw /2. We saw in the first point that gr qw equivalently q. qu p 9r 2 qw 2.
Thus we get that f can have Lipschitz Riemannian gradients only if 1 and p = g. The for all w.
he. or if b and p / 9u/2 and 2 (2)((w Tp² (x(w) x')
D Details for Sec. 3.2 (reduction by bilevel optimization)
D.I Proof of Prop. 3.3
In preparation for the proof of Prop. 3.3, let us first provide a formal proof of the variational representation the squared-TV norm mentioned at the beginning of Sec. 3.2, with a characterization of the set of minimizers. See [Chi17, App. 1] for the rigorous justification of these arguments in the more general context of minimization of convex and positively 1 I-homogeneous integral functionals over the space of signed measures.
23

---

Lemma D.1 ("mp-trick" for the squared TV-norm), We have
1ellbv - -fnn lo(dwo)) inf NEP(W)
(v(dw)) inf (EP(W), faw-RR w SL fn=v
|512 dn-
n(dw)
Moreover the infimum in the third expression is attained at (and only at) q(dw) = [v(dw)]
and the infimum in the fourth expression is attained at (and only at) the same n and f [u(dwr) ||
(((cu)
Proof. The infimum in the third expression is the value of a convex constrained minimization problem, whose Lagrangian is L(m; X) |=[" + x dn - The dual optimality condition implies Ve E supp(n), À = du value ||./|Ty.
án (w)² , so the infinimum attained at 17(ww) = |n(dw]] with optimal TMllyy
The optimality condition for the infimum in the fourth expression follows directly from the one for the third expression and from the constraint fn = V.
Proof of Prop. 3.3. By the lemma above,
2 / |f| dn
inf inf Jx(n) = NEP(W),JWWR NEP(W)
G(fn) +
G(fn) + A
2 / 1 dn
inf vEM(W) PEP(W), f-W+3 sL f/=
inf
inf G(v) + X vEM(W)
/ 151 dn
inf nEP(W), J.W+3 st fn
> 1ul G(v)
inf PEM(W)
inf rEM(W)
G2(z).
Hence the equality of the optimal values. The claimed characterization arg min Jx terms of arg min Gx follows from the characterization of the minimizers of the inner minimization nEPWw, f.-+R Sw 1sf dn in the third line, which is given by the lemma above.
st. fn=
Furthermore, Jx is convex as the partial minimization of (1, - G(v) + $ M which is jointly convex.
D.2 Proof of the explicit form of the two-timescale SDE (3.4)
For ease of reference, we recall here the two-timescale SDE (3.4):
dr - F dwf - V nr Fx,2
E (rd (ri,w))d)
N
V N,
Eズ= 8 (rf (r), wi)dt + /2 dd
By (C.2) with b and 1,
(x w) r (G(hhu](w) + X |r|
V, (x w) (w)) +r
so
and
Fx w) = TVG hpp (w).
Finally, by definition h 1 N
N 8/ (oo
= N Σi م Sug. Hence the second part (3.4).
24

---

D.3 Proof of Prop. 3,4 (Jx satisfies PO, P1 and P2)
Simplifying the expression of the bilevel objective. throughout our analyses of the bilevel problem (3.3).
The following expressions will be useful
Proposition D.2. We have that Jx(n) - G( + //2 [ da where fn is the unique solution of the fixed-poite point equation
Vw E W. fn(w)= = G' Gnifnll(w)).
(D.1)
Furthermore,
A
[bb](w) - 1fp] (w).
(D.2)
Proof.
Consider the optimization problem defining Jx (n(n), for fixed
min G(fn) + X / | dn. /E
This problem is convex since G is, and strongly convex in (,, thanks to the term in X. So there
exists a unique solution which we denote by fn E LWW,, and it is characterized by the first-order optimality condition:
G' n) n Af, 1 0 in M(W).
Now let fn - - which is defined over all of W. Then fn satisfies the fixed-point equation (D.1) by construction. Conversely, for any solution 9n of (D.1), its restriction to supp(n) viewed as an element D. of 2 (WV must in particular satisfy G' G" + Aōn" = in M(W),, and so G. = Jm. and 9,, X G" [.,]]
Gf = -m
Furthermore, by differentiability of G then 17 1 fn continuous (in the total variation sense). So in turn, 77 + (w) the unique solution of (D.1) is continuous for each w (in the total variation sense). So by the envelope theorem, since for any fixed the first variation of n G(fn) + + If dn is w + (f + [f(w) J'Di((w) = G") + If (2))
2
1 2X
= 1n(w)) = - GG' (w),
which is precisely Eq. (D.2).
We remark that the above manipulations rely crucially on the fact that the optimization problem (1.1) over signed measures and not just non-negative measures otherwise we would additionally need constrain f 2 0-, and on the regularization term being ||~|| instead of ||
Preliminary estimates.
Lemma D.3. Under Assumption for any v E M(W), we have sup
[G'(v](o)| S 2LgG(v). wEW
Proof. We follow the proof technique of [GGGM21, Appendix D). Let wo E W and d 1 Ip (= By mean-value theorem there exists E (0,1) such that G(v") - G(v) = G 0/2 =v)]d[/ v), and so
inf S G(L') s ((v) / d'vja(u
Lo V - u1
- u)
2
G(v) ('v) 1
G'(u)(wo) = G(v) 1
G'vy((u))?
To
2To
2T.0
Hence, since G is non-negative by Assumption 1, 1
Vu E W, G'[v[(w)² < G(v) inf G(v)
2Lo
25

---

Lemma D.4. Under Assumption 1, let 17 E P(W) and let fn as in (D.1). Then
sup w
Tup1ss1 - X (2L022(1)
and for each i {1,2},
L full $ X2
B, X (2LoJA(n)++
V
sup
WEW
Moreover, Jx(1) S G(0) for all 1) E P(W).
Proof. For the first inequality, by definition [f = -Afn for all w E so
X² (w,()) = IG' & 2LoG(fan) 2 (a(5mn) A / 15 ( 21) 2
where the first inequality follows from Lem. D.3. For the second part, by Assumption 1, Vv E EMW), sup,, | I, 1ullll
Bi, so
A f = |
- B, + Li 11 = B, + Ls |
1
S B, + L [fr S B₁ + Li x (TnJJ(())
W
by the first part of the lemma.
Finally, the uniform bound (n) follows by taking f = 0 in the infimum defining Jx: Jx(1) serz, G(fn) + + I dn S G(O).
Lemma D.5. Under Assumption 1, Jx P(W) - R is weakly continuous and Vn,n E P(W), (J) Jx(x') BW2(n,n'))
where B = (LL
(L, + )
Proof. For any E P(W), letting Jn as in (D.1), we have [w)](w)) = - |5r] (w) so
Jibi (( V (
[V][]]())] S Asup fyl sup I
W
W
1
( 20) B1) B<
S X 2L.C(O))
by Lem. D.4, uniformly in n E P(W) and w E W. So by Lem. D.8 below, we have (Jx(T) - JA(7') S (W2(n.r)) for all n. E P(W). Moreover W2 metrizes weak convergence, so JA is weakly continuous.
Lemma D.6. Under Assumption 1, let w W and n E P(W). Leth W R and suppose that
Vi E W, Xh(w) + / (nn)n] w' "o""dn(w")))))) (/w"""")">"==>>>>""""""""" = (,""
Then supwew [h(w)] S (1+ 12) 1e 2L.G(O)
Alternatively, suppose that there exists E Te W with || such that
Vu W, Xk(w)+ + / /G" w" od")dn(w" (/"""w"""(w)) - Vw [G" w (') ))) m.
4 Then sup [k(w)] S (1 =+) 4 2(0
B
26

---

Proof. Let g (,,)) + (W) the operator
(w))(w)= / GnMni) w" "dn(n"(" (h(w"))
G is well-defined as bounded operator, since Assumption I implies that |G" (w,,") S Ly- Note that (w,, w") is symmetric in w and w", and that by convexity of G. G" Gf w") 2 0 for all w, w". Consequently, is symmetric positive-semi-defiite oper- ator from (WW) to itself.
On the other hand, let Vi(+) = LG" w/) w')s (w)). By Lem. D.4 we have
1VVill23 S sup w
[G"[fmlll sup lf Lo ā 22 V1. |Vl & sup
wxw
W
Also let V2(+) (s', , V [G" [f w (') Then by Lem. D.4.
[Vell2; S sup V/2 v sup w ",w"
V G* Gf w sup ful + sup Wxw
|G" sup || fill
< い x (L
V2Ln +
L
B
X
(1 2 4 2Lo + LB X = V2.
Le
V2LG(0)
Denote by the restriction of supp(1) viewed as element of L2WW). Then, denoting by id the identity operator on (WW), we may rewrite the assumption as (Aid +g)h = V3 for j = 1 or 2, and so
|M dn = =11b] 12
-1(Aid+9))VV > V\
since G is positive-semi-define and X > 0. Thus for any w E W, we get the point wiitt-wise bound
Ah(w) = VV())) - / (nf(w")"" (ar (we w wh"h""")
A [l((w)] S [V](ww] + ddn(u")) | (f w (")l((((")) [h(w"))
$ VV + [[ 112g |1bll
< VV, + Lo V
Lemma D.7. Under Assumption 1, let n, E P(W) and let fmr fw in (D.1), Then there exist constants H, H' dependent only on x-1,G(0) and Lo, Ln, B1, La such that
sup If -frl HW2(n.n)) W
and
sup wew
IV Ffr S H'We(n.r')
Proof. For each w E W, we denote the first variation of 17 + (n((w) by w' + that this quantity uniformly bounded. 8 By definition, for any W and 17 P(W) and w' W,
(f.(w) oyddwrj Let us show
(w) G'Ifmn)())= =0 (
G² G"(lm) w (''') + (G"(mM((u 85n()) =
on(wr)
("((("")
x ōn(wr) + + G"I w m'dd""") 8n(ww))
(
85((w") - (f w/ (n('').
(D.3)
'The rigorous proof that the first variation "n(dw") follow from the same derivations for the uniform bound, omit it here.
(w) is well-defined for all w, wo E W and n E P(W) would
27

---

So by Lem. D.6 applied to w, w' and 77- Let us now show that
6f(() "y(w")
we indeed have that ffg(w) (y(ww")
bounded by constant uniformly
85
sup sup sup weW TEP(W) w'EW
w Öny(dw")
H
for constant H depending only - Lo, L1, B1, G(O). Indeed, it suffices show that for any H. Now, starting from (D.3) which 17 and differentiating with respect to w/ in the direction s', we get that holds for all w',
s' E Tw W such that I|s'll.r 1, ( s' V au ōn(dwr))
(f.(ww
(fn(w) 8n(w)) w
+ 6" w (nd("")))
V 5fn("")
V
ōn((w")
(8', V [G" [nn]](w, w/ f ())
and so h(w) = (s', V W
dlm(w) By(dww)
satisfies the conditions of Lem. D.6, which proves the claim.
Next let us show that
sup sup sup sup V 8 (s, F WEW SET. (WW) W'EW 1=1_11
S
HH'
87j(dw")
for a constant H' depending only on X-1, Lo, I,, B1, G(0) and T.z. Indeed, starting from (D.3) and differentiating with respect to w' in the direction s', and differentiating with respect to in the direction s, we get
(s, (w))))) ón((w))
+ + V m")n""")) ( s', V.. 5fn(w") = (s, (s. ve {(ss, w. [G" w w')f })..
s' V
ōm((w") / w/
and so
V 688,V (.
S V V of G" [(r('')) + V G" | ('n'''))
8rj(dw")
4 (w="
sup ""EW
G. (n)(W,, w")
sup w'EW
En((w)))
(L 2L
Bi
To (0) + Tn
+
+ La
H =
H
X
by Assumption 1.
Now fix E W. By Lem. D.8 below applied to F(n) fn()),,
we have that
V w ff (w) 5r7"(dw'
(n((w) - (x)( sup sup ("EP(W) w'EW
We(n.n) S HW2(1.))
Likewise, fix any w E W and let = (fr
E T. Then by Lem. D.8 below
V/,((v] E
applied to F(n) (s. v (w)) m
(w) f)(w)] = Vy (s, (w)) H' Lemma D.8. Let W a compact Riemannian manifold and F P(W) R such that Vn E P(W), Vw W, |V B.
Then
Vn. E P(W), [F()) F(/) < BWi(n, BWe(n.7')
28

---

Proof. For any y E W, pose (Ee(z, 1)) the constant speed length- minimizing geodesic in W interpolating between I and y. Also pose (z v) d example if W Rd [z y)
de Eo( v) (s W for any 0. For 0 x) and (z, y) ) =
x for all A
Let the optimal coupling between 17, n in the W sense, and for all 0 E [0, 1], no = (Eo)a7 the pushforward measure of by Eg. Note that for any 0 E [0, 1),
d F(ne)=
S,v F'lnold (Dane)
dl
and that
dl I odno de
d
d /m/w (< (1)()))))))2 d
Vp: W + R,
= Imw
do oZio (a, v))dr(H..))
Il (Ea(z,)) v) (e 8) Eo(r.)) 2o))))
dy(x,,). v).
(The interchange of de and f.wxw on the second line can be justified by the dominated convergence theorem assuming that yp has bounded C norm, which is the case of F' [ns] by assumption.) So by Cauchy-Schwand inequality,
d F(1e) I/mw (Eb(a,)) v). VF (olooo((( dr(x,y)
$ J/ |E (ا
do (1w))
||V (To w)) (1 dr(a, v)
\VF'nll SI 1Z v) ( dr(2,))
sup sup r'P(W)
<B. 1/ dist(a, y)rr(a, v) = (Wi(n,n))
by definition of the geodesic (Eol 1))oe and by definition of the optimal coupling 7. Finally,
[F()) - 1/¹ db F(16) do < sup 8c(0.1]
d8 F(w)
< (n.(n.'').
Proof of the Proposition.
Proof of Prop. 3.4. We first check (PO). The fact that Jx is convex is given by Prop. 3.3. Moreover, let any 3 > let us check that Jx,9 == Jx -- (-Ir) has a minimizer. Indeed, Jx,s weakly continuous as shown in Lem. D.5, and non negaaivee any set of probability measures on W is tight, i.e., any sequence in P(W) has weakly convergent subsequence. So we conclude by the direct method of calculus of variations: let sequence (1n), such that (%,) + p(w) Jx,s and extract a weakly convergent subsequence with limit 7so; then by weak continuity Veec minimizer of Jx,8-
Since W is compact then
We now show that JA satisfies (P1). Recall from (D.2) that (b))(w) - 1/11 (w) with fn -1 G' over W. Let us show the first condition for (P1):
Vn E P2(W), Yo E W,
V J())ssss) SA
SET. 11
for some A < 0o, where 2 denotes the Riemannian Hessian. We have V((l((( = -
2 K[)](1))) = (v T f((w)
29

---

and so, for all E TwM W such that || - 1,
(( s |lo |2 + A
L2 2 2L + B2 x + (AP 22L + Bi )
S 2v2L,G(0)
by Lem. D.4.
Let us now check the second condition for (PI), namely that
Yw E W, Vn,o// E P2(W), [[V][[]] - MMMI, & A W2(n.'))
for some A < co. Indeed,
V [n] JAWII
V Jn
JN V X V +10, fo Vvrll)
< X supsup 1
sup W
sup sup 7/"
sup Sn Stc)
7"' W
w
W
L1 G() H'
B +
< X
G()
HWz(n,
= (1,(1.17
X2
X
by Lem. D.4 and Lem. D.7.
We now turn to the proof of (P2) with the quantitative bound on the local LSI constant. Let 7 E P(W). By the first part of Lem. D.4, we directly have that
(w) [bb((w) = 2 (w) Lo
(2(1).
In particular, by the Holley-Stroock Stroock bounded perturbation argument [HS86], the proximal Gibbs measure v) == e (-- (nl, /Z satisfies LSI with constant 0s; = o, exp (- - (J))))
Finally, we turn to the proof of the bound on the uniform LSI constant along the MFLD trajectory (me)c2.. Given the bound on the local LSI constants, suffices to show that
Vn E P(W), (T) S G(O) and AI 2 0, (x(T:) Jx - (nolT).
The first bound was shown in Lem. D.4. For the second bound, note that J*(n) + - (TilT) decreases with t, since MFLD is precisely the Wasserstein gradient flow for + (2(n) + 8 and H (1) and H (n)[) differ by constant. So, since relative entropy is non-negative,
JA(n) & Jx(n)) + 3 (ndl) s (x() + + ----H (nolT)
for all > as desired.
E Details for Sec. (global convergence by annealing)
The following preliminary lemma allows to control the effect of entropic regularization, using box-kernel smoothing technique similar to [Chi22a].
Lemma E.1. Let W a d-dimensional compact Riemannian manifold and denote by the uniform probability measure over W. Let P(W) - I. and 17" P(W), and suppose that there exist constants A, B >> 0 such that
tn s.t. Wi(n.n") SA J(n) - J(n)) < BW (n.
Denote Jz= I g- H (-I7), for any 3 Then
Be + d
log ( +
log C
min (:W19)99)
Jz(1) < -(TT) + 0ce5mi((1A))
inf
B
where C' = linf infoces1 dd. T (Hw'; dwstw w') S E})]
30

---

Proof. The proof is adapted from [Chi22a]. It is based on constructing an e-smoothed version of nj" Le. a measure 7 which admits density w.r.t. J while being close 1j" an appropriate sense.
Let any < 6 min(1, (1, A}. Given w E W, define the probability measure Ve.w (dw') as the uniform probability measure
the geodesic ball (e(w) == {wo W; dist w/) S c). In other words, doww (w') = 1w 1WWE (.,) Then, let 7 (dw, dww' = vf" (e, (dw') E P(W x W), and
(wB)
let 17 (dw') Swew "e "elww, du') its second marginal.
One can then verify that
drlc (w') = / d'je,w (w )n' (dw) = dr
1(w' E Be(w)) (dw).
wEW
dr
weW
r(B[(w))
Moreover there exists a positive constant such that T(B)(w)) 2 C--led for all $ 1 [GV79, Theorem 3.3]. As consequence,
H (nelT)= =
dr f dr, (w') log dn (w') S sup - log (B:( I dllo((//)) + log C.
weW
Furthermore, by definition of the coupling Yes we have W WT,,,,77 I Wooner 1²) S 6 A.
Therefore, by assumption [(nc) [(1') & BW (n,, 17" Be, and so
(:Wi JSA Jz(n) & Jn(n.) = J (m)) + 3 H (n+lT)
min
I + Be + 8 (d log 1/e) + log C),
and the inequality of the lemma follows by taking the infimum over
E.1 Proof of Prop. 4.1
We state and prove a more precise version of Prop. 4.1 below.
Proposition E.2. Under Assumption I, let > and assume that A
G[O Then X-J
MFLD-Bileve with the temperature schedule Vt. 3r AJ; log (CCV/B AJ3
converges (1 + A)-
multiplicative accuracy in time
2d
(4C//B)
[dLLoG()) AAJ
log (4C//B) )) .log AJS
2JA((n) AJ
(MolT)
TA S a, AJ
log
exp
+
AJ
2 2 C
where C' = max {I, winfwew infosesi e 5 (tw'; distw (w, w') S <t]]]] }
Proof of Prop. E.2. Let (n), the MFLLD-Bievell trajectory with constant inverse temperature parame- ter 3 to be chosen. Denote JA,8 = Jx + 8 (-17). Recall that by Prop. 3.4, Jx.s satisfies a uniformly along the MFLD trajectory with a3 = a exp (=)))))())
by Thm. 2.1, for t,
Ja(m)) < (n)) S inf Jx,s te 29- (Jx,9 - inf Jx,9) S inf Jx +e -2---aa Jx,)(no),
where in the first inequality we used that Jx,8 Jn = 3 (- 2 0. Furthermore, by applying Lem. E. to J = Jx. 1)" arg min J,, A and B / G(O)
4+ 22 + )) the constant from Lem. D.5, we find that
Be + d log
log C
inf JA.S I inf + inf oce51
+
g
B
Taking 3 = 1 for some 2 to chosen, and evaluating at the infimum at e DB we get
d log O'
d
inf Jn I J +
log
3
3B
31

---

where C" = max 11, C}. So in order to guarantee that J2(n) & (1 + A)J3, it suffices to take t such that
J + d+ log C"
d
E -38 ogt (Ja(no) + 3- (100 [T)]T) S (1+
log
B
BB)
8
J2(1o) + + H (NolT)
i.e. t 12
log
Ts5
2
AJ d+Hlog 1
log
d
AR
assuming that is large enough so that the above expression is well-defined. More explicitly, substituting the value of ag and of 3 $s, we have
B
( (HInSC10))
(x + + (NolT)
T. 200
exp
log
AJ7
d+log C²
d log
d
3
3B
sd/B exp
8 1
(oodCo())
J (x + 20
(yol
AB
log
20r
AJ
(1 d- log C' + log s)
Noting that
SAJ
4B AJ
SAJ
log
log log
S
-1
4B
4B
B [1 + d log 0¹ + log 8) S
B
4B
SAJ
d log O² loog
+
AJ
4B
B
4B
AJJ
log C¹ llog
d
8
AJ
4
choose henceforth max
AB (d- log C" + log 11,
1B
}. so that
AJr
B AJ d1+d--log C" + log
AJ
To simplify the final statement, make the assumption that A is small enough that s More explicitly, since we were careful to choose C' 2 1,
4B
(d-l log C" + log AJ;
4B
AJŞ
4B d dd'log log C' + +lg
4B
AJ
AJ < d log C" 4B
1 13
log
AJ
AJ
4B
AJ
AJ
AJ
<
and
log
<
-11
< mimn(1 (1, -
4B
4B
4B
4Be
-1
L1
B1
2l
21
J
Jj
X2
X
4e
2LoLiG(0)
1
2L(0)
= A
<
J
X²
J
X2
Then 8 =
4B
- log C" + +log
48 3 ..3
4d
(
log C" + log
4B
2
4 log C", and
AJ
AJ
AJ
AJJ
3
1 LB)C(0)
Jx(11) + - H (nalT)
T. <
exp
llog
20r
AJ
2d
(4CL/dB)
4dLoG(()
4C1L/d ))
(2Jx(m)) AJ
H (mul))
log
exp
log
log
= Ta.
Or, AJ
AJ
AAJ;
AJ²
2 log C"
Hence the time-complexity upper bound of TA for reaching (1 + m)mulliplicative accuracy.
32

---

Algorithm Annealing of the MFLD.
Require: Functional J P(W) - R. Initialization 7ju So > 0. Schedule K, (Tk)
1: no = 1J0
2: for k = 0. -..... K do
3: Be = 30
4: Run the MFLD with Bg initialized from nă up to Tk.
Oen div(n Viv(nVVV (.(TT + Ani
8x
5: 1io = n..
6: end for
7:
return 1/fx'
E.2 General annealing procedure and its convergence guarantee
The following theorem builds upon and generalizes the idea of (SWON23, Sec. 4.1] to objective functionals I that have a positive optimal value. ensures fast convergence fixed multiplicative accuracy.
Theorem E.3. Let W d-dimensional compact Riemannian manifold, so in particular the uniform measure over W satisfies a, LSI for some ar > O. Let J P(W) R+ convex, suppose that J = min I > 0 and that there exists minimizer 1f", Suppose that there exist constants K1, Cr, A > such that
1. [nlll S (n) for all 17 P(W).
2. (Tr - -((1) & WL (1) all rj E P(W) such that Wi(n. 17") S A.
Fix < < Cr. 3'
min( Let vi the iterates of the annealing procedure of Algorithm I with initialization Bo = d and with the schedule K = [logz(4/(dJ")) and
(r,O2/d
Jso (100 )) (E.1)
To = 2 dlog (2 J& (no)) .0, exp 2 8-1 + log
+
57
2
where C = [infwew info c (ws w') & <e])]]
and the total time-complexity is given by (1330 16 +5 ( )
Then I T S
s T
I
d log
(1,( 8J (5J
ar exp 2x1d (5-1 + log
(8g
(CLC/a
+ +2+
To, (7o)
8J
2
Let us discuss the assumptions of Thm. E.3 and possible generalizations.
Note that the condition 2. of the theorem holds as soon J'(y) W + IR is Cz-Lipschizz for all E P(W), as shown in Lem. D.8, since W, < W₂ S Wo..
The annealing procedure and its convergence guarantee be generalized non-compact manifold W by modifying MFLD to include a confining potential term, as discussed in Sec. A.2.
Condition of the theorem actually holds for any I such that sup,, [T"() w')l s
L 50 and J > with the constant K1
2L
Indeed, one can then show similarly to
Lem. D.3 that
J' S 2L JT) < 2LJ(M) < 2L J J(n)
J
However, when plugging Ki = /2 into the bounds of the theorem, one obtains less favorable dependency of the total time-complexity in J In particular, note that the total complexty guaranteed by the theorem scales exponentially in Kn and polynomially in L/
33

---

The way that the condition 1. of the theorem comes into the proof, is that it allows to guarantee a local LSI constant of J + 5, H at of as = cst e One could similarly formulate an annealing procedure, and state convergence guarantees, tailored to objectives I that satisfy different criteria for LSI, such as the Bakry-Emery curvature- dimension criterion.
The remainder of this subsection is dedicated to proving Thm. E.3.
Proof of Thm. E.3. Fix any 0 < ã v Cv C.min[ Let, for any B > J, = + (-1r).
By condition 1. of the theorem and the Holley- bounded perturbation argument, for any t, k, the proximal Gibbs measure T/) x e [vi), satisfies LSI with the constant
a, exp (-BeK1 (m())
inf a, exp (-Bem1 (-Bem1 a(k).. P20
That is, for any k, Jpp satisfies LSS at n/ for all t 2 O. (To see that a(k) > O, note that for any k, ((H)) I Js, (n S Js. (m6). since H (-17) is non-negative and (v/ ) is a Wasserstein gradient 2 ar exp (- Byri Ja. (nt) > O; but
flow of Jg. and so (<<)) = inf 120 a cxp (-Beni will not make use of this rough bound in the sequel.)
Now let
BR
BR
Tk
log
Ck
20
d
for some a(k) S a<k)) and ox 2 Jo. (mb) min To. to be chosen. Then by Thm. 2.1 applied to Jp. we obtain
To. (m S min Js.. + exp
(Js, (76) - miin Jax)
1B
S min J3,
(Jon
min Toh)) (Jan (16) min JTc)
[ d
d
min J3, kk"
Further, by Lem. F.
log C
Js 5
inf ocegmin(1,A)
Cze+ log
+
+
Be
BR
BR
d
C2 6T
d log
< ()+ +
log
(E.2)
3x
S%
where the last inequality follows by choosing
87 C₁
I min [1, A} since 8 < Ct min(lA))
Then, for all A 2 1 and 2 0,
B&J (vt) Jo. (ot) S Jo. (1 = JrJo, (1T S Tx (00___(1___11 (1 2.8k - (
15-1 H (-IT) N O, that (nt ) is Wasserstein gradient flow for Jap that JR. Js. ( 8 ) (+IT) > 0 since (3x)s is increasing, and that by definition 3x- = 2 3o. So by (E.2),
where used successively that Js. J =BT
C1
(x)) <
(TA=1) k-1
<
+ + 2d
+ 2d + 2 C
65
Cr.
< 22 +0) + 2d
+ 2d + 2 2og C
8J
C1
log C
2d 8-1 + log
+ +2 +
8J
d
34

---

since our choice of S0 = d and K = (lg((1) ensures that 3 0x = 2 Bo S 67
For k = 0 and all t 0, we have more simply (o)(1)) S BoJs. (nk) - Po.Js, (no) dJzt
As result, for all k 2 0 we have
C2 55
log C
1
14 2 0, BaT( S 2d 8 log
+ 2 +
d 2
and so
akk) inf a, exp (-K1 220
2 0r exp (-2xid (6-1 log
CL
log C
I (1o)) (m))) )) = =a()).
1
+ 2 +
+
85
d
2
Moreover, we can choose Ck as
Jo. (16) = Jan (1 S Jon (1 -- S Js. _115-1)s s... S.. by induction,
Jo (1b) min Jap I Ia. (110 =: Tk.
so
Therefore, more explicitly.
3x log T3 2a(k)
Ba log
BE Jop (mn) ar exp 2ndd (8- (5- log
Ct.
log 1
2 J ( )
+ 2
2
55
d
CL
log C 1
= d -oog (2 Jso (no)) af exp 2ndd (8- + log
+ +2+
2 Jo. (1o)
8J
d
since B₁ = 2 3n = dd. d. Note that
K
K
log J&0 (1o)
K
2 log (2t Ja. (no))
Σ
2
122h k 11 log(2)
2
k=0
k=
k=0
(2 1) log J. (1o)
log(2) ((K- - 1)2K + 1)
2
2 log J.z (mo) + -log(2)K2<
1
1
log J&0 (11) +
log
TJ.
5J
5J
1
Ja. (110)
log
8J
8
since K [log2(/(6T"')], hence the announced bound on the total complextyyy Σk=o Th.
Finally, at round K- [log)(4/(67*)), then 3x =2 Bo 2 E
[ 12.3.
]. by (E.2),
d
C1
+ log C
J( s Jsn ( - J + +0)+
log
+
3x 85
3x
log C')
Ct
g
1 + 36
28
28 log
d
5J
which completes the proof.
E.3
Proof of Thm. 4.2
We state a slightly more precise version of Thm. 4.2 below, and prove it as corollary of the more general Thm. E.3. Then Thm. 4.2 follows by choosing 5 = O ((())))))))))) gathering the constants appearing in the bounds, noting that JA (mo) S Jx(m)) + dH (mll () S G(0) + (H))) + d log vol(W).
35

---

Theorem E.4. Under Assumption 1, there exists constants B = poly Ls B,, G(O).>-)) and C dependent only on W such that the following holds. For any ō S J temperature schedule (Be)220 defined by vk S K. E [tx, t4.11.1 = 2 where to = 0 and K = =100g(41(633)) and
B -FLD-Bileeel with the
2Lod ( 8-1 + log
BC
JA. (mTo )):
4-1-tt- = -h-1 dlog (2 J2.0o (170) a0, exp A
+2+ +
8J
2
[BCa/d
achieves (1 + A) multippiicattee accuracy, where A
36 + 28 log
with time- time-complexity
8J2
d log
[J (10) )a7 exp 8J%
(2hg A (xx
(8-1 +-1+log log
[BC/d
J2 (170) )
TA S K++1 I 5J%
+ +2+
6J3,
2
Proof Thmm. 4.2. Let us show that the conditions of Thm. E.3 are satisfied, under Assumption 1, for J Jx- JA is convex and non-negative, and it implied throughout Sec. 4.1 that inf Jx > for the notion of convergence to a fixed multiplicative accuracy to apply (Def. 4.1). The existence of a minimizer 17" is ensured by the weak convexity of JΛ, by similar argument the proof of (PO) in Sec. D.3. We have the condition 1. with λ1 = Lo i.e. [ I Le
X (x(n), by the
first part of Lem. D.4, We also have condition 2. with = oo and Ct = B = (L))
+ / G() + Bi as shown Lem. D.5, since W1 - W S Woo.
Note that annealed MFLle-iileel with the announced temperature annealing schedule (3t)t- precisely corresponds to Algorithm with the schedule (E.1) applied to J = Jxx So the announced time- complexity bound follows directly from the application of Thm. E.3.
F
Details for Sec. 5 (estimates of the local LSI constant)
We begin by presenting the proof of Prop. 5.1, which states that bounding the LSI constant of 7x, leads to a local convergence rate.
Proof of Prop. 5.1. For any 1j E P(W), we denote i((ww) e - where Zy fe 83% dr. First note that for any n. T/ E P(W),
di log dr
(w) + (log Zn - log 2xy)
3 [l] - JbV((())
3 [ - (w)
$ 2 (lf +1frrl) (w) If, fwr
x 2 < 3 2 2Lo 1/22().)))) (i
by Lem. D.4 and Lem. D.7, where H is constant dependent only on A G(O), Lo, L1, B1, Iz.
Now suppose that 77, arg min J2 fx,s satisfies a -LSI. Let s > 0 and no in the d-sublevel set of JA,g, i.e., ?o S == ( inf Jx,s 8)), for some > 0 to be chosen. Denote by (mr) the MFLD trajectory for Jx,g initialized at no. Note that Ss is stable by MFLD since (x,)(1e) decreases with f. So it suffices to show that Jn,8 satisfies (a* ESSL uniformly over Sg.
Choose any 7) E Sá, i.e., such that (n) inf Jx,9 S In particular by Thm. 2.1, holds
g 1 H (nla),)) S JA,,(1) - inf Jx,8 S
Furthermore, since 77,8 satisfies LSI with constant a" then it also satisfies the following Talagrand inequality, as shown in [OV00]:
2 Yny, Wa(n's nx.a) S v H (1
36

---

Then by the inequality noted above, we have
dr
2
log dnxu3
(w) + - HW2(W m. s H H (nlmx,) s H
V
V 38 MVO
for some E I. and so by the Holley- Strolck-ttrockk bounded perturbation argument, f) satisfies LSI with constant a* Mv8 2 a² - for 8 small enough.
F.1
Preliminary estimates for J/ under Assumption 2
Throughout the remainder of this appendix, in the context of Assumption 2, we will use the notations
the Hilbert space H. () with the inner product (F,9) En-ps())g(a
the feature map o W H. given by ((w)(x) x(),,
the symmetric positive-semi-definte operator in H: Kn = f (()))((w)'drd) where denotes adjoint in H.
For any E H. we denote by (h, Vo(w)) (resp. (h, v2 ((w)(w) v² the gradient (resp. Hessian) at w of w + (h, ((w))
The usefuless of these notations is justified by Prop. F.1 below, which gives a simplified expression for Jx and J/
Proposition F.1. Under Assumption 2, letting the Hilbert space H. = (1R and the feature map 6: W + H given by ((w)(x) = (((w,, x)), we have
X (x(n) ? (v, (K, + +Aid)
A
J(b((w) - (d(w), (K, xid))"y)
with K,, f ((w) dnjww,, where denotes adjoint in H. More explicitly, K,, is the integral x') with respect the distribution
operator of the kernel ki (2, x p. i.e.
vh E H (R242), (K,))(z)) = Ex-r (k x' (a(')) in 12
Proof. Under Assumption 2 we have
1 - G(v) -
Ba-o I p((w, x))dv(w) - v(0) 2
1 Hw)dd(w) -m0 1 2
1
so the optimization problem (3.3) defining (x(1). for fixed 17, writes
/ ((w)f(w)dd(()) - n
+ / |sl (w)dn(w))
min SEL((W)
2
This problem is strictly convex thanks to the term in A, and the FOC is Ve, (J (f y, d(w)n(dw)) + >f(w)n(dw) = O. lution of the fixed point equation f(w) - (J (fod y. ((w)) in L,WW). In particular, denoting iu = - (ʃ dhndnn - v). then (n((w) = (()))(n)) d(w)) and, integrating against on.
So the unique minimum fy is so-
I (((())(()dd)(( = d(w) (ww)* hu dn(w)
= -hh, + = = Knnn = (K, + Xid)h, = + h (K, + Aid) 'dy. where a" b= (a, b) and K, = fw ( So the optimal value JA(n) is Jn(n)= / ((w) (ww)
2
+ /// 1sn1 (w)dn(w)
(F.I
H
[ I H
+ A
/ hijo(w) (lw))( h. dr((w)
2 Ahen shn)
(im K, Knnha) 1 (Xhn Whnn + Knhhn)
+ 2
2 ( (N v)
= 2 (y, (K, + A Aid)-'u)
37

---

Further, by applying the envelope theorem on (F.1) (and reasoning similarly to the proof of Prop. D.2 deal with & supp(n), by extending f. E ( into function W - R), we then have
Yw W, [bh](w))
(
- v.
2タ 1fn 1 (w)
A
2
fn(w) ((Aiip((()))
+
A
|
(w)
H
- + A 2,/2m (w) 2 1n
|5 (w) - A 2
X
(my ((w)) )
The characterization of Kn as the integral operator in () of the kernel (2, x²) = Iw (x) (zn)())) follows directly from the definition K, = Iw ((oo)((()""(((" since
Vh E H. Kyhh = /, <(w) (ō(w), h) dn(w),
(K,,)(z))= = Imv ((w)(a) Ex'no ((w) h((')l dn(w)
= =Eg'nup U/x (w)( hix') dn(w)
Exgro (x,
We have the following Wasserstein Lipschitz-conninnityy properties for the bilevel objective func- tional Jx-
Proposition F22. Under Assumption 2, suppose furthermore that sup,. i E 1,, 2, Then for any W Sd and any n, 1/ P(W), it holds
[v'd(w)lx S B; < co for
B1
(Jx Jx(7') -
11yl Wi(n.r))
2B B1
and [\b](()) - J (())((w)
|1 Wi(n.n')
X²
4B,
and [V]][](w)) V>bl'((w)) <
1uyl W1(n.n))
X
4Bo B (Bo + Bi)
and ]/]](w)) v2 [WW
|lll W, ('.
X²
Proof. By Prop. F.I,
[[n[(w) X (d(w), (K, + -id) where K, I ó(w" "(w"))(w") dn(w")
[[]]](()>> -A (δ(w), (K,, + Aid) 1::00'')))))) '/)xx (Vó(w), (K, + Aiddid 4)
(F22
||v s A (]((())) , [v(w)) m. [|(K, + Aid) -
s B 1 (K,, + 2) 1)
$ - BnB 11v/12
since K, is positive-semi definiee by definition and [[ + -) omax mmax((K,, + -iiid
[amin (K, + (ii)) < 1 So by applying Lem. D.8, this shows the first inequality.
Moreover, the first variation of K,, at any 17 is w' + ((w)))(w")" thus by the formula a(x
-X- -0xxxx- for the derivative of matrix inverse.
8
ōn)(w") (Kn, + (id) -KK, + Aid) ((w))((w') (Kn + Aid)
38

---

and so, letting for concision M = (K, + Aid)=
[W w')
= - )x)
(K, + Aid) tooloo . (Kn + xid)) "))
- (o(w), (d()), (lw')(((w))) My)x
= A (Ó(w), Mylu (6(w), ((ww),My).
As result,
V JKbl(,, w') = x (
((Vo(w), My) (o(w), Mbb( + (b(w), My) (VV(w), Mā("
and, using again that || = [| + A) 1--11.pp
'l.pp S _\-1
V Iw, w Il. S ABA 1 || 2B 1 Then applying Lem. D.8 shows the second inequality.
= 2x 22 11
Furthermore, for a fixed w E W. continuing from the expression of Vo (w, w') derived above,
W V v. Mill(w, w')
A (Vo(w'), My) ((Vo(w), My) (6(w), Mx(ww))) + (o(w), My) (Vó(w), Mo(w))) A (d(w'), My)x ((Vo(w), My) (d(w), MV MVp(w))) (x + (o(w), My) (Vo(w), MV (V)(w)))))
SO V V. (w. w'll < 42 B B | æ- and the third inequality follows by applying Lem. D.8 to 1o > ~ (s, VJA (w)) for To arbitrary.
Finally, by differentiating the expression of W V [wlll, w') once more with respect to we get that, for any fixed w E W.
W. v Thlllw w')
A (Ve(w'), My)# : (000)x((222)) ((Vo)(w), My) (6(1), Mox"))) (Vo(), My) (
+ A (Vó(w'), My) ((Vo(w), My) (Vo(w), MM( ( My) (V)
+ A (6(w'), My)x (VV ( (Vo(()) My)
+ A (6(w'), My) (VV(ww) (V((w), MV (w
(6(w), My) (V'6((),
hence Var HM, w < x- | B B (4B2 + 4Bf), and the fourth inequality follows by applying Lem. D.8 to (s. V (w) s),. for Tw W arbitrary.
The following lemma provides explicit upper estimates of the regularity constants Bo, B1, B2 of á appearing in Prop. F.2. in terms of the activation function & and the data distribution p. Lemma F.3. Under Assumption 2, recall that & : W H (R(+1) is defined by ((w)(x)
(((w, 2)), and that yp RR + C2, There exists universal constant > O such that
sup
(1 1 < b
wE&
sup
|V , - 14/1 Na(p),
w EES!
| 1 - (I1""l2440) + (1001l2))))) Na(p)
sup w wEBd
where
Na(p)= = sup (Eep (u,) 1/4
and Af R R RR [[fllee0)) == sup (Eoup x)) ()) weB
11lla51
Note that if p is rotationally invariant, then Ezepl I1( x\) is independent of w, and there exists universal constant such that Na(p) S 1/2 (Exup ||x||* )
39

---

Proof. For the first inequality, we have by definition
sup (1())) sup V Exxp (=)(((x==))) ||,)
For the second inequality, define the orthogonal projector I. = Ta+1 wow Rd+1 + T T fw}- for any w Sd. Then [Vo(w)] (a) = ('((,, ()), I.l so by Cauchy-Schward inequalities,
[Vo(w)lx
sup (11 1
sup sET 11+1
Ezup [f(z) (s. (=(w)()))
Ezup (ss Vo(x)()))) 1/2
Exup [ (1)((11x))) (I x²)² 1/2
sup
sup
sET |1
$ Bxrr (1e0((ws,)))) (()x)) 1/4
(Benpp ((u,))()) 1/4
sup ||u||==1
since ||+|l.. = ||
For the third inequality, the Riemannian Hessian of 6(w) = sp(w, ()) +)): : - R given by
V (w)) (x) = Vur((w, = V ((w, x) [))]] 1l
(o((w, (w, x)) 1m
so similarly by Cauchy-Schwarr inequalities,
Expp [|s" (su)((1,==)) (s, [mi)) + Ezup [Ipe x))) (w, 2)3))//2
[V s sup
TT a =1
s (Ba-p ||0" (1o"((1, (1x)((1,)))())))
(Ramp (Tw (Twos,)²²) 1/4
sup TET 8d
I1 =1
+ EEzp (10(( 1/4
(Eemp [ww, ()) 2/2
Finally, suppose that p is rotationally invariant, and let show that (p) cdd11/2 (Eu ||~||| 1/4 for some universal constant Indeed, for x: ~ p. have that x: and T x/ || are independent and that ~ т. Therefore.
(D( sup Expp ||0|| (u, x/ ||=||) |1
(Bemp [kallt)) (BEwr (u, (1,)))))
sup 1
and sup @1 Eppr (u, xT)² S (d41)= for some universal constant c, which direct conse- quence of the fact that (u, is Lub-Gaussian with sub-Gaussian norm ē/ s/Vd+I + for some uni- versal constant ē [Ver18, Theorem 3.4.6), along with the moment bound for sub-Gaussian random variables [Ver18, Proposition 2.5.2]
Finally, we check rigorously in the following proposition that Assumption 2 with proper additional regularity assumptions on and p. is a special case of Assumption 1.
Proposition E.4. Consider W = Sd and G M(W) + R defined as in Assumption 2. Suppose fur- thermore that Na(p), ||<p| ( -llve'1 [1<p"11140))) oo, where Na(p) and [1-)) are defined in Lem, F33 Then, G and W satisfy Assumption 1.
Proof. The fact that 8d satisfies -I.SI with ar classical and can be found in [BGL114 Sec. 5.7].
40

---

By definition, G(v) wllw ((w)dr(w) - ulli G is non-negative and admits second variations: for any v M(W) and w, w' ES"
G'v](w)= ((w), I o(("')d((w)) - -8) 7
G"Iv[w, w (d(w), o(w'))
and G" [w] w') (Vo(w), d(w'))
(w, w') - (V²o(w), )(w')) V V ~ ("[v(w w') = (Vo(w), VΦ(w ))x-
Consequently, denoting C = supwes
V 12 for i E {0, 1, 2), which are all finite by Lem. F.3,
("Iv](ur w') I C Lo
V V' w') 1l. S CoC1 =: L1
V [") w/ I1, S CoC2 =: L2
V v V w G" [w] w']l $ Cf Iz.
Now for each E (0,1,2). V(D, w. w'), V (uc w') L, 4(u, w) [VGV S Le 1|V vllvv.
Indeed, the right-hand side can be shown by applying the mean-value theorem to g(0) = (s, V'G"lv + 0(v' - v)](w)) over E [0, 1] for each s E (J) Thus, to show the exis- tence of B; < ao such that v(N, w, w²), |/ G' I L, ||lllr + B,, it suffices to check that there exists vp such that llollTv and sup.. V G (rollll <c. Note that for any and w.
V'G'Mu((w) V ((w), /,, o(w')du(w') -
thus V'G'[o[(w) - (V'v(w),,)) u)
and sup [|V'G s C. 1lvllx < co.
Hence the existence of the B, < co is verified. This finishes the verification of Assumption 1.
F.2 Proof of Thm. 5.2
In the single-index setting of Assumption 3, it is intuitive that 8, is minimizer of Jx, for any À 2 and that 77,,9 and 6, are close in certain regimes of /3 and X. For this reason, we will first investigate the properties of [6,] as a proxy of ['],,,,), proving LSI, in Sec. F.2.1. This step uses Lyapunov approach inspired by [MS14; LE23]. We will then prove that these properties carry from J[ over to J [Vx.e), in Sec. F.2.2, thanks quantitative bound on W2(1x,9, 5x) proved in Sec. F.2.3.
show that it amenable to a refined analysis for
Lemma F.5. Under Assumptions 2 and 3, we have
X ( + [1(())1))) (o(v),
Vw E Sd, [(b.[[(w)
2
X
( + [1e112200)) (ar
2
>g((r, w))
for some g: [-1,+1] +
Proof. By Prop. F.1, since y Φ(v),
J((6,] = X
(o(w), (Ks, + - 'o(())
2
41

---

o(v) with eigenvalue [l ((x, v))² |1) it is also an eigenvector of (Ks, Aid) - with eigenvalue (Ils(())] + A)- whence the expression of J'16,] follows.
Since ő(v) is an eigenvector of Ks. = Sw
Moreover, by rotational invariance of p. (xxpp((( ((),, w)) depends only on (e, w), for all w E $4. In other words, there exists 9 such that [' = - >))((,
F.2.1 Lyapunov function analysis for bounding the LSI constant of 8, x e-
Observe that by the assumption o' ci > of Thm. 5.2, [[ -Ag( (v. 1)) has a unique global minimum at v. Moreover, our other assumptions on y will imply that the Riemannian Hessian at optimum V ((oo((()) is positive definite. This motivates us to follow the strategy of [LE23,
Thm. 3.4] for proving LSI for 8, e 9.,,6, Let us first outline the strategy and recall some useful classical notions.
The generator of the Langevin diffusion with invariant measure exp(-8f)7/2#
A ((V..V)
(F.3)
Define U {w distw(w, v) S r for some v E 3 with > 0 to be chosen later. We say : + [1, oo) is Lyapunov function if CW
W - - + bly, for constants 0 > 0 and b 2 0. When proving functional inequalities for a Gibbs measure exp( exp(-8f)) f a typical choice of Lyapunov function is W = exp($(ff - min f)/2), for which the Lyapunov condition writes
f
8² V
S - 0 blu.
(F.4)
2
4
Further, we say probability measure P((d) satisfies local Poincaré inequality U with constant Ky if
[. S
1 |VSIi di, for all smooth f such that [. fdv O. KU
Notice that U has a convex boundary, thus we can use the Bakry-Emery criterion as adapted to manifolds with convex boundaries by [LE23, Proposition B.11] to prove a local Poincaré inequality on U. Specifically, it suffices to have inf Amin (w)) > O.
In summary, Lyapumov condition of the form (F.4), along with control the eigenspectrum of fl)(w,, implies an LSI for e - -f. We record this fact the theorem below, working out the proper dependence on problem parameters for future use.
Theorem F.6. Let E gd. 0 < X < 1 and f gd R of the form f(w) (w,((w, 03) for some increasing function g [-1,,]] + R. Suppose there exist constants Do. D1. Dz, D3, D3 > 0, and E (0, π. T/2) such that if 3 DodA then
3
Vw E Sd,
A
|4
< DiAd
(Lga)
2
4
Vw E E U,
2 Af
|vS12 S - D2
(Lu)
4
Vw E ESd, Amin (V2)(w)) 2 -D3d
(Cgu)
Vw E U, Amia (V2f(w)) 2 D4A
(Cu)
where U = (w E Sd; dist v) S rr. Then (provided that 3 > DodX-1) the probability measure v = exp(-3f)=12 satisfies a-LSI for a constant dependent only on the D, and on r.
Furthermore, the condition B is replaced by B 2 D and if (Lga) is replaced by
1 -
3 4 | < D//133/44
Vw E gd.
(LSaa
2
then (provided that B 2 141 v satisfies a' -LSS for a constant a' dependent only on D/= Dí, D2, D3, DA and on
42

---

Proof. By the Lyapunov criterion for Poincaré inequality [BGL14, Thm. 4.6.2], if the generator L given by (F.3) satisfies the Lyapunov condition <W
W I - b1y for some 0 > 0, b 2 0, U CS and
: gd + R, and if v satisfies local Poincaré inequality on U with constant ky, then v satisfies a Poincaré inequality on S with constant x 2 It
Let us apply this to W = exp( min f)/2). By (Lga) and (Ly),, the Lyapunov condition holds with D2 x and b D1A3(d 1) + D2 x². Moreover, since U has convex boundary (the geodesic in s between any two points in U' remains in U for r < π. 722), by [LE23, Propostion B.11]
satisfies a local Poincaré inequality on with constant
KD 2 Ric, + Admin (V2 f(w)) 2 d -
1 BXD1
where Ricy denotes the Ricci curvature of gd. As a result, v satisfies Poincaré inequality with constant
D2 D. A. 2 CBA,
(F.5)
d- d-1+BXD;
for some constant C depending only on the D., where we used that B 2 DodA
Moreover, by [LE23, Proposition 9. 9.17], if v E P(Sd) satisfies the Poincaré inequality with constant x, and 2 f + Ric, > -K for some K 0 on Sd, then for 3 2 satisfies the LSI with constant TTKK By the assumptions of the theorem, this indeed holds with K = Da Consequently, satisfies LSI with constant
C/(11 D2), which finishes the proof of the first part of the theorem.
The second part, with (Lga) instead of (Lgs), follows by a similar reasoning, except that "D, "",""shin should be replaced by "D' 33/344 in the calculation of (F.5). This still leads to a bound of the form x 2 CMBA provided that B 2 Did X- and the rest of the proof follows without change.
We now verify that J3 satisfies the conditions of Thm. F.6.
Proposition F.7. Under the assumptions of Thm. 5.2, fo = J1,6, satisfies the conditions of Thm. F.6
with Do, ...y D4, r dependent only on Cı, C1 C2. Cz.
Proof. The Riemannian gradient and Hessian of fo = J21ool = -Ag(tu, +) are given by So(w) -Ag' (ww, u))
and 2 fo(w) = ,n (ww, v)) - g' ((w,, o)) (w, 0) Ia+1) Thw
where II, Ia+1 www R Sd {w} for any considering the smooth extension of fo to Rd+1 - TR defined by x + - (n, x)) and using that $d is nub-Riemannian manifold of dd+11 [Bou23, Chap. 5]. In particular since v Iw 1- (w, v)² and Tr II = d,
St. This can be shown by
| fo(w)]] ((w, ((w,,)))) (w, u)
and Afo(w)= = Ir 2 fo(w) = - (o"((ze, u) (1 - (w. (w,u)²) g ((w u)) (20,0) d).
Pose U = fw gd distg (w, v) rr] for some > 0 to be chosen.
Let us verify (T.ge). We have for all w E gd
1
B full = 4 (29 v\) + BDg (ww, u))?) (1 - (w, u)2) + X o o((ww, uw) u)d. 4
Au
2
(R6)
The second term is bounded by 2 Cjd. We can ensure that the first term is non-positive by appropri- ately restricting as follows:
inf [2g" + {² 2 = 2(inf + B2(inf g [-1.]]
=0 20
2Cg 3
2C BAe N
X-1
ct
43

---

Let us verify (Lv). We can upper- opper-bound the first term in (F.6) by negative quantity by restricting B further: by a similar calculation as just above,
1C₂
inf 2g" + L 2 [-1.]]
3
((o 20 - 2g" 3,
2 ²>(o') over [-1,
Then for all w E sd \ U, we have r S distw(w,,)) v) arccos( (w, 0)) - /1 (w. u)?, and so
1 3 V <
A (2SN)(( ((u; (1 (w, 1) マタ ((ue, uu)
2 ~ f0
A o' (wr,, u))
BA
2 0 (o(((w,))) (10, e)² + 2 2( w)d
2
$ 4 o ((ww, u))
A
2B
+ 2/W, uyd
x²
S 4 ((w,, u)).
3Aerr²
S
9x212 2
ct
x
4n72
provided that / 2
2d Acir"
To verify (Cza), simply note that, since
Vo M Ilapp [I = - 1w,
Yw,
V fo(w)) Mlop S >gg((w,, v)) (w. 2) + XC1
+ C)A, g²(s)(1 - x2) + +C1 +Cs++
sup sE(-1,11
and therefore, inf wese Amin (o))) 2 suup || [o((w))))))) -C3 + CI)
Finally, let us verify (Co), Indeed, for any w U. Amiu (V2 fu(w))
min || (u,,)=
~Ag"((w u)) (u, w)² + Ag v((ww, uus v)
2 - ((w, u))1
|| (u, u3² + λe (w, e)
max
= --Alg" 1g"((1, (()((,,))) (1 (w, )) + dei (w. v):
where the bound of the second term follows from (ur, v) 2 0, which can be ensured by taking S 3
Since E U - (w, v) 2 cos(r) 2 - r2, it follows that
Amiin (V2)((w) 2 -- ->
I0"(s) - (2)) + Xei cos
sup cos rSa
2 -
sup Lcos rSe51
V1- s2
dcc cos
C1
= A- sinr + 01 cos r) 2 x
2
for a certain choice of r small enough, dependent only on e₁ and C3.
E.2.2 Lyapunov function analysis for bounding the LSI constant of 17x.8
To prove Thm. 5.2, it only remains to show that the conditions of Thm. F.6 are satisfied for J sl,. instead of [x
Lemma F.8. Under the setting of Assumptions 2 and 3, qa,3 is rotationally invariant except for the direction v, or formally Rv = = Ryy,, "D,ss for orthonormal matrices R. where Rgn denotes the pushforward measure. Moreover, there exists 9n, [-1,]] + R. such that for all w E Sd. JA >gn v)).
Proof. The lemma follows directly from the fact that p is rotationally invariant and that y 6(v). 0
44

---

Lemma E9. Under Assumption 2, suppose furthermore that sup... ('o(w))1 S B, < 00 for
i E 10.1, Then we have, for any n, n' E PP(W).
3 [JA -
Yw E ES"
2 AJAh
|VJ
BB B1 (Br + BB))
|1 + B 2
llvlix Wi(n.')
X2
A3
4B BBB22 + BF)
and |Amin (V22)) - Amin (VJXX[
11 Wi (1(,,')
X2
Proof. By Prop. F.2,
4 Bi (B, + B
1| [nn](w)) - v2 [bll(w)
l1 (1,
op
X²
and Amiin ( (V2 [())(w)) Amin JV S || (x) -2 v² J/ by Weyl's inequality. This shows the second inequality of the lemma.
For the first inequality, we have AV][n](() Ir v (W) and so
I - -
d
d 14 (Bo + B
[bb](w) - v KlW/wo>
ll Wi(,.))
2
2
X2
Moreover, we showed in (F.2) resp. in Prop. F.2 that
B1
4 B?
|| -
llvl
and |V JJV7'1 -
llvll 2 Wi Wi
A
X2
so
B |
B
B
B1 Ilall
4B B |v
I .2
W (1.77)
4
X
2B Bi /
1 Wi(n. (',")),,
which implies the first incquality of the lemma by triangle inequality.
We can now proceed to the proof of Thm. 5.2, thanks to a bound on W2 6,) under Assumption proved in the next section.
Proof Thm. 5.2. For concision, in this proof, we will use the notations 0(.), (()))()))())) S to hide constants dependent only on || (20)) (16)) [1=") Ezup || /d8, c1, C1 C2, C3 and C4.
We established in Prop. F.7 that fo == J [6.] satisfies the conditions (Iga) (Tvr) (Cgt) (Cv) of Thm. F66 with some constants D., r r= 0(1) (in fact only dependent on C1 C1 C2. C3) provided that
3 2 DodA-1 Thus, the first part of the theorem concerning the LSI of 0g x e 90] follows from
Thm. F.6. To prove the second part of the theorem, it suffices to show that ft == J [TX satisfies the
conditions (Iga)) (Lv) (Cg4) (C₂) of Thm. F.6 with some constants Dú D1, D2, Dz. Da, r = O())
By Lem. F.3, there exist constants B, = (() such that sup... Moreover, by Lem. F.12 below, provided that / 2 ((dA), one has
| I B, for (0.
W2(1x.3.5.) & / dv -1 Log(Od-11-1) 12-1 =: W.
Now by the conditions (Iga) (Lvr) (Cza) (Cu) for f = Jo and D; = @(1) (by Prop. E.7), from
Lem. E.9 along with the triangle inequality we have
1
3
Yu gd
Af
|| Ad +dx (dx + 31 1---8-31 ")
2
Vi Edd
AU,
1
AS*
3
||V/>>
2 - D233 + Es (d) -2 + SW
2
Vu E gd.
Amin (v2) 2 - - - ~
Vw E U, A > D&A Es 2
for some constants Eg, E, 0(1). So.
45

---

(Lgg)) for f can be ensured with Dj O(1) provided that ( 3 -3-
1)8x W ( Since already assume that /3 (((A)) this is equivalent to 3A 3W = OAd33//4 ie., /1/41 ---=-=-d=" w O(1).
(Lv) can be ensured with D2 = Ds if 3 is such that
< 3 i.e.,
2
(8- + 1)X- I D2 26, Since we already assume that 3 2 Q(dx)) this is equívalent to
A -W 2 F2 for a certain F2 = (1).
(Cg=) can be ensured with D3 (1)) provided that - W O(1), ie., - ' 0(1).
(Crr) can be ensured with D4 = 24 if E W < 2 X, i... > <
D1
=: F4 = (1))
2E
In summary, since we assume À < 1, we have - 3 < \ and x - v x Hence we will
choose 3 such that 3 +W = O(1) and X >5 v F2 for a certain F3 = ( and this will ensure all four conditions with constants D/, D2, Ds. D. (1). For choices of 3 such that Fg. Now substituting the definition of W, this
3 2 d²A-4, it suffices to have 31/4d-1x -- <
sufficient condition rewrites
S F2 = 2 ...
dx 1-1dd
log
-
log
B dd
S F2.
dd
Since Ye, x > 0, E log x = log x S xτ, then for any & > 0 it suffices to choose / such that
8 1/2 1 9d
EEF + 2
2 P d-e-
dd
Choosing e.g. s = z. we get that a sufficient condition is B 2 [(poly(A-,,d)) ,d)).
Hence we may apply the second part of Thm. F. f"= [((x,s) with constants D2, D2. Ds, Dr O(1), provided that B 2 [(poly(A d)). This concludes the proof of the second part of the theorem.
F.2.3 Bound on Wi (1x,s 5.)
The following lemma shows form of weak coercivity of Jx- Lemma F.10. Under Assumptions 2 and 3, iffurthermore there exist 0₁, C1, C3 C4 0 such that
Vr E +1,, a (') S C1, (1 S Cs. (1 s Cs then there exists a constant ag dependent only on c1, C1, C3, C4 such that vn, Jx(n) - (2(6,) 2 Xag W (n, 8z).
Proof. Since JA is convex,
Jx(n) - J2(6,) 2 / J100,J((1 - 6a)= - / g(vv, w))d(( (w)
f (a(1) - s((1, w))l dn(w).
= Jg
Now let U, = {w E Sd, distgg (w, v) S r for some r 0 to be chosen. We will compute the integral separately on U, and on gd \ Up-
For the part Je. we proceed by second-order Taylor expansion. Namely, for any w E U, \ {u). let I. u such that w cos(O) sin(O)e for some 0 0 S r, since distga(w, v) arccos( ((w, u)) 0. Then (((v, w)) g (coss 0), and
d
d 9cos d²
sin(
dg29(cos 0) = sin( (cos 0) cos(()g) (cos 0)
d3
g 0) = sin( " (cos 0) + 3 sin (f) cos cos(0)g" (0)) (cos 8) + sin(0)g (cos 0).
46

---

Notice that by our assumptions on 9. it is smooth enough at so that sin(f)g (cos 0) + 0 and sin( (cos g f) 0 as 8 + O. Further.
de sup e
dg3 ( 0) S C4 3C3s + C1 =: 6Ms.g
Consequently, by a univariate Taylor expansion with remainder Langrange form around 0, for all 0 < 0 S r, provided that we choose S 2
(1)
we have
g(cos 0) = 9(1) + 0+ 70 g (()1))02
+ 0 (g (3)) for some [0,r) r]
1
1
$ ((1) 29 (1)6
(g cos)
sup [0.r]
6
1
$ (() 29 )'(1)²² +Ms.9 +Ms.00=gg(1 g1)
(0'(1) MMsg8 0²
$ 9(1) 7 o' (1)02
(F.7)
In other words,
Vw E Ur, 2(1) - (1, 2 7 (1) distga u)
[0(1) S U. (a(1) - ((v, (()))) 2 7 (1) /
and so, For the part Someer since g is increasing on [-1, since g' 2 ei > 0, we have
distga (w, u)) dn(w).
J..
Sonv. (o(1) s( dnwww 2 (0(1) - g(cco(r)))] - (U.)]
1 1 (1)re) [1 nUnU)
where the second inequality follows from the Taylor expansion (F.7) above applied to = r.
Thus we showed
(x()) - -2(66) 2 > {(E
[1 [1- (--(()]]]+ + 20 4 / distga (we v) 2)d((()))
(1)
n2 n[U,)] f
clistaa (wotga(w,rr ())dn(())))) }
4
Ur
On the other hand, since distg (v, w) = arccos((v: (w: w)).
W 6w) = / J distga (v. w) dn(ww) + / distga dn(ww
< x - n(Un] + / dister (v, ddy(w.
Hence
(n( (2(6,) 2 \1)
2
sup
min
W (n, 6n)
05r (1(1)
(1(
g'(l)
min
/T2 1
W3 6,)
2.M
2 A C min
C1
/x
W2((1,0) = =30,W3(916b)) 6.).
4
2Ms
Notice that ag only depends C12 C1, C3, Cz.
47

---

We will use the following fact about the surface area of small hyperspherical cap around a pole for bounding Wi (7x,8, 6,). It essentially shows that, for W S", the constant called C in the statement of Lem. E.1 scales with dimension as 2- 2 slVV.
Lemma E11. Fix d 2 2 and E 8d and denote by T the uniform measure on S For any e >> 0, let S. = { w E gd distga (w, v) S <e). There exist universal constants C-, C++ > such that
AA < 4
cI (c/2 <(())) r(()) SC C+ V/Va.
Proof. For w T, the distribution of (w, u) admits a probability density function h(z) (1 22)d/2-1/2/ /Z, where
-2 Z = / (1- 22/d/2 de =
d
= r VT
2' I (d+11
Note that by Gautschi's inequality Vs E (0,
Vz Ve > 0, x . [(++1)
(x + 1)' applied to ((a+s)
s 2 and x d-1, we have 1-1 2
d+1
2
r(2)
2
S0
2π
2T
d+1
S 2
1/
d-1
By definition, since distga (w, v) = arccos((w)) (w,,)), =(S.) = -Seunio) Mz))dz. One can verify
YO<ES & 4. V1- 2 s <cos(c)= S V1
6²
4
So for all 0 < <<< T
7(S.) [ Scooto)
h(z)d== [" h(z)dz
VI-2
1"'
(1 (1-1) (1-1)0/2-1 dt
= z-
(1 - 22d/2211 dz =z-1
VI-23
2vi
1
1
1" / dt
2--
(1 (1 (1-1)d/2-10
-tt= =
2V1 2 J
2/1
1-
= = 1
2
21210/22
< z-
1
d SC4elvVa
2
dv - (T/4)²
for some universal constant C4. In the other direction,
h(2)dz= 7-1 [
(1 - 2222/2212222 d== 2-1 8² (1 tja/2-1 dt
r(S.) 2
VI-2/4
/1-2//
1-22/1 2vT
(1 (1-t) (1-t)d/2-1dss = 2--1 2/11 :d/2-1dt
> z 1
20
2 1-
= -- 2 2d
2/////22 = 7- 7 2 2c(c/2)//Va.
for some universal constants By repeating the same argument with 11 & replaced by v we get get T(S.) 2 c((e11.99)) /Va (e/2g) for some universal constants c', C_.
3.9
The following lemma combines the weak coercivity and weak Lipschitz-continuity of JA by I r-convergence type argument, to show explicit bound on Wi W1 6e). It quantifies the intuitive fact that 77,33 converges weakly to 6, when 8-1 or X t00.
48

---

Lemma F.12. Under Assumptions 2 and 3, if sup
\'o(w))læ s B, < cor {0,1}, and if
3 2 add (BoBi l1ull then
1 1B- (?+log log (BoBI [1vl(x)) log (3 (3-1d))
W2(na, de)
V
ag
where Ö a universal constant and 0g is the constant from Lem, F.10.
Proof. Since Nx.3 arg min Jx.s and JA,8 = J + 8- 1 H (-1r), then for any 1f" P(W),
Jx(1),) S Jx(1x.)) + - (n,,llT) = Jx.a S (A, = Jx(1)) - (1"1T).
Further, we showed in Lem. F.10 that V1,, Jx(n) Jx(6.) 2 Aag W2 so
Aag (7,,,, 6,) & JA(Nx,) (x,) S JA(7") (x + 8 (1)
remains to upper-bound the right-hand side, which we do by choosing as 17" a box-kernel smoothed version of 6 (this part the proof is essentially instantantiation of Lem. E.1). Specifically, let 17" be the uniform measure over the spherical cap S. = w S* distga (w, v) S o for a to be chosen. We showed in Prop. F:2 that
J2(P)- - s BoB 1 (1)
A
where sup |/(((w))) S Bi and by definition
Wi(1f = distga v) dry (w) vol(Sa) / distga (w, v) vol(w) S a.
Moreover by Lem. F11, provided that 0 < a 4.
H (T-TT)= f
dg log dr
vol(S") dnj = log vol(S.)
log ((x)) S log C d doog 0
2
for some universal constant C, and let assume w.l.o.g. that C > 1, so that log C s dlog C. Thus
Bo |Ml σ - -3- - + 3 dloglo2.. 2C.
JA - 6- (n <
Therefore, taking the infimum over 0 <os
BoB1 [1y H
Aag W2 inf 0c05
3-- + 3 2C
B-'d
3 d - -d-3-'d -l'd
+ + 2C
Bo || 2
B-'d 'd (1 log(CC) log( '' + log Bo [1y| (+)
where on the second line we used that the unconstrained infimum of the right-hand side over > 0 is which indeed less than z by assumption. This shows the bound Mi
attained
By B,
W2(T,,,,8.) <
B- (1+ log(2C) Log 'd)) + log BB |1 l
Xag
and the bound announced in the proposition follows by gathering some universal constants into C. G
49

---

F.3 Proof of Prop. 5.3 (examples of activations satisfying the assumptions)
Before presenting the proof, we recall few concepts from the theory of spherical harmonics, and refer to [AH12; FE12] for more details. Let be the uniform probability measure $ The spherical harmonics in dimension
form an orthonormal basis of (T). We denote them by (Yez jk.j. where 2 and & j _ N (d, k). where N (d, 0) = 1 and N [d, k) = 2k+d (k+d-2) for k 2 1 (for k = 0 we have Y01 1). Consequently, any Φ E (,7 can be written as
&
N(d.k)
(6. (z))
k=
j=1
Let Pk.d be the Legendre polynomial (a.k.a. Gegenbauer polynomial) of degree is in dimension d 1, normalized such that Pa.a(1) = Thanks to Rodrigues' formula [AH12, Theorem 2.23]. we can express Legendre polynomials as,
(-1) (1)) (1- (t. 2*T + d/2)
dt (4)
( (1 (1-
We now go over some useful properties of spherical harmonics and Legendre polynomials.
(Addition Formula) We have the following formula which relates Legendre polynomials spherical harmonics [AH12, Theorem 2.9].
(() (1)(1)) = Pa, Vue, v @d.
THecke-Funk Formula) Suppose & (T( is given by o(.) = =(((e,))) for some w e $4. Then [AH12, Theorem 2.22].
I((d+ 1)/2)
(w) [ ((PPL0 (3)la221/2d
(6, Yo)) (P(T)
T(d/2)
(Orthogonality of Legendre Polynomials) Using the addition formula and orthonormality of spherical harmonics, for every k. k 2 O we have,
(Pw,a(w }), Pw ((: = 8K PR.a((w, u))
(d, k)
(Derivative of Legendre Polynomials) For every & 2 j. we have the following identity for derivatives of Legendre polynomials [AH12, Equation (2.89)).
(1) Pj.
where p()) denotes the jth derivative of Pk.ds and
Cjk.d k(k - 1)... (A j + 1)(k 1)(A: d)... (k d d )-2)
(F8)
d(d (d 2) ... (d 2j
Notice that for j > we have p()
We use the tools introduced above to prove the following lemma. Lemma F.13. Suppose p is a spherically symmetric probability measure on Rd+1, Define
[-1,1] + R via g((w, v/) ( x) ((v,)))) for w, v Sd. Then, for every j 1,
(d 1) 3) (d 2j 1) / ||=||3 (00 2) (0() ((0000((0000))0)0)))) (1,,)))dp())
1
((w,
where v0u) denotes the jth derivative of sp
50

---

Proof. We being by introducing the notation (r((w, (((((w, x)). Doing so allows us to only consider functions on g" by conditioning on the norm of input ||=||| Notice that
((w, u)) E
10 [|=||( (w.
- ( s0 ||]||| (u,
| ]
Billll (i u)}] (F19)
||>|||
where
go (ww, o)) = (o) (p,((w, +), ((v, (1) By the Hecke-Fume formula,
(pr((1e,))
= Y = VN(d,k) Yiz(ue),
Othrr
where
I((d+1)/2) +
T(d/2)/T S" ((rt) - 83) (3)
ākrr
Then, by the expansion of (pr((u, 1) in the basis of spherical harmonics,
N(d.k)
pr((w, ())
Σ
OKr (x)()) E (hddk)o
(F.10)
VN(d,k)
j=1
Am0
Via the formula for inner products of Legendre polynomials, we obtain
(r(,w, o)) Σ
Pa,((u, M2m Σ Pedl,Pa
km
As result,
(1w, u)) = Σ P -
(F.11)
k=0
k=
where Gj is given by (F.8). On the other hand, we can directly obtain from (F110),
o ((w, Pld (l)((w,))) =
x)).
Therefore,
0 kir j.k.d
2 Nd
((w, (lu,))
Pr Ph-aad+2((u v)).
N(d+ + 2j, -j)
Moreover, it is straightforward to verify that Cj k) N(d + kj,k-1]
for & 2 j. Therefore,
((w, (P (M))))))))))))) = (d+ + ()(d + +3) ... (d+2) 2j- Σ
ak Ciiakk.kakk-ska (ww,))
k=j
(d + 1d + 3) ... (d + 2ÿ () v)),
where the last identity follows from (F.11), We can now use the fact that x = ripli), and plug the above back into (F.9) obtain
-Bitall / (d (d+ 1) l) ... (d 2j 1) (1 (w, ())) (v, d))
|
((a v)) = Baal
[ (w,,)))
.....dd
||x]| 21
= = / ()dd)((+3) (d dd+22 2j 1) ((x, ((x,)))))
=
which concludes the proof.
51

---

We are now ready to state the proof of Prop. 5.3
Proof of Prop. 5.3. Recall g(w, ut) quently,
( 2(1+11(())1
Let (((w,))) = (o(w), (())) u- Conse-
qq'
g" qa" q
g" 34'9" + 9q""
(A [1
(A+ (((()))))))
(A+ [
We proceed to bound each term separately. By non-nigativity of o. for any r > 0, we have a(w,, v))
R (p((w, a) ((v,==)))
2 (w 2) )o ((), 2) S 5n ((, s r
2 ( inf (0))) PP [[] (1r - S r) (1(0 2) sr]
alsr
2 ( inf p(z))² (1 - (w x)² > - P( x)² > rl)
#|rr
(1- 2
2 ( inf p(z))²
=|gr
(d + 1)r²
where the last inequality follows from Markov inequality along with the fact that B|ax² ] = B[|l=11]
d+1 la+1 for spherically symmetric distributions. Thus, by choosing r = m 2bg//nn we
by
have (()) 2 1 [inf[s]sm ((=)). Furthermore, by the Cauchy-Schward inequality, q((w; o)) S
(l )²} = || Next, we move on to bounding g'. Let ~ T be a uniform random vector on Sd. Then, for any > by Lem. F.13,
q'(uu,, v))
d + E
||>||
xp' (uu,, x) x)) ((v, ]
d+1 E
||a|||
E
[sp' ((w, (P) ((0, x)) [|=][]]]
(inf) (2 (2)))
|| 2 1 { Kw, x) -
EE
Ke, s
Ku
||>|||
d+1
|
[x]|
[inf) St 45' (2)))
2
x2
()
E.
||
1 - TP
(u. xt)²
|1:1|
-P
(v,
d+1
||=||/
||
(inf) (z))² E
72 (d+ 17 )
2 |201|
||==||
1 -
d ++1
Consequently, by choosing r = m 2b,rb2 we obtain o' 2 by (inl sism o(()))² Moreover, by the Cauchy-Schwartz inequality, g' S bz 11 As a result,
b, (inf | ((2)) lmlsm (2))
b2 11 (1 S <
() + (2)
()+ |1 (2)))
Furthermore, by Lem. F.13 and the Cauchy-Schwand inequality,
(d +
b3(d + 1)2
la"1 < d
1
10"1 S (d+ 3)d 5) 11 1240)
Hence,
bg 111 (o)(
14911216))
62 ||pp"1 (2)) ||م1 - b3 I1 |1 "
(A (1
(A + || 22
and
362 11 10'1 11051[2
+ 12 | 1 11 1 (((p)
lg"| <
(A + |1 (2)))))))
which completes the proof.
52

---

F44 Implementation details for Fig. 1
We consider the problem (1.1) where W gd and G is defined as in Assumption 2, where 10, A 10 and
y Rd+1 - is given by a teacher 2NN with 5 neurons defined as follows, The first-layer weights are orthonormal, drawn from the Haar measure, and the second layer weights are drawn i.i.d. from (0, 1.8/a). Its activation is Preacher (z) 24-6, which is the normalized 4th degree Hermite polynomial.
/24
p the empirical distribution of a (covariate) dataset (zt)ign of n 100 training samples,
sampled i..dd from / ((+)
(dd La
i)): with the last coordinate representing bias.
The activation function op of the student 2NN ÿu is the ReLU, p(2)) max(0, 2).
We performed 5 different runs, each corresponding to different teacher network (t) and training dataset (p), and tested all the algorithms considered at each run. So the objective functional G is different for each run, which is why the values shown on the y-axis are offset by G the best value achieved by any the algorithms considered for each run.
For the algorithms using the bilevel formulation, we computed the values and the Wasserstein gradients of JA explicitly by the formulas from Prop. F.1 and (F.2) (the matrix K2 Aid in L3 ~ R" is inverted explicitly).
For the algorithms using MFLD, we used B = 10 We ran the Euler-Maruyama discretization of the noisy particle gradient flow SDE described in Sec. 2 (with an inexact simulation of the Brownian increments described below), using N 1000 particles corresponding to the width of the student 2NN -, and a step size 10 2 for (la) and 10 for (1b). For Wasserstein GF without noise, we used the same discretization but with 8 = 0.
Concerning the initialization of the particles (r", (=)) corresponding to the second resp. first- layer weights of the student network -, the wh are drawn i.i.d. uniformly on Sd, and for the algorithms using the lifting formulation, the ri are drawn i.i.d. from (O, 1).
Note that our simulations of Brownian motion are not exact. To implement MFLD on Sd we simply took gradient steps in R++ with added Gaussian noise, and projected the weights back to the sphere.
thub. com/ com/mousavih/ The code to reproduce this experiment can be found https: h //gi 2024-MFLD WWFFD-bblevee bilevel.
53

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The claims made in the abstract and introduction match the paper's contribu- tions. In particular the three bullet points concluding the introduction summarize the paper's contributions section by section.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The scope and limitations of each optimization dynamics considered is clearly discussed within each section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
54

---

Answer: [Yes]
Justification: For each theorem or proposition or corollary or lemma, be it in the main text in the appendix, the assumptions clearly stated, and all proofs provided.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The contributions this work are theoretical. A numerical illustration given in Fig. 1, for which the implementation details allowing to reproduce the experiment are provided in Sec. F.4.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
55

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
We provide Sec. F.4 full details for the small numerical experiment of Fig. 1, which are sufficient reproduce the experiment. The code used will also be made public at a later date.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nip) cc/ public/guides Seaii///guides//oeeSubbiie
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips.co..c wttppblcc//acacccccccccccccc Meiiddes/Coeee
ssionPolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: The setup of the numerical experiment of Fig. is very simple. Moreover full details are provided in Sec. F.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The purpose the small experiment from Fig. is to compare the qualitative behavior of several algorithms: advantage of MFLD over Wasserstein GF in Fig. la. and advantage of MFLL-Bilevel over MFLDDLiiftig Fig. 1b. This qualitative behavior is clear-cut across the runs, all of which are shown.
Guidelines:
The answer NA means that the paper does not include experiments.
56

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The very small scale of the numerical experiment of Fig. 1 means that any standard laptop or desktop computer can be used reproduce it in, with runtime few minutes.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't ' make it into the paper).
paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https /ttee://nnuuuueee httppppblicc/hhhhccccceeeeeeeeee
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: We have reviewed the Code of Ethics and have not found any deviation of our work from
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
57

---

Justification:
The contributions of this work are theoretical.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
The contributions of this work are theoretical.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
The paper does not existing assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
58

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
59

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
60


Reviews:
[{"review_id": "7oypFjzPR2", "reviewer": "Reviewer_hSwQ", "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.", "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?", "limitations": "Limitations are addressed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and encouraging comments. We address each point in \"Weaknesses\" and \"Questions\" separately.\n- **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.\n- **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:\n    - 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. \n    - Theoretically, the manifold $[-R, R] \\times \\mathcal{W}$ possesses boundaries, which adds significant technical difficulties to the analysis (in particular for time discretization).\n\n    We note that a related modification of the lifting approach is discussed in our answer to Reviewer akPi (4th bullet point).\n- **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.)\n- **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}$. \n\n[GV22] Khashayar Gatmiry and Santosh S Vempala. “Convergence of the Riemannian Langevin algorithm” (2022).\n\n[LE23] Mufan Li and Murat A Erdogdu. “Riemannian Langevin algorithm for solving semidefinite programs” (2023).\n\n[SWN23] Taiji Suzuki, Denny Wu, and Atsushi Nitanda. “Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction” (2023).\n\n[TS24] Shokichi Takakura and Taiji Suzuki. “Mean-field Analysis on Two-layer Neural Networks from a Kernel Perspective” (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "1q0TFfQIpJ", "reviewer": "Reviewer_akPi", "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.\nAs an example, the learning of one neuron is presented at the end of the paper.\n\nIn 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.\nThey observe that the bilevel reduction behaves more favorably compared to the lifting reduction, allowing to obtain stronger and faster convergence.", "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).\n- What is $J$ in line 44?\n- 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.\n- Is there any intuitive insight that could explain, why the bilevel approach works much better than the lifting strategy?\n- 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\".)", "limitations": "I do not see any unadressed limitation.\n\nMoreover, I like how the lifting approach is discussed, despite turning out not to be the way to go.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.\n- 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.\nI 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.\nYet, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed evaluation, and address their questions and comments below.\n- Thank you for the feedback on this part of the introduction, it will be adapted. \n- \"$J$\" on line 44 should be \"$F$\".\n- 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.\n- 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.\n    \n    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.\n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "7uKxk0Ghyo", "reviewer": "Reviewer_3p9w", "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.", "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)$.\n* 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?\n* What is the main issue when trying to extend the results of Section 5 to multi-index models or general target functions?", "limitations": "Discussed in corresponding sections.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n* 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.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful assessment, and address their questions and comments below.\n- **On compactness**: $\\mathcal{W}$ is assumed to be without boundaries. This assumption is missing on line 20. Thank you for pointing this out.\n\n    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.\n- **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.\n- **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. \n- **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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "9kw1BmYTUq", "reviewer": "Reviewer_Hr7F", "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.", "questions": "See *weakness*", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "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?\nAlso, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive evaluation, and address their questions and comments below.\n- 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}}$.\n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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: **Summary** This paper extends mean-field langevin dynamics MFLD to optimization problems over the space of signed measures by two procedures: a lifting approach and a bilevel approach.  The bilevel approach is shown to be superior in various forms.

The reviewers were unanimous in their appreciation of the paper.  We have summarized some of the comments below.

**Strengths**

The bilevel approach is exciting, and the problem setup appears timely and broadly interesting.  Quoting from the reviewers:

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

The authors have considered two approaches (bilevel and lifted), and they have shown how bilevel outperforms the lifted approach.  This is especially significant given the lifted formulation is more standard

>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 well, written, despite being technical:
>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.

Finally, the paper contains additional mathematical novelties which may have extra impact.
> The analysis utilizes novel techniques (namely local LSI bounds) to go beyond the standard results in MFLD.

**Recommendation**

The consensus is that this is a strong acceptance.
Decision: Accept (spotlight)