We present a new extension for Neural Optimal Transport (NOT) training procedure, capable of accurately and efficiently estimating optimal transportation plan via specific regularisation on conjugate potentials. The main bottleneck of existing NOT solvers is associated with the procedure of finding a near-exact approximation of the conjugate operator (i.e., the c-transform), which is done either by optimizing over maximin objectives or by the computationally-intensive fine-tuning of the initial approximated prediction. We resolve both issues by proposing a new, theoretically justified loss in the form of expectile regularization that enforces binding conditions on the learning dual potentials. Such a regularization provides the upper bound estimation over the distribution of possible conjugate potentials and makes the learning stable, eliminating the need for additional extensive finetuning. We formally justify the efficiency of our method, called Expectile-Regularised Neural Optimal Transport (ENOT). ENOT outperforms previous state-of-the-art approaches on the Wasserstein-2 benchmark tasks by a large margin (up to a 3-fold improvement in quality and up to a 10-fold improvement in runtime).

翻译：我们提出了一种神经最优传输（NOT）训练过程的新扩展方法，通过对共轭势能施加特定正则化，能够准确高效地估计最优传输方案。现有NOT求解器的主要瓶颈在于寻找共轭算子（即c-变换）的近乎精确近似，这一过程要么通过优化极大极小目标函数实现，要么依赖于计算密集型的初始近似预测微调。我们通过提出一种理论完备的期望分位正则化损失函数解决了这两个问题，该正则化强制约束对偶势能的学习条件。这种正则化能为可能的共轭势能分布提供上界估计，使学习过程更加稳定，从而消除了额外大规模微调的必要。我们严格论证了所提方法（称为期望分位正则化神经最优传输，ENOT）的优越性。在Wasserstein-2基准任务上，ENOT以显著优势超越了此前最先进方法（质量提升达3倍，运行速度提升达10倍）。