Schrödinger Bridge Problem

The Schrödinger Bridge problem finds the most likely stochastic evolution connecting two observed probability distributions. In the urban context, it answers: given census snapshots of population density at two different times, what is the most probable migration pattern that connects them?

1. Problem Statement

Let $Q$ be a reference stochastic process (e.g., Brownian motion) on $[0, T]$. We observe the marginal distributions at the endpoints:

$$\rho_0(x) = P(x, 0), \qquad \rho_T(x) = P(x, T)$$

The Schrödinger Bridge seeks the path measure $P^*$ that:

$$P^* = \arg\min_{P} \; \text{KL}(P \| Q) \quad \text{subject to} \quad P(x, 0) = \rho_0, \;\; P(x, T) = \rho_T$$

where the Kullback-Leibler divergence between path measures is:

$$\text{KL}(P \| Q) = \int \log\frac{dP}{dQ} \, dP$$

Intuitively, $P^*$ is the stochastic evolution that is closest to the reference process $Q$ while exactly matching the observed marginals. It represents the most likely transport from $\rho_0$ to $\rho_T$.

2. The Schrödinger System

When the reference process $Q$ is Brownian motion with diffusion $\epsilon/2$, the optimal $P^*$ can be characterised by two functions $\varphi(x, t)$ and $\hat{\varphi}(x, t)$ satisfying coupled PDEs:

$$\frac{\partial \varphi}{\partial t} = -\frac{\epsilon}{2}\Delta \varphi + V\varphi \qquad (\text{backward heat equation})$$

$$\frac{\partial \hat{\varphi}}{\partial t} = \frac{\epsilon}{2}\Delta \hat{\varphi} - V\hat{\varphi} \qquad (\text{forward heat equation})$$

The density at any intermediate time is given by the product:

$$\rho(x, t) = \varphi(x, t) \cdot \hat{\varphi}(x, t)$$

For the free case ($V = 0$), the system simplifies. The heat kernel is:

$$K_\epsilon(x, y; t) = \frac{1}{\sqrt{2\pi\epsilon t}} \exp\!\left(-\frac{(x-y)^2}{2\epsilon t}\right)$$

and the Schrödinger system becomes:$\; \varphi(x, 0) \cdot \hat{\varphi}(x, 0) = \rho_0(x)$ and$\varphi(x, T) \cdot \hat{\varphi}(x, T) = \rho_T(x)$, with $\hat{\varphi}(x, t) = \int K_\epsilon(x, y; t) \, \varphi(y, 0) \, dy$.

3. Sinkhorn Algorithm

In the discrete setting with $N$ spatial points, the heat kernel becomes a matrix $\mathbf{K}$ with entries $K_{ij} = K_\epsilon(x_i, x_j; T)$. The Schrödinger Bridge reduces to finding vectors $\boldsymbol{\varphi}$ and $\hat{\boldsymbol{\varphi}}$ such that:

$$\boldsymbol{\varphi} \odot (\mathbf{K} \hat{\boldsymbol{\varphi}}) = \boldsymbol{\rho}_0, \qquad \hat{\boldsymbol{\varphi}} \odot (\mathbf{K}^T \boldsymbol{\varphi}) = \boldsymbol{\rho}_T$$

The Sinkhorn algorithm solves this by alternating projections:

Algorithm: Sinkhorn Iterations

Initialise $\hat{\boldsymbol{\varphi}}^{(0)} = \mathbf{1}$. Then iterate:

$$\boldsymbol{\varphi}^{(k+1)} = \frac{\boldsymbol{\rho}_0}{\mathbf{K} \hat{\boldsymbol{\varphi}}^{(k)}}$$

$$\hat{\boldsymbol{\varphi}}^{(k+1)} = \frac{\boldsymbol{\rho}_T}{\mathbf{K}^T \boldsymbol{\varphi}^{(k+1)}}$$

This converges linearly with rate determined by the second-largest eigenvalue of $\mathbf{K}$. Convergence is guaranteed when $\mathbf{K}$ has strictly positive entries.

The resulting optimal transport plan (coupling) is:

$$\pi^*_{ij} = \varphi_i \, K_{ij} \, \hat{\varphi}_j$$

This is the entropic regularisation of optimal transport, with $\epsilon$ controlling the regularisation strength. As $\epsilon \to 0$, the solution approaches the Monge-Kantorovich optimal transport plan.

4. Connection to Quantum Mechanics

The Schrödinger Bridge has a deep connection to quantum mechanics via Nelson's stochastic mechanics. In Nelson's framework, the Schrödinger equation:

$$i\hbar \frac{\partial \psi}{\partial t} = -\frac{\hbar^2}{2m}\Delta\psi + V\psi$$

