简介这份资源是一套用于计算一维概率分布1-和2-Wasserstein距离的紧凑MATLAB代码面向概率统计、最优传输与机器学习方向的学习者和研究者。它假定输入为离散且均匀的概率分布通过累积分布函数及其伪逆的阶跃函数形式实现距离计算代码已在MATLAB R2017a中测试并针对1-Wasserstein距离与通用p-Wasserstein实现进行了对比后者虽更通用但计算复杂度更高。资源包共3个文件包含1个m脚本、1个md说明文档和1个license许可文件压缩包约4KB体积轻量、便于直接嵌入现有项目。目前已有2492人学习下载读者可借此快速理解Wasserstein距离在均匀离散分布下的简化计算思路掌握从CDF到伪逆的推导与实现方法并可直接复用脚本完成分布间距离度量与结果验证。1. 均匀分布遇上 Wasserstein 距离一段 40 行 MATLAB 代码能解决什么很多人第一次听到 Wasserstein 距离是在做分布对齐或生成模型评估的时候。它比 KL 散度、JS 散度更“讲道理”的地方在于即使两个分布的支撑集完全不重叠它依然能给出一个有意义的数值而不是直接发散或者恒定常数。标题里提到的“均匀分布函数”其实是一个很典型的切入点——两个均匀分布之间的 Wasserstein 距离有解析解拿来验证代码正确性再合适不过。而“1D 和 2-Wasserstein”则分别对应一维情形和二维情形前者有闭式解后者需要求解一个带约束的线性规划或利用高斯近似。这段紧凑的 MATLAB 代码要做的就是让你在几十行内把这两类距离算出来并且能直接嵌入到自己的分布对比流程里。适合谁做信号处理、图像检索、域适应、分布鲁棒优化的从业者尤其是那些不想引入 CVX 或 Python 生态、只想在 MATLAB 里快速验证想法的人。2. 从均匀分布到最优传输Wasserstein 距离到底在算什么2.1 为什么均匀分布是验证 Wasserstein 代码的最佳试金石均匀分布的概率密度函数是常数累积分布函数是线性的分位数函数也是线性的。这三个“线性”叠加在一起使得一维 Wasserstein 距离的积分表达式可以直接手算。假设有两个一维均匀分布 U(a,b) 和 U(c,d)它们的 p-Wasserstein 距离在 p1 和 p2 时都有简洁的封闭形式。如果你写的代码连这两个分布之间的距离都算不对那后面处理经验分布或者高斯混合分布时就更不用谈了。我一般会先用 U(0,1) 和 U(1,2) 做冒烟测试预期结果可以直接用公式算出来对比代码输出误差在 1e-10 量级才算过关。一维情形下p-Wasserstein 距离的定义是W_p(μ,ν) ( ∫_0^1 |F_μ^{-1}(t) - F_ν^{-1}(t)|^p dt )^{1/p}其中 F^{-1} 是分位数函数。对于均匀分布分位数函数就是 a (b-a)t 这种形式代入积分后就是两个线性函数之差的 L^p 范数。当 p2 时积分里是二次函数展开后逐项积分即可。这个推导过程虽然简单但它是理解“为什么 Wasserstein 距离能反映分布的几何差异”的最好例子——两个均匀分布如果位置不同距离就是位置差如果尺度不同距离里就会同时包含位置和尺度的贡献。2.2 一维闭式解与二维线性规划的选型理由一维 Wasserstein 距离之所以有闭式解是因为最优传输计划在实数轴上一定是单调的把源分布的分位数和目标分布的分位数按顺序配对即可。这个性质在高维不成立所以二维及以上通常需要求解一个线性规划问题。但二维有一个特殊情况如果两个分布都是高斯分布那么 2-Wasserstein 距离也有闭式解只涉及均值向量和协方差矩阵的 Bures 度量。对于非高斯的二维分布常见做法是把连续分布离散化成经验分布然后求解一个带等式约束和不等式约束的线性规划。在 MATLAB 里求解线性规划最直接的工具是linprog它属于 Optimization Toolbox。如果你的环境没有这个工具箱也可以自己写一个简单的网络单纯形法或者使用quadprog来求解对偶问题。但考虑到“紧凑代码”这个目标我倾向于用linprog配合稀疏矩阵构造因为传输问题的约束矩阵有很强的结构稀疏存储能显著降低内存占用。对于 n 个源点和 m 个目标点决策变量是 n×m 个传输量约束是行和等于源权重、列和等于目标权重总共 nm 个等式约束。当 nm100 时决策变量有 10000 个等式约束 200 个linprog默认的对偶单纯形法能在一秒内解出来。2.3 用 MATLAB 实现一维 2-Wasserstein 距离的完整步骤先写一维的闭式解函数。输入是两个分布的分位数函数句柄或者直接是样本向量。如果是样本向量先排序然后计算经验分位数。下面这段代码假设输入是两个列向量长度可以不同通过插值统一到相同长度的分位数网格上。function w2 wasserstein1d_2(x, y, nGrid) % x, y: 样本列向量长度可以不同 % nGrid: 分位数网格点数默认 1000 if nargin 3 nGrid 1000; end % 排序 x sort(x(:)); y sort(y(:)); % 在 [0,1] 上均匀取分位点避免 0 和 1 处的边界问题 t linspace(1/(2*nGrid), 1-1/(2*nGrid), nGrid); % 经验分位数函数用线性插值 qx interp1((1:length(x))/length(x) - 0.5/length(x), x, t, linear, extrap); qy interp1((1:length(y))/length(y) - 0.5/length(y), y, t, linear, extrap); % 2-Wasserstein 距离的平方是分位数差平方的积分 diff2 (qx - qy).^2; w2 sqrt(trapz(t, diff2)); end这段代码的关键在于分位点的选取。如果直接用linspace(0,1,nGrid)在 t0 和 t1 处interp1会外插可能产生不稳定的值。我改成取每个小区间的中点这样插值始终在样本范围内数值更稳定。interp1的第三个参数是样本的累积概率位置我用(1:length(x))/length(x) - 0.5/length(x)来近似这是经验分布函数的标准做法。最后用trapz做梯形积分对于线性插值得到的分位数函数梯形积分是精确的。参数说明nGrid控制积分精度对于大多数一维问题1000 已经足够再大只会增加计算时间而不会显著改善结果。如果样本量很小比如少于 50可以把nGrid设成样本量的两倍左右避免插值引入的额外平滑。测试时用x rand(100,1)和y rand(100,1)0.5预期距离在 0.5 附近因为两个均匀分布只差一个位置偏移。2.4 二维离散最优传输的矩阵构造与 linprog 调用二维情形下先把两个分布离散化成带权重的点集。假设源分布有 n 个点坐标矩阵 X 是 n×2权重向量 wx 是 n×1目标分布有 m 个点坐标矩阵 Y 是 m×2权重向量 wy 是 m×1。代价矩阵 C 是 n×m其中 C(i,j) ||X(i,:) - Y(j,:)||^2。传输计划 T 是 n×m 的非负矩阵满足 T*1_m wx 和 T1_n wy。2-Wasserstein 距离的平方就是 sum(sum(T .C))。把 T 按列优先拉直成向量 t长度 nm。等式约束的矩阵 Aeq 是 (nm)×(nm)右端项 beq 是 [wx; wy]。Aeq 的前 n 行对应行和约束每行有 m 个 1后 m 行对应列和约束每行有 n 个 1但位置要错开。用稀疏矩阵构造可以避免生成巨大的稠密矩阵。function w2 wasserstein2d_2(X, wx, Y, wy) % X: n×2 源点坐标 % wx: n×1 源权重sum(wx)1 % Y: m×2 目标点坐标 % wy: m×1 目标权重sum(wy)1 n size(X,1); m size(Y,1); % 代价矩阵欧氏距离平方 C pdist2(X, Y, squaredeuclidean); % 构造等式约束稀疏矩阵 % 行和约束每个源点的流出等于其权重 rows1 repelem((1:n), m); cols1 (1:n*m); vals1 ones(n*m,1); % 列和约束每个目标点的流入等于其权重 rows2 n repmat((1:m), n, 1); cols2 (1:n*m); vals2 ones(n*m,1); Aeq sparse([rows1; rows2], [cols1; cols2], [vals1; vals2], nm, n*m); beq [wx(:); wy(:)]; % 目标函数最小化 sum(T .* C) f C(:); % 变量下界为 0 lb zeros(n*m,1); % 调用 linprog关闭显示 options optimoptions(linprog, Display, off, Algorithm, dual-simplex); [~, fval, exitflag] linprog(f, [], [], Aeq, beq, lb, [], options); if exitflag ~ 1 warning(linprog 未收敛exitflag %d, exitflag); end w2 sqrt(fval); end这段代码的核心是稀疏矩阵的构造。repelem和repmat分别生成行和列约束的行索引sparse一次性组装。注意pdist2需要 Statistics and Machine Learning Toolbox如果没有可以用sum((X - permute(Y,[3,2,1])).^2, 2)手动计算但内存占用会大一些。linprog的dual-simplex算法对这类运输问题通常比interior-point更快因为约束矩阵是全幺模的对偶单纯形法能保持整数解如果权重是整数的话。exitflag为 1 表示收敛其他值需要检查权重是否归一化、是否有负值。参数说明wx和wy必须严格归一化否则beq不满足可行性条件linprog会报错。如果权重里有零可以提前剔除对应的点减少问题规模。对于 nm200 的情况决策变量有 40000 个稀疏矩阵的非零元素约 80000 个现代笔记本上求解时间在几秒到十几秒之间。如果超过这个规模建议改用熵正则化的 Sinkhorn 算法但那就不属于“紧凑代码”的范畴了。3. 避坑与排查那些让 Wasserstein 距离算错的细节3.1 分位数插值在尾部发散现象一维代码在样本量较小时距离值比理论值大很多尤其是两个分布尾部差异大的时候。原因interp1在 t 接近 0 或 1 时依赖外插而经验分位数在尾部本身就不稳定。解决把分位点限制在 [1/(2n), 1-1/(2n)] 内或者改用核密度估计后再积分。我一般直接用中点法牺牲一点理论上的无偏性换取数值稳定性。3.2 linprog 返回 exitflag -2 或 -3现象二维代码报错“问题不可行”或“无界”。原因最常见的是权重向量没有归一化或者wx和wy的和不相等。另一个可能是代价矩阵里有 NaN通常是因为坐标里有缺失值。解决在调用linprog前加一行断言assert(abs(sum(wx)-1)1e-12 abs(sum(wy)-1)1e-12)并用sum(isnan(C(:)))检查代价矩阵。如果权重和不为 1手动除以和。3.3 稀疏矩阵构造时行列索引错位现象二维代码算出来的距离明显偏小甚至接近 0但两个分布明明不同。原因rows2和cols2的对应关系写反了导致列和约束实际约束的是错误的位置。解决用一个小规模例子手动验证比如 n2, m2打印Aeq的全矩阵形式检查每一行是否对应正确的变量。我习惯在调试时用full(Aeq)看一眼确认无误后再改回稀疏。3.4 欧氏距离平方与距离本身的混淆现象2-Wasserstein 距离算出来是理论值的平方。原因代价矩阵用了pdist2的euclidean而不是squaredeuclidean但最后又开了根号。解决记住 2-Wasserstein 的代价是距离的平方所以C必须是平方距离最后sqrt(fval)才得到距离。如果代价用距离本身那算出来的是 1-Wasserstein 距离。这两个别搞混我在早期项目里因为这个 bug 调了一下午。3.5 样本排序后忘记恢复对应关系现象一维代码里如果输入是多维样本的某个投影排序后分位数对了但后续如果需要用传输计划做映射索引就乱了。原因sort只返回排序后的值不返回原始索引。解决用[x_sorted, idx] sort(x)保留索引后续需要把目标分位数映射回原始样本时用idx还原。如果只是算距离不需要这一步但如果要做分布对齐或颜色迁移索引就是必须的。4. 进阶技巧用高斯近似加速二维 Wasserstein 并验证结果二维非高斯分布的线性规划在点数超过 500 时会明显变慢。一个实用的折中是如果两个分布的形状接近高斯直接用 Bures 度量算 2-Wasserstein 距离把均值向量和协方差矩阵代进去。公式是W_2^2 ||μ_x - μ_y||^2 tr(Σ_x Σ_y - 2(Σ_x^{1/2} Σ_y Σ_x^{1/2})^{1/2})在 MATLAB 里矩阵平方根用sqrtm但sqrtm对非对称矩阵可能返回复数。更稳健的做法是先用eig做特征分解再构造对称平方根。下面这个函数封装了高斯近似并和离散线性规划的结果做对比。function [w2_gauss, w2_lp] compare_wasserstein2d(X, Y) % X, Y: 样本矩阵每行一个样本 % 返回高斯近似和离散线性规划两种结果 % 高斯近似 mu_x mean(X, 1); mu_y mean(Y, 1); Sigma_x cov(X); Sigma_y cov(Y); % 对称平方根 sqrt_Sigma_x real(sqrtm(Sigma_x)); term sqrt_Sigma_x * Sigma_y * sqrt_Sigma_x; sqrt_term real(sqrtm(term)); w2_gauss sqrt(norm(mu_x - mu_y)^2 trace(Sigma_x Sigma_y - 2*sqrt_term)); % 离散线性规划均匀权重 n size(X,1); m size(Y,1); wx ones(n,1)/n; wy ones(m,1)/m; w2_lp wasserstein2d_2(X, wx, Y, wy); end这段代码里sqrtm外面套了real是因为数值误差可能让特征值出现极小的负值导致平方根结果带虚部。trace里的项如果出现负数说明协方差矩阵不正定需要检查样本是否退化。高斯近似在样本量大于 200 且分布单峰时和线性规划结果的相对误差通常在 5% 以内。如果误差超过 10%说明分布偏离高斯较远应该以线性规划结果为准。验证方法生成两组二维高斯混合样本一组用mvnrnd生成另一组手动加一个离群簇。分别跑compare_wasserstein2d观察两个数值的差距。如果离群簇权重很小高斯近似依然可用如果离群簇权重超过 20%线性规划结果会明显更大因为传输计划需要把质量搬到远处。这个对比能帮你判断在具体问题里能不能用高斯近似来省时间。我自己的习惯是任何二维 Wasserstein 计算先用高斯近似跑一遍得到一个基准值再用线性规划跑一遍如果两者差距在可接受范围内后续大规模实验就用高斯近似如果差距大就老老实实解线性规划或者改用 Sinkhorn 迭代。这个“先近似后精确”的流程帮我省过很多等待时间也避免过因为近似太粗糙而得出错误结论。希望帮到你。本文还有配套的精品资源点击获取