发表机构
Shanghai Jiao Tong University; Institute of Natural Sciences, Shanghai Jiao Tong University(上海交通大学; 上海交通大学自然科学研究院)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
提出HELLO分层求解器,利用对偶势引导边缘定位,高效求解大规模离散最优传输,在百万点规模上显著提升速度与精度,并支持多种OT变体。
AI 中文摘要
最优传输(OT)在机器学习中用于比较分布和对齐数据集,然而无正则化的离散OT需要求解一个具有二次数量传输变量的线性规划。我们提出HELLO,一种分层求解器,将大规模离散OT转化为边缘定位问题,并利用对偶势引导从粗到细的初始化及层内细化。初始化阶段通过递归子采样层级传播粗对偶势以分配候选边缘。细化阶段则迭代地在每行和每列中插入最大的对偶违反者,直到相对KKT残差满足预定容差,同时通过预算修剪确保线性内存复杂度。对于精确算术细化,我们证明了在符号字典序规则下,算法在全局最优处有限终止。在百万点规模上,HELLO在从个位数到数千的特征维度上,相较于强基线实现了更低的传输目标值和一个数量级的运行时间改进。它还能在单个H100上扩展到每边际128万样本、8192维,峰值GPU内存使用41.6 GiB,同时满足相对KKT残差低于$10^{-6}$。除了标准离散OT,该框架支持一般成对代价,并可作为半离散OT、Gromov--Wasserstein、非平衡OT以及基于OT的流匹配的可扩展平衡OT预言机。
英文摘要
Optimal transport (OT) compares distributions and aligns datasets in machine learning, yet unregularized discrete OT requires a linear program with quadratically many transport variables. We propose HELLO, a hierarchical solver that casts large-scale discrete OT as edge localization and uses dual potentials to guide both coarse-to-fine initialization and within-level refinement. Initialization propagates coarse dual potentials across a recursive subsampling hierarchy to assign candidate edges. Refinement then iteratively inserts the largest dual violators in each row and column until the relative KKT residual meets a prescribed tolerance, while budgeted pruning ensures linear memory complexity. For exact-arithmetic refinement, we prove finite termination at a global optimum under a symbolic lexicographic rule. At the million-point scale, HELLO attains lower transport objectives with order-of-magnitude runtime improvements over strong baselines across feature dimensions from single digits to thousands. It further scales to 1.28 million samples per marginal in 8192 dimensions on a single H100, using 41.6 GiB peak GPU memory while satisfying a full relative KKT residual below $10^{-6}$. Beyond standard discrete OT, the framework supports general pairwise costs and serves as a scalable balanced-OT oracle for semi-discrete OT, Gromov--Wasserstein, unbalanced OT, and OT-based Flow Matching.