Skip to content

optora.divergences.wasserstein

Entropy-regularized Wasserstein (Sinkhorn) divergence between distributions.

SinkhornDivergence

Bases: Divergence

Debiased entropic Wasserstein divergence between discrete distributions.

For discrete distributions p and q (nonnegative tensors that sum to one along their last dimension) sharing a common support with pairwise ground cost cost, the entropic optimal transport cost is

\[ \mathrm{OT}_\epsilon(p, q) = \min_{\pi \in U(p, q)} \langle \mathrm{cost}, \pi \rangle + \epsilon \sum_{ij} \pi_{ij} \big(\log \pi_{ij} - 1\big) \]

where \(U(p, q)\) is the set of transport plans (joint distributions) with marginals p and q, and \(\epsilon\) is the entropic regularization strength. Sinkhorn's algorithm computes the minimizing plan \(\pi\) by alternately updating dual potentials f and g until both marginal constraints hold, and \(\mathrm{OT}_\epsilon\) is then read off as \(\langle \mathrm{cost}, \pi \rangle\) for the converged plan \(\pi = \exp\!\big((f \oplus g - \mathrm{cost}) / \epsilon\big)\). The potential updates are computed with torch.logsumexp (the log-sum-exp trick) rather than by forming the Gibbs kernel \(\exp(-\mathrm{cost} / \epsilon)\) directly: that kernel underflows to exact zero for small \(\epsilon\) or large cost, which silently collapses the whole divergence to zero instead of raising an error, so it is avoided rather than merely guarded against with a larger \(\epsilon\).

Plain entropic OT cost is biased: \(\mathrm{OT}_\epsilon(p, p)\) is not exactly zero for eps > 0. This class instead computes the debiased Sinkhorn divergence

\[ S_\epsilon(p, q) = \mathrm{OT}_\epsilon(p, q) - \tfrac{1}{2}\mathrm{OT}_\epsilon(p, p) - \tfrac{1}{2}\mathrm{OT}_\epsilon(q, q), \]

which removes that self-transport bias so \(S_\epsilon(p, p) = 0\) exactly, as required by the Divergence contract, while still converging to the Wasserstein distance induced by cost as \(\epsilon \to 0\). optora.dro.wasserstein_dro uses this divergence to define Wasserstein-based ambiguity sets.

Attributes:

Name Type Description
cost Tensor

Square, nonnegative pairwise ground cost matrix between the shared support points of p and q, shape (n, n).
epsilon

Positive entropic regularization strength; smaller values approximate the exact Wasserstein distance more closely at the cost of more Sinkhorn iterations to converge.
max_iter

Maximum number of Sinkhorn scaling iterations.
tol

Convergence tolerance on the change in the row dual potential between iterations.
eps

Small positive constant used to clamp p and q away from zero before taking the logarithm, avoiding log(0) without branching.
check_interval

Number of Sinkhorn iterations between host reads of the convergence flag. The potentials are frozen once converged, so the divergence does not depend on this interval.

__init__(cost, epsilon=0.1, max_iter=100, tol=1e-06, eps=1e-12, check_interval=DEFAULT_CHECK_INTERVAL, validate=False)

Initialize the Sinkhorn divergence.

Parameters:

Name Type Description Default
cost Tensor

Square, nonnegative pairwise ground cost matrix between the shared support points of p and q, shape (n, n).

 required
epsilon float

Positive entropic regularization strength.

 0.1
max_iter int

Maximum number of Sinkhorn scaling iterations.

 100
tol float

Convergence tolerance on the change in the row dual potential between iterations.

 1e-06
eps float

Small positive constant used to clamp p and q away from zero before taking the logarithm.

 1e-12
check_interval int

Number of Sinkhorn iterations between host reads of the convergence flag. Raise it to trade redundant frozen iterations for fewer device synchronizations.

 DEFAULT_CHECK_INTERVAL
validate bool

Whether to check that cost is nonnegative. The check reads a reduction over cost on the host, which blocks until the device has produced it, so it is opt-in and off by default to keep construction asynchronous. The shape and hyperparameter checks are metadata-only and always run.

 False

Raises:

Type Description
ValueError

If cost is not a square 2D tensor, if validate is set and cost contains negative entries, or if epsilon, max_iter, tol, eps, or check_interval are not positive.

forward(p, q)

Compute the debiased Sinkhorn divergence of p from q.

Parameters:

Name Type Description Default
p Tensor

Candidate distribution, a nonnegative tensor of shape (n,) that sums to one, indexing cost.

 required
q Tensor

Reference distribution, a nonnegative tensor of shape (n,) that sums to one, indexing cost.

 required

Returns:

Type Description
Tensor

A scalar tensor holding \(S_\epsilon(p, q)\), clamped to be
Tensor

nonnegative to absorb floating-point error near zero.

Raises:

Type Description
ValueError

If the shape of p or q does not match cost.