Statistical Optimal Transport Chapter 7: Metric geometry of the Wasserstein space

Sinho Chewi, Jonathan Niles-Weed and Philippe Rigollet state the arbitrary-P2 Wasserstein geodesic theorem. This is ASTIS's attributed paraphrase and formal proof of its ordered metric-identity component.

Complete source statement (ASTIS restatement)

Let mu0 and mu1 be probability measures on Euclidean space with finite second moments, and let gamma be an optimal quadratic transport plan between them. For pi_t(x,y)=(1-t)x+ty and mu_t=(pi_t)_#gamma, every 0<=s<=t<=1 satisfies W2(mu_s,mu_t)=(t-s)W2(mu0,mu1). ASTIS retains its existing finite-dimensional real inner-product Borel ambient interface. Optimal-plan existence and the remaining geodesic-space conclusions are outside this binding.

displacement interpolation

ordered constant speed

Read the formalized proofs

Each statement and proof below has its own closed Lean disclosure. ASTIS parents, Mathlib calls and external mathematical sources are distinguished in each proof.

ASTIS mathematical exposition

Optimal displacement interpolation has constant speed for arbitrary P2 endpoints

AutoSamplingTheory.TechnicalLemmas.Measure.DisplacementInterpolationConstantSpeed.wassersteinDistance_interpolation_eq_of_le_of_integrable_norm_sq · theorem · Teaching coverage

Statement

Let E be the repository's finite-dimensional real inner-product Borel space. Let mu0 and mu1 be probability measures with finite second moments, and let gamma be a coupling of them attaining the quadratic transport infimum. For 0 <= s <= t <= 1, the 2-Wasserstein distance between the affine pushforwards mu_s and mu_t is exactly (t-s) times the endpoint distance.

