POT库中非平衡Sinkhorn算法的数值稳定性问题分析

问题背景

在Python最优传输工具库POT中，非平衡Sinkhorn算法有两种实现方式：标准版本和稳定化版本。用户发现这两种实现在某些情况下会给出不同的计算结果，特别是在使用吸收(absorbing)机制时。

问题重现

通过以下代码可以重现该问题：

import ot
import numpy as np

# 定义输入分布和成本矩阵
a = [0.3, 0.7]
b = [0.7, 0.3]
M = [[0.0, 1.0], [1.0, 0.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)

# 输出结果比较
print(np.round(Q1, 6))
print(np.round(Q2, 6))

问题本质

该问题本质上源于数值不稳定性，特别是在处理小数值时的对数运算。稳定化版本通过引入额外的对数缩放因子来防止数值下溢，但这种处理方式在实现细节上存在差异。

技术分析

1. 标准实现：直接对u和v向量进行对数运算，保持了原始数值关系
2. 稳定化实现：仅对u和v的最大值进行对数运算，这可能导致信息丢失

解决方案

修改稳定化实现中的对数运算方式，从仅对最大值取对数改为对整个向量取对数：

alpha = alpha + reg * nx.log(u)  # 原为nx.max(u)
beta = beta + reg * nx.log(v)    # 原为nx.max(v)

更深层次理解

在最优传输问题中，数值稳定性是一个常见挑战。非平衡Sinkhorn算法通过引入额外的正则化项来处理质量不守恒的情况，但这增加了数值计算的复杂度。稳定化版本通过缩放因子来防止数值下溢，但需要谨慎处理对数运算的范围。

实际应用建议

1. 对于小规模问题，可以使用标准实现
2. 对于大规模或数值敏感问题，建议使用稳定化版本
3. 比较结果时，应注意设置相同的参数，特别是tau值
4. 结果差异在可接受范围内时，不必过度优化

总结

POT库中的非平衡Sinkhorn算法实现差异反映了数值优化中的常见挑战。理解这些差异有助于用户根据具体问题选择合适的算法变体。该问题的解决方案也展示了数值稳定性处理中的关键细节，对开发类似算法的工程师具有参考价值。