大规模最优传输问题的稀疏化方法
单击上方“图灵人工智能”,选择“星标”公众号
您想知道的人工智能干货,第一时间送达

大规模最优传输问题的稀疏化方法
本文第一作者欧阳夏雪来自中国人民大学,通讯作者为清华大学李梦雨 (https://mengyu8042.github.io/)。其他作者还包括中国人民大学的郑皓,北京师范大学的梁浩贤,北京邮电大学的张静怡,上海财经大学的邱怡轩 (https://statr.me/),以及中国人民大学的孟澄 (https://cheng-bdal.github.io/)。
引言:
2026年初,DeepSeek发布的一篇名为《mHC: Manifold-Constrained Hyper-Connections》的论文引起了广泛的关注。该论文的研究动机之一是其作者注意到残差网络中的超连接(Hyper-Connections)在大规模训练中由于恒等映射属性的破坏,容易产生数值不稳定和信号爆炸等问题。针对这一现象,该论文提出了流形约束超连接(Manifold-constrained hyper-connections)的新架构,利用Sinkhorn算法将残差映射()投影到Birkhoff多胞形(双随机矩阵的流形)上。双随机矩阵是行和与列和均为1的非负方阵,这类矩阵具有范数保持性、复合封闭性以及几何解释性等良好性质,有利于大规模训练的稳定性。该论文应用了特殊情形的Sinkhorn算法,对残差映射做了一步标准化对齐,使这些映射变为双随机矩阵。这样的做法本质上是在求解一个最优传输问题,其目的是找到一个与原始残差映射"最相似"的双随机矩阵。而其中涉及到的Sinkhorn算法,已经在最优传输领域有着十余年的应用。下面将进一步详细介绍最优传输的相关内容。
最优传输(Optimal Transport,OT)提供了一种基于最小传输成本的度量(Wasserstein距离)与转换框架,有效地衡量了两个概率分布之间的相似度。该理论在统计学和机器学习领域都存在着广泛的应用。在统计学中,它不仅提供了具有几何直观的Wasserstein距离,广泛应用于分布比较、非参数假设检验和因果推断中的权重平衡,还能通过传输映射解决诸如单细胞基因组学中常见的"未配对数据对齐"这一经典难题。在机器学习中,最优传输为Wasserstein生成对抗网络(WGAN)提供了稳定训练的损失函数,也为扩散模型、归一化流等生成式模型奠定了分布变换的数学基础。
随着数据规模的快速增长,OT问题所面临的超高计算复杂度已成为制约其广泛应用的关键瓶颈。基于传统线性规划的方式,求解OT问题精确解的计算复杂度通常为 ;近年来虽有部分工作致力于研究精确解的高效求解,但其复杂度也仅能降低至约 的量级。即使通过引入熵正则化项并采用Sinkhorn算法求解近似解,其整体计算复杂度仍然高达 难以满足目前日益增长的超大规模数据集的计算需求。
针对OT问题的计算瓶颈,中国人民大学、北京师范大学、北京邮电大学、上海财经大学和清华大学等机构的学者共同撰写并于WIREs Computational Statistics上发表了综述文章《Sparsification Techniques for Large-Scale Optimal Transport Problems》。该综述介绍了OT问题的稀疏化加速技术——在保持计算精度的前提下,利用OT问题中的矩阵稀疏性,显著降低计算负担,推动Sinkhorn算法应用于更大规模的模型训练问题中。

论文链接(点击阅读原文,下载论文PDF):https://wires.onlinelibrary.wiley.com/doi/10.1002/wics.70056
本综述系统性地梳理了"利用稀疏化方法降低熵正则化OT问题求解负担"的相关研究进展。尽管Sinkhorn算法因其高效性而被广泛应用,但在大规模场景下仍面临着"每轮迭代计算成本高"和"整体算法收敛速度较慢"两大问题。为此,综述指出,现有加速方法可以从两个互补的角度理解:一类方法试图降低每一步迭代的计算复杂度,另一类方法则致力于减少达到收敛所需的迭代次数。这两种思想分别对应于熵正则化OT问题核矩阵(Kernel matrix)的稀疏化与Hessian矩阵的稀疏化。下图概括性地展示了本综述所涵盖的主要方法框架。

背景:
1. 什么是OT
OT问题最早由法国数学家Gaspard Monge于1781年正式提出。二战期间,苏联数学家和经济学家Leonid Kantorovich对该问题进行了系统性的推广与深化。因此,该问题也被称为Monge-Kantorovich运输问题。OT问题可以通过一个直观的物资运输场景来理解:设想存在若干矿山负责开采矿石,并向多个工厂供给原材料,每个工厂对矿石具有给定的需求量。在假设矿山的总产量恰好等于所有工厂的总需求量的前提下,如何规划从矿山到工厂的运输方案,才能使总体运输成本最小?这一过程正是OT问题的核心思想,如下图所示:

我们可以将上述问题形式化为一个数学模型。设有 个物资产地 ,其分布服从 ;同时有 个工厂 ,其需求分布服从 ,分别表示为:
从产地 向工厂 转移单位物资的运输成本记为 ,通常取为 。物资运输问题的目标是在满足供给与需求约束的前提下,寻找一个最优的运输方案,使得总运输成本最小,其数学形式为:
其中 是成本矩阵, 是传输计划, 代表从产地 转移多少物资到工厂 。这个形式的OT问题是一个总是有解、具有强对偶性质的凸优化问题。
然而,虽然原始的OT问题具有良好的理论性质,但在实际计算中面临着严重的复杂度挑战。基于传统线性规划的求解方式,对于两个样本量为 的离散分布(),其计算成本通常高达 。当样本规模 增大时,算法的运行时间和内存方面的开销将会迅速增长,从而显著限制了其在实际场景中的应用。
2. 熵正则化OT和Sinkhorn算法
为了克服传统OT问题在计算上的困难,一个经典而有效的思路就是引入熵正则化项,从而使原始OT问题在数值上变得更易于求解。具体而言,熵正则化OT问题具有如下形式:
其中 是香农熵。通过引入熵正则化,传输计划被迫变得比非正则化情形下更加均匀分散。等价地,该问题也可以写为以下Kullback-Leibler(KL)正则化形式:
这里 是KL散度。正则化参数 控制了解的平滑程度:当 时,问题退化为经典OT,解趋于稀疏但计算困难;当 时,解趋于均匀,但丧失了有意义的传输结构。
熵正则化OT问题是一个严格凸的优化问题,且最优解是唯一的。Marco Cuturi在2013年提出利用Sinkhorn迭代算法高效求解该问题,其每轮迭代仅涉及矩阵-向量乘法,时间复杂度约为 ,极大提升了OT在实际问题中的可扩展性。熵正则化OT的最优解具有一个非常简洁的形式:
其中, 称为Gibbs核矩阵,向量 和 是未知的正值缩放因子,用于调整行和列以满足给定的边际约束。基于这一结构,Sinkhorn算法通过交替地对行和列进行缩放操作,逐步逼近满足边际约束的最优传输计划。
3. 熵正则化OT的对偶形式
熵正则化OT的对偶形式可以写成:
通常将其转化为等价的最小化问题 ,并对对偶变量 做优化。与原始问题相比,该对偶目标函数是一个光滑且可微的函数,其梯度与Hessian矩阵具有显式形式,从而为一阶和二阶优化方法提供了理论基础,并由此发展出多种加速算法。给定 ,可以计算出对应的传输计划:
其梯度和Hessian具有以下显式形式:
注意到,Hessian非对角部分直接由 给出,这意味着Hessian的数值结构与传输计划本身高度相关。当 呈现稀疏或近似稀疏结构时,Hessian也随之具有相似的稀疏性,这一观察正是后续基于Hessian稀疏化加速算法的核心出发点。
方法:
Sinkhorn算法只涉及矩阵乘法,每轮迭代复杂度为 ,且收敛轮数与 有关。尽管Sinkhorn算法将解决原始OT问题的 计算复杂度近似降低到 ,从而使得OT能够在实际中进行使用,但在面对大规模机器学习任务时,Sinkhorn算法仍然面临着计算困难的问题。现有研究指出,引入稀疏化策略是进一步提升计算效率的有效途径,其总体思路可概括为以下两类。
1. 基于核矩阵的稀疏化:这类方法的核心思想是直接降低每次迭代的运算量。通过重要性采样(如Spar-Sink、Spar-GW)或局部敏感哈希等技术,对原始稠密的核矩阵进行稀疏近似,实现了对算法本身的直接加速。 2. 基于Hessian矩阵的稀疏化:这类方法从熵正则化OT问题的对偶形式出发,采用牛顿类二阶优化框架来加速算法整体的收敛性。其核心是通过阈值截断(如SNS算法)或结构化稀疏(如SPLR方法)等技术,对牛顿步中需要求解的稠密Hessian矩阵进行稀疏近似。
1. 核矩阵的稀疏化:近似传输核矩阵以降低单步复杂度
回顾熵正则化 OT 的Sinkhorn形式,最优传输矩阵具有形式(1)。Sinkhorn每轮的主要代价来自于稠密核矩阵与向量的乘法以及存储,因此一个核心想法是:用一个稀疏的 替代原始的 ,把每轮复杂度从 压到接近线性 ,同时仍尽量保持传输结果精度。本篇综述归纳了其中两种代表性方法:
Spar-Sink:随机稀疏,用重要性采样降低方差、保证无偏
Spar-Sink方法的主要出发点是,OT的传输结构与核矩阵的非零/大值结构紧密相关。综述里强调了三个重要观察来支撑这种稀疏采样核矩阵方法的合理性:
1. 与 具有相同的稀疏结构; 2. 传输代价具有求和形式 ,因此稀疏化 可以理解为在这堆求和项里挑选一部分关键项做近似; 3. 从重要性采样角度出发,理想采样概率 应与 成正比,但 是未知的;基于边际约束 和 的存在性,可以推出一个可用的上界,从而得到可实现的采样概率设计。那么 应该怎么选?论文里给出了一个非常简单且重要的形式:
下图来自于本篇综述Figure 3,显示了该算法的总体流程。

LCN-Sinkhorn:把局部稀疏结构和全局低秩结构融合起来,兼顾近邻与远距离相互作用
单纯对核矩阵的元素做随机独立稀疏采样可能丢掉关键的全局几何结构;而单纯做低秩近似虽然抓住全局结构特征,但是在局部细节上可能不够准确,尤其在弱正则或成本矩阵变化剧烈时会面临着极不稳定的问题。LCN-Sinkhorn方法采用了一种两阶段融合的思想:
1. 用局部敏感哈希(locality-sensitive hashing,LSH)做硬阈值截断,专门保留近邻交互(局部稀疏),得到 ; 2. 用Nyström方法对 做低秩近似估计,捕捉全局几何特征,得到 ;
最后用局部修正把两者融为一体,得到一个既有局部精细、又有全局信息的近似核矩阵,其定义为:
其中 表示在 的非零位置上取 的对应元素。下图(节选自本篇综述Figure 5)显示了LCN-Sinkhorn的总体流程。

2. Hessian 矩阵的稀疏化:二阶加速+稀疏近似以降低Newton步代价
相比核矩阵稀疏化是为了让每轮迭代计算代价更少,Hessian稀疏化的逻辑是让算法总体步数更少。其出发点是:由于熵正则化OT的对偶形式是一个光滑且严格凸的函数,理论上使用牛顿法会具有二次局部收敛性质,通常能用远少于Sinkhorn的迭代次数达到高精度的解。但是在对偶空间做牛顿/拟牛顿更新,每次更新都需要求解线性方程组,若直接使用稠密的Hessian,其求解代价通常是 ,反而会比原始Sinkhorn更加高昂。综述指出:Hessian稀疏化方法的目标是将每步求解线性方程组的代价降到 量级,同时尽量保留二阶方法的快速收敛特性。根据式(3),Hessian的非对角块就是 本身,只要 具有近似稀疏结构,那么Hessian也会变得近似稀疏,这正是Hessian稀疏化方法的核心要义。
SNS(Sinkhorn-Newton-Sparse):硬阈值截断Hessian,保留大值元素
SNS的一个核心观察是,当Sinkhorn迭代次数足够多、且正则化系数较小时,Hessian会呈现近似稀疏的结构。于是一个自然的想法是,直接把Hessian中较小的元素置为零,只保留最大的 个元素,得到稀疏近似 。当 时,单步求解牛顿线性方程组的复杂度就变成了 。但它具有一个明显的缺点,那就是 的正定性不能得到保证。
SSNS(Safe and Sparse Newton for Sinkhorn):只对Hessian非对角块做稀疏化,并控制近似误差+保证正定
SSNS旨在满足:1) 必须足够接近 ,即误差可控;2) 必须保持正定性,从而其 Newton方向可解且稳定。作者注意到,在式(3)中,Hessian的对角块 和 对正定性的保证很关键。因此SSNS稀疏化时,保持对角块不动,只在 上做稀疏化。SSNS自适应地删小保大,通过构造一个辅助矩阵 ,计算列方向和行方向的累计和阈值,从 中删除"最小的那一部分质量",得到 。最终稀疏化的Hessian定义为:
当给定阈值 时,SSNS会进行一项行与列的双向筛除来选择 ,其具体策略在下图中给出(节选自综述Figure 6)。这个非对角块的截断操作可以始终避免Hessian的奇异性,在理论方面,SSNS也给出了更强的收敛结论。但它也面临着一些局限性:稀疏程度难以事先控制,导致实际计算成本不易预测;同时,当传输计划 较稠密时,单纯稀疏化可能丢失关键的曲率信息,影响算法收敛与稳定性。

