带动量的退火Sinkhorn:线性内存下带认证的无正则化最优传输
Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory
- MIT(麻省理工学院)
- HEC Paris(巴黎高等商学院)
机构由 AI 辅助整理,请以论文原文为准。
AI总结:
本文提出带动量的退火Sinkhorn(BDRS)方法,用于无正则化离散最优传输,在线性内存中实现原始-对偶认证,并在颜色转移任务中显著加速和降低对偶间隙。
AI中文摘要:
我们刻画了无正则化离散最优传输的Bregman Douglas-Rachford分裂(BDRS),并在线性内存中开发了一个任意时间的原始-对偶认证。我们首先证明,BDRS与使用单次内层Sinkhorn迭代的精确最优传输(IPOT)的热启动不精确近端点方法一致。通过从更新中消除原始传输计划,我们推导出一个等价的对偶形式,揭示BDRS是在隐式逆线性温度调度下的退火Sinkhorn,并带有额外的对数缩放动量项和更冷的核。虽然这解释了BDRS中温度参数作为初始温度的作用,但也将求解器的内存需求从二次降低到线性。利用这一退火视角,我们引入了过松弛BDRS,它在单个递归中结合了退火和过松弛缩放。我们为两种方法推导了一个原始-对偶认证,该认证可以在不构建传输计划的情况下在线性内存中评估,从而提供了一个可计算的停止规则。在DOTmark基准上,将动量与更冷的核相结合,在相同调度下产生的优化间隙远小于退火Sinkhorn。对于$1024\ imes1024$图像之间的像素级颜色转移,BDRS以9倍的速度提升达到了比MDOT-TNT更低的修复传输成本,在24分钟内达到相对对偶间隙$1.59\%$。我们进一步展示了$4238\ imes2365$图像的颜色转移,每幅图像产生1000万像素,隐含传输条目约一百万亿,在单个NVIDIA L40S GPU上,每个方向在35小时内达到最佳相对对偶间隙$2.41\%$和$2.80\%$。
英文摘要:
We characterize Bregman Douglas-Rachford splitting (BDRS) for unregularized discrete optimal transport and develop an anytime primal-dual certificate in linear memory. We first establish that BDRS coincides with warm-started Inexact Proximal point method for exact Optimal Transport (IPOT) using a single inner Sinkhorn iteration. By eliminating the primal transport plan from the updates, we derive an equivalent dual formulation that reveals BDRS as annealed Sinkhorn under an implicit inverse-linear temperature schedule, with an additional log-scaling momentum term and a cooler kernel. While this explains the role of the temperature parameter in BDRS as an initial temperature, it also reduces the solver's memory requirement from quadratic to linear. Utilizing this annealing perspective, we introduce overrelaxed BDRS, which combines annealing and overrelaxed scaling within a single recursion. We derive a primal-dual certificate for both methods that can be evaluated in linear memory without transport plan construction, thus providing a computable stopping rule. On the DOTmark benchmark, combining momentum with the cooler kernel produces substantially smaller optimality gaps than annealed Sinkhorn under the same schedule. For pixel-level color transfer between $1024\times1024$ images, BDRS attains a lower repaired transport cost than MDOT-TNT with a 9$\times$ speed up, reaching a relative duality gap of $1.59\%$ in 24 minutes. We further demonstrate a color transfer with $4238\times2365$ images, yielding 10 million pixels per image and approximately one hundred trillion implicit transport entries, reaching a best relative duality gap of $2.41\%$ and $2.80\%$ within 35 hours in each direction on a single NVIDIA L40S GPU.