Convergence Analysis of the Wasserstein Proximal Algorithm beyond Convexity

Thanks for your careful comments! Below is our reply to your questions and we revised the paper accordingly.

• (On Gronwall inequality) The vanilla Gronwall inequality requires the function to be everywhere differentiable in the interval interior. For our problem, we only have almost everywhere differentiability for the Hopf-Lax semigroup, and the proof of Lemma A.1 (Lemma B.6 in revised version) requires additional monotonicity for the inequality (highlighted in red) to hold. So we added a Remark B.7 to clarify that our inequality extends the classical Gronwall lemma, which requires everywhere differentiability, to the case with almost everywhere differentiability and monotonicity.

• (On existence of minimizer) In the revision, we clarified that $F$ needs to be lower semicontinuous (lsc) in the weak topology and put the assumption of weakly lsc in our Assumption 1 as our regularity assumption. We also added a Lemma B.3 to summarize some classical conditions that guarantees weak lower semicontinuity.

• (On PL definition) We added the references [BV2023, KMV2016] to the PL definition.

• (On subdifferential representation) Thanks for pointing out this technical issue in Lemma A.3! First, we clarify that if $\Theta$ is a compact subset of $\mathbb{R}^d$, then eqn.(18) in Lemma A.3 remains true $\rho_\xi$-a.s. for probability measures supported on compact domains (cf. added Lemma B.5). So the difficult arises primarily when $\Theta = \mathbb{R}^d$ (or unbounded domain). Fortunately, our proof of Theorem 3.2 only requires either one of the followings: (a) Lemma A.3, i.e., the representation of strong subdifferential of the form $(T_{\rho_{\xi}}^{\rho}-I) / \xi$ in terms of Wasserstein gradient $\nabla \dfrac{\delta F}{\delta \rho}(\rho_\xi)$ holds along the proximal trajectory, OR (b) $\nabla \dfrac{\delta F}{\delta \rho}(\rho_\xi)$ is the minimal $L_2(\rho_\xi)$-norm element in the strong subdifferential. Existence (but not uniqueness) of strong subdifferential is guaranteed by [Lemma 10.1.2, (Ambrosio et al., 2005)] when $\rho_\xi\in D(|\partial F|)$. Thus, we impose a slightly stronger assumption in the revision (Assumption 2 (a) or (b) on proximal trajectory). We remark that this new assumption is not restrictive for our main application MFLD. In particular, we show that the MFLD on the standard two-layer architecture with some regularity on the perceptron for either squared (regression) or logistic (classification) loss is a sufficient condition for Assumption 2(b) in Corollary 3.5 (revised version).

• For geodesically strongly convex case (Theorem 3.8 in revised version), we modified our proof to avoid invoking Assumption 2 completely, by heavily utilizing the fact that $(T_{\rho_{\xi}}^{\rho}-I)/\xi$ is a strong subdifferential (Lemma 10.1.2, (Ambrosio et al., 2005)), and provided an iteration rate for the unique minimizer $\rho^*$ under $\mathcal{W}_2$ (in addition to the functional value convergence rate). Similarly as [SKL2020], the $\mathcal{W}_2$ convergence rate on the minimizer is also linear.

• For proof of Theorem 3.8, it suffices to prove the case $\xi=1/L$ because for $\xi<1/L$, a functional that is $(-L)$ geodesically convex is also $-(1/\xi)$ geodesically convex. We apologize for this confusion and revised our proof accordingly.

• We apologize notation typos and corrected them as suggested.

• For small comments, we have resolved all of them.

(Ambrosio et al., 2005) Ambrosio, Gigli, and Savar´e. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.

(Chizat, 2022) Lenaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing, 2022.

So far I was just roughly reading over the revised version (I will have a closer look in the next days). The authors corrected a large number of errors and confusions. However, I already found again a few confusions:

