Convexity along generalized geodesics and the Jordan–Kinderlehrer–Otto step

Let E be a separable Hilbert space and φ an energy on the quadratic Wasserstein space P₂(E). One step of the Jordan–Kinderlehrer–Otto scheme (JKO scheme) for the Wasserstein gradient flow of φ, with time step τ > 0, started at a law μ, chooses a minimizer of

ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ),

that is, a point of the resolvent TauCeti.moreauYosidaResolvent φ τ μ of the metric theory of gradient flows, taken in P₂(E). Ordinary displacement convexity of φ is not the hypothesis that makes this step unique, because ν ↦ W₂ (μ, ν) ^ 2 is not convex along Wasserstein geodesics. It is, however, 2-convex along the generalized geodesics based at μ, the interpolations ν_t = ((1 - t) • y₀ + t • y₁)₊ γ of a three-plan γ of (x, y₀, y₁) whose projections to (x, y₀) and (x, y₁) are optimal couplings out of μ (TauCeti.IsGeneralizedGeodesicPlan.dist_sq_add_le). This file introduces the matching convexity notion for energies and proves that it makes the JKO step unique.

An energy φ is m-convex along generalized geodesics based at μ (TauCeti.GeneralizedGeodesicallyConvex m φ μ) if any two laws ν₀ and ν₁ of its effective domain are joined by the generalized geodesic of some such plan γ along which

φ ν_t ≤ (1 - t) φ ν₀ + t φ ν₁ - (m / 2) t (1 - t) ∫ ‖y₀ - y₁‖² dγ.

The integral is the quadratic cost of the coupling of ν₀ and ν₁ carried by γ; it is at least W₂ (ν₀, ν₁) ^ 2. If φ has this property, then the penalized functional of the JKO step is (m + 1 / τ)-convex along the same curves (TauCeti.GeneralizedGeodesicallyConvex.add_edist_sq_div), so the JKO step is unique as soon as 1 + m τ > 0 (TauCeti.GeneralizedGeodesicallyConvex.subsingleton_moreauYosidaResolvent): the midpoint of the generalized geodesic joining two minimizers would otherwise do strictly better. Since W₂ (ν₀, ν₁) ^ 2 is at most the plan cost, the same convexity holds with W₂ (ν₀, ν₁) ^ 2 when 1 + m τ ≥ 0. If φ is convex along generalized geodesics based at every law, it therefore satisfies the convexity hypothesis TauCeti.PenalizedConvex of the metric theory of gradient flows (TauCeti.penalizedConvex_of_generalizedGeodesicallyConvex), and JKO steps satisfy the discrete variational inequality of TauCeti.Analysis.GradientFlow.MoreauYosida.Convex.

Main definitions

• TauCeti.ConvexAlongGeneralizedGeodesicPlan m φ ν₀ ν₁ γ: the convexity inequality along a fixed three-plan, with the quadratic cost of its endpoint coupling as the correction.
• TauCeti.GeneralizedGeodesicallyConvex m φ μ: the energy φ on P₂(E) is m-convex along generalized geodesics based at μ.

Main results

• TauCeti.IsGeneralizedGeodesicPlan.dist_sq_add_le: the convexity of the squared Wasserstein distance from the base point along a generalized geodesic, as an inequality of real numbers.
• TauCeti.GeneralizedGeodesicallyConvex.add_edist_sq_div: the penalized functional ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ) of an energy that is m-convex along generalized geodesics based at μ is (m + 1 / τ)-convex along them.
• TauCeti.GeneralizedGeodesicallyConvex.eq_of_mem_moreauYosidaResolvent and TauCeti.GeneralizedGeodesicallyConvex.subsingleton_moreauYosidaResolvent: uniqueness of the JKO step when 1 + m τ > 0.
• TauCeti.penalizedConvex_of_generalizedGeodesicallyConvex: convexity along generalized geodesics based at every law is the convexity hypothesis TauCeti.PenalizedConvex of the metric theory, so JKO steps satisfy its discrete variational inequality.

Implementation notes

The JKO step is not given a separate definition: on P₂(E) with its Wasserstein metric it is the resolvent TauCeti.moreauYosidaResolvent φ τ μ, whose penalty edist μ ν ^ 2 / (2τ) is W₂ (μ, ν) ^ 2 / (2τ).

The convexity inequality is stated, as in Ambrosio–Gigli–Savaré, with the cost ∫ ‖y₀ - y₁‖² dγ of the coupling carried by the plan rather than with the smaller W₂ (ν₀, ν₁) ^ 2. For m ≤ 0 this is weaker than the same inequality with W₂ (ν₀, ν₁) ^ 2, for m ≥ 0 it is stronger, and at m = 0 the two agree.

References

