Wasserstein Distance as a Loss Function for Image Morphing

Overview

This project implements a complete Wasserstein-based image morphing system from scratch — no POT, no GeomLoss, no external OT libraries. It transforms an MNIST digit ‘3’ into a digit ‘8’ via the Wasserstein geodesic in the space of probability measures, and contrasts the result against naive pixel blending. The project covers the full pipeline from discrete probability measures and entropic regularization to log-domain Sinkhorn iterations, Kantorovich duality verification, and displacement interpolation on the space of probability measures.

Core implementations covered:

  • Discrete probability measures from grayscale images
  • Squared Euclidean ground cost matrix ($W_2$)
  • Sinkhorn-Knopp algorithm in log-domain (numerically stable)
  • Kantorovich duality verification (primal-dual gap: 2.42e-12)
  • Displacement interpolation (McCann interpolation) for geodesic frames
  • Side-by-side animated visualization (GIF + PNG frames)

Intuitive Explanation

1. What is Wasserstein Distance?

Imagine you have a pile of sand shaped like a mountain, and you want to move all that sand into a new hole shaped like a valley. The Wasserstein distance is simply the cheapest possible way to get the job done. To calculate it, you look at every grain of sand and multiply how much it weighs by how far you have to carry it. If the mountain and the valley are already in the same spot, the distance is zero because you do not have to move anything. If they are far apart, or if the shapes are very different, the distance is high because you are doing more work to shove, drag, and reshuffle the sand into its new home.

This is why it is formally called the Earth Mover’s Distance: it measures the minimum total work required to transport one probability distribution into another. In this project, each MNIST digit is treated as a probability distribution over pixel locations, and the Wasserstein distance quantifies the cost of optimally displacing the ink mass of one digit into the shape of another.

Mathematical Foundations

The Kantorovich Primal Problem

Find a transport plan $\pi$ that minimizes the total cost of moving mass from $\mu$ to $\nu$:

$$\min_{\pi} \sum_{i,j} \pi_{ij} C_{ij}$$

Here $\pi_{ij}$ is the amount of mass transported from source location $i$ to target location $j$, and $C_{ij}$ is the cost of doing so.

Subject to:

$$\sum_{j=1}^{m} \pi_{ij} = \mu_i \quad \forall i \qquad \text{(row marginal)}$$ $$\sum_{i=1}^{n} \pi_{ij} = \nu_j \quad \forall j \qquad \text{(column marginal)}$$ $$\pi_{ij} \geq 0 \quad \forall i, j \qquad \text{(non-negativity)}$$

The row marginal constraint ensures all mass from $\mu$ is fully shipped; the column marginal ensures $\nu$ is exactly satisfied; non-negativity prevents mass from flowing backwards.

The Kantorovich Dual Problem

Find potentials $\phi$ and $\psi$ that maximize the total value:

$$\max_{\phi,\, \psi} \sum_{i} \phi_i \mu_i + \sum_{j} \psi_j \nu_j$$

Subject to:

$$\phi_i + \psi_j \leq C_{ij} \quad \forall i, j$$

The admissibility constraint ensures no pair of potentials over-values any transport link relative to its cost.

Strong Duality: At optimality, the primal and dual objectives are equal:

$$\langle C, \pi^\star \rangle = \langle \phi^\star, \mu \rangle + \langle \psi^\star, \nu \rangle$$

This result is used in this project to numerically verify that the Sinkhorn algorithm has converged to the global optimum of the regularized problem.

Sinkhorn Algorithm

The Sinkhorn algorithm solves the entropy-regularized optimal transport problem, enabling efficient computation of approximate Wasserstein distances on large discrete supports such as image pixel grids.

Entropic Regularization

We solve the regularized problem:

$$\min_{\pi \in \Pi(\mu,\nu)} \langle C, \pi \rangle - \varepsilon H(\pi)$$

where $H(\pi) = -\sum_{ij} \pi_{ij} \log \pi_{ij}$ is the entropy of the transport plan and $\varepsilon > 0$ is the regularization strength.

Adding the entropy term makes the problem strictly convex and yields a unique, smooth solution that can be computed via simple matrix scaling iterations.

Standard Form (Scaling Vectors)

The optimal plan takes the form $\pi = \text{diag}(u) \cdot K \cdot \text{diag}(v)$ where $K = \exp(-C/\varepsilon)$. The Sinkhorn iterations are:

$$u^{(\ell+1)} = \frac{\mu}{K v^{(\ell)}}$$ $$v^{(\ell+1)} = \frac{\nu}{K^\top u^{(\ell+1)}}$$

Each iteration alternately rescales rows then columns of $K$ until the marginal constraints are satisfied. Convergence is indicated when $u$ and $v$ stabilize between iterations.

Log-Domain Form (Numerically Stable)

For small $\varepsilon$, the kernel

$$K = \exp\left(-\frac{C}{\varepsilon}\right)$$

contains entries so close to zero that they underflow to exactly zero in float64, corrupting the scaling vectors. To avoid this, we iterate the potentials