can be rewritten via $\psi = e^{R + iS/\hbar}$ into two real equations. Under the substitution $t \to it$ (Wick rotation), the quantum system becomes exactly the Schrödinger Bridge system with $\epsilon = \hbar/m$.

The analogy is precise:

Quantum Mechanics	Schrödinger Bridge
Wave function $\psi$	Forward potential $\varphi$
Born rule $\|\psi\|^2$	Density $\varphi \hat{\varphi}$
Planck constant $\hbar$	Diffusion $\epsilon$
Propagator	Heat kernel $K_\epsilon$

5. Python: Sinkhorn for Urban Density Transport

We transport a bimodal urban density (two population centres) to a unimodal one (merged city), finding the optimal Schrödinger Bridge connecting them.

Sinkhorn Algorithm for 1D Urban Density Transport

Python

script.py149 lines

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

# --- Spatial grid ---
N = 150
x = np.linspace(-5, 5, N)
dx = x[1] - x[0]

# --- Marginals ---
# rho_0: bimodal (two suburbs)
rho_0 = 0.4 * np.exp(-(x + 1.8)**2 / (2 * 0.4**2)) + 0.6 * np.exp(-(x - 2.0)**2 / (2 * 0.5**2))
rho_0 /= np.sum(rho_0) * dx

# rho_T: unimodal (merged city centre)
rho_T = np.exp(-(x - 0.3)**2 / (2 * 0.8**2))
rho_T /= np.sum(rho_T) * dx

# --- Heat kernel matrix ---
epsilon = 0.8  # diffusion parameter
T_bridge = 1.0
sigma_K = np.sqrt(epsilon * T_bridge)

K = np.zeros((N, N))
for i in range(N):
    K[i, :] = np.exp(-(x[i] - x)**2 / (2 * sigma_K**2))
K *= dx / np.sqrt(2 * np.pi * sigma_K**2)

# --- Sinkhorn iterations ---
max_iter = 200
phi_hat = np.ones(N)  # hat{phi}
phi = np.ones(N)

marginal_errors_0 = []
marginal_errors_T = []

for it in range(max_iter):
    # Update phi
    K_phi_hat = K @ phi_hat
    K_phi_hat = np.maximum(K_phi_hat, 1e-30)
    phi = rho_0 / K_phi_hat

# Update phi_hat
    KT_phi = K.T @ phi
    KT_phi = np.maximum(KT_phi, 1e-30)
    phi_hat = rho_T / KT_phi

# Check marginal errors
    marginal_0 = phi * (K @ phi_hat)
    marginal_T_check = phi_hat * (K.T @ phi)
    err_0 = np.sum(np.abs(marginal_0 - rho_0)) * dx
    err_T = np.sum(np.abs(marginal_T_check - rho_T)) * dx
    marginal_errors_0.append(err_0)
    marginal_errors_T.append(err_T)

# --- Transport plan ---
Pi = np.outer(phi, phi_hat) * K

# --- Intermediate densities via interpolation ---
n_times = 5
times = np.linspace(0, 1, n_times)
rho_interp = []

for t_frac in times:
    if t_frac == 0:
        rho_interp.append(rho_0.copy())
    elif t_frac == 1:
        rho_interp.append(rho_T.copy())
    else:
        # Approximate: interpolate the Schrodinger potentials
        s_fwd = np.sqrt(epsilon * t_frac * T_bridge)
        s_bwd = np.sqrt(epsilon * (1 - t_frac) * T_bridge)

K_fwd = np.zeros((N, N))
        K_bwd = np.zeros((N, N))
        for i in range(N):
            K_fwd[i, :] = np.exp(-(x[i] - x)**2 / (2 * s_fwd**2)) * dx / np.sqrt(2 * np.pi * s_fwd**2)
            K_bwd[i, :] = np.exp(-(x[i] - x)**2 / (2 * s_bwd**2)) * dx / np.sqrt(2 * np.pi * s_bwd**2)

phi_t = K_fwd @ phi
        phi_hat_t = K_bwd @ phi_hat
        rho_t = phi_t * phi_hat_t
        rho_t = np.maximum(rho_t, 0)
        if np.sum(rho_t) * dx > 0:
            rho_t /= np.sum(rho_t) * dx
        rho_interp.append(rho_t)

# --- Plot ---
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Top left: marginals
ax = axes[0, 0]
ax.fill_between(x, rho_0, alpha=0.3, color='#10b981')
ax.plot(x, rho_0, color='#10b981', linewidth=2.5, label=r'$\rho_0$ (two suburbs)')
ax.fill_between(x, rho_T, alpha=0.3, color='#f59e0b')
ax.plot(x, rho_T, color='#f59e0b', linewidth=2.5, label=r'$\rho_T$ (merged city)')
ax.set_xlabel('x', fontsize=12, color='white')
ax.set_ylabel('Density', fontsize=12, color='white')
ax.set_title('Initial and Final Marginals', fontsize=13, color='white')
ax.legend(fontsize=10, facecolor='#1e293b', edgecolor='#10b981', labelcolor='white')
ax.tick_params(colors='white')
ax.grid(True, alpha=0.2)

