Task Memory Sinkhorn Neuralization under varying measure distributions

Optimal Transport (OT) and its important generalization, Unbalanced Optimal Transport (UOT), have emerged as powerful tools in machine learning. However, the Sinkhorn algorithm, a cornerstone solver for both problems, is often hindered by a critical computational bottleneck that arises from its iterative nature. In this paper, we are concerned with neural network approximation of the Sinkhorn algorithm in variable external calling environments. We study approximating and accelerating Sinkhorn in the more general and practical varying-measure-distribution context and propose a new neural network-based Sinkhorn approximation method. In the proposed method, the approximation problem is formulated as approximating convex conjugates determined by task characteristics, thereby enabling successive approximations across trials and epochs. We apply our method to various tasks, including multilabel classification, crowd counting, and knowledge distillation. Extensive experiments on benchmarks such as YFCC-100M, UCF-QNRF, and EURLEX57K demonstrate the efficacy of our method. Our code will be released.