Inexact Bregman Sparse Newton Method for Efficient Optimal Transport¶

会议: NeurIPS 2026
arXiv: 2603.07156
OpenReview: https://openreview.net/forum?id=1bW1n2cWFS
论文: https://arxiv.org/abs/2603.07156
领域: 优化/理论
关键词: 最优传输, Bregman近端点法, 稀疏牛顿法, 半对偶形式, 粗精两阶段求解

一句话总结¶

针对大规模精确最优传输计算复杂度过高且熵正则化数值不稳定的困境，本文提出了非精确Bregman稀疏牛顿法（IBSN），在半对偶空间中结合Hessian稀疏化与可验证停机准则，以二阶超线性收敛加速精确求解。

研究背景与动机¶

最优传输（Optimal Transport, OT）提供了一种具有几何意义的概率测度间距离度量方式，在生成模型、域适应、点云配准及图像处理等领域得到了极为广泛的应用。然而，标准离散最优传输问题本质上是一个大规模线性规划（LP），经典的网络单纯形法和内点法在面对高维或大规模数据时，计算复杂度随样本量呈超立方级激增，难以在可接受时间内处理大规模问题。

为了降低计算门槛，主流方法通过引入负熵正则化转为熵正则化最优传输（EOT），并利用一阶Sinkhorn矩阵缩放算法求解。但Sinkhorn算法仅具有次线性收敛速率，在要求高精度的场景下迭代步数激增；更严重的是，当正则化参数 $\eta$ 趋于0以逼近原OT问题真解时，指数运算极易引发下溢或数值溢出崩溃。近期将EOT作为子问题的Bregman近端点方法（如PINS）引入了牛顿法以加速收敛，但现有方法通常在原始双对偶空间求解规模为 $(m+n) \times (m+n)$ 的线性系统，且要求每个子问题完全精确收敛，造成极其沉重的计算负担；而其他非精确近端算法的内层停机条件又难以在理论与实践中有效验证。

面对这些限制，必须在减少对偶维度、放宽内层解算精度与保持全局收敛之间寻求突破。本文的核心 idea 是通过半对偶变换将对偶变量从 $m+n$ 维消元降至 $n$ 维，并构建带动态阈值与行归一化的正定稀疏Hessian矩阵，配合可验证的Bregman散度非精确停机准则，实现低代价且保证收敛到精确最优传输方案的二阶迭代。

方法详解¶

整体框架¶

