The Sinkhorn algorithm is the state-of-the-art to approximate solutions of entropic optimal transport (OT) distances between discrete probability distributions. We show that meticulously training a neural network to learn initializations to the algorithm via the entropic OT dual problem can significantly speed up convergence, while maintaining desirable properties of the Sinkhorn algorithm, such as differentiability and parallelizability. We train our predictive network in an adversarial fashion using a second, generating network and a self-supervised bootstrapping loss. The predictive network is universal in the sense that it is able to generalize to any pair of distributions of fixed dimension and cost at inference, and we prove that we can make the generating network universal in the sense that it is capable of producing any pair of distributions during training. Furthermore, we show that our network can even be used as a standalone OT solver to approximate regularized transport distances to a few percent error, which makes it the first meta neural OT solver.
翻译:Sinkhorn算法是当前最先进的方法,用于近似离散概率分布之间的熵正则化最优传输(OT)距离。我们证明,通过熵正则化OT对偶问题精心训练神经网络以学习该算法的初始化,可以显著加速收敛,同时保持Sinkhorn算法的理想特性,如可微性和可并行性。我们采用对抗方式训练预测网络,使用第二生成网络和自监督自举损失函数。预测网络具有通用性,即在推理时能够泛化到任意固定维度和成本的分布对,并且我们证明可以使生成网络在训练期间能够产生任意分布对。此外,我们还表明该网络甚至可以作为独立的OT求解器使用,以百分之几的误差近似正则化传输距离,这使其成为首个元神经OT求解器。