Sinkhorn算法Python实战：5行代码计算Wasserstein距离，误差<1%

在机器学习项目中，我们经常需要比较两个概率分布的相似性。传统的KL散度或JS散度虽然计算简单，但无法捕捉分布间的几何结构信息。Wasserstein距离（又称推土机距离）通过考虑"搬运"概率质量的最小成本，成为衡量分布差异的更优选择。然而，经典最优传输问题的高计算复杂度（O(n³ log n)）限制了其应用。本文将介绍基于熵正则化的Sinkhorn算法，用5行Python核心代码实现Wasserstein距离的高效计算，误差可控制在1%以内。

1. Wasserstein距离与最优传输问题

1.1 从推土问题到概率分布

想象这样一个场景：我们需要将M个土堆（每个土堆大小为aₘ）运输到N个坑（每个坑容量为bₙ），从土堆m到坑n的运输成本为c(m,n)。最优传输问题就是寻找总运输成本最小的分配方案。当我们将土堆和坑看作两个概率分布时，这个最小成本就是Wasserstein距离。

数学上，给定两个d维概率向量r和c（即∑rᵢ=1，∑cⱼ=1），以及成本矩阵M∈ℝ⁽ᵈˣᵈ⁾，Wasserstein距离定义为：

d_M(r,c) = min_{P∈U(r,c)} <P,M> 其中U(r,c) = {P∈ℝ⁺⁽ᵈˣᵈ⁾ | P1_d=r, Pᵀ1_d=c}

1.2 计算挑战与熵正则化

直接求解这个线性规划问题对于大规模数据不可行。Cuturi在2013年提出通过熵正则化将问题转化为：

d_M^λ(r,c) = min_{P∈U(r,c)} <P,M> - (1/λ)h(P) 其中h(P)=-∑PᵢⱼlogPᵢⱼ是香农熵

参数λ控制正则化强度：

• λ→0：解趋向均匀分布P=rcᵀ
• λ→∞：恢复原始最优传输问题

2. Sinkhorn算法核心实现

2.1 算法原理

熵正则化后的最优传输问题具有唯一解P⁺=diag(u)Kdiag(v)，其中K=exp(-λM)。Sinkhorn算法通过交替更新标量向量u和v来求解：

u ← r / (Kv) v ← c / (Kᵀu)

2.2 5行Python实现

使用NumPy和POT库，核心计算仅需5行代码：

import numpy as np from ot import sinkhorn def sinkhorn_wasserstein(r, c, M, reg=0.1, max_iter=1000): K = np.exp(-M/reg) # Gibbs核矩阵 u = np.ones_like(r) for _ in range(max_iter): v = c / (K.T @ u) u = r / (K @ v) return np.sum(u * (K * M) @ v) # W距离

2.3 参数选择指南

3. 实战对比：Sinkhorn vs 经典EMD

3.1 实验设置

我们比较两种方法在MNIST数字分布上的表现：

• 经典EMD：ot.emd2
• Sinkhorn：上述实现

# 生成两个MNIST数字的灰度直方图 digits = load_digits().data r = digits[0].astype(np.float64) + 1e-8 r /= r.sum() c = digits[1].astype(np.float64) + 1e-8 c /= c.sum() # 成本矩阵（像素欧氏距离） coord = np.array([[i//8, i%8] for i in range(64)]) M = np.sum((coord[:,None]-coord[None,:])**2, axis=2)**0.5

3.2 结果对比

注意：测试环境为Intel i7-1185G7 CPU，64GB RAM。Sinkhorn算法在保持1%以内误差的同时，实现了数量级的速度提升。

4. 调参技巧与常见问题

4.1 正则化参数λ的选择

λ是精度与速度的权衡：

• 高精度场景：λ=50-100（reg=0.01-0.02）
• 平衡场景：λ=10-20（reg=0.05-0.1）
• 快速近似：λ=5（reg=0.2）

实际建议从λ=10开始，观察结果稳定性：

lambdas = [5, 10, 20, 50] results = {f'λ={l}': sinkhorn_wasserstein(r, c, M, reg=1/l) for l in lambdas}

4.2 数值稳定性处理

指数运算可能导致数值溢出，推荐以下改进：

K = np.exp(-(M - np.min(M))/reg) # 平移确保数值稳定 K = K / np.max(K) # 归一化

4.3 稀疏矩阵加速

当d>1000时，可使用稀疏矩阵运算：

from scipy.sparse import csr_matrix def sparse_sinkhorn(r, c, M_sparse, reg=0.1): K_sparse = csr_matrix(np.exp(-M_sparse.data/reg)) u = np.ones_like(r) for _ in range(1000): Kv = K_sparse.dot(v) u = r / Kv Ktu = K_sparse.T.dot(u) v = c / Ktu return np.sum(u * (K_sparse.multiply(M_sparse)).dot(v))