SPLR(Sparse-Plus-Low-Rank):结构化稀疏+低秩校正
SPLR主要针对 稠密从而导致稀疏策略可能失效的场景,它通过引入低秩修正来补回一些被在稀疏化操作中损失的重要曲率信息。它指定了结构化索引集,再采用硬阈值策略进行选择和稀疏化,最后再加以rank-2修正。和SNS一样,SPLR也是采取硬阈值截断;不同的是,SPLR的稀疏化不是直接对 做阈值,而是只对其中的非对角块 做稀疏化。SPLR通过如下方式构造稀疏索引集 ,定义:
1. 选取 中取值最大的 个元素,其索引记为 ; 2. 定义结构化必选索引集合 ,则 。
于是,SPLR得到了如下稀疏Hessian矩阵,
综述强调:加入 的意义在于可以改善 的条件数,并且给出正定性的保证。为了在 稠密时仍保持足够曲率信息,SPLR在 上补充了一个秩为2的Hessian近似,即 。SNS和SPLR的对比在下图中给出(节选自综述Figure 7)。

总结:
这篇综述系统地回顾了近年来围绕熵正则化OT展开的一系列稀疏化加速算法研究。不同于单纯从工程层面加速Sinkhorn,上述工作回答了一个更接近计算本质的问题:在OT问题中,哪些计算是真正必要的,而哪些又是可以被安全忽略或者近似的。
从这一视角出发,综述将现有方法清晰地划分为两类:核矩阵稀疏化,它关注每一步如何算得更快;Hessian矩阵稀疏化,它关注用更少的步数达到高精度。下图(节选自综述Table 1)显示了各种Sinkhorn加速方法计算复杂度和收敛速度的对比。

这篇综述不仅梳理了当前稀疏化OT方法的技术脉络,也为理解OT的计算瓶颈究竟在哪里、应当如何系统性地突破提供了一个清晰而统一的视角。对于希望将OT应用于大规模数据、复杂结构或高精度任务的研究者而言,这些稀疏化思想已经不再是可选加速技巧,而正在成为算法设计中不可忽视的核心组成部分。

文章精选:
1.图灵奖得主姚期智最新演讲: AI有边界,恰恰是好事
