The Sinkhorn algorithm
The original (Kantorovitch) formulation of discrete optimal transport is as follows. Let \(a \in \R^n\) represents an initial distribution over a set $x_1, \dots, x_m$ of points in $\R^n$. Let $b \in \R^n$ represent a target distribution over a set $y_1, \dots, y_m \in \R^n$ of points. The Kantorovitch problem consists in finding a coupling distribution whose marginal are $a$ and $b$, that is a $m \times m$ that minimizes a transport cost. This writes
\[\begin{equation}\min_{P \in U(a,b)} \sum_{1 \leq i,j \leq n} c(x_i,y_j) P_{i,j}, \end{equation}\]where $c(x_i,y_j)$ is a cost function, $U(a,b)$ is the set of $P \in \R^{m \times m}$ such that $P 1 = a$ and $P^\mathrm{T}1 = b$. The objective function is linear and can be written $\langle C, P \rangle$ where $C_{i,j} = c(x_i, y_j)$. However the coupling constraint is not easily tractable. The Sinkhorn algorithm relies on the following recipes, to approximately solve the issue.
- Entropy regularization.
- Dual formulation.
- Fixed point property of the exponentiated dual variables.
Enropy regularization
Define the entropy function $H(P) = - \sum_{i,j} (P_{i,j} \log(P_{i,j}) - P_{i,j})$, where the last term is added for practicity (also remark that $(x\ln(x) -x)’ = \ln(x)$). We define the entropical regularization of discrete optimal transport as,
\[\min_{P \in U(a,b)}\langle C, P\rangle - \epsilon H(P), \tag{$OT_{\epsilon}$}\]where $\epsilon >0$. As its name suggest, the convex (prove it is!) function $H$ is highest when $P_{i,j}$ is constant, that is when the probability matrix $P$ is the most random. High regularization thus rewards disorder and non-sparsity in $P$. It also induces a strong convexity that provides uniqueness of solutions.
Remark: In passing, we observe that in $OT_{\epsilon}$, \(\langle C, P \rangle - \epsilon H(P) = \epsilon KL(P|\bar K)+cst,\) where $\bar K = (e^{-C_{i,j}/\epsilon})$ and $KL(P|K) = \sum P_{i,j}\ln(P_{i,j}/K_{i,j}) - P_{i,j}+ K_{i,j}$, making the entropic optimal transport a projection of the Gibbs measure $\bar K$ onto the set of couplings.
Dual formulation
In the case of discrete optimal transport, the coupling constraint is just a set of two linear equalities. Previous equation thus reformulates as,
\[\min_{P \in\R^{m \times m}} \langle C, P\rangle - \epsilon H(P) - \max_{f,g \in \R^n} f(P1-a) -g(P^{\mathrm{T}}1-b).\]This Lagrangian provides the dual problem,
\[\max_{f,g \in \R^n} \min_{P \in\R^{m \times m}} \langle C, P\rangle - \epsilon H(P) - f(P1-a) -g(P^{\mathrm{T}}1-b).\]KKT conditions impose that,
- (dual criticality = feasibility) $P\mathbf{1}=a$ and $P^{\mathrm{T}}\mathbf{1} = b$.
- (primal criticality) $\frac{\partial \mathcal{L}}{\partial P_{i,j}} = C_{i,j} - \epsilon \log(P_{i,j}) -f_i - g_j=0$ that is $P_{i,j} = e^{f_i/\epsilon} e^{C_{i,j}/\epsilon} e^{g_j/\epsilon}$, which writes in a matrix fashion, \(P = \mathrm{diag} e^{\odot -f/\epsilon} e^{\odot C/\epsilon} \mathrm{diag} e^{\odot -g/\epsilon},\) where the $\odot$ sign recalls all operations are element-wise.
For simplicity, we replace dual variables $f$ and $g$ by their exponentiated version $u := e^{\odot -f/\epsilon}$ and $v := e^{\odot -g/\epsilon}$, and note that $K:= (e^{C_{i,j}/\epsilon})$ is the Gibbs Kernel of costs $C$. The criticality conditions reformulate as,
\[\begin{cases} \boxed{P = \operatorname{diag}(u)K \operatorname{diag}(v)} \\ P\mathbf{1}=a \\ P^{\mathrm{T}}\mathbf{1} = b \end{cases} \qquad \Rightarrow \qquad \begin{cases} \operatorname{diag}(u)K v = a \\ \operatorname{diag}(v)K u = b. \end{cases} \qquad \Rightarrow \qquad \begin{cases} u = a / Kv & \text{(elementwise)} \\ v = b/Ku & \text{(elementwise)} \end{cases}\]Fixed point property
The unknown $u,~ v$ thus verify the fixed point condition in $\R^n \times \R^n$,
\[\begin{cases} u = a / Kv & \text{(elementwise)} \\ v = b/Ku & \text{(elementwise)} \end{cases}\]The solution can be approached using fixed point iterations. For stability, it is more common to use the Jacobi version of fixed point iterations, in which the updates are made alternatively on $u$ and $v$, resulting in the Sinkhorn-Knopp algorithm for OT.
References: