做跨设备故障识别时碰到过分布偏移我第一反应就是上 MMDMaximum Mean Discrepancy最大均值差异来度量两个特征分布的距离。前段时间在 PyTorch 和 TensorFlow 里各写了一遍 MMD过程中的坑不少这篇把代码和调试经验一起复盘出来。这篇文章不是单纯把公式贴一遍而是把“为什么要这么写”“换成另一个框架要注意什么”“哪些参数会害你调一整天”都讲清楚。适合两类人看一是刚接触域适应、想知道 MMD 到底在算什么的新手二是已经会调接口但想自己实现、遇到梯度或数值问题需要排查的开发者。我会从原理入手给出完整可跑的 PyTorch 与 TensorFlow 代码最后是踩坑记录。1. MMD 原理为什么域适应离不开它1.1 分布差异为什么难度量先看一个实际现象。同一个特征提取网络在 A 设备的数据上训练完测试精度能做到 98%换到 B 设备的数据上直接掉到 70%。网络权重没变数据格式也没变变的是数据分布。A 设备的信号均值在 1 附近B 设备由于传感器型号不同、环境噪声不同信号整体平移到了 2 附近。这种分布偏移在真实项目中非常多见尤其是跨设备、跨时段、跨场景迁移的时候。要解决这个问题第一步就是量化“两个分布的差异”。听起来简单做起来并不容易。最朴素的做法是逐维比较均值把源域特征每个维度的均值算出来目标域每个维度的均值算出来然后求差的范数。但这种方法对分布形状不敏感。举个二维的例子源域样本分布在直线 yx 附近目标域样本分布在直线 y-x 附近。两个分布沿对角线方向完全错开但每个维度的均值却是一样的逐维均值差会得出“分布几乎相同”的错误结论。特征维度更高时这种问题只会更严重。我们需要的是一个能捕捉整体分布形状的度量。MMD 的思路很巧妙不直接比较原始样本而是把样本映射到一个高维特征空间在这个空间里比较均值。这个高维空间用核函数隐式表示既能捕捉非线性结构又不需要真的把无穷维向量算出来。1.2 从一个直观问题到 MMD 的数学定义MMD 的完整定义是两个分布 P 和 Q 在再生核希尔伯特空间中的均值向量之差的范数。写成公式是MMD² E[k(X,X)] E[k(Y,Y)] - 2E[k(X,Y)]其中 X 和 X 是源域中独立同分布的两个样本Y 和 Y 是目标域中的两个样本k 是核函数。直觉上可以这样理解E[k(X,X)] 代表源域内部样本的相似程度E[k(Y,Y)] 代表目标域内部样本的相似程度E[k(X,Y)] 代表跨域样本之间的相似程度。如果两个分布相同那么跨域相似度应当与域内相似度一致MMD 趋近于零如果分布不同域内相似度高于跨域相似度MMD 就是一个正数。这个设计很优雅把一个复杂的分布比较问题转化成了三个期望值的组合。实际计算时我们没有无限多的样本只能用经验估计代替期望。给定 n 个源域样本和 m 个目标域样本经验估计是MMD² ≈ 1/n² ΣΣ k(x_i,x_j) 1/m² ΣΣ k(y_i,y_j) - 2/(nm) ΣΣ k(x_i,y_j)对应到代码里就是算三个核矩阵然后分别求均值。这也是我下面所有实现的基础。1.3 核函数与带宽选择核函数的选择直接决定 MMD 能不能捕捉到分布差异。最常用的是 RBF 高斯核K(x,y) exp(-||x-y||² / (2σ²))其中 σ 是带宽。为什么用高斯核而不是线性核因为线性核算出来的 MMD 本质上还是均值差回到逐维比较的老路。高斯核把样本映射到无穷维空间能捕捉任意阶的统计矩差异。带宽 σ 对结果的影响非常大这可能是 MMD 调参中最关键的一个点。如果 σ 太小高斯核的指数项衰减极快只有距离很近的点之间才有非零相似度。核矩阵会退化成近似单位矩阵所有样本看起来都“互不相同”MMD 对分布的整体偏移失去敏感性梯度也很容易消失。如果 σ 太大所有样本之间的相似度都趋近于 1核矩阵近似全是 1MMD 恒等于 0完全区分不了分布。实践中有一个常见做法用所有样本对之间欧氏距离的中位数作为 σ 的经验值也叫 median heuristic。这个值落在数据点之间距离的中间量级既不会让核矩阵退化成单位阵也不会变成全 1 矩阵。后面我会给出具体代码。单一带宽的 RBF 核有时不够灵活所以出现了多核 MMD用多个不同尺度的 σ 做凸组合。相当于同时用多个放大镜观察分布小 σ 关注局部细节大 σ 关注整体轮廓。多核 MMD 是域适应中效果更稳的选择。2. PyTorch 手写 MMD从核函数到完整计算2.1 距离矩阵的两种写法MMD 的代码核心是核矩阵核矩阵的核心又是两两样本间的欧氏距离平方。在 PyTorch 里距离矩阵有两种主流写法我分别说一下适用场景。第一种是用展开式计算不额外调用 cdistimport torch def compute_dist_sq(x, y): xx (x * x).sum(dim1, keepdimTrue) # [n, 1] yy (y * y).sum(dim1, keepdimTrue) # [m, 1] dist_sq xx - 2 * torch.mm(x, y.t()) yy.t() return torch.clamp(dist_sq, min0.0)这段代码的原理是平方展开公式||x-y||² ||x||² ||y||² - 2x·y。先算每个样本自身的模长平方再算交叉内积最后组合出完整距离矩阵。注意最后一步的 clamp 非常必要浮点运算可能产生 -1e-7 这种微小的负值如果不归零后续取指数会让距离为负的样本对获得大于 1 的相似度破坏语义。第二种写法更简洁直接用 torch.cdistdef compute_dist_sq(x, y): dist torch.cdist(x, y, p2) return dist ** 2一行代码解决但需要注意torch.cdist 在部分 GPU 上的性能不如展开式而且它是先算距离再平方等于多了一次开根号再平方的往返会丢一点精度。矩阵规模不大的时候无所谓但如果你要在一个很大的 batch 上频繁调用 MMD优先使用第一种展开式。2.2 单核 MMD 的完整实现有了距离矩阵RBF 核和单核 MMD 就非常直接了def rbf_kernel(dist_sq, sigma): return torch.exp(-dist_sq / (2 * sigma * sigma)) def mmd_single(x, y, sigma1.0): kxx rbf_kernel(compute_dist_sq(x, x), sigma) kyy rbf_kernel(compute_dist_sq(y, y), sigma) kxy rbf_kernel(compute_dist_sq(x, y), sigma) return kxx.mean() kyy.mean() - 2 * kxy.mean()这里有两个细节值得展开。第一个细节为什么用 mean 而不是严谨的采样估计严格的 MMD 经验估计应当排除 ij 的自配对项即计算 E[k(X,X)] 时X 和 X 要求来自两个独立样本不能是同一个样本自己和自己配对。但深度学习中我们通常在 batch 上计算 MMDbatch 内包含自配对项的均值在期望意义上等价于有偏估计它的梯度方向和无偏估计是一致的而且计算简单、梯度更平滑。实践中绝大多数复现都用这个版本不必纠结。第二个细节损失函数究竟用 MMD 还是 MMD 的平方我建议直接用 MM2 作为损失值。因为反向传播时对 MM2 求梯度等于对核矩阵均值求梯度非常平稳。如果对 MM2 开根号得到真正的 MMD 距离那么在 MM2 接近零时根号函数的梯度会趋向无穷大训练后期极不稳定。所以算法论文里写“最小化 MMD”代码实现里直接最小化 MM2 就好。2.3 多核 MMD 的加权设计多核 MMDMK-MMD的核心是用多个带宽的 RBF 核。最简单有效的实现是等权平均def mmd_multi(x, y, sigma_list[0.1, 0.5, 1.0, 2.5, 5.0, 10.0]): total 0.0 for sigma in sigma_list: total total mmd_single(x, y, sigmasigma) return total / len(sigma_list)严格来说MK-MMD 的权重应当满足凸组合条件并通过核学习算法迭代更新让每个带宽的权重自适应调整。但大量复现经验表明均匀权重已经能取得不错的对齐效果几十个 σ 等比间隔分布在 0.1 到 10 之间就够了。原因在于优化目标是让每个尺度上的分布差异都下降均等对待每个尺度不会产生明显的偏置。如果想更精细可以把权重改成可学习参数用 softmax 归一化sigma_list torch.tensor([0.1, 0.5, 1.0, 2.5, 5.0, 10.0], dtypetorch.float32) weights torch.nn.Parameter(torch.ones_like(sigma_list)) def mmd_learnable(x, y): norm_w torch.softmax(weights, dim0) total 0.0 for sigma, w in zip(sigma_list, norm_w): total total w * mmd_single(x, y, sigmasigma) return total这种方式在复杂任务上会更稳代价是增加了一个权重参数的训练而且权重可能收敛到极端分布导致只有一个 σ 起作用。我的建议是第一版实验先用均匀权重确认 MMD 确实在下滑后再考虑可学习权重。性能上还有一点batch 为 128、特征维度 256 时每个核矩阵是 128x128单核版需要三个矩阵约 196KB 显存可以忽略。但 batch 翻到 1024 时三个 1024x1024 矩阵同时驻留接近 25MB如果再叠加多个 σ显存压力立刻增大。所以多核版更适合中等 batch 场景。3. TensorFlow 实现与 Keras 模型无缝集成3.1 纯 TensorFlow 张量版实现TensorFlow 2.x 下MMD 的张量操作和 PyTorch 思路完全一致API 对应关系需要留意import tensorflow as tf def compute_dist_sq_tf(x, y): xx tf.reduce_sum(tf.square(x), axis1, keepdimsTrue) # [n, 1] yy tf.reduce_sum(tf.square(y), axis1, keepdimsTrue) # [m, 1] dist_sq xx - 2 * tf.matmul(x, y, transpose_bTrue) tf.transpose(yy) return tf.maximum(dist_sq, 0.0) def rbf_kernel_tf(dist_sq, sigma): return tf.exp(-dist_sq / (2 * sigma * sigma)) def mmd_tf(x, y, sigma1.0): kxx rbf_kernel_tf(compute_dist_sq_tf(x, x), sigma) kyy rbf_kernel_tf(compute_dist_sq_tf(y, y), sigma) kxy rbf_kernel_tf(compute_dist_sq_tf(x, y), sigma) return tf.reduce_mean(kxx) tf.reduce_mean(kyy) - 2 * tf.reduce_mean(kxy)个容易踩坑的地方是 tf.transpose(yy)。yy 的形状是 [m, 1]直接参与广播时维度顺序可能不对必须转置成 [1, m] 才能和前面的项对齐。PyTorch 里同样的问题也存在但因为用了 .t()一般不会忘。TensorFlow 里写成 tf.transpose 也简单就是需要对这个细节敏感。3.2 在 Keras 模型中嵌入 MMD 损失TensorFlow 生态里最常见的需求是在 Keras 模型训练中加入 MMD 损失。很多人会尝试把 MMD 直接写成一个自定义损失函数传给 model.compile但这里有个隐藏坑Keras 的 loss 函数接收的是标签和预测值不代表能同时拿到源域特征和目标域特征。而 MMD 恰恰需要两个域的特征同时存在所以更合适的方案是自定义训练循环。# 假设已经定义好 feature_net 和 classifier optimizer tf.keras.optimizers.Adam(1e-3) lambda_ 0.1 tf.function def train_step(src_x, src_y, tgt_x): with tf.GradientTape() as tape: src_feat feature_net(src_x, trainingTrue) tgt_feat feature_net(tgt_x, trainingTrue) pred classifier(src_feat, trainingTrue) clf_loss tf.reduce_mean( tf.keras.losses.sparse_categorical_crossentropy(src_y, pred) ) transfer_loss mmd_tf(src_feat, tgt_feat) total_loss clf_loss lambda_ * transfer_loss grads tape.gradient( total_loss, feature_net.trainable_variables classifier.trainable_variables ) optimizer.apply_gradients(zip(grads, feature_net.trainable_variables classifier.trainable_variables)) return total_loss把 MMD 放在特征层而不是输出层是一个值得强调的经验。特征层的分布偏移直接反映了数据层面的差异通过对齐特征分布可以让下游分类器在新域上沿用源域学到的判别边界。如果放在输出层MMD 会强制两个域的输出概率分布完全一致可能破坏源域已经学到的判别信息导致源域精度也跟着下降。3.3 两种框架的实现差异与对比我把 PyTorch 版和 TensorFlow 版的差异整理成了一张表方便迁移代码时对照操作PyTorchTensorFlow 2.x距离矩阵展开torch.mm(x, y.t())tf.matmul(x, y, transpose_bTrue)转置.t()tf.transpose()数值截断torch.clamp(min0.0)tf.maximum(dist_sq, 0.0)均值tensor.mean()tf.reduce_mean()指数运算torch.exptf.exp自动微分loss.backward()tf.GradientTape除了 API 命名差异两者的数值结果在相同数据、相同 σ、相同随机种子下应当一致到 1e-6 量级。如果发现结果差异明显先检查数据是否经过了不同的归一化再检查距离矩阵里是否漏了 clamp。我遇到过两次“看起来实现一样但数值不同”的情况最后都是因为 float32 下负值截断的时机不同导致的把 clamp 放到最前面就能对齐。4. 实操记录在域适应训练中接入 MMD4.1 用模拟数据快速验证 MMD 的方向性动手写完整代码之前我建议先做一个快速验证分别从两个均值不同的高斯分布中采样看看 MMD 能不能正确地区分它们。torch.manual_seed(0) src torch.randn(200, 64) 1.0 # 源域均值偏移 tgt torch.randn(200, 64) 2.0 # 目标域均值再偏移 print(mmd_multi(src, tgt)) # 应该输出一个正数 noise torch.randn(200, 64) 1.0 print(mmd_multi(src, noise)) # 如果分布接近输出应该明显更小如果第一个输出远大于第二个说明实现正确。这一步能筛掉绝大多数代码错误比如矩阵转置错了、clamp 漏掉、σ 太大导致核矩阵退化成全 1 这类问题。在我自己的实践中这个验证环节为后面省了不少排查时间。模拟数据验证通过后再用真实场景数据训练。我这边通常用一个公开的手写数字分类数据集做域偏移实验把训练集按采集批次分成两组子集一组作为源域另一组作为目标域。目标域没有标签只用来计算 MMD。4.2 网络结构与损失组合用于域适应的网络结构很简单一个两到三层的全连接特征提取器一个单层分类器。关键在于损失函数如何组合total_loss clf_loss lambda_ * transfer_lossloss 权重的选择直接影响最终效果。λ 太小MMD 对特征分布的影响可以忽略目标域精度不会改善λ 太大特征被强烈拉向对齐源域的分类信息被压缩两边精度一起崩。我常用的策略是固定 λ0.1 起步观察目标域精度随训练变化。如果源域精度正常但目标域精度没有回升逐步把 λ 增大到 0.5 或 1.0。如果源域精度开始明显下降说明 λ 过大需要回调。更精细的做法是让 λ 从 0 线性增长到设定值。原因是训练初期分类器还没收敛此时施加强 MMD 会让特征提取器的权重被分布对齐信号主导分类边界得不到充分训练。等训练若干轮之后分类边界基本成形再逐渐增强 MMD 的对齐作用效果更稳定。核心训练循环如下for epoch in range(50): for (src_x, src_y), (tgt_x, _) in zip(src_loader, tgt_loader): feat_src encoder(src_x) feat_tgt encoder(tgt_x) pred classifier(feat_src) clf_loss F.cross_entropy(pred, src_y) transfer_loss mmd_multi(feat_src, feat_tgt) loss clf_loss lambda_schedule(epoch) * transfer_loss optimizer.zero_grad() loss.backward() optimizer.step()第 6 节会再补充一点如果 final MMD 下降缓慢多数时候不是实现问题而是特征归一化和学习率的问题。4.3 训练过程该盯哪些指标训练过程中不要只盯着总 loss 曲线我把日志里每个指标的含义单独说一下源域交叉熵反映分类任务本身是否在正常学习。这个值应该稳步下降。源域精度判断分类器有没有过拟合或欠拟合。如果这个值不上升说明网络结构或学习率有问题。目标域精度这是域适应的最终目标但目标域没有标签只能在保存模型后单独验证。MMD 值反映两个域特征分布的对齐程度。初期 MMD 较大是正常的训练后期应当降到初期的 30% 以下。还有一种需要警惕的情况目标域精度上不去但 MMD 已经降得非常低。这说明特征被过度压缩两个域的特征虽然对齐了但判别信息也被抹掉了。此时应该减小 λ或者把特征维度提高一些让对齐过程中保留更多任务相关的信息。5. 常见问题与排查技巧实录5.1 带宽导致的“假收敛”和“梯度爆炸”很多第一次上手的开发者会遇到一个怪现象两个分布明明明显不同但 MMD 算出来接近 0怎么调都区分不了。这大概率是 σ 太大造成的。前面已经解释过σ 过大会让所有核值接近 1MMD 恒等于 0。反过来σ 太小会让梯度剧烈波动甚至出现 NaN。我的建议是不要手动猜 σ直接写一个 median heuristicdef auto_sigma(x, y): dist_sq compute_dist_sq(x, y) dist torch.sqrt(dist_sq.detach()) med torch.median(dist) return max(med.item(), 1e-3)这个 σ 会自动适配特征尺度的量级。注意两个细节第一detach 掉梯度避免 σ 参与反向传播第二加一个下限防止 batch 内距离中位数为 0尤其是特征被归一化得比较极端时。5.2 训练时 MMD 不下降如果你看到 MMD 在整个训练过程中几乎不动先按顺序检查以下几点。第一检查特征提取器是否在优化器参数列表里。自定义训练循环中很容易只传入分类器参数而漏掉特征提取器导致 MMD 的梯度没有传播到特征层。第二检查有没有误加了 stop_gradient 或 detach。有的实现为了让核矩阵不干扰梯度会在距离矩阵某个位置 detach但如果 detach 加错了位置MMD 就会失去对网络的驱动力。第三检查 λ 是否太小。λ 小到 0.001 量级时MMD 信号被交叉熵完全淹没曲线自然不动。还有一个容易被忽略的原因源域和目标域使用的是同一个特征提取器但 TensorFlow 的 BatchNormalization 层默认在 training 模式下会持续更新均值方差统计量。如果两个域的 batch 交替送入BN 统计量会被两个域轮流拉扯特征分布一直在变动MMD 很难稳定下降。这时可以考虑暂时关闭 BN 的统计量更新或者在特征层之前不放置 BN。5.3 显存占用爆炸MMD 计算需要构建 n×n 的核矩阵batch 翻倍显存占用是翻四倍增长的。如果 batch512 且特征维度较高再乘上多核的 σ 数量显存很容易吃紧。几种缓解手段我都用过。最简单的是减小 batch但 batch 太小会让 MMD 估计方差变大。其次是分块计算把样本拆成多个小块分别计算 MM2 后取平均避免一次性生成完整核矩阵。再高效一点的做法是使用核函数的随机特征近似把核计算转换成有限维特征内积但实现复杂度会上升收益在大多数中等规模任务里并不明显。还有一种工程上很实用的办法对 MMD 做滑动平均估计。每个训练 step 只计算当前 batch 的 MM2同时维护一个 running meanrunning_mmd 0.9 * running_mmd 0.1 * current_mmd.item()这样既不需要额外显存又能得到一个相对稳定的 MMD 监控曲线。5.4 PyTorch 与 TensorFlow 的细节差异坑两套框架我都踩过一些细节坑汇总一下。PyTorch 中torch.mm 只支持二维矩阵乘法如果你的特征张量带着 batch 维和通道维需要先 reshape 成 [n, d] 再计算距离矩阵。TensorFlow 的 tf.matmul 支持带 batch 维的高阶乘法但广播规则更复杂一不小心就会得到形状正确但数字错误的结果。Keras 的另一个典型坑如果你坚持用 model.compile 传自定义 loss且自定义 loss 无法同时拿到两个域的 feature通常的 workaround 是把源域和目标域特征通过一个自定义层拼接后输出多个分支。但这个方案会让代码变得难读而且 debug 困难。我更建议直接用自定义训练循环Surface 清晰也更容易控制 λ 的 schedule。6. 再聊几句MMD 调试的真实感受两套框架写下来我自己的倾向是快速验证效果用 PyTorch手写一个函数挂到 loss 上半小时就能跑通如果项目基建已经在 TensorFlow 上直接用自定义训练循环不要为了 fit() 的便利硬套自定义 loss边界问题容易踩坑。多核 MMD 虽然更稳但第一次跑实验不要上来就用 10 个 σ。3 到 5 个 σ 已经能覆盖大部分需求等确认整个训练流程稳定后再增加尺度数量。σ 的等比间隔建议覆盖数据距离量级的一到两个数量级比如 0.1 到 10 六档足够了。最后分享一个容易忽略的点特征送入 MMD 之前一定要确认两个域的特征尺度一致。如果源域特征的范数普遍在 1 附近而目标域特征的范数在 100 附近即使分布原本接近MMD 也会因为尺度差异而显著拉大。这不是分布不对齐而是数值尺度不匹配。简单地在特征后加一个 L2 归一化往往比调半天 σ 和 λ 更有效。我自己在多个实验里都是靠这一条把目标域精度拉回了正常水平。