Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and Transport Phenomena

本文提出了一种基于可解释人工智能的新方法，能够高效准确地将 Wasserstein 距离归因于数据子群、输入特征或可解释子空间等具体成分，从而深入解析数据集偏移与输运现象。

要点

这篇论文提出了一种名为 WaX（Wasserstein Distances Made Explainable，即“让 Wasserstein 距离变得可解释”）的新方法。

为了让你轻松理解，我们可以把这篇论文的核心思想想象成**“给两群人的差异做体检”**。

1. 背景：我们为什么要比较两群数据？

想象你有两个巨大的仓库，里面装满了成千上万个包裹（这就是数据集）。

• 仓库 A 是去年的包裹。
• 仓库 B 是今年的包裹。

你想知道这两个仓库的包裹分布有什么不同（比如：是不是今年重包裹变多了？是不是来自北方的包裹变少了？）。

在数学界，有一个非常强大的工具叫 Wasserstein 距离（也叫“推土机距离”）。

• 通俗比喻：想象你要把仓库 A 的所有土块（数据点）搬运到仓库 B 的形状。Wasserstein 距离就是计算**“最少需要花多少力气（成本）”**才能把 A 变成 B。
• 问题：以前，我们只能算出“总成本是 100 块”。但这不够！我们不知道这 100 块成本里，是因为**“重箱子变多了”（某个特征变了），还是因为“北方来的包裹变少了”**（某个子群体变了）。就像你只知道修路花了 100 万，但不知道是修桥贵，还是铺路贵。

2. 核心创新：WaX 是什么？

这篇论文的作者说：“别光看总账单，我们要拆解这笔账单！”

他们发明了一种叫 WaX 的方法，利用“可解释人工智能”（XAI）的技术，把那个"100 块的总成本”拆解开来，告诉你：

• 哪几个具体的包裹（数据点）贡献了最大的搬运成本？
• 哪几个特征（比如重量、颜色、产地）是导致成本高的主要原因？

WaX 就像是一个超级侦探，它拿着放大镜，不仅告诉你“路很难走”，还能指着地图说：“看！是因为中间这座桥（瓶颈）太窄了，导致卡车堵在这里，所以运费才这么贵。”

3. 它是如何工作的？（三个步骤）

想象你在玩一个**“层层剥洋葱”**的游戏：

1. 第一步：算出总账（最优运输）
   先像往常一样，算出把仓库 A 变成仓库 B 需要的总力气（Wasserstein 距离）。这时候，计算机已经知道怎么搬运最省力了（这叫“耦合计划”）。

2. 第二步：把数学公式变成“神经网络”
   作者做了一个很巧妙的 trick：他们把计算“搬运成本”的数学公式，强行改写成了一个**“神经网络”**的样子。

   • 比喻：就像把复杂的物理公式，画成了一张电路图。这样，我们就可以用专门分析电路的工具（叫 LRP，层相关性传播）来倒着推。

3. 第三步：倒着推（反向传播）
   从“总成本”开始，顺着电路图往回推：

   • 先推到哪一对包裹（源仓库的一个包裹和目标仓库的一个包裹）最费力气？
   • 再推到哪个特征（比如是“重量”还是“体积”）导致了这一对包裹费力气？
   • 结果：你得到了一份详细的“贡献清单”，告诉你每个特征和每个样本对总差异的贡献有多大。

4. 这个工具能干什么？（三个实际场景）

论文展示了 WaX 在三个真实场景中的大显身手：

场景一：给 AI 模型“排毒”（域适应）

• 问题：你在 A 医院训练了一个 AI 看病，拿到 B 医院用就不准了。因为两家医院的设备不同（数据分布变了），AI 可能偷偷学会了"A 医院的设备特征”而不是“病情”。
• WaX 的作用：它能精准地指出：“嘿！这个 AI 太依赖‘设备型号’这个特征了，这是干扰项，把它删掉！”
• 比喻：就像教学生考试，WaX 告诉老师：“别让学生背‘试卷纸张的颜色’，要让他们背‘知识点’。”这样学生换个考场（新数据集）也能考好。

场景二：观察“时间流逝”（运输现象）

• 问题：想象一群鲍鱼（一种海鲜），一年前和一年后，它们长大了。但鲍鱼群很复杂，有的长得快，有的长得慢，有的变重了，有的变长了。
• WaX 的作用：它能发现：“哦，原来大鲍鱼主要是体重在变，而小鲍鱼主要是长度在变。”
• 比喻：普通的观察只能看到“大家都长大了”。WaX 像是一个**“时间切片显微镜”**，能把不同年龄段、不同生长模式的群体分开看，发现它们各自独特的生长规律。

场景三：找数据集的“潜规则”（数据集差异）

