The Sinkhorn Algorithm

The cost of exact OT

The network simplex finds the exact optimal transport plan in O((N+M)^3 log(N+M)) time. For two histograms with 1000 bins each, that is roughly 8 billion operations per solve. In a training loop that calls solve thousands of times per epoch, exact OT is not practical.

Entropic regularisation cuts this to roughly O(N·M·T) — linear in the number of cells and the number of iterations — by solving a slightly different problem. The solution is no longer exactly optimal, but it is close, and the regularised version is differentiable, which exact OT is not.

Adding an entropy term

The regularised OT objective adds a penalty on the entropy of the plan:

OT_ε(a, b) = min_P  ⟨P, C⟩ − ε · H(P)

where H(P) = −∑_ij P[i,j] log P[i,j] is the entropy of P and ε > 0 is the regularisation strength. Maximising entropy pushes P toward a uniform distribution; the cost term pulls it toward the sparse, diagonal plan. The parameter ε controls the balance.

Sinkhorn iteration

The unique minimiser of the regularised objective has a factored form:

P_ε = diag(u) · K · diag(v)     where K[i,j] = exp(−C[i,j] / ε)

The vectors u and v are found by alternating row and column normalisations — this is the Sinkhorn algorithm:

import numpy as np

def sinkhorn_numpy(C, a, b, reg, n_iter=200):
    K = np.exp(-C / reg)
    v = np.ones(len(b), dtype=np.float64)

    for _ in range(n_iter):
        u = a / (K @ v + 1e-300)
        v = b / (K.T @ u + 1e-300)

    P = np.diag(u) @ K @ np.diag(v)
    return P

N = 16
grid = np.linspace(0, 1, N, dtype=np.float64)
a = np.exp(-0.5 * ((grid - 0.3) / 0.15) ** 2)
b = np.exp(-0.5 * ((grid - 0.7) / 0.15) ** 2)
a /= a.sum(); b /= b.sum()
C = ((grid[:, None] - grid[None, :]) ** 2).astype(np.float64)

P = sinkhorn_numpy(C, a, b, reg=0.05)
print(f"Row error: {np.abs(P.sum(axis=1) - a).max():.2e}")
print(f"Col error: {np.abs(P.sum(axis=0) - b).max():.2e}")

Each iteration enforces one of the two marginal constraints exactly while relaxing the other. Convergence is geometric: the marginal errors drop by a constant factor each round.

Log-domain stability

When ε is small, K[i,j] = exp(−C[i,j]/ε) underflows to zero for large costs. The u and v normalisations then divide by zero. The standard fix is to work in log space throughout, replacing u and v with dual potentials f = ε log u and g = ε log v. torchmatch uses this log-domain implementation:

import torch
import torchmatch
from torchmatch.transport.matrix import Backend

N = 16
grid = torch.linspace(0, 1, N)
a = torch.exp(-0.5 * ((grid - 0.3) / 0.15) ** 2)
b = torch.exp(-0.5 * ((grid - 0.7) / 0.15) ** 2)
a /= a.sum(); b /= b.sum()
C = (grid[:, None] - grid[None, :]) ** 2   # (N, N)

log_plan = torchmatch.transport.matrix.solve(
    C.unsqueeze(0),
    a=a.unsqueeze(0),
    b=b.unsqueeze(0),
    backend=Backend.LOG_SINKHORN,
    reg=0.05,
    n_iter=200,
)
P = log_plan.exp().squeeze(0)
print(f"Transport cost: {(P * C).sum():.4f}")

LOG_SINKHORN is the default backend selected by Backend.AUTO.

What ε controls

The key tradeoff:

• Large ε (e.g. 0.5): the entropy term dominates; P spreads mass everywhere. The plan looks close to the outer product a · bᵀ. The transport cost is above-optimal.
• Small ε (e.g. 0.005): the cost term dominates; P concentrates on the cheapest routes and approaches the exact OT plan. More iterations are needed to converge.

regs = [0.5, 0.1, 0.02, 0.005]

for reg in regs:
    P_r = sinkhorn_numpy(C.numpy(), a.numpy(), b.numpy(), reg=reg, n_iter=500)
    cost_r = (P_r * C.numpy()).sum()
    sparsity = (P_r > 1e-4).mean()
    print(f"ε={reg:5.3f}  cost={cost_r:.4f}  nonzero={sparsity:.1%}")

For most training use cases, reg between 0.01 and 0.1 is a good starting range.

Exact OT as a reference

When you need the true optimal plan — for analysis, for a small problem, or to validate that ε is small enough — the network simplex gives the exact answer:

import torch
import torchmatch
from torchmatch.transport.matrix import Backend

N = 32
grid = torch.linspace(0, 1, N, dtype=torch.float32)
a = torch.ones(N) / N
b = torch.ones(N) / N
C = (grid[:, None] - grid[None, :]) ** 2

plan_exact = torchmatch.transport.matrix.solve(
    C.unsqueeze(0),
    a=a.unsqueeze(0),
    b=b.unsqueeze(0),
    backend=Backend.EXACT_EMD,
)
P_exact = plan_exact.exp().squeeze(0)
print(f"Exact cost: {(P_exact * C).sum():.4f}")
print(f"Non-zero entries: {(P_exact > 1e-6).float().mean():.1%}")

EXACT_EMD runs on CPU only, does not accept reg or n_iter, and returns a sparse plan. For N > a few hundred it becomes slow; the Sinkhorn backends are more practical there.

Sinkhorn divergence

Raw Sinkhorn loss has a self-transport bias: even when a == b, OT_ε(a, a) > 0 because the entropy penalty pushes the plan away from the identity. This is a problem when using the loss as a training objective — the model can never reach zero loss, and the gradient at a == b is non-zero.

Sinkhorn divergence corrects for this by subtracting the self-transport costs:

SD_ε(a, b) = OT_ε(a, b) − ½ OT_ε(a, a) − ½ OT_ε(b, b)

It is zero when a == b, symmetric, and positive otherwise. Use it whenever the loss should be a proper distance — generative model training, distribution matching, point-cloud registration:

import torch
import torchmatch
from torchmatch.transport.matrix import Backend

N = 16
grid = torch.linspace(0, 1, N)
C = (grid[:, None] - grid[None, :]) ** 2

a_base = torch.exp(-0.5 * ((grid - 0.3) / 0.15) ** 2)
a_base /= a_base.sum()
C_t = C.unsqueeze(0)

for offset in [0.0, 0.1, 0.3, 0.5]:
    b = torch.exp(-0.5 * ((grid - 0.3 - offset) / 0.15) ** 2)
    b /= b.sum()
    div = torchmatch.transport.matrix.solve(
        C_t,
        a=a_base.unsqueeze(0),
        b=b.unsqueeze(0),
        backend=Backend.SINKHORN_DIVERGENCE,
        reg=0.05, n_iter=300,
    )
    print(f"offset={offset:.1f}  divergence={div.item():.5f}")
# offset=0.0 → divergence ≈ 0.0
# offset=0.5 → divergence > 0

SINKHORN_DIVERGENCE runs three Sinkhorn solves internally (for (a,b), (a,a), and (b,b)) and returns the scalar divergence, not a plan tensor.

Which backend to use

See also

• Optimal transport: the problem definition, transport plans, and Wasserstein distance.
• Point clouds and shapes: scaling to large point sets without materialising the cost matrix.
• Algorithms: the full derivation of Sinkhorn and log-domain stability.
• Reference: exact parameter names and constraints for each backend.