• L. Ambrosio, N. Gigli, G. Savaré, Gradient Flows in Metric Spaces and in the Space of Probability Measures, 2nd ed., Birkhäuser 2008, §9.2, Definition 9.2.4 and Lemma 9.2.7, and Chapter 4, Assumption 4.0.1.
• R. Jordan, D. Kinderlehrer, F. Otto, The variational formulation of the Fokker–Planck equation, SIAM J. Math. Anal. 29 (1998), 1--17.

def TauCeti.ConvexAlongGeneralizedGeodesicPlan {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] (m : ℝ) (φ : WassersteinSpace 2 E → EReal) (ν₀ ν₁ : WassersteinSpace 2 E) (γ : MeasureTheory.Measure (E × E × E)) : Prop

The energy φ satisfies the m-convexity inequality along the interpolation of the fixed three-plan γ, with endpoint laws ν₀, ν₁. The correction uses the quadratic cost of γ.snd. This predicate alone does not require that γ is a generalized-geodesic plan.

theorem TauCeti.convexAlongGeneralizedGeodesicPlan_def {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {ν₀ ν₁ : WassersteinSpace 2 E} {γ : MeasureTheory.Measure (E × E × E)} : ConvexAlongGeneralizedGeodesicPlan m φ ν₀ ν₁ γ ↔ ∀ t ∈ Set.Icc 0 1, ∀ (ν : WassersteinSpace 2 E), ↑↑ν = generalizedGeodesic γ t → φ ν ≤ ↑(1 - t) * φ ν₀ + ↑t * φ ν₁ - ↑(m / 2 * t * (1 - t) * ∫ (z : E × E), ‖z.1 - z.2‖ ^ 2 ∂γ.snd)

The defining inequality for convexity along a fixed three-plan.

theorem TauCeti.ConvexAlongGeneralizedGeodesicPlan.mono {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] {m m' : ℝ} {φ : WassersteinSpace 2 E → EReal} {ν₀ ν₁ : WassersteinSpace 2 E} {γ : MeasureTheory.Measure (E × E × E)} (h : ConvexAlongGeneralizedGeodesicPlan m φ ν₀ ν₁ γ) (hm : m' ≤ m) : ConvexAlongGeneralizedGeodesicPlan m' φ ν₀ ν₁ γ

Convexity along a fixed three-plan persists with a smaller parameter.

def TauCeti.GeneralizedGeodesicallyConvex {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] (m : ℝ) (φ : WassersteinSpace 2 E → EReal) (μ : WassersteinSpace 2 E) : Prop

The energy φ on P₂(E) is m-convex along generalized geodesics based at μ: any two laws ν₀ and ν₁ with φ ν₀ ≠ ⊤ and φ ν₁ ≠ ⊤ (energy ⊥ is allowed) carry a quadratic generalized-geodesic plan γ from ν₀ to ν₁ based at μ along whose generalized geodesic ν_t, for every t ∈ [0, 1], φ ν_t ≤ (1 - t) φ ν₀ + t φ ν₁ - (m / 2) t (1 - t) ∫ ‖y₀ - y₁‖² dγ. The law ν_t lies in P₂(E) (TauCeti.exists_coe_eq_generalizedGeodesic); the integral is the quadratic cost of the coupling γ.snd of ν₀ and ν₁.

theorem TauCeti.generalizedGeodesicallyConvex_def {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {μ : WassersteinSpace 2 E} : GeneralizedGeodesicallyConvex m φ μ ↔ ∀ ⦃ν₀ ν₁ : WassersteinSpace 2 E⦄, φ ν₀ ≠ ⊤ → φ ν₁ ≠ ⊤ → ∃ (γ : MeasureTheory.Measure (E × E × E)), IsGeneralizedGeodesicPlan 2 γ ↑↑μ ↑↑ν₀ ↑↑ν₁ ∧ ConvexAlongGeneralizedGeodesicPlan m φ ν₀ ν₁ γ

The defining condition of convexity along generalized geodesics.

theorem TauCeti.GeneralizedGeodesicallyConvex.mono {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [NormedSpace ℝ E] {m m' : ℝ} {φ : WassersteinSpace 2 E → EReal} {μ : WassersteinSpace 2 E} (h : GeneralizedGeodesicallyConvex m φ μ) (hm : m' ≤ m) : GeneralizedGeodesicallyConvex m' φ μ

Convexity along generalized geodesics with a parameter m implies it for every smaller parameter.

theorem TauCeti.IsGeneralizedGeodesicPlan.dist_sq_add_le {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E] [CompleteSpace E] [InnerProductSpace ℝ E] {γ : MeasureTheory.Measure (E × E × E)} {t : ℝ} {μ ν₀ ν₁ ν : WassersteinSpace 2 E} (hγ : IsGeneralizedGeodesicPlan 2 γ ↑↑μ ↑↑ν₀ ↑↑ν₁) (ht : t ∈ Set.Icc 0 1) (hν : ↑↑ν = generalizedGeodesic γ t) : dist μ ν ^ 2 + t * (1 - t) * ∫ (z : E × E), ‖z.1 - z.2‖ ^ 2 ∂γ.snd ≤ (1 - t) * dist μ ν₀ ^ 2 + t * dist μ ν₁ ^ 2

The squared distance is convex along generalized geodesics, in P₂(E). Along the generalized geodesic ν_t of a quadratic generalized-geodesic plan γ from ν₀ to ν₁ based at μ, and for t ∈ [0, 1], W₂ (μ, ν_t) ^ 2 + t (1 - t) ∫ ‖y₀ - y₁‖² dγ ≤ (1 - t) W₂ (μ, ν₀) ^ 2 + t W₂ (μ, ν₁) ^ 2.

theorem TauCeti.GeneralizedGeodesicallyConvex.add_edist_sq_div {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E] [CompleteSpace E] [InnerProductSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {μ : WassersteinSpace 2 E} {τ : NNReal} (h : GeneralizedGeodesicallyConvex m φ μ) (hτ : 0 < τ) : GeneralizedGeodesicallyConvex (m + (↑τ)⁻¹) (fun (ν : WassersteinSpace 2 E) => φ ν + ↑(edist μ ν ^ 2 / (2 * ↑τ))) μ

The penalized functional of the JKO step is convex along generalized geodesics. If φ is m-convex along generalized geodesics based at μ, then for a time step τ > 0 the functional ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ) minimized by the JKO step from μ is (m + 1 / τ)-convex along generalized geodesics based at μ, along the same plans.

theorem TauCeti.GeneralizedGeodesicallyConvex.eq_of_mem_moreauYosidaResolvent {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E] [CompleteSpace E] [InnerProductSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {μ : WassersteinSpace 2 E} {τ : NNReal} (h : GeneralizedGeodesicallyConvex m φ μ) (hτ : 0 < τ) (hmτ : 0 < 1 + m * ↑τ) {ν₀ ν₁ : WassersteinSpace 2 E} (h₀ : ν₀ ∈ moreauYosidaResolvent φ τ μ) (h₁ : ν₁ ∈ moreauYosidaResolvent φ τ μ) (htop : φ ν₀ ≠ ⊤) (hbot : φ ν₀ ≠ ⊥) : ν₀ = ν₁

Uniqueness of the JKO step. Let φ be m-convex along generalized geodesics based at μ, and let τ > 0 be a time step with 1 + m τ > 0. Two minimizers ν₀ and ν₁ of ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ), the first of finite energy, are equal.

theorem TauCeti.GeneralizedGeodesicallyConvex.subsingleton_moreauYosidaResolvent {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E] [CompleteSpace E] [InnerProductSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {μ : WassersteinSpace 2 E} {τ : NNReal} (h : GeneralizedGeodesicallyConvex m φ μ) (hτ : 0 < τ) (hmτ : 0 < 1 + m * ↑τ) (hbot : ∀ (ν : WassersteinSpace 2 E), φ ν ≠ ⊥) (hdom : ∃ (ν : WassersteinSpace 2 E), φ ν ≠ ⊤) : (moreauYosidaResolvent φ τ μ).Subsingleton

Uniqueness of the JKO step. Let φ be an energy on P₂(E) that never takes the value ⊥ and is finite somewhere, and that is m-convex along generalized geodesics based at μ. Then for every time step τ > 0 with 1 + m τ > 0, the JKO step from μ, that is, the set of minimizers of ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ), has at most one element.

theorem TauCeti.penalizedConvex_of_generalizedGeodesicallyConvex {E : Type u_1} [NormedAddCommGroup E] [MeasurableSpace E] [BorelSpace E] [SecondCountableTopology E] [CompleteSpace E] [InnerProductSpace ℝ E] {m : ℝ} {φ : WassersteinSpace 2 E → EReal} {τ : NNReal} (h : ∀ (μ : WassersteinSpace 2 E), GeneralizedGeodesicallyConvex m φ μ) (hτ : 0 < τ) (hmτ : 0 ≤ 1 + m * ↑τ) : PenalizedConvex m τ φ

Convexity along generalized geodesics gives penalized convexity. If φ is m-convex along generalized geodesics based at every law of P₂(E), then for every time step τ > 0 with 1 + m τ ≥ 0 it is penalized convex with parameter m: the penalized functional ν ↦ φ ν + W₂ (μ, ν) ^ 2 / (2τ) is (m + 1 / τ)-convex along a generalized geodesic based at μ. So JKO steps satisfy the discrete variational inequality TauCeti.PenalizedConvex.apply_add_le_of_mem_moreauYosidaResolvent and the estimates for minimizing movements that follow from it.