• 问题：你有两个名人照片库（CelebA 和 LFW）。你想看看它们有什么不同。
• WaX 的作用：它发现：
  • 一个主要差异是性别比例（LFW 里男性政治家多，CelebA 里女演员多）。
  • 另一个差异是配饰（LFW 里戴眼镜、打网球的人多）。
  • 还有一个差异是人数（LFW 里有很多双人合影，而 CelebA 多是单人）。
• 比喻：就像两个不同的朋友圈，WaX 能帮你分析出：“哦，A 朋友圈喜欢晒自拍，B 朋友圈喜欢晒聚会和运动。”这能帮你决定训练 AI 时该用哪个数据，或者怎么混合它们。

5. 总结：为什么这很重要？

• 以前：我们只知道“两个数据集不一样”，但不知道为什么不一样，也不知道哪里不一样。
• 现在：有了 WaX，我们可以精准定位差异的来源。
  • 是数据质量问题？（比如某个特征全是噪点）
  • 是群体结构问题？（比如某个子群体消失了）
  • 是特征定义问题？（比如“重量”这个特征在两个数据集里定义不同）

一句话总结：
这篇论文给“比较两个数据集”这件事，装上了一盏探照灯。以前我们只能看到两个山丘离得很远（总距离大），现在 WaX 能照亮山丘上的每一块石头，告诉我们：“看！是因为这块大石头（特征）和那块小石头（样本）的位置不对，才让路变得这么难走。”

这让科学家和工程师能更聪明地处理数据，让 AI 模型更 robust（鲁棒），也能让我们更深刻地理解数据背后的物理或社会现象。

技术摘要

论文技术总结：Wasserstein Distances Made Explainable (WaX)

1. 研究背景与问题 (Problem)

Wasserstein 距离（又称推土机距离，Earth Mover's Distance）是衡量两个概率分布之间差异的强大工具，广泛应用于最优传输（Optimal Transport, OT）、领域自适应、时间序列分析和生物医学等领域。然而，现有的 Wasserstein 距离分析存在以下局限性：

• 缺乏可解释性：传统的分析仅关注距离数值本身或传输计划（Transport Plan/Coupling, $\gamma^*$）。虽然传输计划展示了源分布和目标分布之间的点对点映射，但它难以直接指出哪些具体的输入特征（features）或哪些数据子群（subgroups）导致了分布间的巨大差异。
• 对模型参数不敏感：传输计划往往对 Wasserstein 距离的具体定义（如指数 $p$ 和 $q$ 的选择）不敏感，导致无法捕捉到模型参数变化所反映的关键数据变化（例如，对异常值敏感度的变化）。
• 现有方法的不足：现有的可解释性 AI（XAI）方法主要针对分类或回归模型的预测进行归因，缺乏针对“分布间距离”这一标量输出的系统性归因方法。

核心问题：如何高效、准确地解释 Wasserstein 距离，将其归因到具体的数据点、输入特征或可解释的子空间，从而揭示导致分布偏移（Dataset Shifts）或传输现象（Transport Phenomena）的根本原因？

2. 方法论 (Methodology)

作者提出了一种名为 WaX (Wasserstein distances made explainable) 的新框架，基于**可解释性 AI（XAI）中的层间相关性传播（Layer-wise Relevance Propagation, LRP）**技术。

2.1 核心思想：神经化 - 传播 (Neuralization-Propagation)

WaX 将预计算好的 Wasserstein 距离模型“神经化”为一个等效的神经网络结构，然后利用 LRP 规则进行反向传播，将距离值分解归因到输入层。

1. 神经化 (Neuralization)：

   • 假设源分布 $\mu$ 和目标分布 $\nu$ 是经验分布（由数据点 $x_k$ 和 $y_l$ 组成）。
   • 首先通过最优传输求解器计算最优传输计划 $\gamma^*$ 并固定。
   • 将 Wasserstein 距离 $W_p$ 的计算过程重写为两层神经网络：
     • 第一层：计算源和目标数据点对之间的差异向量及其范数 $z_{kl} = \|x_k - y_l\|_q$。
     • 第二层：根据传输计划 $\gamma^*$ 对差异进行加权并应用 $L_p$ 范数聚合得到 $W_p$。
   • 这种重写使得原本复杂的迭代优化问题转化为一个确定性的前向传播过程。

2. 传播 (Propagation)：

   • 利用 LRP 规则将 $W_p$ 的值反向传播回输入。
   • 步骤 1（实例归因）：将 $W_p$ 分解为每个数据点对 $(x_k, y_l)$ 的贡献 $R_{kl}$。
   • 步骤 2（特征归因）：进一步将 $R_{kl}$ 分解为每个输入特征 $i$ 的贡献 $R_i$。
   • 超参数控制：引入超参数 $\alpha$ 和 $\beta$ 来控制归因的稀疏度（即归因是均匀分布还是集中在少数关键样本/特征上）。作者提出了启发式设置：$\alpha = p$ 和 $\beta = \min(p+2, q)$，以平衡局部性和全局性。

2.2 扩展：基于子空间的解释 (U-WaX)