IBSN整体采用外层Bregman近端点框架与内层半对偶稀疏牛顿法相结合的架构。输入为边际分布向量 $\boldsymbol{a}\in\mathbb{R}^m_}$、$\boldsymbol{b}\in\mathbb{R}^n_{$ 以及代价矩阵 $C\in\mathbb{R}^{m\times n}$，目标是精确求解最优传输矩阵 $X^\star\in\Omega$。

算法在外层迭代中生成一系列以负熵Bregman散度为正则的子问题。首先在双对偶空间运行若干步Sinkhorn迭代快速获得粗解，随后解析消去对偶变量 $\boldsymbol{\zeta}$ 得到仅含 $\boldsymbol{\gamma}\in\mathbb{R}^n$ 的光滑半对偶目标 $\mathcal{L}^k(\boldsymbol{\gamma})$。在内层阶段，根据梯度范数自适应设定截断阈值，对半对偶Hessian矩阵执行稀疏化与行和归一化投影，利用共轭梯度法（CG）求解带谱修正的牛顿方向并配合Armijo线搜索更新对偶变量；最后通过可验证的舍入投影Bregman散度准则提前终止子问题，进入下一个外层近端更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：边际分布 a, b 与代价矩阵 C"] --> B["外层Bregman近端子问题构建<br/>引入上一轮解的Bregman散度"]
    B --> C["半对偶降维与Sinkhorn热启动<br/>解析消去对偶变量 ζ 得 n 维半对偶"]
    C --> D["自适应Hessian稀疏化<br/>基于梯度范数截断与最大行保底归一化"]
    D --> E["稀疏牛顿方向求解与步长搜索<br/>CG求解正定系统 + Armijo回退"]
    E --> F["可验证非精确停机检验<br/>舍入投影Bregman散度判定提前终止"]
    F -->|满足停机容差| G["输出：精确最优传输方案 X^K"]
    F -->|未达容差| D

关键设计¶

1. 半对偶降维与Sinkhorn热启动：消除冗余自由度并压低求解空间

传统EOT双对偶表述包含两个对偶向量 $\boldsymbol{\zeta}\in\mathbb{R}^m$ 和 $\boldsymbol{\gamma}\in\mathbb{R}^n$，导致Hessian矩阵的维度高达 $(m+n)\times(m+n)$。针对这一维度瓶颈，本文利用一阶导数驻点条件 $\frac{\partial Z^k(\boldsymbol{\gamma},\boldsymbol{\zeta})}{\partial\boldsymbol{\zeta}}=\boldsymbol{0}_m$ 显式解出 $\boldsymbol{\zeta}(\boldsymbol{\gamma})$： $$\zeta(\boldsymbol{\gamma})_i = \eta\log a_i - \eta\log\Big(\sum_{j=1}^n X_{ij}^k\exp\{(\gamma_j - C_{ij})/\eta\}\Big)$$ 将其代回目标函数后，优化问题收缩为仅关于 $\boldsymbol{\gamma}\in\mathbb{R}^n$ 的半对偶问题 $\min_{\boldsymbol{\gamma}\in\mathbb{R}^n}\mathcal{L}^k(\boldsymbol{\gamma})$。变量数从 $m+n$ 直接降为 $n$，不仅将存储与CG求解开销从超大矩阵降至 $n\times n$ 级别，更在高度非平衡场景（$m\gg n$）下将Hessian完全约束在低维空间。在每个外层步初始，先运行数步Sinkhorn算法生成粗解并投影至零和正交补子空间 $\boldsymbol{1}_n^\perp$，为后续二阶牛顿迭代提供高质量的热启动初值。

2. 结构自适应Hessian稀疏化：保正定与误差可控的稀疏逼近

精确半对偶Hessian为 $H^k(\boldsymbol{\gamma})=\frac{1}{\eta}(\operatorname{diag}(P^k(\boldsymbol{\gamma})^\top\boldsymbol{a})-P^k(\boldsymbol{\gamma})^\top\operatorname{diag}(\boldsymbol{a})P^k(\boldsymbol{\gamma}))$，虽然在数学上结构稠密，但其核心分量 $P^k(\boldsymbol{\gamma})$ 的绝大多数元素趋于零。为降低矩阵向量乘复杂度，本文提出了带保护机制的截断策略：选定主导行 $i^\star=\arg\max_{i\in[m]}a_i$ 完整保留，其余所有行按阈值 $\rho$ 截断并重新归一化行和为1： $$[P_\rho]_{ij} \leftarrow \frac{P_{ij}\boldsymbol{1}_{\{P_{ij}\ge\rho\}}}{\sum_{k=1}^n P_{ik}\boldsymbol{1}_{\{P_{ik}\ge\rho\}}},\quad \forall i\neq i^\star;\qquad [P_\rho]_{i^\star\cdot}\leftarrow P_{i^\star\cdot}$$ 该结构设计不仅严格保障了稀疏矩阵 $H_\rho$ 在子空间 $\boldsymbol{1}_n^\perp$ 上严格正定（其核空间恒为 $\operatorname{span}\{\boldsymbol{1}_n\}$），避免了牛顿矩阵奇异导致的崩溃，而且建立了严谨的近似误差上界： $$\|H - H_\rho\| \le \frac{6mn\|\boldsymbol{a}\|_\infty}{\eta}\rho$$ 算法进一步令 $\rho = \frac{\eta}{mn}\|\boldsymbol{g}^{k,v}\|$，随着迭代靠近极值点、梯度范数逐渐衰减，截断阈值自适应收缩，使得Hessian近似精度在局部自发恢复，从而保证了二阶方法的局部二次收敛速率。

3. 可验证非精确停机准则：松弛子问题计算负担

如果强行要求每个外层Bregman子问题迭代至机器精度，外层循环的计算代价将极其昂贵。本文引入Bregman散度投影停机条件，在保证整体全局收敛的前提下对内层子问题实施早停： $$D_\phi(\mathcal{P}_\Omega(X^{k,v+1}), X^{k,v+1}) \le \min\{m, n\}\cdot\mu_k,\quad \text{其中 } \mu_k=\max\left\{\frac{10^{-4}}{(k+1)^2}, 10^{-11}\right\}$$ 为避免每步显式求解精确投影，算法采用近线性复杂度的Greenkhorn舍入过程实现可行域投影 $\mathcal{P}_\Omega$，并在检查散度前先用轻量级的梯度范数滤波 $\|\nabla\mathcal{L}^k(\boldsymbol{\gamma}^{k,v+1})\| < \mu_k$ 过滤无效评估。理论证明只要容差序列可和（$\sum_k \mu_k < \infty$），该非精确序列就严格全局收敛至原OT线性规划的真解 $f^\star$。

实验关键数据¶

主实验¶

实验在合成数据（Uniform和Square代价矩阵，维度 $m=n\in\{1000, 5000, 10000\}$）及真实图像基准（MNIST、Fashion-MNIST $784\times 784$、DOTmark $32\times 32$ 和 $64\times 64$）上展开，统一以相对KKT残差 $\Delta_{\text{kkt}}<10^{-11}$ 或Gurobi高精度目标间隙（Gap）为终止准则，基线对比了PINS、HOT、IBSink、IPOT及ExtraGrad等主流算法。

| 任务设定 | 维度 (m=n) | 指标 | IBSN (本文) | PINS (二阶SOTA) | IBSink (一阶近端) | 表现分析 | |---|---|---|---|---|---| | Synthetic: Square | 1,000 | 运行时间 (s) / 间隙 | 24.68 / 1.07e-8 | 49.32 / 1.07e-8 | 138.50 / 1.07e-8 | 耗时减少约 50.0%，二阶优势显著 | | Synthetic: Square | 5,000 | 运行时间 (s) / 间隙 | 457.10 / 1.49e-7 | 986.30 / 1.49e-7 | 2540.20 / 1.49e-7 | 较PINS加速 2.16×，较一阶法提速数倍 | | Synthetic: Square | 10,000 | 运行时间 (s) / 间隙 | 2205.74 / 1.59e-7 | 4830.15 / 1.59e-7 | >5000 (超时) | 规模扩大时半对偶稀疏牛顿扩展性突出 | | DOTmark Shapes | 1,024 (32×32) | CG总耗时 (s) / CG步数 | 0.37 / 360.3 | 1.45 / 763.2 | - (一阶无CG) | 求解牛顿方向所需CG迭代缩减超过 52.8% | | DOTmark Shapes | 4,096 (64×64) | CG总耗时 (s) / 迭代步数 | 6.26 / 14 | 28.85 / 12 | - | CG方向计算耗时降低 78.3% |

消融实验¶

为明确Hessian稀疏化策略在牛顿阶段的核心贡献，本文设立了消融对照组IBN（即移除Hessian截断与稀疏化、保留半对偶牛顿架构的变体），在合成数据集上测量CG方向求解时间与整机运行效率。

实验配置	数据集与规模	Newton迭代数	CG总耗时 (s)	算法总时间 (s)	最终间隙 (Gap)
IBN (无稀疏化稠密Hessian)	Uniform (m=n=1000)	177	28.38	33.47	3.18e-10
IBSN (完整模型: 稀疏Hessian)	Uniform (m=n=1000)	176	3.86 (-86.4%)	11.04 (-67.0%)	3.19e-10
IBN (无稀疏化稠密Hessian)	Square (m=n=5000)	301	803.53	1101.43	1.49e-07
IBSN (完整模型: 稀疏Hessian)	Square (m=n=5000)	301	74.37 (-90.7%)	457.10 (-58.5%)	1.49e-07
IBN (无稀疏化稠密Hessian)	Square (m=n=10000)	300	2654.52	4303.08	1.59e-07
IBSN (完整模型: 稀疏Hessian)	Square (m=n=10000)	300	267.15 (-89.9%)	2205.74 (-48.7%)	1.59e-07

关键发现¶

CG方向计算开销实现数量级缩减：消融数据显示，Hessian稀疏化策略在不同规模下均将CG求解时间压缩了 86%–91%，且未引起牛顿迭代次数的增加或精度损失，彻底打破了二阶法矩阵向量乘瓶颈。
破解非精确求解中段瓶颈：一阶非精确方法（IBSink）在容差序列 $\mu_k$ 逐渐收紧的中间阶段会遭遇严重的迭代步数停滞，而IBSN凭借二阶二次收敛特性，在容差严格时仍能迅速跃出瓶颈区。
非平衡场景与维度鲁棒性：在球面地理人口传输测试（$m\gg n$）中，半对偶机制使得牛顿系统完全独立于源样本维度 $m$ 的膨胀，IBSN在高度非平衡传输任务中的提速比随 $m$ 增长呈线性上升趋势。

亮点与洞察¶

半对偶消元的维度压制：直接通过KKT一阶条件解析剔除一组对偶变量，将原本双对偶下的鞍点/高维牛顿系统收缩至单边分布空间，在非平衡最优传输中尤其具备降维打击优势。
保留主导行的正定子空间构造：稀疏化直接置零会导致Hessian失去强凸性或陷入奇异，本文通过强制保留最大边际行 $i^\star$ 的设计，以极简形式在数学上严格锁定了子空间 $\boldsymbol{1}_n^\perp$ 内的最小正特征值下界。
自适应阈值动态退火机制：令截断阈值 $\rho$ 与梯度范数自适应绑定，在远离极值时激进稀疏化以换取极致速度，在接近收敛时自动退火至高保真Hessian，兼具粗搜速度与二次精敛特性。

局限与展望¶

极端规模下的线搜索开销上升：消融分析显示，当维度扩展至 $10000\times 10000$ 时，CG耗时占比大幅下降，而矩阵指数计算、Armijo线搜索与舍入投影散度的计算逐渐成为整体耗时的主要部分，未来有待开发免线搜索的二阶信赖域或拟牛顿加速版本。
稠密代价矩阵的高维内存瓶颈：当前算法依然需要在内存中实例化或索引大小为 $m\times n$ 的完整代价矩阵，对于百万级超大规模测度配准，仍需探索低秩核分解或分块流式加载技术。
正则化参数的手动选取经验性较强：外层参数 $\eta$ 的衰减调度与初始选取仍依赖网格调优，若能建立自适应参数动态调节策略将进一步提高实用性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙融合半对偶降维、行保底稀疏Hessian与非精确Bregman停机准则，理论体系严密闭环。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖不同规模合成数据、标准真实图像集与非平衡球面实验，消融与收敛分析翔实。
写作质量: ⭐⭐⭐⭐⭐ 符号系统严谨，理论定理完备，算法各环节设计意图清晰。
价值: ⭐⭐⭐⭐⭐ 为大规模精确最优传输计算提供了极具实用性与理论保障的二阶加速基准工具。