线性成本计算Wasserstein-p距离的离散图像表示算法

5190使用线性成本计算图像之间的Wasserstein-p距离0Yidong Chen 1 , 2 Chen Li 1 Zhonghua Lu 1 , �01 中国科学院计算机网络信息中心 2 中国科学院大学0chenyidong@cnic.cn, lichen@sccas.cn, zhlu@cnic.cn0摘要0当将图像表示为离散测度时，由于解决相应的Kantorovich问题的复杂性，计算它们之间的Wasserstein-p距离是具有挑战性的。在本文中，我们提出了一种新的算法，通过将最优传输（OT）问题限制在一个子集上来计算离散测度之间的Wasserstein-p距离。首先，我们定义了限制的OT问题，并证明了限制问题的解收敛到Kantorovich的OT解。其次，我们提出了SparseSinkhorn算法来解决限制问题，并提供了一个多尺度算法来估计子集。最后，我们在CUDA上实现了所提出的算法，并展示了在时间和内存需求方面的线性计算成本。我们计算Wasserstein-p距离，估计传输映射，并在大小从64×64到1920×1200的彩色图像之间进行颜色转换。（我们的代码可在https://github.com/ucascnic/CudaOT找到）01. 引言0最优传输（OT）问题的悠久历史可以追溯到Monge（1781年）、Toloston（1930年）、Kantorovich（1942年）和Brenier（1991年）的开创性工作[6]。到目前为止，这个理论为研究提供了一个肥沃的土壤，与凸性[18]、偏微分方程[8]、经济问题[13]、机器学习[3]和计算机视觉[28]有着深刻的联系。OT理论定义了概率测度上的一种距离，称为Wasserstein距离。在大多数情况下，Wasserstein距离无法通过解析计算，而且在计算上比Lp距离更昂贵。随着问题规模的增长，计算Wasserstein距离和解决最优耦合的计算工作量成为许多应用的瓶颈。因此，需要高效的算法来计算Wasserstein距离。解决离散OT问题的直接方法是使用基于线性规划的算法，如0匈牙利方法[26]、拍卖算法[4]和网络单纯形[31]通常具有数值稳定性。它们的局限性在于算法的复杂性，尤其是解决大问题时的内存需求。然而，基于近似的方法在原始公式中添加了一个熵项，并使目标成为一个强凸函数，这有助于开发基于迭代的算法，如Sinkhorn算法[41]、Greenkhorn算法[2]和ScreeningSinkhorn算法[1]。Sinkhorn算法为原始OT问题提供了一种高效可扩展的近似方法，并且易于并行化。然而，对于大问题来说，存储密集矩阵的成本很高，并且由于核的某些元素变得太小而无法存储为正数，而变为null，出现数值问题。与Sinkhorn相比，Greenkhorn的收敛性更好，计算成本更低，因为它只需要在每次迭代中更新特定的行或列，但对于大离散OT问题仍然存在数值问题。解决Sinkhorn算法的数值问题的一种有效方法是在对数域上进行迭代[25,40]，这样可以保证每次迭代期间的核是有界的。缺点是它需要在每次迭代中进行许多额外的指数运算。稳定稀疏算法在每次迭代期间生成稳定的核的思想在[37]中进行了研究，不仅消除了Sinkhorn算法的数值问题，还解决了大点云之间的离散OT问题。相关的缩放算法在[10]中研究了计算不平衡OT问题。稀疏算法的局限性在于在迭代过程中生成稀疏核需要相当长的时间。用于测量近似的多尺度方案在[30]中研究了解决半离散OT问题，并在解决离散OT问题（如Shielding方法[36]）的解决时间和内存需求方面显示出显著改进，该方法对于计算Wasserstein-2距离非常有效，但计算Wasserstein-1距离的收敛速度非常慢。有µ =n�i=1µiδxi,ν =m�j=1νjδyi,(1)n�i=1µi =m�j=1νj = 1,µi, νj ≥ 0,(2)Lc(µ, ν) :=minπ∈Π(µ,ν) E(π) := ⟨c, π⟩.(3)H(π) := −minπ∈Π(µ,ν) Eε(π) := ⟨c, π⟩ − εH(π),(4)max(α,β)∈(Rn,Rm) Jε(α, β) := ⟨α, µ⟩ + ⟨β, ν⟩ − ε⟨eα/ε, Keβ/ε⟩,π∗ε = diag(exp (α/ε))Kdiag(exp (β/ε)).(5)5200同时，也有基于梯度的算法[14,24]用于解决大规模的最优输运问题，其复杂度比Sinkhorn算法低。这些算法对于半离散的最优输运问题[7,27]效果良好，但对于大规模稠密离散最优输运问题仍然存在数值问题，因为梯度在迭代过程中会变得接近于零。本文的重点是在一个受限的子空间上开发基于Sinkhorn的算法，以减少迭代过程中的时间和内存需求。我们的工作既补充了理论工作，也补充了实践工作。通过提供证明，保证了限制在子空间上的非负代价函数的最优输运问题收敛到原始的非正则化最优输运问题。我们提出了多尺度稀疏Sinkhorn（M3S）算法，并在CUDA上实现了它，用于求解具有线性代价的大规模稠密离散最优输运问题，以及计算Wasserstein-p距离和在大小为1920×1200的图像之间进行颜色转换。本文的主要贡献如下。首先，我们引入了受限正则化最优输运问题，并证明了它收敛到原始非正则化最优输运问题。其次，基于收敛的保证，我们提出了M3S算法来计算大规模稠密离散最优输运问题。第三，我们提供了所提算法的CUDA实现，用于计算Wasserstein-p距离和在图像之间进行颜色转换，图像大小可达到1920×1200。本文的其余部分组织如下。在第2节中，我们回顾了非正则化和正则化的离散最优输运。第3节和第4节包含了我们的主要贡献，即受限正则化离散最优输运的收敛性和M3S算法。第5节介绍了所提算法的数值结果。第6节总结了本文。02. 离散最优输运和熵正则化0符号说明。我们用  R +  表示非负实数，用  [ n ]  表示整数集合 { 1 , . . . , n }。标准的欧几里得内积用  �∙ ,  ∙�  表示。概率单纯形用  Σ n  表示，即  { a i  ∈  R +  : � n i=1 a i  = 1 } 。对于离散的有限空间  Z  （通常是  X ，  Y  和  X  ×  Y  ），我们用  | Z |表示其基数。对于矩阵  A  ∈  R n × m0其支撑集定义为  spt A  =  { ( i, j ) | A ij  >  0 } .  坐标  r i (A )  和  c j ( A )  表示  A  的第 i 行和第 j 列的和. 对于  a, b∈  R n ，运算符  ⊙ 和  �  表示逐点乘法和除法，例如 a  ⊙b  ∈  R n ，  ( a  ⊙  b ) i  :=  a i  ∙  b i ，其中  i  ∈  [ n ] .函数 exp  和  log  通过逐点应用到所有分量上扩展到  R n：  exp( a ) i  := exp( a i ) . 我们用  1 n  和  0 n  表示  R n中全1和全0的向量.  diag( a ) 表示以向量  a为对角线的对角矩阵.0对角线。02.1. 离散最优输运0给定离散概率测度  µ  和  ν ，使得0集合  Π( µ, ν ) :=  { π  ∈  R n × m + | π 1 m  =  µ, π T 1n  =  ν }  称为  µ  和  ν  之间的耦合。设  c  ∈  R n × m +是一个代价矩阵，其中从  X  中的  x i  到  Y  中的  y j的质量转移的代价为  c ij 。离散最优输运问题定义为0我们关注代价矩阵  c是非负稠密矩阵的一般情况。在这种情况下，我们称（3）为稠密问题。当  m  =  n 时，该问题是线性规划问题，其最佳理论复杂度为O ( n 5 2  )  [ 29 ]。02.2. 熵正则化0定义1 (熵函数) 对于离散测度 π ∈ R n × m + ，我们定义π 的熵 H ( π ) 为0i,j ( π ij log π ij  −  π ij) ,0注1 如果对于某些 i ， j ，使得 π ij = 0 ，则将 π ij log π ij的值定义为0。由于当 z → 0 时，lim z → 0 z log z =0，熵函数 H ( π ) 是连续的。0在 R n × m + 上定义的函数。0离散最优OT问题的熵正则化定义为0其中 ε > 0是正则化参数。通过应用Fenchel-Rockafellar对偶理论([34]第6节)，对应的对偶问题给出0其中 K ∈ R n × m + 是具有条目 K ij := exp( − c ij /ε )的Gibbs核。问题(4)的最优耦合 π � ε 可以描述为u(l+1)=µ ⊘ (Kv(l)),v(l+1) = ν ⊘ (KT u(l+1)),minπ∈ΠN (µ,ν) ENε (π) :=�(i,j)∈Nπijcij,(6)ΠN (µ, ν) := {π ∈ Rn×m+|�{j|(i,j)∈N }πij = µi, i ∈ [n],�{i|(i,j)∈N }πij = νj, j ∈ [m],πij = 0, (i, j) ∈ N0 \ N}.minπ∈ΠN (µ,ν) ENε (π) :=�(i,j)∈N�πijcij − ε(πij log1πij + πij)�.(7)5210通过更新变量 u ( l ) 和 v ( l )，结合(5)和边际条件，引出了一种自然而著名的Sinkhorn算法0其中初始值 v (0) 给定为 v (0) = 1 m 。当 ε → 0时，正则化OT问题(4)的解收敛到非正则化问题(3)[11]，并且在[37]中对 E ε ( π ) 和 J ε ( α, β )的原始-对偶间隙进行了良好的估计。然而，在实践中，存储核 K所需的内存变得太大而无法忽略。例如，如果我们使用Sinkhorn算法计算两个256×256的灰度图像之间的Wasserstein-p距离，核的大小为2^16×2^16。以双精度存储核需要32GB内存，这是非常高要求且低效率的。03. 在子集上的最优输运0我们提出了在子集上限制的正则化最优输运问题，并证明了在一定条件下，受限最优输运的解收敛到非正则化问题(3)，这启发我们通过使用更少的内存快速准确地计算Wasserstein距离和最优耦合。0定义2 (受限OT问题) 设 c ∈ R n × m +是一个成本矩阵，N 0 是基础集合，即 c 的条目的索引，即N 0 = [ n ] × [ m ]。如果 N � N0，则受限制的非正则化最优输运问题定义为0其中受限耦合 Π N ( µ, ν ) 定义为0对于ε > 0，限制在N上的正则化OT问题定义为0问题(6)和(7)在集合 Π N ( µ, ν )非空时是可行的。显然，如果我们能够估计一个合适的子集0如果集合N的基数|N|较小且保证问题的可行性，那么问题可以更快地解决并且占用更少的内存。关键是，在什么条件下问题是可行的，以及受限正则化问题是否收敛到非正则化问题。我们陈述并证明了以下定理。0定理1令π�为非正则化问题（3）的所有最优解中熵最大的解。0π� = arg min π {−H(π) : π ∈ Π(µ, ν), �c, π� = Lc(µ, ν)}0如果spt π� � N，则（i）min π ∈Π(µ,ν) ENε(π)是可行的。0（ii）方程（7）的唯一解ˆπ�(ε)收敛到π�，即0当ε → 0时，ˆπ�(ε) = π�。（8）0限制是针对ˆπ�(ε)的每个元素进行的。（iii）min π ∈ Π(µ,ν)ENε(π)收敛到非正则化Kantorovich问题的解。0化Kantorovich问题的解。0当ε → 0时，min π ∈ Π(µ,ν) ENε(π) = Lc(µ,ν)。（9）0如果我们将原始的正则化问题（4）限制在一个稀疏子集N上，使得spt π� �N，那么解决问题并得到一个近似解就足够了。对于这样一个子集N和一个矩阵A ∈ Rn×m，我们用rN(A) :=0{ j | ( ∙ ,j ) ∈N } A ( ∙ , j )  ∈  R n and  c N ( A) := �0分别是Rm和Rn。在子集N上限制的Sinkhorn迭代给出了算法1。0算法1 RestrictedSinkhorn ( N , K, µ, ν, ε )01:  k  ←  0 2:  v(0) ←1/m 3:  repeat 4: if  kis even  then05: u(k+1)i ← µi/rNi(Kv(k)) for all i ∈ [n]06: v(k+1) ← v(k)08: v(k+1)j ← νj/cNj(KTu(k+1)) for all j ∈ [m]09: u(k+1) ← u(k)010: end if012: π(k)ij ← u(k)i exp(−cij/ε) v(k)j for (i, j) ∈ N013:  until  || rN(π(k)) − µ ||2 + || cN(π(k)) − ν ||2 < δ2014:  α(k) ← ε log u(k),  β(k) ← ε log v(k)015:  输出  π(k),  α(k),  β(k)Kij(ε, α, β) = exp�KNij :=� Kij(i, j) ∈ N,0(i, j) ∈ N0 \ N.(11)5:(π(l+1), α(l+1), β(l)) ← RestrictedSinkhorn(N (l),KN (l), µ, ν, ε(l))6:ε(l+1) ← ε(l)/(1 + σ)(l+1)10:(l+1) ← N (l+1) ∪ (i, j)∆ ≜inf{π∈V :⟨c,π⟩>⟨c,π∗⟩}⟨c, π⟩ − ⟨c, π∗⟩.(12)and θ < ( s5220Sinkhorn算法和RestrictedSinkhorn算法之间唯一的区别是我们在一个子集上进行迭代，而不是整个域[23]。RestrictedSinkhorn算法输出两个额外的变量u(k)和v(k)，它们在找到一个新的子集N满足spt π* �N中起到关键作用。我们稍后会讨论它们。以下定理保证了RestrictedSinkhorn算法的收敛性。（证明见附录）。0定理2 令N � N0为一个子空间，满足spt π* � N。算法1输出一个满足||rN(π) − µ||2 +||cN(π) − ν||2 < δ2的矩阵π。0经过O(ρh(δ)^-2 log(h/s))次迭代，其中h = �0(i,j) ∈ N πij, ρ = max{||rN(π)||∞, ||cN(π)||∞}，s =min{(i,j) ∈ N, πij > 0} πij。0推论1 令N � N0为一个子空间，满足spt π* �N。算法1在满足||rN(π) − µ||1 + ||cN(π) − ν||1 <δ2的情况下，经过O(Nρh(δ)^-2log(h/s))次迭代输出一个满足||rN(π) − µ||1 + ||cN(π) − ν||1 <δ2的矩阵π，其中N = max{m, n}。0推论1中的额外因子N是必要的，因为我们需要进行N次迭代将l2界限转化为l1界限。算法1中的停止准则||rN(π(k)) −µ||2 + ||cN(π(k)) − ν||2 < δ2可以调整为||rN(π(k)) − µ||1+ ||cN(π(k)) − ν||1 <δ，这遵循了[2]的思路并且是一个更强的l1近似。以下定理证明了这个调整的合理性。0定理3 令N � N0为一个子空间，满足spt π* �N。算法1在满足停止准则||rN(π(k)) − µ||1 + ||cN(π(k)) −ν||1 < δ的情况下，经过O((δ)^-2log(h/s))次迭代输出一个满足||rN(π) − µ||1 + ||cN(π) −ν||1 < δ的矩阵π。0算法1在一个子集上进行标准Sinkhorn迭代，这会导致当ε→0时的数值问题。为了避免这种情况，我们定义了以下稳定核函数来替代原始核函数。替代生成稳定迭代，从数学上讲与原始Sinkhorn算法等价，但在实践中有显著改进。0定义3（稳定核函数[37]）对于  ε >  0 ，  α  ∈  R n0and  β  ∈  R m ，与代价矩阵c相关的稳定核函数K ( ε, α, β)定义为0ε ( α i  +  β j  −  c ij ) � ，( i, j )  ∈ N 0 . (10)0定义4（受限核函数）对于一个空间N � N 0 ，受限核函数KN 定义为0为了避免迭代过程中的数值问题，我们打算从选择一个合适的子空间N开始进行迭代，然后将子集细化为一个“更小”的子集，这在算法2中总结。我们将在下一节中展示如何通过多尺度方案生成子空间N。0算法2 SparseSinkhorn ( X ,  Y  ,  N ,  α ,  β ,  ε ,  θ )01:  初始化  l  ←  0 ,  ε (0)02:  N (0) ← N ,  α (0) ←  α, β (0)04: 通过(11)生成核函数  K N  ( l )08: 对于  ( i, j )  ∈ N ( l )09: 如果  K N  ( l ) ij ( ε ( l +1) , α ( l +1) , β ( l +1) )  > θ  则011: 结束012: 结束循环014:  直到  ε ( l ) < ε 15:  输出  π ( l ) ,α ( l ) , β ( l )0缩放步骤  ε ( l +1) ←  ε ( l ) / (1 +  σ )  (算法20步骤12）生成一个逐渐减小的序列{ ε ( l ) }，收敛到零（ σ>  0）。对于正参数ε l ，算法2的步骤5在子集N ( l)上执行RestrictedSinkhorn算法，输出一个耦合π � ( l+1)和两个对偶变量α ( l +1)和β ( l+1)。在截断步骤（算法2的步骤7-12）之后，通过使用参数θ截断核函数生成新的子集。截断核函数生成新的子集的思想受到[37]的启发。然而，我们应该指出，我们的算法通过在旧的子集N ( l )上截断核函数来生成新的子集N ( l+1)，而不是在整个空间X ×Y上截断，这显著降低了计算成本。根据[42]的思想，我们提供了一个理论分析来验证所提出的截断步骤当| X |  =  | Y|  =  n时的有效性。0定义5 令V为集合Π( µ, ν)的顶点，受限OT的次优间隙∆定义为0定理4 令ε (0) < ∆0δ ) (1+ σ ) ，算法2的步骤7-12输出一个子集N ( l+1)，使得spt π � � N ( l +1)，其中s = min { ( i,j ) ∈N  ,π ij04. 通过多尺度方案估计子集0算法2通过在受限子空间上执行Sinkhorn迭代来接近非正则化的OT问题5230N。剩下的问题是，如何为算法2提供一个合适的N。我们采用多尺度方案，通过使用前一尺度的解来确定子集N。此外，前一尺度的解为当前尺度提供了初始化，从而显著减少了内存使用量。0定义6（分层划分[38]）对于离散集合X，分层划分是一个有序的元组 (X(0), ..., X(K-1))，其中X(0) = {{x} : x ∈ X} 是将 X划分为单例集合的平凡划分，并且子单元通过合并来自子单元的上一级的单元构建。对于 k ∈ {1, ..., K-1} 和任意 x i ∈0ˆxi ∈ ˆX ˆxi，并且我们称 ˆxi 为 x i 的子节点。0对于离散测度 µ = ∑ni=1 µi δxi，其多尺度测度方案是元组(µ(0), ..., µ(K-1))，可以从细粒度 (k = 0) 到粗粒度 (k =K-1) 的尺度进行分解，设置0µ(k) = 0i ∈ J(k) µ(k)i δx(k)i, X(k) = {x(k)i, i ∈ J(k)}. (13)0从 µ(0) = µ（和 X(0) = {X}）开始，我们提取 X(k) =∪i∈J(k)C(k)i 的支持的聚类，我们用 X(k+1) = {x(k+1)i, i∈ J(k+1)}表示相应的聚类中心。接下来，我们通过收集每个聚类中的质量来计算权重 µ(k+1)i = µk(C(k)i)。令 X ∈ Rn 和 Y ∈Rm 为具有分层划分 (X(0), ..., X(K-1)) 和 (Y(0), ..., Y(K-1))的向量，其中 X(k) = {x(k)i, i ∈ J(k)X} 和 Y(k) = {y(k)j, j ∈0J(k)Y}。对于成本矩阵 c ∈ Rn×m，成本映射 C 给出 C: X ×Y → Rn×m，其中 C(xi, yj) = cij。成本矩阵的分层划分(c(0), ..., c(K-1)) 是0c(k)ij = min(x,y) C(x, y), (i, j) ∈ J(k)X × J(k)Y, (14)0(x, y) ∈ children(x(k)i), children(y(k)j).0对于 k ∈ {0, ..., K-1}。图1（左侧）显示了两个离散测度 X和 Y的三级多尺度结构（每个测度有八个元素，测度（X，Y）有64个元素）。分层划分是通过合并来自上一级的相邻测度创建的。第一层（k =0）表示原始测度，（x(0)i，y(0)i）是第二层（k =0）中（x(1)i，y(1)i）的子节点。第三层只有四个元素。每个元素是来自第二层的测度（x(1)i，y(1)i）的组合。我们进一步说明了如何将多尺度方案与SparseSinkhorn算法结合起来解决受限OT问题（4）。我们从粗层开始通过计算最优耦合0调用SparseSinkhorn算法。然后生成新的耦合并估计下一层的子集。然后再次调用SparseSinkhorn算法，并计算新的耦合。循环继续，直到达到第一层 k =0。最后，生成子集 N(0) 并计算最优耦合π(0)。该过程在算法3中总结。请参考图1（右侧）的具体示例。0算法3 多尺度SparseSinkhorn01: 构建多尺度结构 { (Xk, µ(k)) } K-1 k=0 , { (Y(k), ν(k)) } K-1 k=0 和 { c(k) } K-1 k=002: α(K-1), β(K-1) ← 003:  N(K-1) ← J(K-1)X × J(K-1)Y04: 对于 k = K-1, ..., 0 do 5: (π(k), α(k), β(k)) ←SparseSinkhorn(X(k), Y(k), N(k), α(k), β(k), ε(k), θ)08: 对于 (i, j) ∈ spt π(k) 执行09: 对于 (x(k-1)i, y(k-1)j) ∈ (children(x(k)ˆi),0对于所有（x（k）ˆi，y（k）ˆj）的子节点（x（k）ˆi，y（k）ˆj）do010：N（k-1）←N（k-1）∪（i，j）011：（α（k-1）i，β（k-1）j）←（α（k）ˆi，β（k）012：结束循环013：结束循环014：结束if015：结束循环16：输出π（0）0注2：使用多尺度方案计算OT问题首先在[30]中提出并在[21，27，37]中进行了研究。我们的算法与现有方法的区别在于方法[27，30]仅使用粗略解作为热启动，但不为下一层提供新的子集。此外，我们根据集合sptπ（k）计算子集N（l-1），而多尺度算法[37]计算集合X×Y中的子集，与我们相比，估计子集的时间更长。05.数值实验0本节介绍数值实验，以展示M3S算法在计算时间和内存需求方面的性能。此外，我们计算彩色图像之间的Wasserstein-p距离，并在1920×1200图像之间传输颜色。所有报告的运行时间都是在具有64G-B内存、2.7GHz英特尔XeonE5-2697处理器和RTX（2080Ti）GPU的计算机上获得的。(xi(1),yj(1))(xi(2),yj(2) )(xi(0),yj(0))(0)(2)(1)(0)(1)1024409616384655362621440.10.31310301003001000| X |Solving time (s) FastEMD Network simplex Sinkhorn GPU M2S M3S M2S(keops) M3S(keops)5240（最优耦合）0第k=2层0调用SparseSinkhorn0调用SparseSinkhorn0调用SparseSinkhorn0生成子集0初始化子集0生成子集0获取最优耦合0获取最优耦合0获取最优耦合0第k=1层0第k=0层0（1）�0图1.左：离散测度（X，Y）的三级分层划分。右：多尺度方案与SparseSinkhorn算法的结合。粉色区域表示子集{N（l）}。算法从第三层开始，k = 3。子集N（2）0在整个空间中初始化，并调用SparseSinkhorn算法获取最优耦合π（2）。使用耦合的支撑生成子集N（1），然后调用SparseSinkhorn算法获取最优耦合π（1）。最后，我们生成N（0）并在第一层获取最优耦合π（0）。（x（3）i，y（3）i）的附加测度是其子测度（x（2）i，y（2）i）。05.1.性能分析0我们首先考虑n×n灰度图像之间的OT。成本矩阵c∈Rn2×n2由c（x，y）=|x-y|p计算，即像素位置之间的lp距离。实验在DOT基准[39]上进行测试，图像大小范围从32×32到512×512，即X和Y的基数范围从210到218，完全耦合空间的维度范围从220到236。我们考虑两个大小相等的图像之间的传输，即|X|=|Y|。我们首先计算lp距离为p∈[1，2.5]的OT问题。对于p=1，我们计算Earth Mover'sDistance。我们设置θ=10-8用于细化新的子集，δ=10-5用于估计误差界。0M3S算法的性能与我们的GPU实现的Sinkhorn算法进行了比较，并与CPU实现的CPLEX网络单纯形法[12]、FastEMD[32]、由POT库[19]实现的Greenkhorn算法[2]进行了比较。此外，我们还将M3S与由[17]实现的多尺度Sinkhorn（M2S）进行了比较。我们还为两种算法都配备了keops后端[9]并设置了相同的参数进行比较。图2（右侧）比较了计算Wasserstein-1距离的不同算法的内存使用情况，即我们计算了Earth MoverDistance（EMD）。M3S算法的内存需求呈O（n）增加。Sinkhorn算法的内存需求呈O（n^2）增加，并且在我们的设备（RTX 2080ti）上内存不足。0图2.不同算法的运行时间和内存需求。右：FastEMD的求解时间随O(n^3)的阶数增长。Sinkhorn算法显示出O(n^2)的时间复杂度。M2S算法的时间复杂度为O(nlog(n))，而M3S的时间复杂度为线性增长。左：M2S和M3S算法都显示出线性内存需求，而M3S算法需要更少的内存。当配备keops后端时，M3S和M2S在GPU上分配相同的缓冲区，随着问题规模的增长，M3S运行速度更快。0当问题规模大于128×128（2^14）时，FastEMD和传输网络单纯形方法的内存需求呈O(n^3)增长。在所有测试中，FastEMD方法使用的内存最小，相比Sinkhorn算法和网络单纯形方法，但计算大小为256×256的两个图像之间的EMD仍然需要约30.2GB的内存。当N增加到超过2^36时，所提出的算法显示出线性内存需求来计算EMD。图2（左）显示了通过不同算法计算EMD的求解时间。在所有测试中，5250图4.通过M3S算法计算30个彩色图像（1920×1200）之间的Wasserstein-2距离。我们计算图像之间的Wasserstein距离并构建距离矩阵。我们使用多维缩放重新构建距离矩阵，并获得每个图像中心点的三维坐标。根据它们的坐标，图像在R3中呈现。我们以一幅图像为参考，将其他图像投影到该图像的平面上。我们的工作是在3D空间中测量彩色图像之间的距离的扩展（参见[35]，图6）。0图5. 1920×1200图像之间的颜色转换。我们通过提出的算法在RGB空间中计算每个像素从源图像到参考图像的最优传输映射。0FastEMD的时间需求呈O(n^3)增长，而所提出的算法的时间需求呈线性增长。计算大小为512×512的两个图像之间的EMD平均需要28.90秒。其他方法在观察时间（1200秒）内无法收敛。我们通过比较不同p的成本c(x, y) =|x−y|p的平均求解时间来进行进一步研究。从图3（左）可以看出，求解时间随不同p变化。当p在1到1.4范围内时，我们观察到求解大小为256×256的两个图像之间的OT平均需要7秒，而当p接近2.0时，求解时间减少。图3（右）显示了内存使用情况。0在GPU上存储稀疏矩阵K(x,y)，它占用了主要的内存资源。对于不同的p，我们分配的内存随问题规模线性增长。当图像大小为64×64时，我们在GPU上为OT问题分配64MB内存，并且随着图像大小增加到128×128，内存分配增加到256MB的四倍。表1报告了不同图像大小n的相对误差|Wp−W�p|/W�p。成本函数由c(x, y) =|x−y|p给出。通过FastEMD获得精确解W�p，它计算了未正则化问题（3）的解。1.01.21.41.61.82.02.22.40.3131030100Solving time/sExponent p n = 512 n = 128 n = 256 n = 6440961638465536262144416642561024409664MB for n = 64 1GB for n = 256 0.25GB for n = 128 Memory allocated for kernel (MB)X p = 1 p = 1.2 p = 1.5 p = 1.8 p = 2.0 p = 2.54GB for n = 512 Relative error (10−4)Image size3264128256512p10.1340.4923.561.27−1.20.1570.08226.232.34−1.50.1040.1320.09780.0522−1.70.1130.1800.8630.411−20.1090.8030.06750.9230.460��5260图3.成本c(x, y) = |x−y|p的求解时间和内存使用情况。0表1.M3S和FastEMD计算的Wp和W�p之间的相对误差。'-'表示FastEMD内存不足，因此无法进行相对误差比较。05.2.彩色图像之间的Wasserstein距离0Wasserstein-p距离是图像的有用度量[35]，因为它量化了直观的图像相似性概念，特别是用于量化不同尺寸图像之间的距离。在本节中，我们通过计算大小为1920×1200的彩色图像的Wasserstein-2距离来浏览彩色图像数据库。我们将彩色图像视为图像的RGB空间中的点。将基准图像映射到RGB空间，然后计算所有图像之间的距离。为了在R3空间中可视化图像的相对位置，我们使用多维缩放（MDS）技术[5]，将一组图像嵌入到欧几里德空间中作为点。3DMDS嵌入的可视化可以以感知直观的方式组织和改进最近邻查询的结果。用户可以通过使用提出的算法计算Wasserstein-2距离快速导航到感兴趣的图像空间部分。图4显示了30个不同彩色图像之间的Wasserstein-2距离，距离矩阵W30×30是由W(i, i) = 0构成的，每个条目W(i,j)是图像i和j之间的Wasserstein-2距离。通过将矩阵嵌入到R3空间中，获得每个图像的三维坐标。距离越小，两个图像越相似。Wasserstein距离是浏览大尺寸彩色图像数据库的导航。05.3. 大尺寸图像之间的颜色转换0颜色转换在计算机图形学和计算机视觉领域受到了广泛关注。颜色转换的目的是重新着色给定的图像或0通过导出与另一幅图像作为参考图像之间的映射来对图像进行修改[15]。用户可以通过选择参考图像来修改原始图像，使得原始图像获得该参考图像的调色板。在颜色转换中计算最优传输映射已经在[16, 20, 22,33]中进行了研究。然而，随着图像尺寸的增大，计算成本也大幅增加。在本节中，我们通过使用提出的算法计算OT映射在图像之间进行颜色转换。我们将u: Ωu � Z2 → Σ �R，其中Ωu是u的像素网格，Σ是量化的RGB颜色空间。我们将Xi = (Ui) ∈ R30用于指定空间组件(xi ∈ Ωu)和颜色组件(Ui ∈Σ)。源测度µX由µX: X → 表示0i ∈ IXµiδXi(X)，目标测度νY类似。Nx是图像中所有像素的数量。两个测度X和Y之间的正则化OT问题如下所示0π�ε = arg min π ∈ Π(µ,ν) �C,π� − εH(π). (15)0传输成本取为Cij = ||Ui −Vj||22，即RGB空间的l2范数的平方。为了保持核的稀疏性，点云被归一化到[0,1]3。我们定义了一个一对一的最优耦合映射T，从最优耦合π�ε中得到0T(Xi) = Yj, j ∈ arg min(πε)ij. (16)0通过计算最优耦合和OT映射，我们在图5中展示了六个不同季节之间的颜色转换示例。OT映射(16)定义了一个像素到像素的颜色转换，避免了平滑传输映射的错误颜色伪影以及颜色对比度的损失。06. 结论0我们提出了受限正则化OT问题，并证明了受限最优解的收敛性。基于理论保证，我们引入了M3S算法来解决子集中的正则化OT问题。所提出的算法在计算OT问题方面具有线性时间和内存需求的成本。所提出的算法足够准确，可以计算Wasserstein-p距离，并比较彩色图像的RGB云特征。通过计算RGB颜色空间之间的最优传输映射，实现了尺寸为1920×1200的图像之间的颜色转换。致谢：本研究得到了中国国家自然科学基金（项目编号61873254）和广西科技重大专项（桂科AA18118054）的支持。5270参考文献0[1] Mokhtar Z. Alaya, Maxime Berar, Gilles Gasso, and AlainRakotomamonjy. Screening sinkhorn algorithm for regularizedoptimal transport. In Advances in Neural Information ProcessingSystems, pages 12169–12179, 2019. 10[2] Jason Altschuler, Jonathan Niles-Weed, and PhilippeRigollet. Near-linear time approximation algorithms foroptimal transport via sinkhorn iteration. In Advances inNeural Information Processing Systems, pages 1961–1971,2017. 1, 4, 60[3] Martin Arjovsky and L´eon Bottou. Towards PrincipledMethods for Training Generative Adversarial Networks.arXiv:1701.04862, 2017. 10[4] D. P. Bertsekas. The auction algorithm: A distributedrelaxation method for the assignment proble