Sinkhorn Algorithm

Sinkhorn solves entropy-smoothed transport by alternately rescaling matrix rows and columns, giving fast GPU-friendly approximate couplings.

Alternate row and column scalings converge a kernel to the required margins

Why Does This Exist?

Exact transport linear programs scale cubically and resist GPUs, yet ML needs transport daily: matching batches, aligning embeddings, averaging distributions. Adding a little entropy turns the hard program into matrix scaling with a unique solution reached by alternating normalization. Divide rows to hit source margins, divide columns to hit target margins, repeat. Each sweep is a batched multiply, so thousands of small problems solve side by side on accelerators.

This page covers the iteration itself: the kernel, the updates, and what the smoothing costs. The linear algebra underneath lives in linear systems, and the cost language lives in complexity analysis.

The one idea you need first

The Gibbs kernel $K_{ij} = \exp(-C_{ij}/\epsilon)$ turns costs into affinities, where $C_{ij}$ is the cost of moving mass from $i$ to $j$ and $\epsilon$ controls smoothing. Cheap routes get weight near 1, expensive ones near 0. The optimal smoothed coupling has factored form $\pi = \text{diag}(u)K\text{diag}(v)$ , where $u$ and $v$ are scaling vectors that enforce the margins.

Think of It Like This

Balancing a spreadsheet two ways

A budget table must hit row totals and column totals at once. Fix the rows and the columns break. Fix the columns and the rows break. So you alternate: scale each row to its target, then each column, then rows again.

Each pass cuts the remaining error, and the table settles on the unique sheet satisfying both sides. Sinkhorn is that alternating fix with multiplicative corrections instead of additive ones.

Where the analogy stops: the kernel encodes costs exponentially, so tiny smoothing makes entries underflow without log-domain care.

How It Actually Works

Start from $v = 1$ , where $1$ means a vector of ones. Then iterate:

u \gets \mu/(Kv), \qquad v \gets \nu/(K^T u)

In plain words: $u$ rescales rows so $Kv$ hits the source margins $\mu$ , then $v$ rescales columns so $K^T u$ hits the target margins $\nu$ . Each sweep costs $O(nm)$ for $n$ sources and $m$ targets. Convergence is linear with a rate tied to the kernel contraction. Small $\epsilon$ sharpens the plan toward exact transport but slows and destabilizes the iteration. Log-domain updates on potentials $f$ and $g$ with softmin operations fix the underflow.

Worked example

Take $\mu = \nu = (0.5, 0.5)$ , costs $C = [[0,1],[1,0]]$ , and $\epsilon = 0.5$ . Then $K = [[1, e^{-2}],[e^{-2},1]] \approx [[1, 0.135],[0.135,1]]$ , because $e^{-2} \approx 0.13534$ .

Init $v = (1,1)$ : $Kv = (1.135, 1.135)$ . Then $u = (0.5/1.135, 0.5/1.135) \approx (0.44, 0.44)$ .
Then $K^T u \approx (0.5, 0.5)$ , so $v = (1.0, 1.0)$ . Symmetry converges in one step.
Plan $\pi_{ij} = u_i K_{ij} v_j$ : diagonal $0.44$ , off-diagonal $0.44 \times 0.135 \approx 0.06$ . Row sums are $0.44 + 0.06 = 0.5$ , and column sums match.
Cost is $2 \times 0.06 \times 1 = 0.12$ . Exact transport would put $0.5$ on the diagonal with cost $0$ . The $0.12$ gap is the smoothing price at $\epsilon = 0.5$ .

Code

The point is the alternating rescaling, not the formula retyped. This runs the iteration on the worked fixture and prints the plan the hand trace found.

1import numpy as np2
3
4def sinkhorn(mu: np.ndarray, nu: np.ndarray, C: np.ndarray, eps: float, iters: int = 100) -> np.ndarray:5    K = np.exp(-C / eps)  # Gibbs kernel turns costs into affinities6    v = np.ones_like(nu)7    for _ in range(iters):8        u = mu / (K @ v)  # match row margins9        v = nu / (K.T @ u)  # match column margins10    return np.diag(u) @ K @ np.diag(v)11
12
13C = np.array([[0.0, 1.0], [1.0, 0.0]])14print(np.round(sinkhorn(np.array([0.5, 0.5]), np.array([0.5, 0.5]), C, 0.5), 3))15# -> [[0.44 0.06] [0.06 0.44]] matching the hand computation

Diagonal $0.44$ and off-diagonal $0.06$ agree with the hand-derived $u = (0.44, 0.44)$ and $v = (1, 1)$ .

Watch Out For

Running Sinkhorn in raw space with tiny epsilon

Entries $\exp(-C/\epsilon)$ underflow to zero in float32 once the exponent drops below about $-88$ , producing division by zero and NaN plans. Small $\epsilon$ triggers it first, exactly when you wanted sharper answers.

Symptom: NaNs appear as $\epsilon$ shrinks, while larger $\epsilon$ runs fine. Fix: switch to log-domain softmin updates on $f$ and $g$ , or stabilize with periodic rescaling once $\epsilon$ drops under the cost median.

Stopping after a fixed 10 iterations always

Loose tolerance leaves margin errors that downstream barycenters amplify into visible artifacts. Fixed counts ignore problem difficulty: easy symmetric cases converge in one step, while small $\epsilon$ cases need hundreds.

Symptom: rows sum to $0.48$ instead of $0.5$ , and averages drift. Fix: check marginal violation $\|r - \mu\|_1$ and iterate to $10^{-6}$ or your loss tolerance instead of trusting a count.

The Quick Version

Sinkhorn solves entropic transport by alternating row and column rescaling to meet both margins.
The Gibbs kernel encodes costs exponentially, while scaling vectors $u$ and $v$ enforce the margins.
Each iteration is matrix-vector products, which is why it runs well on GPUs.
On the worked $2 \times 2$ case, $u = (0.44, 0.44)$ and $v = (1, 1)$ give diagonal $0.44$ and cost $0.12$ versus exact cost $0$ .
Small smoothing sharpens plans but needs log-domain arithmetic to avoid underflow.
Verify marginal errors instead of trusting a fixed iteration count.