Wasserstein距离与分布鲁棒优化：让模型从容应对数据分布偏移

简介面向机器学习、优化算法与电力系统交叉领域研究者的分布鲁棒优化项目代码包特别适合正在处理风电、光伏等新能源出力不确定性的电力系统研究人员也适合希望系统掌握推土机距离这一度量工具的算法工程师和研究生。资源围绕该距离展开先阐明其如何衡量真实分布与经验分布之间的差异再讲解如何构造Wasserstein球来限定不确定集并给出球半径的确定方法、风光出力样本的标准化流程以及借助数学工具把分布鲁棒约束转化为易求解形式的完整思路。压缩包共4个文件其中Python主程序为核心实现txt为依赖与使用说明inscode和html分别用于在线配置与结果预览整体仅6KB轻量紧凑、注释清晰便于直接运行和进一步扩展。已有110人学习。通过学习可以直接复用该距离计算、模糊集构建与鲁棒优化求解相关代码为电力系统鲁棒调度、储能配置和新能源消纳等实际场景提供可复用的实验基础是一份兼顾理论与工程落地的紧凑代码资料。 做机器学习这几年最让人头疼的不是模型不够复杂而是训练时表现完美的模型上线后遇到数据分布一波动效果立刻垮掉。分布偏移是常态而不是例外尤其在金融风控、供应链预测、工业传感这类场景里。今天要聊的Wasserstein距离和分布鲁棒优化就是专门针对这种问题的工具箱。Wasserstein距离负责度量分布之间的差异分布鲁棒优化则把这个差异直接写进训练目标让模型在“最坏情况”的分布偏移下依然不掉链子。这篇文章从原理讲到代码实现再把项目落地时踩过的坑一并列出来希望对正在做鲁棒性优化的朋友有实际帮助。1. 先看清问题标准模型在面对分布偏移时有多脆弱1.1 经验风险最小化的隐含假设我们平时训练模型绝大多数都是最小化经验风险[ \min_{\theta} \frac{1}{n} \sum_{i1}^{n} \ell(\theta; z_i) ]其中 (z_i(x_i,y_i)) 是训练样本。这个目标隐含了一个假设训练样本是从某个固定分布 (P) 中独立同分布采样的而我们希望模型在同一个分布 (P) 上的期望风险也能很小。问题就出在这里——当你把模型部署到线上数据来源变了或者环境发生了微小变化训练分布 (P) 和实际分布 (Q) 就不一致了。此时 ERM 训练出来的模型没有对这种偏移做任何防护效果暴跌几乎是必然的。我之前做过一个传感器故障预测项目模型在离线数据上 AUC 有 0.92上线一周后跌到 0.83。排查了数据质量、特征一致性最后才发现是传感器标定参数变了导致特征分布偏移。这种情况不是偶然而是所有依赖历史数据的系统的通病。1.2 一个直观的小实验给你一个简单的例子感受一下。生成一组一维回归数据真实关系是 (y2x\epsilon)训练集里加入 5% 的异常点把 (y) 人为加一个很大的偏移。用普通最小二乘拟合再用一个对离群点稳健的方法拟合你会发现最小二乘的系数估计被那 5% 的异常点拉偏了很多。这个现象的本质是 ERM 对所有训练样本一视同仁而异常点或者分布尾部样本对梯度的影响可能被放大。分布鲁棒优化的思路正好相反它假设测试分布不是固定的而是在训练分布附近的一个“不确定集合”内浮动然后针对集合里最坏的那个分布去优化模型。这种“悲观主义”的训练策略恰恰能提升模型在真实世界中的生存能力。2. Wasserstein距离从“搬土问题”到分布度量的数学工具2.1 最优传输视角Wasserstein距离也叫推土机距离Earth Movers Distance它的直觉很有意思想象你有两堆土形状分别代表两个概率分布要把一堆土搬成另一堆土最小的运输成本就是两个分布之间的 Wasserstein 距离。数学上两个分布 (P) 和 (Q) 之间的 p 阶 Wasserstein 距离定义为[ W_p(P,Q)\left(\inf_{\pi \in \Pi(P,Q)} \int |x-y|^p , d\pi(x,y)\right)^{1/p} ]其中 (\Pi(P,Q)) 是所有以 (P) 和 (Q) 为边缘分布的联合分布集合。这个定义初看有点抽象但你可以把 (\pi) 理解成一个运输计划(\pi(x,y)) 表示从 (x) 点运到 (y) 点的“土量”目标是最小化总运输成本。对于经验分布Wasserstein距离的计算可以写成一个线性规划问题。假设有两个经验分布分别有 (n) 和 (m) 个样本点代价矩阵 (C) 中 (C_{ij}|x_i-y_j|^p)那么[ W_p^p(\hat{P}n,\hat{Q}m)\min{\pi} \sum{i,j} C_{ij} \pi_{ij} ]约束是 (\pi \mathbf{1}_m \frac{1}{n}\mathbf{1}_n)(\pi^T \mathbf{1}_n \frac{1}{m}\mathbf{1}m)且 (\pi{ij} \ge 0)。2.2 为什么选Wasserstein而不是KL散度做分布鲁棒优化选择哪个距离来定义“不确定集合”非常关键。KL 散度和 Total Variation 距离虽然常见但有一个致命缺陷当两个分布的支撑集不重叠时KL散度会变成无穷大Total Variation 也无法提供连续的度量。而 Wasserstein 距离对支撑集的差异不敏感即使两个分布一个在 (x0) 附近一个在 (x10) 附近只要代价函数定义合理距离就是一个有限值。这一点在维数灾难下有更实际的意义。高维空间中真实数据分布往往集中在低维流形上。KL散度对这种“流形偏移”非常敏感哪怕一点点偏移都会导致散度爆炸。Wasserstein 距离则能优雅地处理这种几何结构上的微小变化因此更适合作为分布鲁棒优化中的“距离标尺”。下表是我在实际对比中的感受距离度量是否度量几何结构对非重叠支撑集计算难度对离群点敏感性KL散度否无穷大低高Total Variation否能算低中Wasserstein-1是有限中高中Wasserstein-2是有限中高中2.3 模糊集把分布偏移框在一个球里有了 Wasserstein 距离就可以构造所谓的模糊集ambiguity set[ \mathcal{B}_\varepsilon(\hat{P}_n) { Q : W_p(Q, \hat{P}_n) \le \varepsilon } ]这个集合里的所有分布都可以当作“可能的真实分布”。半径 (\varepsilon) 代表我们对分布偏移的容忍程度。分布鲁棒优化就是在这些分布上最小化最坏情况下的期望损失。注意这里 (\varepsilon) 的选择是有讲究的。理论上随着样本量 (n) 增大经验分布 (\hat{P}_n) 会以 (\sqrt{d/n}) 的量级收敛到真实分布 (P)所以 (\varepsilon) 可以设成这个收敛速度的某个倍数。但在实际项目中我更倾向于把 (\varepsilon) 当成一个超参数来调用交叉验证确定这和调正则化系数是一个逻辑。3. 分布鲁棒优化把“最坏情况”写进训练目标3.1 一般形式与对偶问题分布鲁棒优化的标准形式是[ \min_{\theta} \sup_{Q \in \mathcal{B}_\varepsilon(\hat{P}n)} \mathbb{E}{Z \sim Q}[\ell(\theta; Z)] ]直接求解这个 min-max 问题是很难的。但幸运的是在相当一般的条件下这个内层最大化可以转化为一个对偶问题。核心结果是这样的——当 Wasserstein 距离的代价函数是度量或者满足特定条件时有[ \sup_{Q: W_p(Q,\hat{P}_n) \le \varepsilon} \mathbb{E}_Q[\ell(\theta; Z)]\inf_{\lambda \ge 0} \left{ \lambda \varepsilon \frac{1}{n} \sum_{i1}^{n} \sup_{z \in \mathcal{Z}} \left[ \ell(\theta; z) - \lambda c(z, z_i) \right] \right} ]这个对偶形式把对分布的优化转换成了对每个样本点 (z_i) 的“局部扰动”优化加上一个关于 (\lambda) 的标量优化。(\lambda) 可以理解为对扰动幅度的惩罚系数它与模糊集半径 (\varepsilon) 是鞍点关系。3.2 从对偶到正则化的桥有一个经常被忽视的洞察Wasserstein DRO 在很多情况下和正则化是等价的。以线性回归为例如果损失是平方损失代价函数用欧氏距离平方那么对偶问题在某些条件下会退化成一个带 L2 正则的岭回归。这一点很关键因为它解释了为什么 DRO 能提升泛化性能它不只是“数据增强”或者“对抗训练”它其实是通过在分布层面做扰动实现了比普通正则化更精细的模型约束。正则化是对参数空间做约束而 DRO 是对数据分布做约束后者更贴近实际问题。3.3 为什么 DRO 比对抗训练更“温柔”对抗训练如 FGSM、PGD通常是在输入样本上做有界的对抗扰动扰动方向是让损失增大的方向。这种扰动是“极端点”而且它不考虑扰动后样本的分布特征。Wasserstein DRO 则不同它允许的扰动是在一个概率分布球内也就是说扰动后的样本整体上仍然呈现出与训练分布相近的结构只是在一些局部区域发生了质量转移。这更符合真实世界中“分布偏移”的形态。4. 项目代码用 Python 实现 Wasserstein 距离与 DRO4.1 环境准备这个项目我主要用 Python 3.10核心依赖是numpy、cvxpy、scipy。cvxpy用来求解 Wasserstein 距离的线性规划scipy.optimize用来求解 DRO 对偶问题。pip install numpy scipy cvxpy4.2 实现 Wasserstein 距离计算Wasserstein 距离的线性规划实现看起来简单但有几个容易出错的细节。以下是我在项目中使用的版本import numpy as np import cvxpy as cp def wasserstein_distance(X, Y, p2): 计算两组样本之间的Wasserstein距离 X: (n, d) 第一组样本 Y: (m, d) 第二组样本 p: 距离阶数 n, d X.shape m Y.shape[0] # 代价矩阵样本间距离的p次方 diff X[:, None, :] - Y[None, :, :] if p 1: C np.sqrt((diff ** 2).sum(-1)) elif p 2: C (diff ** 2).sum(-1) else: C np.sqrt((diff ** 2).sum(-1)) ** p # 传输计划变量 pi cp.Variable((n, m), nonnegTrue) # 边际约束两个经验分布都是均匀权重 p_marginal np.ones(n) / n q_marginal np.ones(m) / m constraints [ pi np.ones(m) p_marginal, pi.T np.ones(n) q_marginal ] prob cp.Problem(cp.Minimize(cp.sum(cp.multiply(C, pi))), constraints) prob.solve(solvercp.OSQP, verboseFalse) if p 2: # 如果是Wasserstein-2距离需要对结果开根号 return np.sqrt(prob.value) return prob.value这里有几个实操要点。第一代价矩阵 (C) 的维度是 ((n,m))如果样本量超过几千直接构建线性规划会比较吃力建议用分批估计或者采样近似。第二OSQP求解器对中小规模问题非常快但如果你要精确到小数点后很多位可以换成CLARABEL或者MOSEK如果有 license。第三两个经验分布都默认是均匀权重如果你的样本有权重记得修改边际约束。4.3 实现 Wasserstein DRO 线性回归推导对偶目标这一步是整个项目的核心。对于线性回归损失 (\ell(\theta; (x,y))(y-\theta^T x)^2)代价 (c(z,z_i)|x-x_i|^2(y-y_i)^2)。内层最大化问题可以通过 KKT 条件求得闭式解[ \sup_{z} \left[ (y-\theta^T x)^2 - \lambda(|x-x_i|^2(y-y_i)^2) \right] \frac{\lambda (y_i-\theta^T x_i)^2}{\lambda - (1|\theta|^2)} ]这个闭式解成立的条件是 (\lambda 1|\theta|^2)。于是 DRO 目标变成[ \min_{\theta, \lambda} \left{ \lambda \varepsilon \frac{1}{n} \sum_{i1}^{n} \frac{\lambda (y_i-\theta^T x_i)^2}{\lambda - (1|\theta|^2)} \right} ]对应的 Python 实现import numpy as np from scipy.optimize import minimize def dro_linear_regression(X, y, eps): Wasserstein DRO 线性回归平方损失 平方代价 X: (n, d) 特征 y: (n,) 标签 eps: Wasserstein 模糊集半径 n, d X.shape # 用最小二乘解作为初始值 theta0 np.linalg.lstsq(X, y, rcondNone)[0] lambda0 1.0 theta0 theta0 1.0 def objective(vars): theta vars[:d] lam vars[d] s 1.0 theta theta # 保证lambda 1 ||theta||^2否则目标无界 if lam s: return 1e12 r y - X theta inner lam * np.mean(r ** 2) / (lam - s) return lam * eps inner # 约束lambda - (1 theta^T theta) 0 cons ({ type: ineq, fun: lambda vars: vars[d] - (1.0 vars[:d] vars[:d]) }) result minimize( objective, np.r_[theta0, lambda0], methodSLSQP, constraintscons, options{maxiter: 200} ) return result.x[:d]4.4 对比实验ERM vs DRO写一个合成数据实验来对比# 生成数据真实theta [2.0, -1.5] np.random.seed(42) n 200 X np.random.randn(n, 2) theta_true np.array([2.0, -1.5]) y X theta_true 0.1 * np.random.randn(n) # 混入5%离群点 outlier_idx np.random.choice(n, sizeint(0.05 * n), replaceFalse) y[outlier_idx] np.random.randn(len(outlier_idx)) * 10 # ERM 和 DRO theta_erm np.linalg.lstsq(X, y, rcondNone)[0] theta_dro dro_linear_regression(X, y, eps0.05) # 在另一个偏移测试集上评估 X_test X 0.2 * np.random.randn(n, 2) y_test X_test theta_true 0.1 * np.random.randn(n) erm_test_error np.mean((y_test - X_test theta_erm) ** 2) dro_test_error np.mean((y_test - X_test theta_dro) ** 2) print(fERM test MSE: {erm_test_error:.4f}) print(fDRO test MSE: {dro_test_error:.4f})实验结果在我的场景下通常是 DRO 的测试误差比 ERM 低 15%~30%尤其在离群点比例较高且测试分布有偏移的情况下。注意DRO 的训练误差通常会略高于 ERM因为它刻意牺牲了一部分训练集拟合度来换取分布鲁棒性这是正常的。5. 调参与避坑经验5.1 半径 ε 怎么选这是整个 DRO 模型最核心的超参数。ε 太小模型退化成 ERM鲁棒性提升有限ε 太大模型会过度保守把所有样本都当作潜在异常点导致欠拟合。实践中我一般按以下思路取如果数据量 n 在几千到几万可以先算一下经验分布到几个参考分布的 Wasserstein 距离观察量级再在这个量级的 0.1 倍到 1 倍之间网格搜索。如果数据量很大可以按 (\varepsilon \propto \sqrt{d/n}) 的理论速率来缩放再乘一个常数。用验证集的 worst-slice 误差比如按误差排序取最差的 10% 的样本计算 MSE来选 ε比用平均误差更合理。5.2 求解器与数值问题的坑这里有一个我踩过很多次的坑对偶目标函数里那个 (\lambda - (1|\theta|^2)) 的分母在优化过程中很容易跑到零附近导致目标函数值爆炸。这就是为什么代码里加了约束 (\lambda 1|\theta|^2)。但 SLSQP 在边界附近可能不稳定我的做法是把约束改成 (\lambda \ge 1|\theta|^2 0.1)加一个小的安全余量。另外如果你的业务场景数据维度很高比如上千维直接求解线性规划形式的 Wasserstein 距离几乎不可行。这时候可以考虑用 Sinkhorn 距离加熵正则的最优传输计算复杂度大幅下降适合大规模应用。用随机梯度法估计对偶问题每次只用一小批样本计算内层 sup 的近似值。5.3 DRO 与其他鲁棒方法的对比我经常被问到一个问题既然有 dropout、L2 正则、对抗训练这些方法为什么还要用 DRO我的理解是L2 正则约束的是参数范数它可以缓解过拟合但对分布偏移没有直接的建模。对抗训练在输入空间上做扰动但扰动的分布特性不强容易出现“看起来增强了换个方向偏移就失效”的情况。DRO 则是在概率分布球内做优化它对“整体分布”的变化更敏感也更贴近真实系统的失效模式。当然这三者可以结合使用在项目中我通常会在 DRO 目标函数里再加一个很小的 L2 正则项效果往往更好。5.4 常见问题速查表问题可能原因解决方案目标函数出现 inf/NaNλ 接近边界 1||θ||²加安全余量或重新参数化 λDRO 结果和 ERM 几乎一样ε 设置过小调大 ε或改用最差分组误差来选训练时间过长每次迭代都求解线性规划用 Sinkhorn 近似或对偶随机梯度高维数据效果不佳距离度量在稀疏空间中失去区分度先降维或改用加权代价函数验证集整体误差上升ε 过大过度保守减小 ε或在 DRO 目标中加正则项6. 适用场景与选型建议6.1 金融风控金融行业是 DRO 最适合的土壤之一。市场状态切换、用户行为变化、政策调整都会导致数据分布发生显著偏移。我在信用评分模型中使用 Wasserstein DRO 的经验是它在客群迁徙、宏观经济波动导致的违约率上升场景下效果比普通评分卡模型稳定得多。具体做法是把样本按时间窗口切分用当前窗口与历史窗口之间的 Wasserstein 距离来度量分布漂移程度然后动态调整 ε。6.2 供应链与库存管理供应链中的需求预测天然受到促销、季节、突发事件影响需求分布经常在一夜之间改变。传统的报童模型newsvendor model假设需求分布已知而分布鲁棒版本则在 Wasserstein 球内考虑最坏需求分布决策者不需要准确估计需求分布只需要定一个合理的 ε。这在需求数据稀疏但有历史偏差的场景下特别实用。6.3 自然语言处理在 NLP 中Wasserstein DRO 可以用于抵抗域偏移。比如情感分析模型在一个领域的评论上训练投放到另一个领域的评论上时词汇分布会发生明显变化。此时用 Wasserstein 距离来约束词嵌入层面的分布变化比简单地做域对抗训练更可控。不过说实话NLP 场景中 DRO 的落地难度比数值型场景高因为文本的“距离”定义需要额外的语义功夫。6.4 什么时候不该用 DRO也不是所有场景都适合用 DRO。如果你的数据分布本身就很稳定或者你有充足的在线学习机制来不断更新模型DRO 带来的鲁棒性收益就会被这种动态更新抵消反而增加了计算成本。另外如果样本量特别小少于几百Wasserstein 距离估计本身方差很大ε 很难调准这个时候 DRO 的优势不大更多的收益来自贝叶斯方法或简单的正则化。我个人在实际项目中最大的体会是Wasserstein 距离不仅仅是构造 DRO 模糊集的工具它本身就是一个非常优秀的分布漂移监测指标。与其把 DRO 当成“吃了就能变强”的银弹不如先把它当成一套诊断方法定期计算当前数据分布与训练分布的 Wasserstein 距离一旦距离超过阈值就是模型需要重新训练或者调整的信号。这样用起来比单纯追求在某一个测试集上指标提升要踏实得多。最后再分享一个小细节。DRO 目标里的 ε 和正则化系数一样需要在每次数据量变化时重新校准。数据量翻倍经验分布更接近真实分布(\varepsilon) 就应该按 (\sqrt{d/n}) 的比例缩小。如果忽略这一点模型会因为 ε 相对过大而变得过度保守。这个坑我在三个月里踩了好几次希望你能避开。本文还有配套的精品资源点击获取