\[\mu_r=((x,y)\mapsto(1-r)x+ry)_\#\gamma,\qquad W_2(\mu_s,\mu_t)=(t-s)W_2(\mu_0,\mu_1),\qquad 0\le s\le t\le1.\]

All objects and hypotheses

E retains the existing NormedAddCommGroup, real InnerProductSpace, FiniteDimensional, Borel, SecondCountableTopology, StandardBorelSpace and Nonempty assumptions used by the interpolation and exact triangle APIs.
mu0 and mu1 are probability measures, expressed by IsProbabilityMeasure instances, and their real squared norms are integrable. Absolute continuity with respect to volume is not assumed.
gamma is supplied and is a quadratic-optimal coupling: it has marginals mu0 and mu1 and its ENNReal quadratic cost equals the transport-cost infimum. The theorem does not prove that gamma exists.
The parameters satisfy 0 <= s <= t <= 1. The ordered formula uses t-s rather than an absolute value.
The output is an equality in ENNReal. Finite second moments are used to prove that the endpoint Wasserstein distance is below infinity before additive cancellation.

Mathematical proof

1. Build probability laws along the interpolation

The optimal coupling has probability mass because its first marginal is the probability law mu0. Mapping this coupling by the measurable affine point map gives a probability measure mu_s at time s and similarly mu_t at time t. No density or absolute-continuity argument is used.

\[\mu_r=(\pi_r)_\#\gamma,\qquad \pi_r(x,y)=(1-r)x+ry,\qquad \gamma( E\times E)=1\Longrightarrow\mu_r(E)=1.\]

Corresponding Lean step

Transport.isProbabilityMeasure_of_isCoupling_left supplies the probability instance for gamma, and isProbabilityMeasure_displacementInterpolation transports it through the measurable affine map.

2. Use the canonical two-time coupling for the upper bound

Push gamma forward by the pair of affine maps at s and t. This is a coupling of mu_s and mu_t, and its displacement is exactly (t-s)(y-x). Quadratic optimality at the endpoints therefore gives the sharp upper bound for each of the three ordered subintervals [0,s], [s,t], and [t,1].

\[W_2(\mu_s,\mu_t)\le(t-s)W_2(\mu_0,\mu_1).\]

Corresponding Lean step

wassersteinDistance_interpolation_le packages the measurable two-time coupling, exact quadratic-cost scaling, endpoint optimality and monotonicity of the positive square root. Endpoint rewrites use displacementInterpolation_zero and displacementInterpolation_one.

3. Compare the three pieces with the endpoint distance

The exact Wasserstein triangle inequality along mu0, mu_s, mu_t, mu1 bounds the endpoint distance by the sum of the three segment distances. Replacing the prefix and suffix by their sharp upper bounds produces a lower bound on the middle segment after cancellation. The coefficients s, t-s and 1-t are nonnegative and sum to one.

\[W_2(\mu_0,\mu_1)\le W_2(\mu_0,\mu_s)+W_2(\mu_s,\mu_t)+W_2(\mu_t,\mu_1),\qquad s+(t-s)+(1-t)=1.\]

Corresponding Lean step

WassersteinTriangleExact.wassersteinDistance_triangle is applied twice. interpolation_coefficients_sum_one converts the coefficient partition into an ENNReal equality.

4. Justify extended-real cancellation using finite moments

Cancellation in ENNReal is invalid when the common prefix or suffix is infinite. The endpoint probability and finite-second-moment hypotheses feed the existing independent-product-coupling lemma, proving the endpoint W2 distance is finite. Multiplying by finite real coefficients keeps the prefix and suffix finite, so the two additive cancellations are sound. Combining the resulting lower bound with the canonical-coupling upper bound gives equality.

\[\int\|x\|^2\,d\mu_0<\infty,\quad\int\|y\|^2\,d\mu_1<\infty\Longrightarrow W_2(\mu_0,\mu_1)<\infty\Longrightarrow W_2(\mu_s,\mu_t)=(t-s)W_2(\mu_0,\mu_1).\]

Corresponding Lean step

WassersteinFiniteSecondMoment.wassersteinDistance_lt_top_of_integrable_norm_sq proves hd. ENNReal.mul_lt_top and ENNReal.add_le_add_iff_left/right justify the two cancellations, and le_antisymm closes the equality.

Lean statement · wassersteinDistance_interpolation_eq_of_le_of_integrable_norm_sq

The theorem quantifies over the existing ambient E, endpoint measures mu0 and mu1 with probability instances, an optimal quadratic coupling gamma, explicit Integrable squared-norm hypotheses, and ordered s,t in the unit interval. It returns the exact ENNReal Wasserstein-distance identity for the two displacementInterpolation pushforwards.

Braces mark parameters Lean can infer; square brackets request structures such as a measurable space or probability measure. Named hypotheses are mathematical premises, not facts established by this declaration. Section parameters are described in the mathematical hypotheses above; the module link retains their exact source context.

theorem wassersteinDistance_interpolation_eq_of_le_of_integrable_norm_sq
    {γ : Measure (E × E)} {μ₀ μ₁ : Measure E}
    [IsProbabilityMeasure μ₀] [IsProbabilityMeasure μ₁]
    (hμ₀ : Integrable (fun x : E => ‖x‖ ^ 2) μ₀)
    (hμ₁ : Integrable (fun x : E => ‖x‖ ^ 2) μ₁)
    (hγ : IsQuadraticOptimalCoupling γ μ₀ μ₁)
    {s t : ℝ} (hs0 : 0 ≤ s) (hst : s ≤ t) (ht1 : t ≤ 1) :
    WassersteinSpace.wassersteinDistance
        (displacementInterpolation γ s)
        (displacementInterpolation γ t) =
      ENNReal.ofReal (t - s) *
        WassersteinSpace.wassersteinDistance μ₀ μ₁

Exact module and namespace context

Lean proof · wassersteinDistance_interpolation_eq_of_le_of_integrable_norm_sq

The proof reuses the existing sharp interpolation upper bound and exact triangle inequality. Its only endpoint-domain change is to call wassersteinDistance_lt_top_of_integrable_norm_sq directly instead of unpacking the stronger P2ac predicate. The historical wassersteinDistance_interpolation_eq_of_le declaration and module remain unchanged. The focused compatibility test installs probability instances from its P2ac hypotheses and derives the same signature from this general theorem.

theorem wassersteinDistance_interpolation_eq_of_le_of_integrable_norm_sq
    {γ : Measure (E × E)} {μ₀ μ₁ : Measure E}
    [IsProbabilityMeasure μ₀] [IsProbabilityMeasure μ₁]
    (hμ₀ : Integrable (fun x : E => ‖x‖ ^ 2) μ₀)
    (hμ₁ : Integrable (fun x : E => ‖x‖ ^ 2) μ₁)
    (hγ : IsQuadraticOptimalCoupling γ μ₀ μ₁)
    {s t : ℝ} (hs0 : 0 ≤ s) (hst : s ≤ t) (ht1 : t ≤ 1) :
    WassersteinSpace.wassersteinDistance
        (displacementInterpolation γ s)
        (displacementInterpolation γ t) =
      ENNReal.ofReal (t - s) *
        WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
  let μs : Measure E := displacementInterpolation γ s
  let μt : Measure E := displacementInterpolation γ t
  letI : IsProbabilityMeasure μs := by
    dsimp [μs]
    exact isProbabilityMeasure_displacementInterpolation hγ.1 s
  letI : IsProbabilityMeasure μt := by
    dsimp [μt]
    exact isProbabilityMeasure_displacementInterpolation hγ.1 t

  have hupper :
      WassersteinSpace.wassersteinDistance μs μt ≤
        ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    simpa [μs, μt] using wassersteinDistance_interpolation_le (E := E) hγ hst

  have h0s :
      WassersteinSpace.wassersteinDistance μ₀ μs ≤
        ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    have h := wassersteinDistance_interpolation_le (E := E) hγ hs0
    rw [displacementInterpolation_zero hγ.1] at h
    simpa [μs] using h

  have ht1' :
      WassersteinSpace.wassersteinDistance μt μ₁ ≤
        ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    have h := wassersteinDistance_interpolation_le (E := E) hγ ht1
    rw [displacementInterpolation_one hγ.1] at h
    simpa [μt] using h

  have htri_left :=
    WassersteinTriangleExact.wassersteinDistance_triangle μ₀ μs μ₁
  have htri_right :=
    WassersteinTriangleExact.wassersteinDistance_triangle μs μt μ₁
  have hchain :
      WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
        WassersteinSpace.wassersteinDistance μ₀ μs +
          WassersteinSpace.wassersteinDistance μs μt +
            WassersteinSpace.wassersteinDistance μt μ₁ := by
    calc
      WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
          WassersteinSpace.wassersteinDistance μ₀ μs +
            WassersteinSpace.wassersteinDistance μs μ₁ := htri_left
      _ ≤ WassersteinSpace.wassersteinDistance μ₀ μs +
            (WassersteinSpace.wassersteinDistance μs μt +
              WassersteinSpace.wassersteinDistance μt μ₁) := by
        exact add_le_add le_rfl htri_right
      _ = WassersteinSpace.wassersteinDistance μ₀ μs +
            WassersteinSpace.wassersteinDistance μs μt +
              WassersteinSpace.wassersteinDistance μt μ₁ := by
        simp [add_assoc]

  have hchain' :
      WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
        ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          WassersteinSpace.wassersteinDistance μs μt +
            ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    exact hchain.trans (add_le_add (add_le_add h0s le_rfl) ht1')

  have hcoeff := interpolation_coefficients_sum_one hs0 hst ht1
  have hpartition :
      WassersteinSpace.wassersteinDistance μ₀ μ₁ =
        ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
            ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    calc
      WassersteinSpace.wassersteinDistance μ₀ μ₁ =
          1 * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by simp
      _ = (ENNReal.ofReal s + ENNReal.ofReal (t - s) + ENNReal.ofReal (1 - t)) *
            WassersteinSpace.wassersteinDistance μ₀ μ₁ := by rw [hcoeff]
      _ = ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
            ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
              ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
        simp [add_mul, add_assoc]

  have hcancelInput :
      ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
            ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
        ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          WassersteinSpace.wassersteinDistance μs μt +
            ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ := by
    rw [← hpartition]
    exact hchain'

  have hd : WassersteinSpace.wassersteinDistance μ₀ μ₁ < ∞ :=
    WassersteinFiniteSecondMoment.wassersteinDistance_lt_top_of_integrable_norm_sq
      μ₀ μ₁ hμ₀ hμ₁
  have hprefix :
      ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ < ∞ :=
    ENNReal.mul_lt_top ENNReal.ofReal_lt_top hd
  have hsuffix :
      ENNReal.ofReal (1 - t) * WassersteinSpace.wassersteinDistance μ₀ μ₁ < ∞ :=
    ENNReal.mul_lt_top ENNReal.ofReal_lt_top hd

  have hcancelSuffix :
      ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
        ENNReal.ofReal s * WassersteinSpace.wassersteinDistance μ₀ μ₁ +
          WassersteinSpace.wassersteinDistance μs μt := by
    exact (ENNReal.add_le_add_iff_right hsuffix.ne).mp
      (by simpa [add_assoc] using hcancelInput)
  have hlower :
      ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁ ≤
        WassersteinSpace.wassersteinDistance μs μt :=
    (ENNReal.add_le_add_iff_left hprefix.ne).mp hcancelSuffix

  change WassersteinSpace.wassersteinDistance μs μt =
    ENNReal.ofReal (t - s) * WassersteinSpace.wassersteinDistance μ₀ μ₁
  exact le_antisymm hupper hlower

end

end DisplacementInterpolationConstantSpeed
end Measure
end TechnicalLemmas
end AutoSamplingTheory

Exact module and namespace context

Source assumptions versus formal assumptions

Source	Actual Lean	Difference kind	Why it matters
mu0 and mu1 belong to W2=P2(R^d)	IsProbabilityMeasure mu0 and mu1 plus Integrable (fun x => norm x squared) for both endpoints	same	The source's P2 membership is expanded into probability normalization and finite real second moments. Absolute continuity is absent on both sides.
Euclidean space R^d	Finite-dimensional real inner-product Borel space with SecondCountableTopology, StandardBorelSpace and Nonempty	API-limitation	The existing ASTIS affine-map, measure-map and exact Wasserstein triangle interfaces use these ambient classes. Every Euclidean R^d source instance satisfies them; no stronger endpoint regularity is introduced.
gamma is an optimal plan in Gamma(mu0,mu1)	IsQuadraticOptimalCoupling gamma mu0 mu1	same	The predicate records both marginals and exact attainment of the ENNReal quadratic transport-cost infimum. Existence is an input in this declaration.
0 <= s <= t <= 1 and mu_r=(pi_r)_#gamma	Explicit hs0, hst and ht1 inequalities and displacementInterpolation gamma r	same	The affine pushforward is definitionally the source path. Ordered times make ofReal(t-s) the exact nonnegative coefficient.
A finite real W2 equality	An equality in ENNReal, with endpoint finiteness derived before cancellation	API-limitation	ASTIS's Wasserstein distance is ENNReal-valued so infinite costs remain representable. The finite-moment hypotheses prove the displayed values are finite; the equality preserves the source constant without totalized real conversion.
Theorem 7.6 also concludes W2 is a geodesic space and constructs a path for every endpoint pair	Only the ordered metric identity for a supplied gamma	same	This binding is intentionally one proof obligation of the source theorem. The existence and packaging obligations remain separately red and are not inferred from the conditional theorem.

Exactly the ordered metric identity conditional on a supplied optimal coupling. No optimal-plan existence, path continuity, uniqueness, arbitrary-time absolute-value wrapper, Benamou-Brenier formula, full geodesic-space instance, or chapter completion. Relative to Log-Concave Sampling Theorem 1.3.23, this is an explicitly attributed endpoint-domain generalization; its old P2ac theorem remains a corollary.

Encoder–denoiser: accepted · equivalent-after-elaboration

Detected semantic differences

domains: The source writes R^d, while Lean uses the coordinate-free finite-dimensional real inner-product Borel interface. Specialization to R^d is exact, and the extension uses no structure absent from finite-dimensional Euclidean space. — Independent source review classification: domain-clarification; nonblocking.
assumptions: Source membership in P2 is unbundled in Lean into probability instances and integrability of squared norm. This is the same endpoint condition, not an absolute-continuity assumption and not a strengthening beyond P2. — Independent source review classification: notation-resolution; nonblocking.
conclusion: Lean states W2 in ENNReal and writes ENNReal.ofReal(t-s). Endpoint second moments prove finiteness, so this is equivalent to the source's finite real-valued equality. — Independent source review classification: notation-resolution; nonblocking.
scopes: The declaration proves the ordered metric-identity component conditional on gamma. It does not prove optimal-plan existence or package the full geodesic-space conclusion; those exclusions are explicit in both the packet and candidate exposition. — Independent source review classification: domain-clarification; nonblocking.

A generalization is not a source correction. Proposed missing conditions require separate independent repair review. No proposed repair silently changes the original theorem.

Scope and omitted-condition boundaries

This is exactly the ordered metric-identity component of Statistical Optimal Transport Theorem 7.6, conditional on a supplied optimal plan; it does not prove optimal-plan existence or the complete geodesic-space assertion.
Relative to Log-Concave Sampling Theorem 1.3.23, removing endpoint absolute continuity is a genuine generalization. That source's P2ac theorem and IsWassersteinGeodesic interface remain unchanged and follow as a compatibility corollary.
No uniqueness, continuity equation, Benamou-Brenier identity, optimal map, path continuity, arbitrary-time symmetry wrapper, or completion of either source chapter is claimed.

Source and reuse

ASTIS parents called

Mathlib API called (external library)

ENNReal.mul_lt_top
ENNReal.add_le_add_iff_left
ENNReal.add_le_add_iff_right
le_antisymm

Mathematical sources

Chewi, Niles-Weed and Rigollet, Statistical Optimal Transport, Theorem 7.6, printed pp.209-210/PDF pp.215-216 — Direct source for arbitrary P2 endpoints, the optimal-plan affine pushforward and the constant-speed identity. ASTIS formalizes only the ordered metric identity conditional on the plan.
Chewi, Log-Concave Sampling, Theorem 1.3.23 and Definition 1.3.25, printed p.30/PDF p.42 — The existing P2ac source theorem is the compatibility specialization. It is not the source for removing absolute continuity.

ASTIS prose is not a quotation or a source-equivalence certificate. Definitions and aliases are explained as constructions, not counted as new mathematical proofs.

Metric geometry of the Wasserstein space

Source → theorem map → reusable Lean nodes

Source audit

Upstream alignment

Frontier Cells

Graph placement

Mathematical orientation

Section source map

Arbitrary-P2 displacement interpolation has constant speed

Complete source statement (ASTIS restatement)

displacement interpolation

ordered constant speed

Read the formalized proofs

Optimal displacement interpolation has constant speed for arbitrary P2 endpoints

Statement

All objects and hypotheses

Mathematical proof

1. Build probability laws along the interpolation

2. Use the canonical two-time coupling for the upper bound

3. Compare the three pieces with the endpoint distance

4. Justify extended-real cancellation using finite moments

Source assumptions versus formal assumptions

Detected semantic differences

Scope and omitted-condition boundaries

Source and reuse

ASTIS parents called

Mathlib API called (external library)

Mathematical sources

Which proof edges are actually covered?