# Top right: Sinkhorn convergence
ax = axes[0, 1]
ax.semilogy(marginal_errors_0, color='#10b981', linewidth=2, label='Marginal 0 error')
ax.semilogy(marginal_errors_T, color='#f59e0b', linewidth=2, label='Marginal T error')
ax.set_xlabel('Iteration', fontsize=12, color='white')
ax.set_ylabel('L1 Error', fontsize=12, color='white')
ax.set_title('Sinkhorn Convergence', fontsize=13, color='white')
ax.legend(fontsize=10, facecolor='#1e293b', edgecolor='#10b981', labelcolor='white')
ax.tick_params(colors='white')
ax.grid(True, alpha=0.2)

# Bottom left: intermediate densities
ax = axes[1, 0]
colors_t = ['#10b981', '#06b6d4', '#8b5cf6', '#f59e0b', '#ef4444']
for i, (t_val, rho_t) in enumerate(zip(times, rho_interp)):
    ax.plot(x, rho_t, color=colors_t[i], linewidth=2, label=f't = {t_val:.2f}')
ax.set_xlabel('x', fontsize=12, color='white')
ax.set_ylabel('Density', fontsize=12, color='white')
ax.set_title('Density Evolution (Schrodinger Bridge)', fontsize=13, color='white')
ax.legend(fontsize=9, facecolor='#1e293b', edgecolor='#10b981', labelcolor='white')
ax.tick_params(colors='white')
ax.grid(True, alpha=0.2)

# Bottom right: transport plan
ax = axes[1, 1]
extent = [x[0], x[-1], x[-1], x[0]]
im = ax.imshow(Pi, extent=extent, aspect='auto', cmap='viridis')
ax.set_xlabel('Destination x', fontsize=12, color='white')
ax.set_ylabel('Origin x', fontsize=12, color='white')
ax.set_title(r'Optimal Transport Plan $pi^*$', fontsize=13, color='white')
ax.tick_params(colors='white')
cbar = plt.colorbar(im, ax=ax, fraction=0.046)
cbar.ax.tick_params(colors='white')

fig.patch.set_facecolor('#0f172a')
for ax_row in axes:
    for ax in ax_row:
        ax.set_facecolor('#0f172a')
plt.tight_layout()
plt.savefig('output.png', dpi=130, bbox_inches='tight', facecolor='#0f172a')
plt.show()

print(f"Final marginal errors: rho_0={marginal_errors_0[-1]:.2e}, rho_T={marginal_errors_T[-1]:.2e}")
print(f"Transport plan mass: {np.sum(Pi) * dx**2:.4f} (should be ~1.0)")
print(f"Epsilon (diffusion): {epsilon}, Sinkhorn iterations: {max_iter}")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

6. Properties and Urban Applications

Key properties of the Schrödinger Bridge:

Entropic regularisation: As $\epsilon \to 0$, the bridge approaches the deterministic optimal transport (Monge-Kantorovich) solution. Larger $\epsilon$ gives more diffuse transport.
Uniqueness: The Schrödinger Bridge is unique when $\rho_0, \rho_T > 0$ everywhere.
Time symmetry: The bridge from $\rho_0$ to $\rho_T$ is the time-reversal of the bridge from $\rho_T$ to $\rho_0$.
Urban interpretation: The Schrödinger Bridge gives the most likely population flow between two census snapshots, accounting for random individual mobility decisions. The parameter $\epsilon$ encodes the "temperature" of mobility decisions.

Connection to Wasserstein Distance

The cost of the Schrödinger Bridge is related to the Wasserstein-2 distance:

$$\epsilon \cdot \text{KL}(P^* \| Q) = \frac{1}{2}W_2^2(\rho_0, \rho_T) + \mathcal{O}(\epsilon)$$

This establishes the Schrödinger Bridge as a regularised version of optimal transport, computationally tractable via the Sinkhorn algorithm with $O(N^2)$ per iteration.

← Fokker-Planck Equation Part VI: Network Analysis →

Quantum Mechanics	Schrödinger Bridge
Wave function \(\psi\)	Forward potential \(\varphi\)
Born rule \(\|\psi\|^2\)	Density \(\varphi \hat{\varphi}\)
Planck constant \(\hbar\)	Diffusion \(\epsilon\)
Propagator	Heat kernel \(K_\epsilon\)