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.
- Fiexed 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. We define the entropical regularization of discrete optimal transport as, \(\min_{P \in U(a,b)}\langle C, P\rangle - \epsilon H(P),\) 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.
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) $P1=a$ and $P^{\mathrm{T}}1 = b$.
- (primal criticality) $\frac{\partial \mathcal{L}}{\partial P_{i,j}} = C_{i,j} - \epsilon \log(P_{i,j}) -f_i - g_j$ 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.