Non-geodesically-convex optimization in the Wasserstein space

Authors: hoang  phuc
hau  luu
hanlin  yu
bernardo  williams
petrus  mikkola
marcelo  hartmann
kai  puolamäki
arto  klami
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 14:12:20

Abstract: We study a class of optimization problems in the Wasserstein space (the space of probability measures) where the objective function is nonconvex along generalized geodesics. Specifically, the objective exhibits some difference-of-convex structure along these geodesics. The setting also encompasses sampling problems where the logarithm of the target distribution is difference-of-convex. We derive multiple convergence insights for a novel semi Forward-Backward Euler scheme under several nonconvex (and possibly nonsmooth) regimes. Notably, the semi Forward-Backward Euler is just a slight modification of the Forward-Backward Euler whose convergence is---to our knowledge---still unknown in our very general non-geodesically-convex setting.

Main paper: 
 introduction: 
 Sampling and optimization are intertwined. For example, the (overdamped) Langevin dynamics, typically considered a sampling algorithm, can be considered as gradient descent optimization where a suitable amount of Gaussian noise is injected at each step. There are also deeper connections. At the limit of infinitesimal stepsize, the law of the Langevin dynamics is governed by the Fokker-Planck equation describing a diffusion over time of probability measures. In the seminal paper [39] , Jordan, Kinderlehrer, and Otto reinterpreted the Fokker-Planck equation as the gradient flow of the functional relative entropy, a.k.a. Kullback-Leibler (KL) divergence, in the (Wasserstein) space of finite secondmoment probability measures equipped with the Wasserstein metric. The discovery connects the two fields and encourages optimization in the Wasserstein space, even conceptually, as it directly gives insight into the sampling context. Studies in continuous-time dynamics [21, 12, 66, 30] seem natural and enjoy nice theoretical properties without discretization errors. Another line of research studies discretization of Wasserstein gradient flow by either quantifying the discretization error between the continuous-time flow and the discrete-time flow [39, 67, 27, 23, 28] or viewing discrete-time flows as iterative optimization schemes in the Wasserstein space [65, 26, 70, 11] where the primary focus is on (geodesically) convex optimization problems.Nonconvex, nonsmooth optimization is challenging, even in Euclidean space, quoting Rockafellar [63] : "In fact the great watershed in optimization isn't between linearity and nonlinearity, but convexity and nonconvexity." The landscape of nonconvex problems is mostly underexplored in the Wasserstein space. In the sampling language, it amounts to sampling from a non-log-concave and possibly non-log-Lipschitz-smooth target distribution. Recently, Balasubramanian et al. [9] advocated the need for a sound theory for non-log-concave sampling and provided some guarantees for the unadjusted Langevin algorithm (ULA) in sampling from log-smooth (Lipschitz/Hölder smooth) densities. These results are preliminary for the ULA (and its variants) with a specific class of densities (smooth). Theoretical understandings of other classes of algorithms and densities are needed.We approach the subject through the lens of nonconvex optimization in the space of probability distributions and pose discretized Wasserstein gradient flows as iterative minimization algorithms. This allows us to, on the one hand, use and extend tools from classical nonconvex optimization and, on the other hand, derive more connections between sampling and optimization.We study the following non-geodesically-convex optimization problem defined over the space P 2 (X) of probability measures µ over X = R d with finite second moment, i.e., ∥x∥ 2 dµ(x) < +∞, min µ∈P2(X)EQUATIONwhere F : X → R is a nonconvex function which can be represented as a difference of two convex functions G and H, E F (µ) := F (x)dµ(x) is the potential energy, and H : P 2 (X) → R ∪ {+∞} plays a role as the regularizer which is assumed to be a convex function along generalized geodesics.Why difference-of-convex structure? Nonconvexity lies at the difference-of-convex (DC) structure F = G -H, where G and H are called the first and second DC components, respectively. F being nonconvex implies E F being non-geodesically-convex in general. First, the class of DC functions is very rich, and DC structures are present everywhere in real-world applications [60, 46, 47, 1, 22, 56, 58] . Weakly convex and Lipschitz smooth (L-smooth or simply smooth) functions are two subclasses of DC functions. Furthermore, any continuous function can be approximated by a sequence of DC functions over a compact, convex domain [8] . We also remark that many nonconvex functions admit quite natural DC decompositions, for example, an L-smooth function F has the following splittings: F (x) = α∥x∥ 2 -(α∥x∥ 2 -F (x)) and F (x) = (F (x) + α∥x∥ 2 ) -α∥x∥ 2 whenever α ≥ L/2. Second, DC functions preserve enough structure to extend convex analysis. Such structure is key in classic DC programming [60] and in our Wasserstein space analysis with optimal transport tools.Context Many problems in machine learning and sampling fall into the spectrum of problem (1) . For example, refer to a discussion in [65] that inspired our work. The regularizer H can be the internal energy [4, Sect. 10.4.3] . Under McCann's condition, the internal energy is convex along generalized geodesics [4, Prop. 9.3.9]. In particular, the negative entropy, H (µ) = log(µ(x))dµ(x) if µ is absolutely continuous w.r.t. Lebesgue measure, +∞ otherwise, is a special case of internal energy satisfying McCann's condition. In the latter case, F(µ) = D KL (µ∥µ * ) + const where µ * (x) ∝ exp(-F (x)), the optimization problem reduces to a sampling problem with log-DC target distribution. In Bayesian inference, the posterior structure depends on both the prior and likelihood. If the likelihood is log-smooth, it exhibits the aforementioned DC splittings. Logpriors, often nonsmooth to capture sparsity or low rank, typically also have explicit DC structures [48, 32, 25] . In the context of infinitely wide one-layer neural networks and Maximum Mean Discrepancy [51, 7, 21] , let µ * be the optimal distribution over a network's parameters, k be a given kernel, the regularizer is then the interaction energy H (µ) = k(x, y)dµ(x)dµ(y) and F (x) = -2 k(x, y)dµ * (y). In general, H is not convex along generalized geodesics and F is nonconvex but not necessarily DC. When the kernel has Lipschitz gradient, we can adjust both H and F as H (µ) = k(x, y)+α∥x∥ 2 +α∥y∥ 2 dµ(x)dµ(y) and F (x) = -2 k(x, y)dµ * (y)-2α∥x∥ 2 for some α > 0 making H generalized geodesically convex and F concave (hence DC); Appx. A.2.Our idea is to minimize (1) in the space of probability distributions by discretization of the gradient flow of F, leveraging on the JKO (Jordan, Kinderlehrer, and Otto) operator (2) . In the previous work [70] , this has been done with the Forward-Backward (FB) Euler discretization, but it lacks convergence analysis. Recently, Salim et al. [65] did some study on FB Euler, but their results do not apply here because F is nonconvex and possibly nonsmooth. Further leveraging on the DC structure of F and inspired by classical DC programming literature [60] , we subtly modify the FB Euler to give rise to a scheme named semi FB Euler that enjoys major theoretical advantages as we can provide a wide range of convergence analysis. Regarding the name, "semi" addresses the splitting of the potential energy. The scheme can be nevertheless reinterpreted as FB Euler. solution, if exists, is denoted by T ν µ and called the optimal (Monge) map. Monge's problem can be ill-posed, e.g., no such T ν µ exists when µ is a Dirac mass and ν is absolutely continuous [5] . By relaxing Monge's formulation, Kantorovich considers min γ∈Γ(µ,ν) X×X ∥x -y∥ 2 dγ(x, y), where Γ(µ, ν) denotes the set of probabilities over X × X whose marginals are µ and ν, i.e, γ ∈ Γ(µ, ν) iff proj 1# γ = µ, proj 2# γ = ν where proj 1 , proj 2 are the projections onto the first X space and the second X space, respectively. Such γ is called a plan. Kantorovich's formulation is well-posed because Γ(µ, ν) is non-empty (at least µ × ν ∈ Γ(µ, ν)) and the arg min element actually exists (see [5, Sect. 2.2] ). The set of optimal plans between µ and ν is denoted by Γ o (µ, ν). In terms of random variables, any pairs (X, Y ) where X ∼ µ, Y ∼ ν is called a coupling of µ and ν while it is called an optimal coupling if the joint law of X and Y is in Γ o (µ, ν).In P 2 (X), the min value in Kantorovich's problem specifies a valid metric referred to as Wasserstein distance, W 2 (µ, ν) = ( X×X ∥x -y∥ 2 dγ(x, y)) 1/2 for some, and thus all, γ ∈ Γ o (µ, ν). The metric space (P 2 (X), W 2 ) is then called the Wasserstein space. In P 2 (X), beside the convergence notion induced by the Wasserstein metric, there is a weaker notion of convergence called narrow convergence: we say a sequence{µ n } n∈N ⊂ P 2 (X) converges narrowly to µ ∈ P 2 (X) if X ϕ(x)dµ n (x) → X ϕ(x)dµ(x) for all ϕ ∈ C b (X).Convergence in the Wasserstein metric implies narrow convergence but the converse is not necessarily true. The extra condition to make it true is m 2 (µ n ) → m 2 (µ). We denote Wasserstein and narrow convergence by Wass --→ and narrow ---→, respectively.If µ ∈ P 2,abs (X), ν ∈ P 2 (X), Monge's formulation is well-posed and the unique (µ-a.e.) solution exists, and in this case, it is safe to talk about (and use) the optimal transport map T ν µ . Moreover, there exists some convex function f such that T ν µ = ∇f µ-a.e. Kantorovich's problem also has a unique solution γ and it is given by γ = (I, T ν µ ) # µ where I is the identity map. This is known as Brenier theorem or polar factorization theorem [18] .Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We briefly mention our contribution in the abstract (propose a new scheme named semi Forward-Backward Euler and derive multiple convergence insights) and provide a detailed contribution paragraph in Sect. 1 where each claim is pointed to the corresponding theorem.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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 can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
