用神经网络加速多分布运输计算,速度提升数十倍。
Neural Estimation for Scaling Entropic Multimarginal Optimal Transport
- 用小批量训练的神经网络替代传统算法,降低计算复杂度。
- 在相同精度下,处理样本量和分布数均显著增加,速度提升数量级。
- 适合大规模机器学习中的多分布对齐任务,可直接嵌入流水线。
多边际最优传输(MOT)是建模多个分布间交互的强大框架,但其应用受限于高计算开销。熵正则化通过多边际Sinkhorn算法实现计算加速,但时间复杂度随数据集大小n和边际数k呈O(n^k)增长,对多数机器学习任务仍不可行。本文提出神经熵正则多边际最优传输(NEMOT),采用小批量训练的神经网络,将计算复杂度从数据集规模转移到小批量大小,实现显著提速。我们给出了非渐近误差界,证明NEMOT的精度保证。数值实验表明,相比现有方法,NEMOT在速度上取得数量级提升,可处理更多样本与边际。该方法可无缝集成至大规模机器学习流程中,拓展了熵正则MOT在多边缘数据任务中的实际应用能力。
原文摘要 · Abstract (English)
Multimarginal optimal transport (MOT) is a powerful framework for modeling interactions between multiple distributions, yet its applicability is bottlenecked by a high computational overhead. Entropic regularization provides computational speedups via the multimarginal Sinkhorn algorithm, whose time complexity, for a dataset size $n$ and $k$ marginals, generally scales as $O(n^k)$. However, this dependence on the dataset size $n$ is computationally prohibitive for many machine learning problems. In this work, we propose a new computational framework for entropic MOT, dubbed Neural Entropic MOT (NEMOT), that enjoys significantly improved scalability. NEMOT employs neural networks trained using mini-batches, which transfers the computational complexity from the dataset size to the size of the mini-batch, leading to substantial gains. We provide formal guarantees on the accuracy of NEMOT via non-asymptotic error bounds. We supplement these with numerical results that demonstrate the performance gains of NEMOT over Sinkhorn's algorithm, as well as extensions to neural computation of multimarginal entropic Gromov-Wasserstein alignment. In particular, orders-of-magnitude speedups are observed relative to the state-of-the-art, with a notable increase in the feasible number of samples and marginals. NEMOT seamlessly integrates as a module in large-scale machine learning pipelines, and can serve to expand the practical applicability of entropic MOT for tasks involving multimarginal data.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。