ENOT：用于神经最优传输快速准确训练的分位数正则化
机器学习
2024-10-21 v4 人工智能
摘要
我们提出了一种新的神经最优传输 (NOT) 训练方法，能够通过对偶 Kantorovich 势的特定正则化，准确且高效地估计最优传输计划。现有 NOT 求解器的主要瓶颈在于寻找共轭算子（即 c-变换）的近似精确解的过程，这通常需要通过优化非凸极大极小目标或对初始近似预测进行计算密集的精细调整来实现。我们通过提出一种新的、具有理论依据的分位数正则化 (expectile regularisation) 损失形式解决了这两个问题，该损失对对偶势的学习过程施加了约束条件。这种正则化提供了可能共轭势分布的上界估计，使学习过程稳定，完全消除了对额外广泛精细调整的需求。所提出的方法称为分位数正则化神经最优传输 (ENOT)，在既定的 Wasserstein-2 基准任务上大幅超越了之前的最先进 (SOTA) 方法（质量提升高达 3 倍，运行时间提升高达 10 倍）。此外，我们展示了 ENOT 在不同成本函数下于图像生成等不同任务中的性能，证明了所提算法的鲁棒性。OTT-JAX 库包含了我们的 ENOT 算法实现：https://ott-jax.readthedocs.io/en/latest/tutorials/ENOT.html
引用
@article{arxiv.2403.03777,
title = {ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport},
author = {Nazar Buzun and Maksim Bobrin and Dmitry V. Dylov},
journal= {arXiv preprint arXiv:2403.03777},
year = {2024}
}