Rethinking Initialization of the Sinkhorn Algorithm

While the optimal transport (OT) problem was originally formulated as a linear program, regularizing it with an entropic penalty has been favored by practitioners in many recent applications, where that regularization is seen as beneficial from both computational and statistical perspectives. The Sinkhorn fixed-point algorithm isthe most popular approach to solve that regularized problem, and, as a result, multiple attempts have been made to reduce its runtime using, e.g., annealing in the regularization parameter, momentum or acceleration in the iterates. The premiseof this work is that initialization of the Sinkhorn algorithm has received comparatively little attention, possibly due to two preconceptions: since the regularized OT problem is convex, it may not be worth crafting a good initialization, since any isguaranteed to work; secondly, because the outputs of the Sinkhorn algorithm are often differentiated in end-to-end pipelines, a data-dependent initialization would bias Jacobian estimates obtained when unrolling iterations. We challenge this conventional wisdom, and show that data-dependent initializers result in dramatic speed-ups, without affecting the correctness of Jacobian maps, as long as those are recovered using implicit differentiation. Our initializations rely on simple closed-forms for exact or approximate OT solutions, using known results in the 1D, Gaussianor GMM settings. These initializations can be used with little to no prior tuning, and result in consistent speed-ups for a variety of OT problems.