A Linear Complexity Algorithm for Optimal Transport Problem with Log-type Cost
Ziyuan Lyu, Zihao Wang, Hao Wu, Shuai Yang
TL;DR
This work tackles accelerating entropy-regularized optimal transport with log-type transport costs by introducing Fast Sinkhorn for Log-type Costs (FSL). FSL leverages a polynomial kernel expansion and series rearrangement to transform the costly matrix-vector products in Sinkhorn iterations into a linear-time procedure when $(ML)^d \ll N$, with $L=1/\varepsilon$ and $M$ the polynomial degree. The authors extend the framework to practical problems: (i) differentiable Sinkhorn ranking with log-type costs, and (ii) the far-field reflector and refractor costs in 2D, demonstrating substantial speedups while maintaining accuracy (transport plans nearly identical to the standard Sinkhorn solution). These contributions enable scalable OT computations for large-scale ranking and optical-design tasks and offer a pathway to higher-dimensional and spherical extensions.
Abstract
In [Q. Liao et al., Commun. Math. Sci., 20(2022)], a linear-time Sinkhorn algorithm is developed based on dynamic programming, which significantly reduces the computational complexity involved in solving optimal transport problems. However, this algorithm is specifically designed for the Wasserstein-1 metric. We are curious whether the preceding dynamic programming framework can be extended to tackle optimal transport problems with different transport costs. Notably, two special kinds of optimal transport problems, the Sinkhorn ranking and the far-field reflector and refractor problems, are closely associated with the log-type transport costs. Interestingly, by employing series rearrangement and dynamic programming techniques, it is feasible to perform the matrix-vector multiplication within the Sinkhorn iteration in linear time for this type of cost. This paper provides a detailed exposition of its implementation and applications, with numerical simulations demonstrating the effectiveness and efficiency of our methods.
