PythonOT/POT项目中不平衡Sinkhorn算法的数值稳定性问题分析

概述

在PythonOT/POT项目中，用户报告了一个关于不平衡Sinkhorn算法实现的问题。具体表现为sinkhorn_stabilized_unbalanced函数与sinkhorn_knopp_unbalanced函数在吸收(absorbing)机制启用时会产生不同的计算结果，而当不启用吸收机制时，两者结果一致。

问题背景

Sinkhorn算法是计算最优传输问题的一种有效方法，通过引入熵正则化来近似求解。在不平衡最优传输问题中，算法需要处理质量不守恒的情况，即输入分布的总质量不等于输出分布的总质量。POT库提供了多种实现变体，包括标准版本和数值稳定版本。

问题现象

用户提供的示例代码清晰地展示了这个问题：

import ot
import numpy as np
a=[.3, .7]
b=[.7, .3]
M=[[0., 1.], [1., 0.]]
reg = 0.01
reg_m = 100
Q1 = ot.sinkhorn_unbalanced(a, b, M, reg, reg_m, reg_type='entropy')
Q2 = ot.unbalanced.sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type='entropy', tau=1000)

理论上，Q1和Q2应该给出相同的结果，但实际上却出现了差异。

技术分析

数值稳定性机制

数值稳定版本的Sinkhorn算法通过引入对数域计算(log-domain)来避免数值下溢问题。在标准实现中，当数值变得非常小时，直接计算会导致精度损失。稳定版本通过保持计算在对数域中进行来缓解这个问题。

吸收机制的影响

吸收(absorbing)是数值稳定算法中的一种技术，用于处理极端小的数值情况。当数值低于某个阈值(tau)时，算法会进入吸收状态，采用不同的计算策略。用户观察到的差异正是发生在吸收机制被触发时。

潜在修复方案

用户提出了一个修改建议：将原来的nx.max(u)改为直接使用u。这个修改确实可以消除差异，但需要仔细评估其对数值稳定性的影响：

alpha = alpha + reg * nx.log(u)  # 修改后
beta = beta + reg * nx.log(v)    # 修改后

深入理解

1. 正则化参数的作用：reg参数控制熵正则化的强度，较小的值会使问题更接近原始最优传输问题，但也更容易出现数值不稳定。

2. 质量松弛参数：reg_m控制质量不守恒的惩罚强度，较大的值强制更严格的质量守恒。

3. 对数域计算：稳定版本通过在对数域中操作，避免了小数值的乘法运算，转而使用加法运算，这在数值上更稳定。

解决方案建议

1. 参数调整：首先尝试调整tau参数，找到一个平衡点，既保持数值稳定性又获得合理结果。

2. 算法选择：根据问题规模选择合适的算法变体，对于小规模问题，标准版本可能足够。

3. 实现验证：检查两种实现的一致性测试，确保在非吸收情况下结果一致。

结论

这个问题揭示了数值算法实现中的常见挑战：在追求数值稳定性的同时保持算法的理论正确性。POT库提供了多种实现选项，用户应根据具体问题的数值特性选择最合适的变体。对于大多数实际应用，数值稳定版本提供了更好的鲁棒性，尽管在极端参数下可能与标准版本存在微小差异。