subdifferential calculus in the wasserstein space: 
 Apart from being a metric space, (P 2 (X), W 2 ) also enjoys some pre-Riemannian structure making subdifferential calculus on it possible. Let us have a picture of a manifold in mind. Firstly, the tangent space [4] of P 2 (X) at µ is Tan µ P 2 (X) := {∇ψ : ψ ∈ C ∞ c (X)} L 2 (X,X,µ), where the closure is w.r.t. the L 2 (X, X, µ)-topology. Intuitively, for ψ ∈ C ∞ c (X), I + ϵ∇ψ is an optimal transport map if ϵ > 0 is small enough [43] , so ∇ψ plays a role as "tangent vector".Let ϕ : P 2 (X) → R∪{+∞}, we denote dom(ϕ) = {µ ∈ P 2 (X) : ϕ(µ) < +∞}. Let µ ∈ dom(ϕ), we say that a map ξ ∈ L 2 (X, X, µ) belongs to the Fréchet subdifferential [15, 43] ∂ - F ϕ(µ) if ϕ(ν) -ϕ(µ) ≥ sup γ∈Γo(µ,ν) X×X ⟨ξ(x), y -x⟩dγ(x, y) + o(W 2 (µ, ν)) for all ν ∈ P 2 (X), where the little-o notation means lim s→0 o(s)/s = 0. If ∂ - F ϕ(µ) ̸ = ∅, we say ϕ is Fréchet subdifferentiable at µ. We also denote dom(∂ - F ϕ) = {µ ∈ P 2 (X) : ∂ - F ϕ(µ) ̸ = ∅}. Similarly, we say that ξ ∈ L 2 (X, X, µ) belongs to the (Fréchet) superdifferential ∂ + F ϕ(µ) of ϕ at µ if -ξ ∈ ∂ - F (-ϕ)(µ). In other words, ∂ - F (-ϕ)(µ) = -∂ + F ϕ(µ). We say ϕ is Wassertein differentiable [15, 43] at µ ∈ dom(ϕ) if ∂ - F ϕ(µ) ∩ ∂ + F ϕ(µ) ̸ = ∅.We call an element of the intersection, denoted by ∇ W ϕ(µ), a Wasserstein gradient of ϕ at µ, and it holds ϕ(ν) -ϕ(µ) = X×X ⟨∇ W ϕ(µ)(x), y -x⟩dγ(x, y) + o(W 2 (µ, ν)), for all ν ∈ P 2 (X) and any γ ∈ Γ o (µ, ν). The Wasserstein gradient is not unique in general, but its parallel component in Tan µ P 2 (X) is unique, and this parallel component is again a valid Wasserstein gradient as the orthogonal component plays no role in the above definitions, i.e., if ξ ⊥ ∈ Tan µ P 2 (X) ⊥ , it holds X×X ⟨ξ ⊥ (x), y -x⟩dγ(x, y) = 0 for any ν ∈ P 2 (X) and γ ∈ Γ o (µ, ν) [43, Prop. 2.5] . We may refer to this parallel component as the unique Wasserstein gradient of ϕ at µ.A function ϕ : P 2 (X) → R ∪ {+∞} is called proper if dom(ϕ) ̸ = ∅, while it is called lower semicontinuous (l.s.c) if for any sequence µ n Wass --→ µ, it holds lim inf n ϕ(µ n ) ≥ ϕ(µ).We next recall (a simplified version of) generalized geodesic convexity. Definition 1. [65] Let ϕ : P 2 (X) → R ∪ {+∞}. We say ϕ is convex along generalized geodesics if ∀µ, π ∈ P 2 (X), ∀ν ∈ P 2,abs (X), ϕ((tTµ ν + (1 -t)T π ν ) # ν) ≤ tϕ(µ) + (1 -t)ϕ(π), ∀t ∈ [0, 1]. The curve t → (tT µ ν + (1 -t)T π ν )# ν (called a generalized geodesic) interpolates from π to µ as t runs from 0 to 1. The definition says that ϕ is convex along these curves. If µ ∈ P 2,abs (X) and ν = µ, the curve is a geodesic in (P 2 (X), W 2 ). If the definition is relaxed to the class of geodesics only, we say that ϕ is convex along geodesics.An important characterization of Fréchet subdifferential of a geodesically convex function is that we can drop the little-o notation in its definition in Sect. 2.3 [4, Sect 10.1.1] . As a convention, for a geodesically convex function ϕ, the Fréchet subdifferential ∂ - F will be simply written as ∂.First-order optimality conditions Let ϕ : P 2 (X) → R∪{+∞} be a proper function. We call µ * a Fréchet stationary point ofµ * ∈ P 2 (X) is a global minimizer of ϕ if ϕ(µ * ) ≤ ϕ(µ), ∀µ ∈ P 2 (Xϕ if 0 ∈ ∂ - F ϕ(µ * ).Fréchet stationarity is a necessary condition for local optimality. In other words, if µ * is a local minimizer, it is a Fréchet stationary point (Lem. 5 in Appendix). In addition, if ϕ is Wasserstein differentiable at µ * , ∇ W ϕ(µ * )(x) = 0 µ * -a.e. [43] . When ϕ is geodesically convex, Fréchet stationarity is a sufficient condition for global optimality (Lem. 6 in Appendix).3 Semi Forward-Backward Euler for difference-of-convex structuresQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We mentioned the expensiveness of the method, especially since the JKO operator is not scalable at least for now. See Sect. 5.Guidelines:• The answer NA means that the paper has no 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 a separate "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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
wasserstein gradient flows: different types of discretizations: 
 To neatly present the idea of minimizing F via discretized gradient flow, we first assume for a moment that F is infinitely differentiable and H is the negative entropy. See also a discussion in [65] .We wish to minimize (1) in the space of probability distributions. A natural idea is to apply discretizations of the gradient flow of F, where the gradient flow is defined (under some technical assumptions [39] ) as the limit η → 0 + of the following scheme with some simple time-interpolationEQUATIONStraightforwardly, given a fixed η > 0, (2) gives back a discretization for this flow known as Backward Euler. On the other hand, if F is Wasserstein differentiable (Sect. 2.2), the Forward Euler discretization reads [70] µ n+1 = (I -η∇ W F(µ n )) # µ n , which is reinterpreted as doing gradient descent in the space of probability distributions. These are optimization methods that work directly on the objective function F itself. However, the composite structure of F (a sum of several terms) can also be exploited. One such scheme is the unadjusted Langevin algorithm (ULA), where it first takes a gradient step w.r.t. the potential part, then follows the heat flow corresponding to the entropy part [70] : ν n+1 = (I -η∇F ) # µ n , and µ n+1 = N (0, 2ηI) * ν n+1 , where * is the convolution. This ULA is "viewed" in the space of distributions (Eulerian approach), a more familiar and equivalent form of the ULA from the particle perspective (Lagrangian approach) goes likex n+1 = x n -η∇F (x n ) + √ 2ηz k where z k ∼ N (0, I).The ULA is known to be asymptotically biased even for Gaussian target measure (Ornstein-Uhlenbeck process). To correct this bias, the Metropolis-Hasting accept-reject step [62] is sometimes introduced. Metropolis-Hasting algorithm [52, 36] is a much more general framework that works with quite any proposal (e.g., a random walk) whose convergence analysis is based on the Markov kernel satisfying the detailed balance condition. This convergence framework is different from what is considered in this work: we are more interested in the underlying dynamics of the chain. Metropolis-Hasting algorithm is indeed another story.In optimization, for composite structure, Forward-Backward (FB) Euler and its variants are methods of choice [59, 10] . The corresponding FB Euler for F will take the gradient step (forward) according to the potential, and JKO step (backward) w.r.t. the negative entropy (FB Euler) ν n+1 = (I -η∇F ) # µ n , and µ n+1 ∈ JKO ηH (ν n+1 ).This scheme appears in [70] without convergence analysis, and later on [65] derives non-asymptotic convergence guarantees under the assumption F being convex and Lipschitz smooth.In this work, as F is nonconvex and nonsmooth, the theory in [65] does not apply, and the convergence (if any) of (3) remains mysterious. The DC structure of F can be further exploited. In DC programming [60] , the forward step should be applied to the concave part, while the backward step should be applied to the convex part. We hence propose the following semi FB Euler (semi FB Euler) ν n+1 = (I + η∇H) # µ n , and Additionally, we also offer convergence guarantees when H is nonsmooth.EQUATIONOur goal is to minimize the non-geodesically-convex functional F(µ) = E F (µ)+H (µ) over P 2 (X), where F = G -H is a DC function. We make Assumption 1 throughout the paper: Assumption 1.(i) The objective function F is bounded below.(ii) G, H : X → R are convex functions and have quadratic growth.(iii) H : P 2 (X) → R ∪ {+∞} is proper, l.s.c, and convex along generalized geodesics in (P 2 (X), W 2 ), and dom(H) ⊂ P 2,abs (X).(iv) There exists η 0 > 0 such that ∀η ∈ (0, η 0 ), JKO η(E G +H ) (µ) ̸ = ∅ for every µ ∈ P 2 (X).Note that Assumption 1(iv) is a commonly-used assumption to simplify technical complication when working with the JKO operator [4, 15, 65] . First, it follows from Assumption 1(iii), dom(F) ⊂ P 2,abs (X). By analogy to DC programming in Euclidean space, we call µ * ∈ dom(F) a critical point of F = H + E G -E H if ∂(H + E G )(µ * ) ∩ ∂E H (µ * ) ̸ = ∅.We allow H to be non-differentiable in some derivations, meaning that ∂H (convex subdifferential [54] ) contains multiple elements in general. We first pick a selector S of ∂H, i.e., S : X → X, such that S(x) ∈ ∂H(x). By the axiom of choice (Zermelo, 1904, see, e.g., [37] ), such selection always exists. However, an arbitrary selector can behave badly, e.g., not measurable. We shall first restrict ourselves to the class of Borel measurable selectors (see Appx. A.1 for an existence discussion). Assumption 2 (Measurability). The selector S is Borel measurable.We recall the semi FB scheme (4) but for nonsmooth F as follows: start with an initial distribution µ 0 ∈ P 2,abs (X), given a discretization stepsize 0 < η < η 0 , we repeat the following two steps:ν n+1 = (I + ηS) # µ n ◁ push forward step; µ n+1 = JKO η(E G +H ) (ν n+1 ) ◁ JKO step.Well-definiteness and properties: Given µ n ∈ P 2 (X), it follows from Lem. (4) that ν n+1 ∈ P 2 (X).The two generated sequences are then in P 2 (X). Moreover, it follows from Assumption 1 that {µ n } n∈N are in P 2,abs (X), so are {ν n } n∈N using Lem. 9 by noting that I + ηS is subgradient of a strongly convex function x → (1/2)∥x∥ 2 + ηH(x).4 Convergence analysis 4.1 Asymptotic analysis Lemma 1 (Descent lemma). Under Assumptions 1 and 2, let {µ n } n∈N be the sequence of distributions produced by semi FB Euler starting from some µ 0 ∈ P 2,abs (X) with 0 < η < η 0 . Then it holdsF(µ n+1 ) ≤ F(µ n ) -1 η X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x), ∀n ∈ N.Lem. 1 shows that the objective does not increase along semi FB Euler's iterates. Proof of Lem. 1 is in Appx. A.3. By using Lem. 1, we establish asymptotic convergence for semi FB Euler as follows.For the asymptotic convergence analysis, we need the following assumption on H. some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
non asymptotic analysis: 
 To measure how fast the algorithm converges, we need some convergence measurement. First, for proximal-type algorithms in Euclidean space, the notion of gradient mapping G η (x n ) is usually used (see, e.g., [33, 55] and [38, Eq. ( 5)]) and we measure the rate ∥G η (x n )∥ 2 → 0. In analogy as in Euclidean space, we define the Wasserstein (sub)gradient mapping as follows G η (µ) :=1 η I -T JKO η(E G +H ) ((I+ηS) # µ) µ, and we measure the rate of ∥G η (µ n )∥ 2 L 2 (X,X,µn) → 0. Theorem 2 (Convergence rate: Wasserstein (sub)gradient mapping). Under Assumptions 1, 2, let {µ n } n∈N be the sequence of distributions produced by semi FB Euler starting from some µ 0 ∈ P 2,abs (X) with 0 < η < η 0 . Then it holds minn=1,N ∥G η (µ n )∥ 2 L 2 (X,X,µn) = O(N -1 ).Proof of Thm. 2 is in Appx. A.5. This theorem holds without requiring G and H to be differentiable.Next, if H is twice differentiable with uniformly bounded Hessian, we can derive a stronger convergence guarantee based on Fréchet stationarity (see Sect. 2.4). In other words, we evaluate the rate of dist (0,∂ - F F(µ n )) := inf ξ∈∂ - F F (µn) ∥ξ∥ L 2 (X,X;µn) → 0. Assumption 4. H ∈ C 2 (X) whose Hessian is bounded uniformly (H is then L H -smooth).Theorem 3 (Convergence rate: Fréchet subdifferentials). Under Assumptions 1, 4, let {µ n } n∈N be the sequence of distributions produced by semi FB Euler starting from some µ 0 ∈ P 2,abs (X) with0 < η < η 0 , then min n=1,N dist (0, ∂ - F F(µ n )) = O N -1 2 .Proof of Thm. 3 is in Appx. A.6.Fast convergence can be obtained under isoperimetry, e.g., log-Sobolev inequality (LSI). There are certain connections between LSI in sampling and the Łojasiewicz condition in optimization allowing linear convergence. In nonconvex optimization in Euclidean space, analytic and subanalytic functions are a large class satisfying Łojasiewicz condition [44, 14] . Subanalytic DC programs are studied in [45] . In the infinite-dimensional setting of the Wasserstein space, the Łojasiewicz condition should be regarded as functional inequalities [12] .Assumption 5 (Łojasiewicz condition in the Wasserstein space). Assume that F * is the optimal value of F, and assume there exist r 0 ∈ (F * , +∞], θ ∈ [0, 1), and c > 0 such that for all µ ∈ P 2 (X),F(µ) -F * < r 0 ⇒ c (F(µ) -F * ) θ ≤ ∥ξ∥ L 2 (X,X,µ) , ∀ξ ∈ ∂ - F F(µ), where the convention 0 0 = 0 is used. θ ∈ [0, 1) is called Łojasiewicz exponent of F at optimality. Remark 2. If H is the is negative entropy, F ∈ C 2 (X)whose Hessian is bounded uniformly, then F is Wasserstein differentiable at any µ ∈ P 2,abs (X) and [43, Prop. 2.12, E.g. 2.3] ∇ W F(µ) = ∇µ µ + ∇F , where, by abuse of notation, the probability density function of µ is still denoted by µ. We have∥∇ W F(µ)∥ 2 L 2 (X,X,µ) = ∇µ(x) µ(x) + ∇F (x) 2 dµ(x) = µ(x) ∇ log µ(x) µ * (x) 2dx, where µ * ∝ exp(-F ). On the other hand, F(µ)-F * = D KL (µ∥µ * ). The log-Sobolev inequality with parameterα > 0 inequality reads [57] D KL (µ∥µ * ) ≤ 1 2α FI(µ∥µ * ) := 1 2α µ(x) ∇ log µ(x) µ * (x) 2dx, where FI(µ∥µ * ) is the relative Fisher information of µ with respect to µ * . Therefore, log-Sobolev inequality is a special case of Łojasiewicz condition with r 0 = +∞, c = √ 2α, and θ = 1/2. In another case, when the objective function is the Maximum Mean Discrepancy, under some regularity assumption of the kernel, it holds [7] 2(F(µ) -F(µ * )) ≤ ∥µ * -µ∥ Ḣ-1 (µ) × ∥∇ W F(µ)∥ 2 dµ(x)where ∥µ * -µ∥ Ḣ-1 (µ) is the weighted negative Sobolev distance. This is "nearly" the Łojasiewicz condition, with a caveat that ∥µ * -µ∥ Ḣ-1 (µ) may be unbounded. Nevertheless, assuming the boundedness of this term along the algorithm's iterates is sufficient for convergence. Theorem 4. Under Assumptions 1, 4 and Assumption 5 with parameters (r 0 , c, θ). Let {µ n } n∈N be the sequence of distributions produced by semi FB Euler starting from some sufficiently warm-up µ 0 ∈ P 2,abs (X) such that F(µ 0 ) < r 0 and with stepsize 0 < η < η 0 , then(i) if θ = 0, F(µ n ) -F * converges to 0 in a finite number of steps; (ii) if θ ∈ (0, 1/2], F(µ n ) -F * = O M M +1 n where M = 2(η 2 L 2 H +1) c 2 η ; (iii) if θ ∈ (1/2, 1), F(µ n ) -F * converges sublinearly to 0, i.e., F(µ n ) -F * = O n -1 2θ-1 . Proof of Thm. 4 is in Appx. A.7. Remark 3.In the usual sampling case, i.e., H is the negative entropy, and under log-Sobolev condition, r 0 = +∞. Therefore, µ 0 can be arbitrarily in P 2,abs (X). In the general case, however, a good enough starting point (i.e., F(µ 0 ) < r 0 ) is needed to guarantee we are in the region where Łojasiewicz condition comes into play. In such a case, F(µ n ) -F * = D KL (µ n ∥µ * ) where µ * (x) ∝ exp(-F (x)) is the target distribution (see Rmk. 2), so Thm. 4 provides convergence rate of {µ n } n∈N to µ * in terms of KL divergence and this convergence is exponentially fast if θ ∈ (0, 1/2]. Theorem 5. Under the same set of assumptions as in Thm. 4, the sequence {µ n } n∈N is a Cauchy sequence under Wasserstein topology. Furthermore, as the Wasserstein space (P 2 (X), W 2 ) is complete [5, Thm. 2.2], every Cauchy sequence is convergent, i.e., there exists µ * ∈ P 2 (X) such that µ n Wass --→ µ * . The limit distribution µ * is indeed the global minimizer of F. In addition:(i) if θ = 0, W 2 (µ n , µ * ) converges to 0 in a finite number of steps; (ii) if θ ∈ (0, 1/2], W 2 (µ n , µ * ) = O M M +1 n , where M = 1 + (2(η 2 L 2 H +1)) 1 2θ (1-θ)η 1-θ θ c 1 θ ; (iii) if θ ∈ (1/2, 1), W 2 (µ n , µ * ) = O n -1-θ 2θ-1.Proof of Thm. 5 is in Appx. A.8. This theorem provides convergence to optimality in terms of Wasserstein distance. Remark 4. If H is the negative entropy and θ = 1/2, under some technical assumptions on F (e.g., continuously twice differentiable), LSI implies Talagrand inequality [57] (in optimization, known as Łojasiewicz implies quadratic growth [40] ), meaning that KL divergence controls squared Wasserstein distance, so fast convergence under KL divergence implies fast convergence under Wasserstein distance. 
practical implementations: 
 The push-forward step ν n+1 = (I + η∇H) # µ n is rather straightforward: if Z are samples from µ n then Z + η∇H(Z) are samples from ν n+1 . On the other hand, to move from ν n+1 to µ n+1 we have to work out the JKO operator. Recent advances [53, 2] propose using the gradient of an input-convex neural network (ICNN) [6] to approximate the optimal Monge map pushing ν n+1 to µ n+1 , which we briefly describe as follows. This approach is inspired by Brenier theorem asserting that an optimal Monge map has to be the (sub)gradient field of some convex function. Therefore, one can "parametrize" µ ∈ P 2,abs (X) as µ = ∇ψ # ν n+1 for some convex function ψ. We then write the JKO objective asEQUATIONWhile the two last terms (potential energy and squared Wasserstein distance) in ( 5) can be handled efficiently by the Monte Carlo method using samples from ν n+1 , the first term H might be complicated as it possibly involves the (unavailable) density of ν n+1 . We remark that the easy case would be H being another potential energy or an interaction energy. In such a case, Monte Carlo approximations are again readily applicable. The tricky case would be H being the negative entropy that requires the density of ν n+1 . Fortunately, we have the following change of entropy formula: for any T :X → X diffeomorphic, any ρ ∈ P 2,abs (X), it holds -H (T # ρ) = -H (ρ) + X log | det ∇T (x)|dρ(x).Therefore, ( 5) can be written as (up to a constant that does not depend on ψ)X -log det ∇ 2 ψ(x) + G(∇ψ(x)) + 1 2η ∥x -∇ψ(x)∥ 2 dν n+1 (x).Note that this entropy formula can be extended naturally to the case of general internal energy [2] , which means we can also handle this general case. Let us now consider the entropy case for simplicity.We can leverage on a class of input convex neural networks [6] ψ θ (x) (θ is the neural network's parameters, x is the input) in which x → ψ θ (x) is convex. Optimizing over θ can then be solved effectively by standard deep learning optimizers (e.g. Adam). The complete scheme is given in Alg.1. We also remark that [2] further proposes fast approximation for log det ∇ 2 ψ and [31] leverages on the variational formula of the KL to propose an even faster scheme. Nevertheless, these schemes are generally expensive. For illustrative purposes, we adopt the vanilla version of [53] . The iteration complexity is thus cubic in d, linear in the size of the ICNN, and linear in the iteration count k [53] .Input: Initial measure µ 0 ∈ P 2,abs (X), discretization step size η > 0, number of steps K > 0, batch size B.for k = 1 to K do for i = 1, 2, . . . do Draw a batch of samples Z ∼ µ 0 of size B; Ξ ← (I + η∇H) • ∇ x ψ θ k • (I + η∇H) • ∇ x ψ θ k-1 • . . . • (I + η∇H)(Z); W 2 2 ← 1 B ξ∈Ξ ∥∇ x ψ θ (ξ) -ξ∥ 2 ; U ← 1 B ξ∈Ξ G(∇ x ψ θ (ξ)); ∆H ← -1 B ξ∈Ξ log det ∇ 2 x ψ θ (ξ). L ← 1 2η W 2 2 + U + ∆H .Apply an optimization step (e.g., Adam) over θ using ∇ θ L. end for θ k+1 ← θ. end for 6 We perform numerical sampling experiments from non-log-concave distributions: the Gaussian mixture distribution and the distance-to-set-prior [61] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We provide our code/results in a GitHub repo mentioned in Sect. 6. We also instruct how to reproduce our experiments in README.md.Guidelines:• The answer NA means that paper does not include experiments requiring code. 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 a 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. 
discussion and related work: 
 We first narrow down our discussion on FB Euler and its variants in the Wasserstein space. When H is the negative entropy, Wibisono [70] provides some insightful discussion on how FB Euler should be consistent (no asymptotic bias) because the backward step is adjoint to the forward step, hence preserves stationarity. However, no convergence theory is presented for FB Euler in the Wasserstein space in [70] . Recently, Salim et al. [65] provide convergence guarantee for FB Euler within the following setting: H is convex along generalized geodesics, F is Lipschitz smooth and convex/strongly convex. This setting remains a "convex + convex" structure, while ours has a "convex + concave" structure. A natural extension of our work would be a full-fledged study of DC programming in the Wasserstein space F = G -H where G and H are convex along generalized geodesics and H is not necessarily potential energy E H . This problem possesses another implementation challenge regarding the Wasserstein gradient of H.Other works that bear a tangential relation to ours involve the use of forward-only Euler and ULA. Arbel et al. [7] study Wasserstein gradient flows with forward Euler for maximum mean discrepancy. This objective function is also nonconvex (specifically, weakly convex). Durmus et al. [26] analyze the ULA from the convex optimization perspective. Vempala et al. [67] show that LSI and Hessian boundedness suffice for fast convergence of the ULA where "fast" is understood as fast to the biased target since ULA is a biased algorithm. Balasubramanian et al. [9] analyze the ULA under quite mild conditions: log-density is Lipschitz/Hölder smooth. Bernton [11] studies the proximal-ULA also under the convex assumption, where the difference to the ULA is the first step: gradient descent is replaced by the proximal operator. Similar to ULA, proximal-ULA is asymptotically biased. To address nonsmoothness, another line of research utilizes Moreau-Yosida envelopes to create smooth approximations of the ULA dynamics [29, 50] . This approach is also applicable to certain classes of non-log-concave distributions [50] and is more of a flavour of discretization error quantification.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We report in the Numerical illustration section 6 the mean curves in bold and individual run curves shown faintly in the background playing the same role as error bars. We provide sufficient information in Appendix C.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean. Justification: This is standard machine learning research without any sensitive or humanrelated data. We do not see any negative impact on the society. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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] Justification: This is theoretical research on an optimization problem defined in the Wasserstein space. We do not see any societal impact that our work may bring about.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 negative 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
conclusion: 
 We propose a new semi FB Euler scheme as a discretization of Wasserstein gradient flow and show that it has favourably theoretical guarantees that the commonly used FB Euler does not yet have if the objective function is not convex along generalized geodesics. Our theoretical analysis opens up interesting avenues for future work. Given the ubiquity of nonconvexity, we hope that the idea can be reused in various contexts, such as with different optimal transport cost functions, different base spaces in the Wasserstein space [42] , or submanifolds of the Wasserstein space (e.g., Bures-Wasserstein [42, 24] ). 
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 (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: the paper poses no such risks. It is a theoretical paper.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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: [Yes]Justification: We explicitly mention in Sect. 5 that we use the deep learning approach proposed in [53] to compute the JKO. We also acknowledge that we leverage the code from [53] in Section C and in README.md in the GitHub repo; the license information is explicitly mentioned in Section C.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • 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. 
 
 Appendix: Lemma 2 (Transfer lemma). [3, Sect. 1] Let T : R m → R n be a measurable map, and µ ∈ P(R m ), then T # µ ∈ P(R n ) and f (y)d(T # µ)(y) = (f • T )(x)dµ(x) for every measurable function f : R n → R, where the above identity has to be understood that: one of the integrals exits (potentially ±∞) iff the other one exists, and in such a case they are equal. Consequently, for a bounded function f , the above integrals exist as real numbers that are equal. Lemma 3. [4, Rmk. 6.2.11] Let µ, ν ∈ P 2,abs (X), then T µ ν • T ν µ = I µ-a.e. and T ν µ • T µ ν = I ν-a.e. Theorem 6 (Characterization of Fréchet subdifferential for geodesically convex functions). [4, Section 10.1] Suppose ϕ : P 2 (X) → R ∪ {+∞} is proper, l.s.c, convex on geodesics. Let µ ∈ dom(∂ϕ) ∩ P 2,abs (X), then a vector ξ ∈ L 2 (X, X, µ) belongs to the Fréchet subdifferential of ϕ at µ if and only ifϕ(ν) -ϕ(µ) ≥ X ⟨ξ(x), T ν µ (x) -x⟩dµ(x) ∀ν ∈ dom(ϕ).Lemma 4. Let H : X → R be a convex function having quadratic growth and ξ be a measurable selector of ∂H, i.e., ξ(x) ∈ ∂H(x) for all x ∈ X. Then, for all µ ∈ P 2 (X),EQUATIONIn other words, ξ ∈ L 2 (X, X, µ) for all µ ∈ P 2 (X).Proof. Since ξ(x) ∈ ∂H(x), by tangent inequality for convex functions, we haveH(y) ≥ H(x) + ⟨ξ(x), y -x⟩, ∀y ∈ X.By picking y = x + ϵξ(x) for some ϵ > 0, we getEQUATIONSince H has quadratic growth, for some a > 0,H(x + ϵξ(x)) -H(x) ≤ |H(x + ϵξ(x))| + |H(x)| ≤ a ∥x + ϵξ(x)∥ 2 + 1 + a(∥x∥ 2 + 1) ≤ 2a + a∥x∥ 2 + a(∥x∥ + ϵ∥ξ(x)∥) 2 ≤ 2a + a∥x∥ 2 + 2a(∥x∥ 2 + ϵ 2 ∥ξ(x)∥ 2 ) = 2a + 3a∥x∥ 2 + 2aϵ 2 ∥ξ(x)∥ 2 .Combining with (7) , it holdsϵ(1 -2aϵ)∥ξ(x)∥ 2 ≤ 2a + 3a∥x∥ 2 .By choosing 0 < ϵ < 1/(2a), we obtain∥ξ(x)∥ 2 ≤ 2a ϵ(1 -2aϵ) + 3a ϵ(1 -2aϵ) ∥x∥ 2 .Therefore, ∥ξ(x)∥ 2 has quadratic growth and -as a consequence -(6) holds for any µ ∈ P 2 (X).Lemma 5. Let ϕ : P 2 (X) → R ∪ {+∞} be a proper function. Let µ * be a local minimizer of ϕ, then µ * is a Fréchet stationary point of ϕ.Proof. There exists r > 0 such that ϕ(µ * ) ≤ ϕ(µ) for all µ ∈ P 2 (X) : W 2 (µ, µ * ) < r. It follows that lim inf µ Wass --→µ * ϕ(µ) -ϕ(µ * ) W 2 (µ, µ * ) ≥ 0, so 0 ∈ ∂ - F ϕ(µ * ),+ E G )(µ * ) ∩ ∂E H (µ * ) ̸ = ∅.Proof. Since µ * is a local minimizer of F, there exists r > 0 such thatEQUATIONLet ξ be a measurable selector of ∂H. Thanks to Lemma 4, ξ ∈ L 2 (X, X, µ * ). According to [5, Prop. 4 .13], ξ ∈ ∂E H (µ * ). It follows from Thm. 6 thatEQUATIONFrom ( 8) and ( 9), for µ ∈ B(µ * , r),EQUATIONProof. Item a. is trivial from the definition of Fréchet subdifferential. For item b., from item a., we first see that,EQUATIONOn the other hand, we apply item a. for U + V and -V to obtainEQUATIONFrom ( 11) and ( 12), we derive (10) . Lemma 10. [65] Let G : P 2 (X) → R ∪ {+∞} be proper and l.s.c. Suppose that G is convex along generalized geodesics. Let ν ∈ P 2,abs (X), µ, π ∈ P 2 (X). If ξ ∈ ∂G(µ), then [19] .X ⟨ξ • T µ ν (x), T π ν (x) -T µ ν (x)⟩dν(x) ≤ G(π) -G(µ). X ∞ such that ∇ψ * | B : B → X ∞ isClaim [19] :EQUATIONA sketch of the claim proof goes as follows. For the first equality in ( 15), first we haveσ(τ ∪ {∞}) ⊂ σ(τ ∞ ) because (1) τ ⊂ τ ∞ and (2) {∞} = X ∞ \ X ∈ σ(τ ∞ ) as X, X ∞ ∈ τ ∞ . On the other hand, σ(τ ∞ ) ⊂ σ(τ ∪ {∞}) because, again, of the construction of τ ∞ : let U ∈ τ ∞ , if U ∈ τ then U ∈ σ(τ ∪ {∞}), otherwise U = (X \ S) ∪ {∞} for some compact set S ⊂ X. As X = R d , S is closed, so (X \ S) ∈ τ implying U ∈ σ(τ ∪ {∞}).For the second equality in (15) , it is straight forward to verify thatG := σ(τ ) ∪ {V ∪ {∞} : V ∈ σ(τ )} is a sigma-algebra in X ∞ . Since τ ∪ {∞} ⊂ G , it holds σ(τ ∪ {∞}) ⊂ G . Conversely, as σ(τ ) ⊂ σ(τ ∪ {∞}) and -consequently - V ∪ {∞} ∈ σ(τ ∪ {∞}) for all V ∈ σ(τ ), it holds G ⊂ σ(τ ∪ {∞}).We conclude that (∇ψ* | B ) -1 : X → X is a Borel (w.r.t. τ ) selector of ∂ψ. As a consequence, S := (∇ψ * | B ) -1 -I is a Borel measurable selector of ∂H.The notion of Maximum Mean Discrepancy (MMD) between two distributions is introduced in [34] . Given a kernel k : X × X → R and denote by K its reproducing kernel Hilbert space. Given two distributions µ and ν, the maximum mean discrepancy between them are defined asD MML (µ, ν) = ∥f µ,ν ∥ K , where f µ,ν (z) = k(x, z)dν(x) -k(x, z)dµ(x), ∀z ∈ X,where f µ,ν is called the witness function. Given some target (or optimal) distribution µ * , we seek to optimize F(µ) = 1 2 D MML (µ, µ * ) 2 . Note that F(µ) admits the following free-energy expression [7] F(µ) = F (x)dµ(x) + 1 2 k(x, x ′ )dµ(x)dµ(x ′ ) + CwhereF (x) = -k(x, x ′ )dµ * (x ′ ), C = 1 2 k(x, x ′ )dµ * (x)dµ * (x ′ ).In general, F is not convex along generalized geodesics. Rather, it exhibits some weakly convex structure [7] that falls within the DC spectrum as detailed in Subsection A.2.3.See also [7, 65] . We include the discussion here for completeness.Consider a one-hidden-layer neural networkf (x) = 1 n n i=1 ψ(x, z i )where ψ(x, z i ) = w i σ(⟨x, θ i ⟩), z i = (θ i , w i ) ∈ Z is the parameters (in and out) associated with i-th hidden neuron, n is the number of hidden neurons, σ is the activation function.When the number of hidden neurons n tends to infinity, f can be written as f (x) = ψ(x, z)dµ(z) for some distribution µ over the parameter space Z. Given data (x, y) ∼ p(x, y) and assuming quadratic loss, the optimization problem reads min µ∈P2(Z)L(µ) := E (x,y)∼p(x,y) y -ψ(x, z)dµ(z).We assume that the model is well-specified, i.e., there exists µ * ∈ P 2 (Z) such thatEQUATIONExpanding L, we getL(µ) = -2 E (x,y)∼p(x,y) (yψ(x, z))dµ(z) + E x∼p(x) (ψ(x, z)ψ(x, z ′ ))dµ(z)dµ(z ′ ) + Cwhere C does not depend on µ.Naturally, we define the kernelk(z, z ′ ) = E x∼p(x) (ψ(x, z)ψ(x, z ′ )) ∀z, z ′ ∈ Z.Using the assumption that the model is well-specified ( 16), we obtaink(z, z ′ )dµ * (z ′ ) = E x∼p(x) (ψ(x, z)ψ(x, z ′ ))dµ * (z ′ ) = E x∼p(x) ψ(x, z)ψ(x, z ′ )dµ * (z ′ ) = E x∼p(x) ψ(x, z) ψ(x, z ′ )dµ * (z ′ ) = E x∼p(x) ψ(x, z)E y∼p(y|x) (y)= E (x,y)∼p(x,y) (yψ(x, z)) .Therefore, L can be written asL(µ) = -2 k(z, z ′ )dµ * (z ′ ) dµ(z) + k(z, z ′ )dµ(z)dµ(z ′ ) + C.This is exactly the MMD setting in Sect. A.2.1.Let k be a kernel whose gradient is Lipschitz continuous, i.e., for some L > 0,∥∇k(x, y) -∇k(x ′ , y ′ )∥ ≤ L ∥x -x ′ ∥ 2 + ∥y -y ′ ∥ 2 1 2 , ∀x, x ′ , y, y ′ ∈ X,which can be expressed equivalently as∥∇ x k(x, y) -∇ x k(x ′ , y ′ )∥ 2 + ∥∇ y k(x, y) -∇ y k(x ′ , y ′ )∥ 2 ≤ L 2 ∥x -x ′ ∥ 2 + ∥y -y ′ ∥ 2 .Let µ * be some target distribution and consider the free-energy functionalF(µ) = k(x, y)dµ(x)dµ(y) -2 k(x, y)dµ * (y)dµ(x).Let α > 0, we can rewrite F as followsF(µ) = α∥x∥ 2 + α∥y∥ 2 + k(x, y) dµ(x)dµ(y) -2 α∥x∥ 2 + k(x, y)dµ * (y) dµ(x).As k is Lipschitz smooth w.r.t. (x, y), x → k(x, y)dµ * (y) is also Lipschitz smooth. Indeed,∇ x k(x, y)dµ * (y) -∇ x k(x ′ , y)dµ * (y) = (∇ x k(x, y) -∇ x k(x ′ , y))dµ * (y) ≤ ∥∇ x k(x, y) -∇ x k(x ′ , y)∥ 2 dµ * (y) 1 2 ≤ L ∥x -x ′ ∥ 2 dµ * (y) 1 2 = L∥x -x ′ ∥.Next, as a standard result, if f is an L-smooth function, x → (α/2)∥x∥ 2 ± f (x) are convex whenever α ≥ L. Therefore, for α ≥ L, W (x, y) := α∥x∥ 2 + α∥y∥ 2 + k(x, y) is convex and F (x) = -2 α∥x∥ 2 + X k(x, y)dµ * (y) is concave. From [4, Prop. 9.3.5], the interaction energy corresponding to W is generalized geodesically convex.Since I + ηS is a subgradient selector of a convex function, the optimal transport between µ n and ν n+1 is given byEQUATIONand between µ n+1 and ν n+1 [4, Lem. 10.1.2]EQUATIONSince E H is convex along generalized geodesics [4, Prop. 9.3.2] and S is a subgradient of E H at µ n [5, Proposition 4.13], by Lem. 10 it holds, for any ν ∈ P 2,abs (X),E H (µ n+1 ) ≥ E H (µ n ) + X ⟨S • T µn ν (x), T µn+1 ν (x) -T µn ν (x)⟩dν(x).By choosing ν = ν n+1 (note that ν n+1 ∈ P 2,abs (X)),EQUATIONwhere the second equality uses ( 17) and the last one uses Lem. 3.On the other hand, since E G + H is convex along generalized geodesics, by applying Lem. 10 forEQUATIONwhere the last equality uses Lem. 3.By adding (20) and (19) side by side,EQUATIONA.4 Proof of Theorem 1Let µ * ∈ P 2 (X) be a cluster point of {µ n } n∈N . There exists a subsequence µ n k--→ µ * . It holdslim inf k→∞ F(µ n k ) = lim inf k→∞ (H (µ n k ) + E F (µ n k )) = lim inf k→∞ H (µ n k ) + E F (µ * ) ≥ H (µ * ) + E F (µ * ),since H is l.s.c. and E F is continuous w.r.t. Wasserstein topology. Therefore, H (µ * ) < +∞, which further implies that µ * ∈ P 2,abs (X).We haveEQUATIONWe observe thatT µn+1 νn+1 • T νn+1 µnis a (possibly non-optimal) transport pushing µ n to µ n+1 , by the optimality of T µn+1 µnEQUATION)By Lem. 1 and ( 22), (23) ,EQUATIONNote that F is bounded below (Assumption 1), telescoping (24) gives usEQUATIONIn particular, W 2 (µ n , µ n+1 ) → 0. This together withµ n k Wass --→ µ * implies µ n k +1 Wass --→ µ * .Under Assumption 3, S = ∇H and S is continuous. Next, recall ν n k +1 = (I + ηS) # µ n k , we showEQUATIONThus, let f be a continuous and bounded test functional in X, by using transfer lemma 2,EQUATIONEQUATIONsince S is continuous. So ν n k +1 narrow ---→ ν * . We go one step further and prove that ν n k +1 actually converges to ν * in the Wasserstein metric. This boils down to showing convergence in second-order moments, i.e.,m 2 (ν n k +1 ) → m 2 (ν * ), which is equivalent to showing that X ∥x + ηS(x)∥ 2 dµ n k (x) → X ∥x + ηS(x)∥ 2 dµ * (x).On the other hand, ψ(x) := ∥x+ηS(x)∥ 2 has quadratic growth (follows from Lem. 4) and µ n k → µ * in the Wasserstein metric, so [3, Prop. 2.4 ]lim k→∞ X ∥x + ηS(x)∥ 2 dµ n k (x) = X ∥x + ηS(x)∥ 2 dµ * (x) Therefore, ν n k +1 → ν * in Wasserstein metric.To proceed further, we need the following theorem stating that the graph of the subdifferential of a geodesically convex function is closed under the product of Wasserstein and weak topologies. Theorem 8 (Closedness of subdifferential graph). [4, Lemma 10.1.3] Let ϕ be a geodesically convex functional satisfying dom(∂ϕ) ⊂ P 2,abs (X). Let {µ n } n∈N be a sequence converging in Wasserstein metric to µ ∈ dom(ϕ). Let ξ n ∈ ∂ϕ(µ n ) be satisfyingEQUATION)and converging weakly to ξ ∈ L 2 (X, X, µ) in the following sense:EQUATION)Then ξ ∈ ∂ϕ(µ).As a side note, we need the notion of weak convergence in the above theorem because -unlike subdifferentials in flat Euclidean space -each ξ n lives in its own L 2 (X, X, µ n ) space.Back to our proof, for item (29) , we show thatEQUATIONWe proceed as follows to prove (31) . We first show that sup n∈N m 2 (µ n ) < +∞. By contradiction, by assuming sup n∈N m 2 (µ n ) = +∞, we can extract a subsequence {µ n k } k∈N such thatEQUATION)By compactness assumption, there further exists a subsequence {µ n k i } i∈N such that µ n k i converges (in Wasserstein metric) to some µ * * ∈ P 2 (X), implying that lim i→∞ m 2 (µ n k i ) = m 2 (µ * * ), which contradicts (32) . Therefore, sup n∈N m 2 (µ n ) < +∞. We next show that sup n∈N m 2 (ν n ) < +∞. Indeed, as ∥S∥ 2 has quadratic growth,∥S(x)∥ 2 ≤ c(∥x∥ 2 + 1)for some c > 0. So m 2 (ν n+1 ) = X ∥x∥ 2 dν n+1 (x) = X ∥x∥ 2 d(I + ηS) # µ n (x) = X ∥x + ηS(x)∥ 2 dµ n (x) ≤ 2 X ∥x∥ 2 dµ n (x) + 2η 2 X ∥S(x)∥ 2 dµ n (x) ≤ (2 + 2cη 2 )m 2 (µ n ) + 2η 2 c, implying that sup n∈N m 2 (ν n ) < +∞. This in conjunction with G having quadratic growth im- plies that sup n∈N |E G (ν n )| < +∞. Furthermore, inf n (E G + H )(µ n ) > -∞EQUATION)The finiteness as in (31) then follows fromsup n∈N |E G (ν n )| < +∞, sup n∈N -(E G + H )(µ n ) <+∞ as proved and sup n∈N H (ν n ) < +∞ as assumed.We next prove that there is a subsequence of {µ n k } k∈N such thatEQUATION)We consider the sequence of optimal plans between µ n and ν n as followsρ n = (I, T νn µn ) # µ n , ∀n ∈ N.We observe that m 2 (ρ n ) = X×X (∥x∥ 2 + ∥y∥ 2 )dρ n (x, y) = X ∥x∥ 2 dµ n (x) + X ∥y∥ 2 dν n (y) < +∞, so ρ n ∈ P 2 (X ×X) for all n ∈ N. Since sup n∈N m 2 (ν n ) < +∞EQUATION)We can see that proj 1# ρ * = µ * , proj 2# ρ * = ν * . We further show that ρ * ∈ P 2 (X × X), or equivalently,EQUATIONLet C ∈ N, thanks to the narrow convergence in (35) , we haveX×X min{∥x∥ 2 + ∥y∥ 2 , C}dρ n k i +1 (x, y) → X×X min{∥x∥ 2 + ∥y∥ 2 , C}dρ * (x, y).Furthermore,X×X min{∥x∥ 2 + ∥y∥ 2 , C}dρ n k i +1 (x, y) ≤ sup n∈N m 2 (µ n ) + sup n∈N m 2 (ν n ) := M < +∞.Passing to the limit, we getX×X min{∥x∥ 2 + ∥y∥ 2 , C}dρ * (x, y) ≤ M for all C ∈ N.Sending C to ∞ and applying Monotone Convergence Theorem we derive (36) .Back to the main proof, since {ρ n k i +1 } i∈N is a sequence of optimal plans, its limit, ρ * is also optimal [3, Proposition 2.5]. Therefore,ρ * = (I, T ν * µ * ) # µ * .Moreover, asm 2 (ρ n k i +1 ) = m 2 (µ n k i +1 ) + m 2 (ν n k i +1 ) → m 2 (µ * ) + m 2 (ν * ) = m 2 (ρ * ),we haveEQUATIONIndeed, (x, y) → ζ(x) proj i (y) where proj i is the projection into the i-th coordinate is continuous and has quadratic growth since ζ(x) is bounded and proj i (y) is linear.Since 37) holds. Consequently, (34) also holds by noticing thatEQUATIONX xζ(x)dµ n k j +1 (x) → X xζ(x)dµ * (x) as j → ∞.By the closedness of subdifferential graph of ∂(E G + H ) (Thm. 8), we obtainEQUATIONOn the other hand, by the definition of ν * in (26), we get T ν * µ * = I + ηS. Together with (38) ,S ∈ ∂(E G + H )(µ * ). Noting that S ∈ ∂E H (µ * ), it holds S ∈ ∂(E G + H )(µ * ) ∩ ∂E H (µ *) and we conclude µ * is a critical point of F.We see that∥G η (µ n )∥ 2 L 2 (X,X;µn) = 1 η 2 X x -T JKO η(E G +H ) ((I+ηS) # µn) µn (x) 2 dµ n (x) = 1 η 2 W 2 2 (µ n , JKO η(E G +H ) ((I + ηS) # µ n )) = 1 η 2 W 2 2 (µ n , µ n+1 ).On the other hand, it follows from (25) of the proof of Thm. 1 that min i=1,N W 2 2 (µ n , µ n+1 ) = O(N -1EQUATION)We then have the following evaluations:EQUATION.By transfer lemma 2,EQUATIONOn the other hand, by using the trivial identityEQUATIONwhere the last equality uses (I + η∇H) • T µn νn+1 = I ν n+1 -a.e. The Hessian of H is bounded uniformly, ∇H is Lipschitz, let's say ∥∇H(x)-∇H(y)∥ ≤ L H ∥x-y∥ for all x, y ∈ X. We continue evaluating (41) as followsEQUATIONwhere the third equality uses (42) .From ( 40), (41) , and (43), we deriveEQUATION.On the other hand, by telescoping Lem. 1, we obtain∞ n=1 X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x) < +∞.Therefore,N -1 n=0 dist (0, ∂ F F(µ n+1 )) ≤ 2(η 2 L 2 H + 1) η N -1 n=0 X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x) 1 2 ≤ 2(η 2 L 2 H + 1) η N N -1 n=0 X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x) 1 2 ≤ 2(η 2 L 2 H + 1)N η +∞ n=0 X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x) 1 2We derivemin n=1,N dist (0, ∂ F F(µ n )) = O 1 √ N .A.7 Proof of Theorem 4Since H ∈ C 2 (X) whose Hessian is uniformly bounded, recall from (39) that∂ - F F(µ n+1 ) = ∂(E G + H )(µ n+1 ) -∇H ∋ T νn+1 µn+1 -I η -∇H.Since F(µ 0 )-F * < r 0 and the sequence {F(µ n )} n∈N is not increasing (Lem. 1), F(µ n )-F * < r 0 for all n ∈ N. Łojasiewicz condition impliesEQUATIONwhere the last inequality follows from ( 41) and ( 43) and L H is the Lipschitz constant of ∇H.Combining with Lem. 1, we derivec (F(µ n+1 ) -F * ) θ ≤ 2(η 2 L 2 H + 1) √ η (F(µ n ) -F(µ n+1 )) 1 2orEQUATIONWe then use the following lemma [71, Lem. 4 ].Lemma 11. Let {s k } k∈N be a nonincreasing and nonnegative real sequence. Assume that there exist α ≥ 0 and β > 0 such that for all sufficiently large k,EQUATIONThen (i) if α = 0, the sequence {s k } k∈N converges to 0 in a finite number of steps;(ii) if α ∈ (0, 1], the sequence {s k } k∈N converges linearly to 0 with rate β β+1 ;(iii) if α > 1, the sequence {s k } k∈N converges sublinearly to 0, i.e., there exists τ > 0:s k ≤ τ k -1 α-1for sufficiently large k.Compared to [71, Lem. 4], we have dropped the assumption s k → 0 in Lem. 11 because this assumption is vacuous, i.e., it can be induced by (47) and nonnegativity of {s k } k∈N .We now apply Lem. 11 for s k = F(µ k ) -F * using ( 46) to derive the followings(i) if θ = 0, F(µ n ) -F * converges to 0 in a finite number of steps; (ii) if θ ∈ (0, 1/2], F(µ n ) -F * converges to 0 linearly (exponentially fast) with rate F(µ n ) -F * = O M M + 1 n where M = 2(η 2 L 2 H + 1) c 2 η ; (iii) if θ ∈ (1/2, 1), F(µ n ) -F * converges sublinearly to 0, i.e., F(µ n ) -F * = O n -1 2θ-1 .A.8 Proof of Theorem 5Cauchy sequence.By replacing n := n -1 in (45) and rearrangingEQUATIONIt follows from Lem. 1 and (48) thatEQUATIONSince the function s :R + → R, s(t) = t 1-θ is concave if θ ∈ [0, 1), tangent inequality holds s ′ (a)(a -b) ≤ s(a) -s(b).Note that s ′ (t) = (1 -θ)t -θ , the above inequality further impliesEQUATIONFrom ( 49) and ( 50)X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x) ≤ 2(η 2 L 2 H + 1) (1 -θ)c X ∥T µn-1 νn (x) -T µn νn (x)∥ 2 dν n (x) 1 2 × (F(µ n ) -F * ) 1-θ -(F(µ n+1 ) -F * ) 1-θor equivalently,EQUATIONwhere r n := X ∥T µn νn+1 (x) -T µn+1 νn+1 (x)∥ 2 dν n+1 (x). By telescoping (51) from n = 1 to +∞ we obtain+∞ n=1 r n √ r n-1 < +∞.On the other handEQUATIONwe derive 22), ( 23) of the proof of Thm. 1,+∞ n=0 √ r n < +∞. From (r n ≥ W 2 2 (µ n , µ n+1 ), we obtain +∞ n=0 W 2 (µ n , µ n+1 ) < +∞.or, in other words, {µ n } n∈N is a Cauchy sequence under Wasserstein topology. The Wasserstein space (P 2 (X), W 2 ) is complete [5, Thm. 2.2], every Cauchy sequence is convergent, i.e., there existsµ * ∈ P 2 (X) such that µ n Wass --→ µ * .We prove that µ * is actually an optimal solution of F by showing F(µ * ) = F * . Indeed, firstly, as G and H have quadratic growth, it holdsE G (µ n ) → E G (µ * ), E H (µ n ) → E H (µ * ).On the other hand,F * = lim n→∞ F(µ n ) = lim inf n→∞ F(µ n ) = lim inf n→∞ H (µ n ) + E G (µ * ) -E H (µ * ) ≥ H (µ * ) + E G (µ * ) -E H (µ * ) = F(µ * )since H is l.s.c. The equality has to occur, i.e., F * = F(µ * ), due to the optimality of F * .n } n∈N . (i) If θ = 0From item (i) of Thm. 4, there exists n 0 ∈ N such that F(µ n ) = F * for all n ≥ n 0 . It then follows from (24) that µ n0 = µ n0+1 = µ n0+2 = . . ., which further implies that µ n = µ * for all n ≥ n 0 .EQUATION)where the last inequality uses triangle inequalityW 2 (µ i , µ * ) ≤ N -1 n=i W 2 (µ n , µ n+1 ) + W 2 (µ N , µ * )and lets N → ∞ with a notice that µ N Wass --→ µ * .From ( 51) and ( 52),EQUATIONTelescope (54) for n = i to +∞,EQUATIONwhere the last inequality uses (45) . Since r i → 0 as i → ∞, r i < 1 for i sufficiently large. It follows from ( 55) that: for i sufficiently larges i ≤ M √ r i-1 = M (s i-1 -s i )whereM = 1 + (2(η 2 L 2 H + 1)) 1 2θ(1 -θ)ηEQUATION)Rewriting as ) implies: for all i sufficiently large,EQUATIONs i ≤ M r 1-θ 2θ i-1 = M (s i-1 -s i ) 1-θ θwhere M is the same as in (56) .Applying Lem. 11(iii),s i = O i -1-θ 2θ-1, which implies (by ( 53)) W 2 (µ i , µ * ) = O i -1-θ 2θ-1.Following a similar approach to the semi FB Euler's implementation outlined in Alg. 1, we present a practical implementation of FB Euler for the sampling context in Alg. 2.Input: Initial measure µ 0 ∈ P 2,abs (X), discretization stepsize η > 0, number of steps K > 0, batch size B. for k = 1 to K do for i = 1, 2, . . . do Draw a batch of samples Z ∼ µ 0 of size B;Ξ ← (I -η∇F ) • ∇ x ψ θ k • (I -η∇F ) • ∇ x ψ θ k-1 • . . . • (I -η∇F )(Z); W 2 2 ← 1 B ξ∈Ξ ∥∇ x ψ θ (ξ) -ξ∥ 2 ; ∆H ← -1 B ξ∈Ξ log det ∇ 2 x ψ θ (ξ). L ← 1 2η W 2 2 + ∆H .Apply an optimization step (e.g., Adam) over θ using ∇ θ L. end for θ k+1 ← θ. end forWe perform numerical experiments in a high-performance computing cluster with GPU support. We use Python version 3.8.0. We allocate 8G memory for the experiments. The total running time for all experiments is a couple of hours. Our implementation is based on the code of [53] (MIT license) with the DenseICNN architecture [41] .Consider a target Gaussian mixture of the following form:π(x) ∝ exp(-F (x)) := K i=1 π i exp - ∥x -x i ∥ 2 σ 2 .We writeF (x) = -log K i=1 π i exp - ∥x -x i ∥ 2 σ 2 = -log K i=1 π i exp - ∥x∥ 2 + ∥x i ∥ 2 -2⟨x, x i ⟩ σ 2 = -log K i=1 π i exp - ∥x∥ 2 σ 2 × exp - ∥x i ∥ 2 -2⟨x, x i ⟩ σ 2 = ∥x∥ 2 σ 2 -log K i=1 π i exp - ∥x i ∥ 2 -2⟨x, x i ⟩ σ 2 convexwhich is DC. Note that the convexity of the second component is thanks to (a) log-sum-exp is convex and (b) the composite of a convex function and an affine function is convex.Experiment details We set K = 5 and randomly generate x 1 , x 2 , . . . , x 5 ∈ R 2 . We set σ = 1.The initial distribution is µ 0 = N (0, 16I). We use η = 0.1 for both FB Euler and semi FB Euler. We train both algorithms for 40 iterations using Adam optimizer with a batch size of 512 in which the first 20 iterations use a learning rate of 5 × 10 -3 while the latter 20 iterations use 2 × 10 -3 . For the baseline ULA, we run 10000 chains in parallel for 4000 iterations with a learning rate of 10 -3 .We run the above experiment 5 times where x 1 , x 2 , . . . , x 5 are randomly generated each time. Let π be the original prior, Θ be the constraint set that we want to impose, and the distance-to-set prior [61] is defined by, for some ρ > 0,π(θ) ∝ π(θ) exp - ρ 2 d(θ, Θ) 2 ,that penalize exponentially θ deviating from the constraint set.Given data y, using the this distance-to-set prior, the posterior readsπ(θ|y) ∝ L(θ|y)π(θ) exp - ρ 2 d(θ, Θ) 2 .where L(θ|y) is the likelihood.The structure of π(θ|y) depends on three separate components: the original prior, the likelihood, and the constraint set Θ. In the ideal case, π(θ) and L(θ|y) are given in nice forms (e.g., log-concave), and Θ is a convex set. As a fact, if Θ is a convex set, θ → d(θ, Θ) 2 is convex, making the whole posterior log-concave. If Θ is additionally closed, d(θ, Θ) 2 is L-smooth.However, whenever Θ is nonconvex, the function θ → d(θ, Θ) 2 is not continuously differentiable. This is induced by the Motzkin-Bunt theorem [16, Thm. 9.2.5] asserting that any Chebyshev set (a set S ⊂ X is called Chebyshev if every point in X has a unique nearest point in S) has to be closed and convex.On the other hand, d(θ, Θ) 2 is always DC regardless the geometric structure of Θ:EQUATIONNote that the supremum of an arbitrary family of affine functions is convex.Therefore, the log-DC structure of the whole posterior only depends on whether the original prior and the likelihood are log-DC, which is likely to be the case.Distance-to-set prior relaxed von Mises-Fisher In directional statistics, the von Mises-Fisher distribution is a distribution over unit-length vectors (unit sphere). It can be described as a restriction of a Gaussian distribution in a sphere. By using the distance-to-set prior, we can relax the spherical constraint asπ(θ) ∝ exp(-F (θ)) := exp -κ (θ -µ) ⊤ (θ -µ) 2 × exp - ρ 2 dist(θ, S) 2where S denote the unit sphere in some R d space. By the DC structure (57) of the distance function, F is DC with the following compositionF (θ) = κ ∥θ -µ∥ 2 2 + ρ 2 ∥θ∥ 2 - ρ 2 sup x∈S -∥x∥ 2 + 2⟨x, θ⟩ := G(θ) -H(θ).We also note that H is not continuously differentiable because S is nonconvex. Furthermore, ρ proj S (θ) ∈ ∂H(θ) where proj S (θ) is the projection of θ onto S, which can be computed explicitly in this case.Experiment details We consider a unit circle in R 2 with centre µ = (1, 1.5). We set κ = 1 and ρ = 100. The initial distribution is µ 0 = N (0, 16I). We use η = 0.1 for both FB Euler and semi FB Euler. We train both algorithms for 40 iterations using Adam optimizer with a batch size of 512 in which the first 20 iterations use a learning rate of 5 × 10 -3 while the latter 20 iterations use 2 × 10 -3 .NeurIPS Paper Checklist38th Conference on Neural Information Processing Systems (NeurIPS 2024).Our code is available at https://github.com/MCS-hub/OW24 


Decision: The paper takes a stab at DC optimization in the Wasserstein space and proposes a semi Forward-Backard Euler scheme (section 3, equation 4). All the reviewers like the work especially as it offers a fresh perspective on the problem. Hence, we have decided to accept the work. Please incorporate all the asked changes. 

One reviewer had comments on whether the developments are complete and suggested some explorations. Although it would not be possible to accommodate such detailed explorations in the final version, I would still encourage the authors to provide a short commentary on the future directions / limitations.

Reviews: review: paper_summary: paper_summary: This paper studies minimization algorithms for functionals on the Wasserstein space with the following difference of convex functions (DC) structure:
$$
\min_{\mu \in \mathcal{P}_2(X)} \mathcal{F}(\mu) := \int (G(x) - H(x)) d\mu(x) + \mathcal{H}(\mu),
$$ where $G$ and $H$ are convex functions and $\mathcal{H}$ is convex along (generalized) geodesics in the Wasserstein space. The main motivation for this problem is that it is the sampling analog to difference of convex functions problems in optimization, giving a problem of intermediate difficulty between convex and true non-convex. The authors also mention an application to maximum mean discrepancy.

They study a modified forward Euler sceme which they refer to as "semi FB Euler", where given a current iterate $\mu_n$ they update as
$$
\mu_{n + 1/2} := (I + \gamma \nabla H) \mu_n,
$$ and then
$$
\mu_{n + 1} := JKO_{\gamma(H+ E_G)}(\mu_{n + 1/2}).
$$ A similar scheme was considered in [1,2] but with the difference that the function $G$ was included in the first step but not in the second; the motivation for this difference is that it more closely aligns with DC programming. One benefit of this modification is that $\mu_{n + 1/2}$ is pushed forward by an optimal transport map, which makes the problem more tractable. 

The main focus of the work are various convergence results for the above scheme, including an asymptotic result as well as results on the norm of a gradient analog and the distance between $0$ and the sub-differential. The authors also consider convergence under a gradient domination condition, and conclude with some numerics. 

[1] Wibisono, Andre. "Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem." Conference on Learning Theory. PMLR, 2018.

[2] Salim, Adil, Anna Korba, and Giulia Luise. "The Wasserstein proximal gradient algorithm." Advances in Neural Information Processing Systems 33 (2020): 12356-12366.
strength_weakness: strength_weakness: - The Łojasiewicz condition is not verified to hold in any settings other than their the standard case of KL minimization.

- The problem is somewhat far from practice, firstly because the JKO step is very challenging to implement in practice, and secondly because the setting the authors study is only motivated by one application, to maximum mean discrepancy functionals. They do mention a connection to two-layer neural networks in the "Context" section on page 2 but this was impossible to understand from the little they wrote.

- The proofs don't yield explict constants in their results, which is important because there could be hidden dimension-dependence.
questions: questions: - The thorough introduction to Wasserstein space is much appreciated, but perhaps it could be largely moved to an appendix so that more space is left in the main text for more discussion of applications and high-level ideas of the proofs.

- Please give more space to the connection to two-layer neural networks mentioned in the "Context" section on page 2.

- Since it is a general fact (see, e.g. [1]) that a gradient domination condition implies that $d(x, x*)^2 <= C(F(x) - F(x*))$, I feel that Theorem 5 should either be removed, or stated as a corollary in an appendix and just mentioned as a brief remark in the main text. 

[1] Otto, Felix, and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality." Journal of Functional Analysis 173.2 (2000): 361-400.
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_summary: This work focuses on the optimization in the Wasserstein space of a functional which is a sum of an internal energy (convex along a generalized geodesic) and of a potential energy, whose potential is a difference of convex functions. To solve such problem, the authors propose to generalize the semi Forward-Backard Euler scheme to the Wasserstein space, and study it theoretically, showing the convergence in general settings. Finally, they demonstrate the convergence of the algorithm on two toy objectives.
strength_weakness: strength_weakness: This is mainly a theoretical work. Thus, the following weaknesses on the experiment section and on the method are minor in my opinion.
- As underlined in the experiment section, the scheme involves to compute the JKO operator, which is computationally costly as it requires neural networks to solve it
- The experiments focus on 2D toy examples
questions: questions: The focus of the paper is on functionals obtained as a sum of a potential and an internal energy. Could the analysis be done with more general DC functionals (e.g. with interaction energies.). I guess yes since the MMD is presented in the introduction as being such a functional. It would also have been nice to test the semi Forward-Backward scheme on the MMD, as it is not geodesically convex, and does not converge well without tricks such as adding noise (e.g. for Gaussian kernel).
Assumption 2 supposes that the sublevel sets are compact w.r.t the Wasserstein topology. Are there easy examples of functionals which satisfy this (e.g. those used in the experiment section)?

Typos:
- In the abstract, I found the sentence line 3 "When the regularization term is the negative entropy, the optimization problem becomes a sampling problem where it minimizes the Kullback-Leibler divergence between a probability measure (optimization variable) and a target probability measure whose logarithmic probability density is a nonconvex function." weird as it is not specified that the objective function is a potential.
- Line 14: "can be considered gradient descent": lack a word?
Rating: 7
limitations: limitations: The authors have acknowledged the limitations in the main text
paper_summary: paper_summary: This work, of a theoretical nature, considers the problem of minimizing a functional $$\mathcal{F}$$ over the space of probability measures of the form

$$\mathcal{F}(\mu) = \int (G(x) - H(x)) d\mu(x) + \mathcal{H}(\mu)$$

where $G,H$ are **convex** potentials and $\mathcal{H}$ is (typically) the negative entropy (the work actually extends to functionals $\mathcal{H}$ that are convex along generalized geodesics). 

This setting substantially differs from standard optimization in the space of probability measure in that $\mathcal{F}$ is not convex along generalized geodesics, the standard assumption to establish convergence of gradient flows in this spaces, following the seminal work of Ambrosio, Gigli and Savaré. Nonetheless, the "difference of convex functions" (DC) structure enables the derivation of a two-step scheme to minimize $\mathcal{F}$. Namely, given a current iterate $\mu_n$, one first computes $\nu_{n+1}$ using a **forward** (explicit) Euler scheme _using only the concave term_ $-H$, then obtain $\mu_{n+1}$ using a **backward** (implicit) iterate on the convex (along generalized geodesics) term $\mu \mapsto \braket{G,\mu} + \mathcal{H}(\mu)$. 

This approach is referred to as a _semi_ forward-backward (FB) Euler scheme, seen as a variation of the known FB scheme (which would process similarly, but performing an explicit iterate on $G-H$ and an implicit iterate on $\mathcal{H}$). 

In contrast to the standard FB scheme, this semi-FB scheme that leverages the DC structure enables the derivation of strong theoretical guarantees: existence of accumulation points that are critical points of $\mathcal{F}$, and under additional standard assumptions ($\sim$ PL inequality), convergence rate toward the global minimum of $\mathcal{F}$.
strength_weakness: strength_weakness: 1. While I understand that the paper is mostly theoretical, the "from theory to numeric" side may be enhanced (typically, putting some of the material of Part B, in particular Algorithm 2, to the main body). 

2. Still targeting numerical applications, the approach seems somewhat limited in that the numerical experiments, according to the appendix, have been run on a fairly powerful hardware and yet took a couple of hours to run, while they remain at the proof-of-concept level on toy datasets. 

3. Though this does not diminishes the intrinsic quality of the paper, the paper relies on technical proofs deferred to the appendix and I could not proofread all of them. I read the proof of Lemma 1 and Theorem 1, which seem correct to me, and quickly check other proofs but cannot guarantee the correctness of the theorems (though they look sounded)---which are nonetheless the main (if not unique) contributions of the work. I understand that this is not the authors fault, but I wanted to stress that point in my review. 

**Minor aspects:**
- It may be better to use $\tau$ instead of $\gamma$ to denote the step-size in JKO schemes, since $\gamma$ is often used to denote a transport plan. 
- I believe it's better to avoid starting sentences with mathematical symbol. That makes the document a bit easier to read. See for instance line 117 "continuous w.r.t. $\mathcal{L}^d$. $\mu$-a.e. stands..." is a bit confusing.
questions: questions: 1. Genuine question: is there any convergence in the regime $\gamma \to 0$ toward a limit curve $(\mu_t)_t$, and if so does it coincides with the usual Wasserstein gradient flow of $\mathcal{F}$? Said differently, is it true that the (time-continuous) Wasserstein gradient flow of the KL for a log-DC target distribution converges globally? (As a consequence of the fact that the sequence $(\mu_n)_n$ you build do converge globally in that case)
limitations: limitations: NA
paper_summary: paper_summary: This paper introduces and analyzes an optimization scheme, termed "semi-FB Euler", for a particular minimization problem over the Wasserstein space $P_2(R^d)$: $\min_\mu \mathcal{F}(\mu) = \int (G-H) d\mu + \mathcal{H}(\mu)$, where $G$ and $H$ are convex functions over $R^d$ and $\mathcal{H}$ is a functional that is convex over generalized geodesics. Namely, "semi-FB Euler" is defined by equation (4) of the paper: starting from some $\mu_0$,
$$\begin{cases}
\nu\_{n+1} = (I + \gamma \nabla H)\_{\sharp} \mu_n, \\\\
\mu\_{n+1} \in \arg\min\_{P\_2(R^d)} \int G d\mu + \mathcal{H}(\mu) + \frac{1}{2\gamma} W\_2^2(\cdot, \nu\_{n+1}).
\end{cases}$$

(The second step, called a JKO step, is assumed solvable for the theoretical parts of the paper.)
Convergence to critical points of $\mathcal{F}$ is established under mild regularity assumptions, both asymptotically (Theorem 1) and with non-asymptotic rates (Theorems 2, 3). Under a Lojasiewicz inequality assumption, convergence to a minimizer is established with non-asymptotic rates (Theorems 4, 5).

Numerical illustrations with synthetic data are provided, in which an input-convex neural network is used to implement the JKO step.
strength_weakness: strength_weakness: The proposed "semi-FB Euler" corresponds to a forward-backward splitting scheme of $\mathcal{F}$ into $\mathcal{F}_1(\mu) = \int (-H) d\mu$ and $\mathcal{F}_2(\mu) = \int G d\mu + \mathcal{H}(\mu)$ in the Wasserstein optimization geometry. (Incidentally, the name "semi-FB Euler" can be misleading, as it is really a forward-backward scheme.)
So the setting considered in this paper is very similar to that of "The Wasserstein proximal gradient algorithm" ([56] Salim, Korba, Luise, 2020), which corresponds to the case where $H$ is concave and $G=0$. In fact, it seems that the only way that convexity of $G$ comes into play in the proofs, is to ensure convexity of $\mathcal{F}_2$ over generalized geodesics. (Please do correct me if this is incorrect.)

With this in mind, the contributions of this paper are to study what happens when $-H$ is concave in $\mathcal{F}_1(\mu) = \int (-H) d\mu$, as opposed to convex in [56]. It seems that concavity of $-H$ is used for:
1. the existence of a Borel measurable selector $S$ for the subdifferential $\partial H$ (section A.1), but it is not essential to the paper;
2. the descent lemma (Lemma 1);
3. quadratic growth of $\|S\|^2$ under quadratic growth of $H$ (Lemma 4 and proof of Thm 1 line 706);
4. it seems that convexity of $H$ is implicity used in the very last step of the proof of Thm 1 (line 735). (Incidentally, please make the last step of the proof of Thm 1 more explicit.)

In summary, compared to the setting of [56], which corresponds to $-H$ being smooth and convex, this work shows that if $-H$ is concave, then the forward-backward algorithm (with a good choice of the selector $S$ when $H$ is non-differentiable) converges for any choice of step-size $\gamma$.

I find this theoretical contribution interesting, but I am not sure if it is sufficiently significant. Besides, there is also the weakness, acknowledged by the authors (line 362) that the proposed algorithm relies on a JKO step which is computationally expensive. Hence my overall rating for this submission.
questions: questions: - Please address the points in "Weaknesses".
- The items 1.-4. in "Weaknesses can also be guaranteed if, instead of $-H$ concave, one assumes $H$ is smooth (which is Assumption 5) and $\gamma$ is small enough (see [56, Thm 5] for item 2 and [56, Lemma 2] for item 4). Could your results be extended to the case where $H$ is smooth instead of convex? If yes, then this could be another way to show theoretical guarantees for forward-backward schemes outside of the convex case (in the spirit of your conclusion, line 393).
- What are the "interesting avenues for future work" mentioned in the conclusion?
limitations: limitations: The authors adequately addressed the limitations.

Comments: title: Thanks for your response
Thanks
comment: Thank you for the response.

We will incorporate those suggestions into the final version of the paper.
Thank you again for a detailed response.

I stand by the points made in my two previous comments, so I keep my updated rating of 3 (Reject). I would like to reiterate my opinion that the technical core of the work has undeniable value, but the submission would need substantial modifications, as was discussed.
With regard to the Lojasiewicz condition, it is interesting that it can almost be satisfied in the MMD case, but I would emphasize that the reason the log-Sobolev inequality is so important is that it holds uniformly. In particular, the gap between a PL inequality at some points versus a PL inequality everywhere can be very large and difficult to bridge. Also, I still don't find the connection to MMD particularly convincing, so this setting is not completely convincing anyways.

With regard to the fact that gradient domination implies quadratic growth: sure, the case where $(f(x) - f(x_*))^\theta \leqslant C\|\nabla f(x)\|$  is not explicitly proved in [1], but the proof can be quite easily modified to handle that case. At a minimum, you should add citations to [1] and [2]. Note also that using their results you could likely simplify your proofs and make the constants more explicit. 

With regard to your remaining points, thanks for your responses. Overall, my opinion of the work has not substantially changed, so I leave my score as-is.

[1] Otto, Felix, and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality." Journal of Functional Analysis 173.2 (2000): 361-400.
[2] Karimi, Hamed, Julie Nutini, and Mark Schmidt. "Linear convergence of gradient and proximal-gradient methods under the polyak-łojasiewicz condition." Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16. Springer International Publishing, 2016.
We thank the reviewer for the response.


***$\bullet$ It is indeed preferable to clarify from the outset***: we will clarify in the abstract that we consider the DC structure.

***$\bullet$ I would recommend to pick a side clearly***: we would pick the general side that considers all geodesically convex regularizers because (1) we want to cover the case of Maximum Mean Discrepancy that has value in the context of deep learning theory, (2) our work is more about optimization in the Wasserstein space, motivated by sampling but should not be restricted to that area. 

About the implementation, apart from the negative entropy regularizer, we can discuss the following regularizers given in the following forms: (a) the interaction energy, (b) another (non-smooth) potential, and (c) general internal energy. We also note that the negative entropy considered in our paper is (probably) already the trickiest regularizer. It requires computing $E_{x \sim \rho_t}(\log(\rho_t(x)))$. Although samples from $\rho_t$ can be simulated, we cannot use the vanilla idea of Monte Carlo approximation because the function inside the expectation involves the probability density of $\rho_t$, which is unknown. The key technique making this possible is the change of entropy formula, saying that the change in the entropy can be computed via the determinant of the Jacobian of the transformation [Mokrov21] (line 828 in our work). For other types of regularizers, such as interaction or potential (or their summation), the implementation is straightforward (via Monte Carlo) since the expectation no longer involves the probability density. It is also worth noting that: in the case of the potential regularizer, the JKO operator reduces to the proximal operator in the Euclidean space [Wibisono18], and it is known that many proximal operators in the Euclidean space have closed-form solutions [Parikh14]. 

On the other hand, the change of entropy formula can be naturally extended to the case of general internal energy $E_{x \sim \rho_t}(f(\rho_t))$ thanks to the change of variable formula [Lemmas 3.2, 3.4, A-Melis22]. This is pretty much in the same spirit as normalizing flows, another active ML research area.

*[Mokrov21] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J. M., \& Burnaev, E. (2021). Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34, 15243-15256.*

*[A-Melis22] Alvarez-Melis, D., Schiff, Y., \& Mroueh, Y. (2022). Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research.*

*[Wibisono18] Wibisono, A. (2018). Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Conference on Learning Theory (pp. 2093-3027). PMLR.*

*[Parikh14] Parikh, N., \& Boyd, S. (2014). Proximal algorithms. Foundations and trends® in Optimization, 1(3), 127-239.*


***$\bullet$ This fact should be promoted and discussed much more***: We will discuss more DC splittings for Bayesian posterior sampling as suggested. 

***$\bullet$ Perhaps you could motivate the practicality of your method by empirically comparing your method vs. traditional sampling methods on target measures which are DC but non-smooth?***: we can compare with traditional sampling methods like ULA in the revised version. We agree with the reviewer's concern regarding the JKO scalability, as we also already explicitly admitted it in the main text of our submission (Line 362). Our work was mainly inspired by [36, Salim, Korba, Luise - NeurIPS2020], another purely theoretical work. We believe theoretical works have their own merits in the long run as they provide insight and connections that can guide practical research. The JKO operator, not only studied in the ML community but primarily by the PDE community, has demonstrated that purely theoretical objects can have far-reaching impacts, with the paper by Jordan, Kinderlehrer, and Otto in 1998 being seminal. Notice that even though our work is theoretical in nature, we have provided two concrete examples of DC decomposition, demonstrating that our approach can be implemented in practice.
Thank you for the detailed response!
- It is indeed preferable to clarify from the outset (abstract or title) that the structure you are considering is DC, in my opinion. 
- I would recommend to pick a side clearly: either focus on sampling applications, with $\mathcal{H}$ being the entropy, or implement the extension to all Wasserstein-geodesically-DC functionals. The presentation in the submission is somewhat muddled, as general Wasserstein-geodesically-convex $\mathcal{H}$ is considered but all the examples and illustrations are for the entropy. 
- I was not aware of the fact that some "natural" Bayesian inference problems have a DC (with explicit splitting) but non-smooth structure, thank you for pointing this out. This fact should be promoted and discussed much more, in my opinion, as it shows that for certain settings of practical interest, DC-type considerations are in some sense unavoidable. This is much more convincing to me than the abstract argument that "more functions satisfy DC".
- I do not find the theoretical contributions of your work sufficient to bypass the issue of practical performance. The implementability of the JKO step is indeed a difficult question. Perhaps you could motivate the practicality of your method by empirically comparing your method vs. traditional sampling methods on target measures which are DC but non-smooth?

Looking forward to seeing your thoughts on these suggestions!
We thank the reviewer for the response.

***The title of the submission "Non-geodesically-convex optimization in the Wasserstein space" indicates...***

We agree that smoothness (L-smooth) is the go-to condition when there is no convexity. As we understand, the reviewer advocates for this class of L-smooth nonconvex functions. We also assume we are talking about the nonconvexity of $F$ in the potential part, $\mathcal{E}_F$ (inducing the nonconvexity along Wasserstein geodesics). We emphasise is that the class of L-smooth  functions is just a subclass of DC functions thanks to the well-known result: if $F$ is L-smooth, then it admits the following DC decompositions: (1)
$F(x) = (F(x) + (\eta/2) \Vert x \Vert^2) - (\eta/2) \Vert x \Vert^2$
or (2)
$F(x) = (\eta/2) \Vert x \Vert^2 - ((\eta/2) \Vert x \Vert^2 - F(x))$
whenever $\eta \geq L$.

In other words, our results readily apply to the class of $L$-smooth functions raised by the reviewer. Nevertheless, we can explicitly mention in the abstract that the nonconvex structure we are studying is DC.

***given the approximation result you cited...***

The result we cited advocates the richness of the class of DC functions since it is a universal approximator for any continuous function over a compact region.

Obtaining explicit DC decomposition for a given function has been studied thoroughly in the DC programming literature. It is known that many objective functions have natural DC decomposition (we are talking about the $F$ part in the potential) (see, e.g., [Nouiehed19, Lethi05]). We can add a paragraph to discuss the explicit decomposition in the revised version of the paper. Note that, in Appendix C.1 and C.2 in our submission, we already provided two concrete examples of the DC decomposition for the log of the Gaussian mixture and also von Mises-Fisher with distance-to-set prior, demonstrating that it can be done in practice. We can point out DC structures of other types of non-log-concave densities, including the posterior, as suggested.

*[Nouiehed19] Nouiehed, M., Pang, J. S., \& Razaviyayn, M. (2019). On the pervasiveness of difference-convexity in optimization and statistics. Mathematical Programming, 174(1), 195-222.*

*[Lethi05] H. A. Le Thi, \& T. Pham Dinh (2005). The DC (difference of convex functions) programming and DCA revisited with DC models of real world nonconvex optimization problems. Annals of operations research, 133, 23-46.*

***Another argument in favor of studying DC functions on the Wasserstein space, instead of smooth ones, might be that your analysis holds under quite a few assumptions***: This is indeed why we focus on DC class (containing L-smooth class), and we will make this advantage even more clear in the revised version. We already made an attempt to explain that in the paragraph **Why difference-of-convex structure** and will make it more pronounced.

***(i) It is still unclear to me that the algorithm you analyze (using a proximal step in Wasserstein space) has hopes of becoming efficiently implementable***: our work is mainly theoretical, in the same vein as in ([56], Salim, Korba, Luise - NeurIPS2020).

***(ii) The example setting you promoted most, namely sampling from a log-DC measure, deserves more discussion***: It is possible to discuss the DC structure of posterior distributions more extensively in the revised version. The splitting of the posterior depends on two terms: log-likelihood and log-prior. If we assume that the likelihood is nice enough, i.e., log-L-smooth, it is already DC with known splitting as above (of course, other structures can also be considered). The log-prior is normally not smooth to model sparsity or low rank (LASSO, group LASSO, SCAD, Capped-$\ell_1$, PiL, etc). Most of them have some explicit DC structure [Lethi15].

*[Lethi15] Lethi, H. A., Dinh, T. P., Le, H. M., \& Vo, X. T. (2015). DC approximation approaches for sparse optimization. European Journal of Operational Research, 244(1), 26-46.*


***My recommendation to the authors would be to exploit their (technically interesting!) analysis into a full-fledged study of DC optimization in the Wasserstein space***: this is a promising future research direction. However, our current paper is already technically large (26 pages or so), and integrating all of this into the current version makes it overly heavy. Moreover, we did not study the suggested class initially because we did not find any practical problems with the concave part given in a more general form than the potential energy.
Thank you for clarifying what your contributions are. However, this clarification makes the significance of this work less convincing than I initially thought.

The title of the submission "Non-geodesically-convex optimization in the Wasserstein space" indicates that 1. your audience is one that is interested in optimization questions; 2. you consider objective functionals which are non-convex, and the go-to condition in the non-convex optimization literature is smoothness. However your submission is actually about Difference-of-Convex optimization in the Wasserstein space, which is quite different, both in its application cases and in the kind of theory involved for the analysis.

I understand the general appeal of the DC assumption, given the approximation result you cited ("any continuous function $f$ can be approximated by a sequence of DC functions $g-h$ over a compact, convex domain"). But a) this approximation result does not say how to choose the splitting $g-h$ in a usable way for optimization [please correct me if incorrect], and besides b) you do not prove such an approximation result for the Wasserstein space.