$$f = \varepsilon \log u,\qquad g = \varepsilon \log v$$

directly:

$$f_i' = \varepsilon\log\mu_i - \varepsilon\,\mathrm{LSE}_j\!\left(\frac{g_j - C_{ij}}{\varepsilon}\right)$$ $$g_j' = \varepsilon\log\nu_j - \varepsilon\,\mathrm{LSE}_i\!\left(\frac{f_i' - C_{ij}}{\varepsilon}\right)$$

The Log-Sum-Exp operator is defined with numerical stabilization via:

$$\mathrm{LSE}_k(x_k) = M + \log\left(\sum_k \exp(x_k - M)\right), \qquad M = \max_k x_k$$

This formulation avoids numerical underflow because the kernel $K$ is never formed explicitly. Instead, all computations are performed in the log domain, where the exponential terms remain well-scaled even when $\varepsilon$ is very small.

Displacement Interpolation

Displacement interpolation (McCann interpolation) defines a geodesic path in the space of probability measures by physically moving mass along the optimal transport trajectories — not fading one image into another.

For $t \in [0, 1]$, define the interpolation map $T_t : \mathbb{R}^2 \times \mathbb{R}^2 \to \mathbb{R}^2$:

$$T_t(x, y) = (1 - t)\,x + t\,y$$

The interpolated measure $\mu_t$ is the push-forward of the optimal transport plan $\pi^\star$ under this map:

$$\mu_t = (T_t)_{\#}\, \pi^\star$$

Geometrically: every particle of mass at pixel location $x$ travels in a straight line toward its matched location $y$, arriving at the intermediate position $T_t(x,y)$ at time $t$. The intermediate measure is a coherent, solid object — not a superposition of two transparent images.

Comparison with Naive Interpolation

Naive pixel blending linearly interpolates pixel intensities directly:

$$I_t = (1-t)\,I_{\mu} + t\,I_{\nu}$$

This produces a double-exposure effect — two transparent digits occupying the same space simultaneously at $t = 0.5$. It preserves no spatial geometry and moves no mass. Displacement interpolation, by contrast, encodes a coherent redistribution of mass through space: structure is moved, not mixed.

Results

Findings

The log-domain Sinkhorn algorithm converged to the global optimum of the regularized problem, verified by a primal-dual duality gap of 2.42 × 10⁻¹² — at floating-point precision for float64. Marginal constraint errors remained below 1 × 10⁻⁶, confirming that the computed transport plan is feasible. The displacement interpolation frames produce a geometrically coherent morphing sequence, while the naive pixel blending baseline exhibits the expected ghosting artifact at intermediate timesteps.

Verification Table

Quantity Value
Primal objective $\langle C, \pi \rangle - \varepsilon H(\pi)$ −0.95284338
Dual objective $\langle \phi, \mu \rangle + \langle \psi, \nu \rangle$ −0.95284338
Duality gap 2.42 × 10⁻¹²
Marginal constraint error < 1 × 10⁻⁶
Strong duality satisfied ✓

The duality gap of 2.42 × 10⁻¹² confirms that the log-domain Sinkhorn algorithm converged to the global optimum of the regularized problem.

Figures & Animations

Wasserstein Morphing

Left: Wasserstein geodesic via displacement interpolation. Right: naive pixel blending. Both panels animate across the same 10 timesteps $t \in \{0.0, 0.1, \ldots, 1.0\}$.

Key Insights

1. Why Pixel Blending Fails

Pixel blending fails because it treats an image as a collection of independent light bulbs that simply dim or brighten in place. This ignores the spatial geometry of the mass distribution entirely, resulting in a ghosting effect where the first digit disappears while the second one fades in. At $t = 0.5$, you see two transparent, overlapping images occupying the same space simultaneously — a double exposure, not a transformation.

In contrast, Wasserstein distance treats the image as a physical pile of ink. By computing the optimal transport plan, it finds the most efficient way to physically displace each grain of mass from its source location to its target location. The intermediate frames preserve the geometric structure of the digit — mass is moved, not mixed. This is what the push-forward $(T_t)_{\#} \pi^\star$ encodes mathematically: a coherent redistribution of mass through space, not a pointwise average of intensities. Blending destroys structure; displacement interpolation preserves it.

Implementation Notes

  • No external OT libraries. Sinkhorn is implemented from scratch in pure NumPy to expose every numerical detail explicitly.
  • Cost matrix normalization. $C$ is normalized to $[0, 1]$ before the Sinkhorn iteration to prevent $-C/\varepsilon$ from producing values that underflow even before LSE stabilization.
  • Log-domain throughout. The kernel $K = \exp(-C/\varepsilon)$ is never explicitly materialized. All computation stays in log-space.
  • Epsilon = 0.1 used for the final transport plan. Smaller $\varepsilon$ sharpens the plan but requires more iterations.

Dependencies

numpy
scipy
matplotlib
torchvision   # mnist download only