为了处理更复杂的异质传输现象，作者提出了 U-WaX，将输入空间分解为多个正交子空间（代表不同的“概念”或“子偏移”）。

• 通过优化一个正交矩阵 $U$，最大化子空间内的传输统计量（如尾部统计量），从而解耦不同的传输模式。
• 这使得分析不仅能指出“什么特征”重要，还能指出“哪些特征组合”在特定的子群中导致了分布偏移。

2.3 理论性质

• 守恒性 (Conservation)：归因分数的总和严格等于原始的 Wasserstein 距离值 ($\sum R_i = W_p$)。
• 梯度等价性：在特定超参数设置下，LRP 归因等价于梯度计算，保证了数学上的严谨性。
• 模型无关性：适用于各种 Wasserstein 变体（如 Sinkhorn 正则化距离）和 Minkowski 距离。

3. 主要贡献 (Key Contributions)

1. 首创 Wasserstein 距离的可解释性框架：填补了 XAI 文献中关于“分布间距离归因”的空白，首次系统性地解决了如何解释 Wasserstein 距离的问题。
2. 提出 WaX 算法：基于神经化 - 传播范式，能够高效、准确地计算数据点、输入特征及子空间的归因分数。
3. 理论保证与高效性：证明了方法的守恒性，并展示了其计算效率（仅需一次前向传播，无需像 Occlusion 方法那样重复计算多次距离）。
4. 广泛的适用性验证：在多个数据集（表格数据、时间序列、图像）和多种 Wasserstein 参数设置下进行了验证，证明了其优于基线方法。
5. 实际应用案例：
   • 领域自适应：识别并剪枝领域特异性特征，提升分类器鲁棒性。
   • 现象洞察：解构复杂的传输现象（如生物衰老过程、流体瓶颈效应）。
   • 数据集差异分析：揭示大型数据集（如 CelebA vs LFW）之间细微的语义和人口统计学差异。

4. 实验结果 (Results)

4.1 解释忠实度 (Faithfulness)

• 使用对称相关性增益 (SRG) 指标评估。该指标衡量移除/保留被解释为“重要”的特征后，Wasserstein 距离的变化程度。
• 结果：WaX 在几乎所有测试的 Wasserstein 模型（$W_1, W_2, W_{10}$ 等）和数据集上，SRG 分数均显著高于基线方法（MeanShift, Occlusion, Coupling）。
• 对比：
  • 优于 MeanShift：MeanShift 仅关注均值差异，忽略了方差和复杂分布形状。
  • 优于 Occlusion：Occlusion 计算昂贵且仅能捕捉局部效应，无法处理非线性特征交互；WaX 在计算效率上快 10-90 倍，且归因更准确。
  • 优于 Coupling：传统的传输计划对模型参数（如 $p, q$）不敏感，而 WaX 能敏锐捕捉参数变化带来的关键特征变化。

4.2 传输现象表征 (Characterization of Transport)

• 在时间序列数据（空气质量、电力消耗）和病理图像数据上，WaX 生成的归因向量与真实传输方向（Ground Truth）的余弦相似度最高。
• 证明了 WaX 能捕捉到均值移动之外的复杂分布变化（如方差变化、非线性扭曲）。

4.3 实际用例表现

• 领域自适应：在 Office-Caltech10 基准测试中，利用 WaX 剪枝特征后，分类器在目标域上的准确率显著提升，且在稀疏特征表示下表现优于现有的 FeatureOT 方法。
• 生物衰老分析：U-WaX 成功将海胆（Abalone）的衰老过程解耦为不同的子群，揭示了不同体型海胆在体重增长上的非线性规律，这是传统聚类方法无法做到的。
• 数据集差异：在 CelebA 和 LFW 数据集对比中，U-WaX 不仅发现了性别比例差异，还识别出了“佩戴眼镜”、“网球主题”、“多人合影”等具体的语义子偏移。

5. 意义与影响 (Significance)

• 工具化：WaX 为研究人员提供了一个交互式工具，用于验证传输模型的有效性、选择超参数以及开发更先进的基于 Wasserstein 的模型。
• 深度洞察：它超越了传统的统计分析，能够揭示数据分布变化背后的物理机制或语义原因（如流体瓶颈、生物生长规律、数据集偏差）。
• 鲁棒性提升：通过识别并移除导致分布偏移的特定特征，WaX 直接帮助构建更鲁棒的机器学习模型，特别是在医疗等对偏差敏感的领域。
• 未来方向：该方法为解释更复杂的传输模型（如 Gromov-Wasserstein, Sliced Wasserstein）奠定了基础，推动了可解释性 AI 从“单点预测”向“分布比较”的范式转变。

总结：这篇论文通过引入 XAI 技术，成功将黑盒的 Wasserstein 距离计算转化为可解释的归因分析，不仅提升了理论深度，还在解决实际的领域自适应、科学发现和数据处理问题上展现了巨大的实用价值。