Another argument in favor of studying DC functions on the Wasserstein space, instead of smooth ones, might be that your analysis holds under quite few assumptions. But, particularly in a venue such as NeurIPS, a theoretical analysis only has value when it applies (or there are reasons to think it will apply in the future) to (i) an algorithm _and_ (ii) a setting of practical interest. Now
- (i) after reading your answer to Reviewer 7EW7, it is still unclear to me that the algorithm you analyze (using a proximal step in Wasserstein space) has hopes of becoming efficiently implementable. If, as you say, methods proposed 2 years ago allow for faster computation, then it is up to you to implement them.
- (ii) The example setting you promoted most, namely sampling from a log-DC measure, deserves more discussion. Sampling problems of this form typically arise in Bayesian inference, where the target measure is a posterior likelihood which can be quite involved. It is not at all clear how to choose a splitting of it (or rather of its log) in a tractable way.

For these reasons, I downgrade my rating from 4 to 3. Nonetheless I remain open to further discussion. My recommendation to the authors would be to exploit their (technically interesting!) analysis into a full-fledged study of DC optimization in the Wasserstein space. (I realize that this is not easy to incorporate into the current submission unfortunately.)

PS1: I understand that your analysis can be extended to the case where the objective $\mathcal{F}$ is of the form $\tilde{\mathcal{G}} - \tilde{\mathcal{H}}$ where $\tilde{\mathcal{G}}, \tilde{\mathcal{H}}$ are general geodesically convex functionals, instead of $\tilde{\mathcal{G}}$ being linear ($\tilde{\mathcal{G}} = \mathcal{E}_{-H}$ and $\tilde{\mathcal{H}}= \mathcal{H} + \mathcal{E}_G$ in the paper). My comments above implicitly assume this extension.

PS2, regarding the name semi-FB Euler: I do not have a strong opinion on the matter, but I would like to point out that choosing such a name does not contribute at all to convey the non-triviality of the method, and that this name would be quite awkward in the context of the natural extension mentioned in PS1 (with general $\tilde{\mathcal{G}} - \tilde{\mathcal{H}}$).
Thank you for taking time answering my review :-)