• Regarding the Gronwall inequality: You are basically employing the integral form of the Gronwall lemma. It is contained in most standard text books, Wikipedia, or in the Encyclopedia of Mathematics (Thm 3 of https://encyclopediaofmath.org/wiki/Gronwall_lemma). It just requires $u$ to be measurable and non-negative (no differentiability assumption). I disagree with the claim that Lemma B.6 extends the existing forms of the Gronwall inequality in any direction. If you add continuity to the assumptions (you apply it to the Moreau-Yoshida approximation which is continuous), then the statement is e.g. given in Appendix B, inequality j of the text book "Partial Differential Equations" of L. C. Evans.

• I am a bit confused by the newly introduced Assumption 2. It says that you assume that either Assumption 2 (a) or Assumption 2 (b) is fulfilled. But Assumption 2 (a) obviously implies Assumption 2 (b). So wouldn't it be sufficient just to state Assumption 2 (b)?

We want to express our sincere appreciation for your careful review!

(On Assumption 2)

Thank you so much for pointing out this confusion! We only kept Assumption 2b in our new paper version as you suggested.

(On Gronwall inequality)

Thank you for the references to the integral form of Gronwall inequality! However, from our understanding, it is still different from our lemma to the best of our knowledge. We compare our lemma with the integral form of Gronwall inequality in the following,

(1) Following the similar proof strategy of our Lemma B.6, we can show

If $u(t)$ is a increasing positive function, and $\frac{d}{dt}u(t)\geq u(t)c(t)$ almost everywhere in $(0,\xi)$, where $c(t)>0$ for $(0,\xi)$, we have,

we can see that the inequality direction differs from Gronwall inequality under this setting. We believe this already shows some differences.

(2) If $u(t)$ is a decreasing positive function, and $\frac{d}{dt}u(t)\leq u(t)c(t)$ almost everywhere in $(0,\xi)$, where $c(t)<0$ for $(0,\xi)$, we seemly have some "Gronwall-type" inequality in this case,

And in this case, by monotonicity, we actually have

which is very similar to the integral form of Gronwall lemma.

However, without monotonicity, we will not have the equation above. On the other hand, in the equation above, $c(t)<0$ which is different from (J) equation in Appendix B of Evans' book and Theorem 2 in Encyclopedia, where $c(t)>0$ is assumed.

(3) In conclusion, (J) equation in Appendix B of Evans' book and Theorem 2 in Encyclopedia can control that $u(\xi)$ at most increases "exponentially". However, our Lemma B.6 can guarantee that $u(\xi)$ at least shrinks "exponentially" (or at least increases "exponentially" by (1)).

Thank you so much for your feedback and appreciate your careful reviewing again!

• (On Remark 3.3) We highly agree that current illustration in Remark 3.3 is not straightforward and thank you for pointing this out. We will mainly modify Remark 3.3 from the following two aspects as you suggested.

  • We directly state Lemma B.5: If $\Theta$ is compact, then we have

    $$\nabla\frac{\delta F}{\delta \rho}(\rho_\xi)= ({T_{\rho_\xi}^{\rho}-\mathbf{id}})/{\xi}$$

    Therefore, Assumption 2 holds.

  • For $\Theta=\mathbb{R}^d$, we explain why ”$(T_{\rho_\xi}^{\rho}-\mathbf{id}) / \xi$ is a strong subdifferential of $F$ at $\rho_\xi$“ is important: If $\Theta=\mathbb{R}^d$, we can not directly have $\nabla\frac{\delta F}{\delta \rho}(\rho_\xi)= ({T_{\rho_\xi}^{\rho}-\mathbf{id}})/{\xi}$ as in compact case (we will add some discussion in Lemma B.5). However, by [Lemma 10.1.2, (Ambrosio et al., 2005)], $\rho_{\xi}\in D(|\partial F|)$ and $(T_{\rho_\xi}^{\rho}-\mathbf{id}) / \xi$ is a (not necessarily unique) strong subdifferential. If we can prove $\nabla\dfrac{\delta F}{\delta \rho}(\rho_\xi)$ is also a strong subdifferential with minimum $L_2(\rho_\xi)$-norm under some conditions, then Assumption 2 holds.

• (On PL inequality) If we understand your question correctly, you want to have an example that doesn't contain that many conditions as in Corollary 3.5. The example in Corollary 3.5 is indeed involved. We provide an abstract version of it here, also see (Chizat, 2022):

  Assume the functional can be written as $F_{\tau}=R(\rho)+\tau \int \rho log\rho$. If $F_\tau$ satisfies the following two properties, then the PL inequality holds,

  (1) $R$ is convex in the linear sense. (Not geodesic convexity)

  (2) Uniform LSI: see Definition C.1 in our paper, also see [Assumption 3, (Chizat, 2022)]. Basically, it requires every proximal Gibbs distribution to satisfy LSI with a constant.

  Under these two conditions, we can prove a PL inequality, see (2) in the proof of Corollary 3.5 or [Section 3, (Chizat, 2022)].

• (On Gronwall inequality) In our Lemma B.6, we proved this "Gronwall" inequality for a monotone function. As you pointed out, we believe the proof is already somewhere. However, we kept the proof in our paper for completeness.

We are sorry for the format issues: we will clean the references, and solve the format issues related to the "norm" and Theorem 3.4.

If you have further questions, please don't hesitate to let us know! Thank you again for your constructive comments!

(Ambrosio et al., 2005) Ambrosio, Gigli, and Savar´e. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.

(Chizat, 2022) Lenaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing, 2022.

I was reading through parts of the revised paper. The authors have addressed several of the comments and I raised my score.

Some comments and concerns remain:

• Most readers won't have any chance to understand the meaning of Remark 3.3. Please explain, why it is interesting that $(T-id)/\xi$ is a strong subdifferential. Avoid referencing future claims (Cor 3.5 and 3.6). Why does the existence of the first variation implies its specific form (it would be better to directly state Lemma B.5 instead of writing a sentence which cannot be understood at this point).

• In order to distinguish PL over LSI: Can you give some specific examples (despite the KL divergence) where PL holds?

• Regarding the Gronwall inequality: I see that my reference was not correct due to the varying sign of the C. But I am sure that a form of the Gronwall inequality for continuous, almost-everywhere differentiable $u$ already exists (for instance the Wikipedia proof never uses the strict everywhere differentiabiltiy and has no sign restriction on C) and would do the job in your case (Your modification (1) can also be done for the standard proof of Gronwall). However, since I believe that Lemma B.6 is correct (just not new) and it is not considered as a main contribution of this paper (neither by authors nor by reviewers), I will just ignore this issue for this review.

• Even though it does not influence my score, the authors definitely should invest significantly more efforts in proofreading and formatting their paper. E.g.:

  • As other reviewers pointed to that, the reference list still does not look appropriate (references which have neither journal nor arxiv identifiers, wrong capitalization etc.)

  • stuff like using \left and \right for the norm lines throughout the paper (e.g. in Assumption 2)

  • The proof of Thm 3.4 does not fit the margins

• (On the implementation of WPA) To compute $T^{n+1}$ that maps $\rho_n$ to $\rho_{n+1}$, each step we will train a small neural network to minimize the loss of the right hand side of (17), and thus can approximate $T^{n+1}$. In our case where an entropy regularization is present, we use a neural network to compute the map $T^{n+1}$. Since the particle number is finite, $H(\rho)=\int \rho \log\rho$ is not well-defined and cannot be estimated directly. We instead compute the map $T^{n+1}$ via a change of variable formula, see [Example 2, (Yao et al. 2024)]. For details of the change of variable formula, we refer to [Theorem 1, (Mokrov, et al., 2021)]. We also remark that the training of a transport map is closely connected to the proximal point algorithm in Euclidean space:

  $(\theta_{j}^{n+1})=argmin_{T}\ \dfrac{1}{N}\sum_{i=1}^{N}l(\dfrac{1}{m}\sum_{j=1}^m \varphi(\theta_j^{n+1},x_i),y_i)+\dfrac{\lambda}{m}\sum_{j=1}^{m}\Vert \theta_j^{n+1}\Vert^2 +\dfrac{1}{2m\eta}\sum_{j=1}^{m}\Vert \theta_j^{n+1}-\theta_j^n\Vert^2$

  which is applicable when there is no entropy regularization term. In the noise-free case, we might not need to introduce a neural network to compute $T^{n+1}$. Instead, we can directly compute the particle position at next iteration. However introducing a neural network allows us to learn a map for each step, which provides certain flexibility. For instance, if the number of particles changes, we need to compute (*) for each step from scratch if we don't compute a map for each step. By contrast, if we already trained a neural network, then we are able to reuse the trained neural network for each step.

• (On proximal sampler) For KL divergence, (Chen et al., 2022) proposed a proximal sampler. At every step, it requires to implement an independent approximate RGO for each particle. If we want to increase the number of sampled particles, we need to repeat the implementation of approximate RGO for newly added particles. Even though the smoothness condition (or semiconvexity of $f$) is not listed in [Theorem 3, (Chen et al., 2022)], the chain rule used in the main lemma (Lemma 12) implicitly requires geodesic semiconvexity (for KL divergence it requires $f$ to be semiconvex), see (Section 10.4.E, (Ambrosio et al., 2005)). And also, only if $f$ is smooth, then approximate RGO can be implemented "effectively", for example by gradient descent. In addition, we want to stress here that all of the bound established in (Chen et al., 2022) are only based on the setting of sampling from a fixed function.

  Under LSI condition and semiconvexity, WPA and proximal sampler share the same convergence rate for sampling from a fixed function.

• (On Unbiased Sampler and the application to MFLD) First, we do not cover the particle discretization error, which is a challenging problem for proximal gradient methods. (Kook et al., 2024) and some other works introduce various methods to analyze the error induced from time-discretization error and particle-discretization error, for "Langevin" type algorithms. We have added the discussion in our appendix. However, to the best of our knowledge, (Kook et al., 2024) mostly focused on a general framework for analyzing the error for mean-field sampling (specifically, Langevin-type algorithm). There is no new algorithm proposed. However, our paper focus on the time-discretization error of Wasserstein proximal algorithm when geodesic convexity is missing and leave the particle discretization error to future.

• (On experimental setup) For the experiments, choosing $d=2$ is because the higher the dimension is, the more particle we need to well approximate the two time-discretizations, Langevin algorithm (assume infinite particle) and proximal algorithm. For example, the experiments in (Chen et al., 2024) also uses $d=2$ to demonstrate the influence of particle number.

(Yao et al., 2024) "Wasserstein proximal coordinate gradient algorithms." JMLR, 2024.

(Chen et al., 2022) Chen, Yongxin, et al. "Improved analysis for a proximal algorithm for sampling." Conference on Learning Theory. PMLR, 2022.

(Xu et al., 2024)Xu C, Cheng X, Xie Y. Normalizing flow neural networks by JKO scheme[J]. Advances in Neural Information Processing Systems, 2024, 36.

(Kook et al., 2024) Kook, Yunbum, et al. "Sampling from the mean-field stationary distribution." The Thirty Seventh Annual Conference on Learning Theory. PMLR, 2024.

(Ambrosio et al., 2005) Luigi Ambrosio, Nicola Gigli, and Giuseppe Savar´e. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.

(Chen et al., 2024) Chen F, Lin Y, Ren Z, et al. Uniform-in-time propagation of chaos for kinetic mean field Langevin dynamics[J]. Electronic Journal of Probability, 2024, 29: 1-43.

(Mokrov et al., 2021) Petr Mokrov, et al. Advances in Neural Information Processing Systems, 2021.

Thank you for the clarification on the implementation of WPA by training neural networks. For proximal sampler, we can control the error of implementing RGO by convex optimization or rejection sampling, but the error of implementing WPA is still unknow, which could be a limitation. However, this seems to be a new perspective for sampling. In this case, I will raise my score to 6.

Thank you so much for your careful reviewing. Some similar questions were also raised by Reviewer jLpS. So we answer those questions first.

• (On Gronwall inequality) The vanilla Gronwall inequality requires the function to be everywhere differentiable in the interval interior. For our problem, we only have almost everywhere differentiability for the Hopf-Lax semigroup, and the proof of Lemma A.1 (Lemma B.6 in revised version) requires additional monotonicity for the inequality (highlighted in red) to hold. So we added a Remark B.7 to clarify that our inequality extends the classical Gronwall lemma, which requires everywhere differentiability, to the case with almost everywhere differentiability and monotonicity.

• (On existence of minimizer) In the revision, we clarified that $F$ needs to be lower semicontinuous (lsc) in the weak topology and put the assumption of weakly lsc in our Assumption 1 as our regularity assumption. We also added a Lemma B.3 to summarize some classical conditions that guarantees weak lower semicontinuity.

• (On subdifferential representation) Thanks for pointing out this technical issue in Lemma A.3! First, we clarify that if $\Theta$ is a compact subset of $\mathbb{R}^d$, then eqn.(18) in Lemma A.3 remains true $\rho_\xi$-a.s. for probability measures supported on compact domains (cf. added Lemma B.5). So the difficult arises primarily when $\Theta = \mathbb{R}^d$ (or unbounded domain). Fortunately, our proof of Theorem 3.2 only requires either one of the followings: (a) Lemma A.3, i.e., the representation of strong subdifferential of the form $(T_{\rho_{\xi}}^{\rho}-I) / \xi$ in terms of Wasserstein gradient $\nabla \dfrac{\delta F}{\delta \rho}(\rho_\xi)$ holds along the proximal trajectory, OR (b) $\nabla \dfrac{\delta F}{\delta \rho}(\rho_\xi)$ is the minimal $L_2(\rho_\xi)$-norm element in the strong subdifferential. Existence (but not uniqueness) of strong subdifferential is guaranteed by [Lemma 10.1.2, (Ambrosio et al., 2005)] at proximal $\rho_\xi$. Thus, we impose a slightly stronger assumption in the revision (Assumption 2 (a) or (b) on proximal trajectory). We remark that this new assumption is not restrictive for our main application MFLD. In particular, we show that the MFLD on the standard two-layer architecture with some regularity on the perceptron for either squared (regression) or logistic (classification) loss is a sufficient condition for Assumption 2(b) in Corollary 3.5 (revised version).

Next we address other questions:

• "the authors write that they assume that the strong subdifferential exists for every $\rho\in\mathcal{P}_2^a(\Theta)$''; this will never be the case; you do not mean this, so please fix."

Reply: We apologize for the confusion. We clarified that our result relies on the fact that strong subdifferential at the proximal $\rho_\xi\in\mathcal{P}_2^a(\Theta)$.

• "Lem A.3: Do you need some smoothness assumptions on $\rho$?"

Reply: Yes, [Lemma 10.4.2, (Aambrosio, et al. 2005)] requires smoothness assumption we ignored before. Please see our detailed reply in (On subdifferential representation).

• "Lem A4: This is exactly Prop. 3-1 , 3.3 (Ambrosio 2013)."

Thanks. This part is not our contribution, and thus we put it to appendix and cited (Ambrosio 2013).

• "Def B.1: reformulate, what is defined here Proof of Thm 3.2: line 316: why you have an additional $\xi$ in the second summand? line 333: the last equality should be an inequality The rest appears to be correct."

We reformulated Def B.1 (now Def C.1 in revised version). The additional $\xi$ and the last equality is a typo. We fixed this.

• "Cor 3.4: what do you mean by φ(x,y) ''is uniformly in y'' (uniformly what?)"

We have corrected this part in Corollary 3.8 in our new version, also we refer to Proposition 5.1 (Chizat, 2022) for details.

We also cleaned up the reference list for capitals.

(Ambrosio et al., 2005) Ambrosio, Gigli, and Savare. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.

(Chizat, 2022) Lenaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing, 2022

For Reviewer ppSk

• (Connection to Euclidean case) Under convexity and PL inequality in Euclidean space, a linear convergence of proximal algorithm can be proved, which is a classical result. However, there is no strict proof to show that under merely PL inequality, the linear convergence can be achieved. Our proof strategy actually can provide a strict proof for linear convergence of proximal under merely PL inequality in Euclidean space. For Euclidean case, with the strong lemma from (Ambrosio, 2013), the main challenge will only be how to obtain a "Gronwall" lemma for function that is only almost everywhere differentiable, see Lemma B.6 and Remark B.7.

• (Key Challenges in Wasserstein space) In Wasserstein space, the analysis becomes more difficult. Except for the aforementioned generalization of the "Gronwall" lemma, there is an additional main challenge. In Wasserstein space, the strong subdifferential might not exist everywhere and might not be unique. However, in Euclidean case where we have $0=\theta_{\xi}-\theta+\xi\nabla f(\theta_{\xi})$ for $\theta_\xi=argmin f(\tilde\theta)+\dfrac{\Vert\theta-\tilde\theta\Vert^2}{2\xi}$, when we assume $f$ is differentiable everywhere. In Wasserstein space we don't have that $\dfrac{T_{\rho_\xi}^{\rho}-I}{\xi}=\nabla\dfrac{\delta F}{\delta \rho}(\rho_\xi)$ directly, see Remark 3.3. Solving this problem is generally quite technical and might involve some other conditions (such as smoothness of the cost functional, Assumption 3 in (Chizat, 2022)), which are not needed for the Euclidean case.

• (On Shaper bound for strongly convex objective) The main ingredient comes from the proof of Theorem 3.4 (paper in new version), where we utilized the Lemma 3.1 which connect the time derivative of Hopf-Lax formula (a.k.a. Moreau envelope) and the distance between the proximal $\rho_\xi$ and $\rho$. Specifically, we are able to compute an integral over time $\xi$, which allows us to have a better bound. In this way, for every $\rho_\xi,$ where $\xi\in[0,\eta]$, we can utilize the definition of strong convexity along geodesics.

  However, the classical proof for strongly convex objective, only utilizes strongly convexity along geodesics for $\xi=\eta$ once, which leads to a weaker bound.

  Again, if we want to extend a faster rate for strongly convex objective in Euclidean space, it is much easier than Wasserstein space case and straightforward as strong convexity implies PL inequality in Euclidean space.

• (Other cases of PL inequality) We know one example from (Chizat, 2022) on Kernel Maximum Mean Discrepancy. Under [Proposition 5.2, (chizat, 2022)], the objective functional satisfies PL type inequality.

• (Extension to KL inequality) Thank you for your suggestion! We believe it is possible to extend our analysis to Wasserstein KL inequality setting, and there is already some research on continuous time dynamics under Wasserstein KL inequality (Blanchet, et al., 2018).

(Chizat, 2022) Lenaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing, 2022. URL https://arxiv.org/abs/2202.01009

(Ambrosio et al., 2013) Calculus and heat flow in metric measure spaces and applications to spaces with Ricci bounds from below. Inventiones mathematicae, 195(2):289–391, February 2013. URL http://dx.doi.org/10.1007/s00222-013-0456-1.

(Blanchet, et al., 2018) Blanchet A, Bolte J. A family of functional inequalities: Łojasiewicz inequalities and displacement convex functions[J]. Journal of Functional Analysis, 2018, 275(7): 1650-1673.