5. 进阶应用场景

5.1 Wasserstein GAN中的使用

在WGAN中，Sinkhorn距离可作为判别器的损失函数：

# PyTorch实现示例 import torch class SinkhornLoss(nn.Module): def __init__(self, reg=0.1): super().__init__() self.reg = reg def forward(self, r, c, M): K = torch.exp(-M/self.reg) u = torch.ones_like(r) for _ in range(100): v = c / (K.t() @ u) u = r / (K @ v) return torch.sum(u * (K * M) @ v)

5.2 点云配准

匹配两个点云分布时，Sinkhorn算法比ICP更鲁棒：

def align_point_clouds(source, target): # source/target: [n,3]和[m,3]点云 M = torch.cdist(source, target) # 成对距离矩阵 P = sinkhorn(torch.ones(n)/n, torch.ones(m)/m, M) aligned = P @ target # 加权平均 return aligned

5.3 跨域图像匹配

匹配不同域（如素描↔照片）的特征分布：

def domain_adaptation(feats_A, feats_B): # feats_A/B: [n,d]和[m,d]特征矩阵 M = 1 - cosine_similarity(feats_A, feats_B) P = sinkhorn(np.ones(n)/n, np.ones(m)/m, M) return P # 传输矩阵可用于特征对齐

6. 与其他方法的对比

6.1 计算效率对比

在d=1000的分布上测试：

6.2 质量对比

在CIFAR-10生成任务中：

7. 工程优化技巧

7.1 GPU加速

使用CuPy或PyTorch实现GPU加速：

import cupy as cp def gpu_sinkhorn(r, c, M, reg=0.1): M_gpu = cp.array(M) K = cp.exp(-M_gpu/reg) u = cp.ones_like(r) for _ in range(1000): v = c / (K.T @ u) u = r / (K @ v) return cp.sum(u * (K * M_gpu) @ v).get()

7.2 批处理实现

同时计算多个分布对的距离：

def batch_sinkhorn(R, C, M, reg=0.1): # R: [b,n], C: [b,m], M: [n,m] K = torch.exp(-M/reg) U = torch.ones_like(R) for _ in range(100): V = C / (K.t() @ U) U = R / (K @ V) return torch.sum(U * (K * M) @ V, dim=1)

7.3 自动微分支持

Sinkhorn迭代本身可微，可直接嵌入神经网络：

# PyTorch自动微分示例 r = torch.rand(64, requires_grad=True) c = torch.rand(64, requires_grad=True) M = torch.rand(64,64) distance = sinkhorn_wasserstein(r, c, M) distance.backward() # 可计算梯度

8. 数学原理深入

8.1 Sinkhorn定理

对于任意正矩阵A，存在对角矩阵D₁,D₂使得D₁AD₂是双随机矩阵（行列和均为1）。Sinkhorn迭代就是通过交替行列归一化寻找这些对角矩阵。

8.2 对偶问题视角

原始问题等价于求解对偶问题：

max_{f,g} fᵀr + gᵀc - (1/λ)∑exp[λ(fᵢ + gⱼ - Mᵢⱼ)]

Sinkhorn迭代实际上是对偶坐标上升法。

8.3 收敛性证明

可以证明：

• 迭代过程单调降低原始目标
• 线性收敛率：‖u⁽ᵏ⁾ - u*‖ ≤ C/(1+λ)ᵏ
• 对固定λ，迭代复杂度O(n²/ε)

9. 变体与扩展

9.1 不平衡最优传输

当∑rᵢ ≠ ∑cⱼ时，添加边际违反惩罚：

def unbalanced_sinkhorn(r, c, M, reg=0.1, tau=0.5): K = np.exp(-M/reg) u, v = np.ones_like(r), np.ones_like(c) for _ in range(1000): u = (r/(K@v)) ** tau # 添加tau控制 v = (c/(K.T@u)) ** tau return np.sum(u*(K*M)@v)

9.2 低秩近似

对于超大规模问题，使用低秩分解：

def lowrank_sinkhorn(r, c, U, V, reg=0.1): # U: [n,k], V: [m,k], M ≈ UVᵀ K = np.exp(-(U@V.T)/reg) ... # 标准Sinkhorn迭代

10. 实际案例：图像风格迁移

10.1 颜色分布匹配

将目标图像的颜色分布迁移到源图像：

def color_transfer(source, target): # 将图像转换为3D颜色点云 src_pixels = source.reshape(-1,3) tgt_pixels = target.reshape(-1,3) # 计算颜色距离矩阵 M = np.sum((src_pixels[:,None]-tgt_pixels[None,:])**2, axis=2) # 计算最优传输 P = sinkhorn(np.ones(len(src_pixels)), np.ones(len(tgt_pixels)), M) # 应用传输矩阵 transferred = P @ tgt_pixels return transferred.reshape(source.shape)

10.2 结果对比