跳转至

Explanation of Dynamic Physical Field Predictions Using WassersteinGrad: Application to Storm Surges

会议: NeurIPS 2026
arXiv: 2604.22580
代码: https://github.com/meteofrance/py4cast
领域: 时间序列
关键词: 动态物理场, 特征归因, 最优传输, 瓦瑟斯坦重心理论, 自回归天气预报

一句话总结

针对自回归天气预报等动态物理场模型在梯度平滑归因时产生的空间几何位移与相错误，本文提出了 WassersteinGrad，利用熵正则化 2-Wasserstein 重心提取扰动归因图的几何共识，在保持高忠实度与稀疏度的同时显著抑制归因漂移。

研究背景与动机

深度神经网络在天气预报与地球系统模拟等高风险物理动力学任务中，预测性能已逐步比肩乃至超越传统数值物理求解器。然而，神经网络固有的黑盒属性给气象运营决策的信任构建带来了巨大阻碍。在严苛的气象预报流水线中，业务专家必须明确获知高空流场或地表气压等哪些动力学特征驱动了特定区域的极端天气演变。为了揭示高维连续网格输入对预报目标的贡献，以显著图与输入梯度乘积为代表的基于梯度的特征归因方法成为了首选。然而，深度模型强烈的非线性极易引发“梯度破碎”（Shattered Gradients Problem），导致原始输入梯度在视觉上呈现出高频、散乱且难以解释的斑图。

在传统计算机视觉中，SmoothGrad 及其变体通过在输入中叠加各向同性高斯白噪声、并对多次扰动得到的梯度图执行逐点算术平均，能够极其有效地过滤局部高频抖动。然而，当这种逐点平滑策略被直接移植到连续几何结构和具有强平流迁移特性的动态物理场时，却暴露出根本性的失效机制。气象场并非离散物体的简单拼贴，其本质由流体动力学偏微分方程主导。实验表明，输入的微小随机扰动不会在特征归因图上产生平稳的振幅噪声，而是引发了归因质量的结构性空间迁移——归因重心从真实物理起因位置被系统性“挖除”并搬运到相邻网格。这种漂移在气象检验领域被称为典型的“相错误”（Phase Error），即预测对了现象本身却偏离了空间位置。在自回归多步滚动预测中，各预测步引入的几何畸变逐层累加，若依旧强行采用逐点算术平均，不同扰动下空间错位的归因高亮区域会相互抵消并虚假平滑，彻底破坏了物理因果的局部性。

本文的核心洞察在于：既然输入扰动在深层架构与自回归反向传播中引发的是空间概率测度的几何形变而非逐点独立噪声，归因的聚合就必须脱离欧氏空间的数值平均，转向度量空间支撑集几何距离的最优传输理论。本文的核心 idea 是：将扰动输入得到的梯度归因图映射为空间概率测度，基于熵正则化 2-Wasserstein 重心求解所有变形归因图的几何测地平均，从而在物理连续空间中提取具有空间相干性与几何共识的稳健可解释表征。

方法详解

整体框架

WassersteinGrad 旨在为高维多通道自回归物理场预报模型提供物理自洽且空间聚焦的特征归因。给定输入物理场状态序列与预测目标标量（如巴黎区域未来累积降水均值），整个归因流水线包含四个阶段：通道与目标区域投影、单通道局部高斯扰动采样、反向传播生成扰动梯度测度族、以及基于卷积 Sinkhorn 算法的 Wasserstein 重心聚合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入物理场状态<br/>x ∈ R^(H×W×C)"] --> B["扰动特征切片与自回归前向<br/>单通道高斯加噪 ε_i ~ N(0, σ^2 I)"]
    B --> C["目标标量提取与反向传播<br/>G_i = ∇_x Y_target"]
    C --> D["测度映射：归一化绝对梯度<br/>μ_i = |G_i| / Σ|G_i|"]
    D --> E["几何共识：熵正则化 2-Wasserstein 重心<br/>卷积 Sinkhorn 迭代求解 μ*_λ"]
    E --> F{"输出变体路由"}
    F -->|直接输出测度| G["WGBary：纯几何共识归因图"]
    F -->|掩码基准梯度| H["WGBary×Grad：保持符号与稀疏性"]

关键设计

1. 梯度位移诱导机制分析：揭示 CNN 池化与 Transformer 注意力切换对反向传播的几何推移

针对现有 XAI 理论普遍将梯度扰动假设为空间平稳幅度噪声的缺陷，本文在机理层面上剖析了输入噪声转化为空间几何位移的内在架构根源。在卷积神经网络（CNN）中，高斯扰动在经过卷积核线性传播后转化为预激活高斯分布，一方面通过 ReLU 激活引起死神经元意外放电，开辟虚假梯度路径；更关键的是在空间下采样最大池化（Max Pooling）中，局部窗口内的微小噪声尖峰便会改变 argmax 选中的网格坐标，导致反向传播梯度物理性地跳转到错误位置。

在 Transformer 注意力层中，由于自注意力算子不具备全局 Lipschitz 连续性，且注意力交互能量在均场相互作用粒子系统视角下存在多个局部极大值与亚稳态（Metastability），初始输入的扰动会驱使 token 测度流向不同的几何聚类吸引子。由于前向注意力矩阵 \(A = \mathrm{softmax}(QK^\top / \sqrt{d_k})\) 直接以乘法因子形式嵌在反向传播梯度的链式法则中： $$ \nabla_{O_{\mathrm{in}}}\mathcal{Y}{\mathrm{target},i} = A^{\top}\Delta $$ 注意力路由的拓扑重排直接导致反向传播得到的特征归因图发生空间平移。在多步自回归预测中，每一步反向传播均叠加独立的几何扭曲，导致梯度质心在物理网格上产生大幅伪平移。}W_{V}^{\top} + \frac{1}{\sqrt{d_{k}}}\Delta_{S}KW_{Q}^{\top} + \frac{1}{\sqrt{d_{k}}}\Delta_{S}^{\top}QW_{K}^{\top

2. 归因测度映射：将符号梯度变换为严格空间概率测度

为了使最优传输（Optimal Transport, OT）理论能够应用于具备正负符号且量纲任意的原始梯度图，算法必须将空间梯度连续场转化为满足质量守恒与非负性约束的离散概率测度。对于第 \(i\) 次带噪输入反向传播得到的单通道梯度切片 \(G_i \in \mathbb{R}^{H \times W}\)，算法对其绝对值实施 \(\ell_1\) 归一化： $$ \mu_i(u) = \frac{|G_i(u)|}{\sum_{v \in \Omega} |G_i(v)|}, \quad \forall u \in \Omega \subset \mathbb{Z}^2 $$ 该映射操作不仅严格保证了测度总质量为 1，而且完整保留了正向驱动与负向抑制特征在空间网格上的绝对重要性密度，为后续在几何支撑集上度量测度间搬迁代价铺平了道路。

3. 熵正则化 Wasserstein 重心聚合：卷积 Sinkhorn 提取几何共识

为了从 \(N\) 个空间位移的扰动测度族 \(\{\mu_1, \dots, \mu_N\}\) 中提取出剔除相误差的几何共识，算法摒弃传统的逐点算术平均 \(\frac{1}{N}\sum \mu_i\)，转向 2-Wasserstein 空间中的 Fréchet 均值求解。引入熵正则化后，优化目标转化为求解平滑运输代价下的重心测度 \(\mu_\lambda^*\)： $$ \mu_\lambda^* = \arg\min_{\mu \in \mathcal{P}(\Omega)} \sum_{i=1}^N W_{2,\lambda}^2(\mu, \mu_i) = \arg\min_{\mu \in \mathcal{P}(\Omega)} \sum_{i=1}^N \left( \min_{\pi_i \in \Pi(\mu, \mu_i)} \langle M, \pi_i \rangle + \lambda \sum_{u,v} \pi_{i}(u,v)(\log \pi_{i}(u,v) - 1) \right) $$ 其中代价矩阵 \(M_{uv} = \|u - v\|_2^2\) 显式编码了二维连续物理网格的欧几里得空间距离。借助离散网格上高斯核的可分离卷积性质，该问题通过卷积 Sinkhorn 迭代算法在 GPU 上快速并行收敛，使得原本规模为 \(\mathcal{O}((HW)^2)\) 的线性规划问题降低到与卷积平滑同等的复杂度，计算时间仅为传统 SmoothGrad 的 1.96 倍。

4. 双模式输出策略：平衡几何相干性与符号稀疏性

为了满足不同物理预报场景的解释需求，本文设计了两种下游归因呈现形式： - WGBary：直接将计算得到的重心测度 \(G_W = \mu_\lambda^*\) 作为最终归因图。它最大程度反映了特征空间分布的几何不确定性与整体结构，彻底消除了空间破碎感，但牺牲了梯度的正负符号信息； - WGBary×Grad：将重心测度 \(\mu_\lambda^*\) 作为空间加权注意力掩码作用于无扰动的基准梯度 \(G_{\text{base}} = \nabla_{x_{\text{in}}} \phi([f^T(x)]_{c_{\text{out}}})\) 上，即 \(G_W = \mu_\lambda^* \times G_{\text{base}}\)。该设计由 Wasserstein 重心锁定可靠的几何空间支撑区域，同时利用基准梯度恢复特征贡献的符号方向与高动态反差，极大提升了空间归因的稀疏度。

实验关键数据

主实验

实验基于法国气象局高分辨率气象基准数据集 TITAN（AROME 区域模型 2.5 km 分辨率，法国本土 \(512 \times 640\) 网格，21 个物理状态通道）与通过业务气象学家验收的 UNetRPP 混合架构开展。测试任务针对巴黎区域地表降水预测，归因于 250 hPa 高空纬向风（急流引导作用）。在 100 场极端降水天气事件上，分别评测了单步（\(t+1\)）与 5 步自回归（\(t+5\)）下的稀疏性（Gini 指数 \(\uparrow\)）、忠实度（ROAD-AUC \(\uparrow\)）以及幅度与结构鲁棒性（\(\text{LLE}_{\ell_2} \downarrow\) 和 \(\text{LLE}_{\cos} \downarrow \times 10^3\)）。

方法 单步 \(t+1\) Sparsity \(\uparrow\) 单步 \(t+1\) Faith. \(\uparrow\) 单步 \(t+1\) \(\text{LLE}_{\cos} \downarrow\) 自回归 \(t+5\) Sparsity \(\uparrow\) 自回归 \(t+5\) Faith. \(\uparrow\) 自回归 \(t+5\) \(\text{LLE}_{\cos} \downarrow\)
BaseGrad \(0.547 \pm .001\) \(0.772 \pm .010\) \(3.5 \pm 0.07\) \(0.531 \pm .002\) \(0.714 \pm .019\) \(4.2 \pm 0.10\)
IntegratedGrad \(0.628 \pm .004\) \(0.810 \pm .011\) \(2.9 \pm 0.10\) \(0.632 \pm .004\) \(0.739 \pm .018\) \(6.3 \pm 0.10\)
SmoothGrad \(0.539 \pm .002\) \(0.809 \pm .012\) \(2.1 \pm 0.05\) \(0.526 \pm .001\) \(0.757 \pm .017\) \(2.6 \pm 0.06\)
VarGrad 0.871 \(\pm .004\) \(0.809 \pm .012\) \(2.3 \pm 0.06\) \(0.710 \pm .006\) \(0.740 \pm .013\) \(3.0 \pm 0.10\)
WGBary (本文) \(0.340 \pm .002\) 0.815 \(\pm .011\) 0.1 \(\pm 0.01\) \(0.330 \pm .002\) 0.762 \(\pm .016\) 0.2 \(\pm 0.01\)
WGBary×Grad (本文) \(0.849 \pm .003\) \(0.800 \pm .010\) \(3.6 \pm 0.08\) 0.750 \(\pm .004\) \(0.760 \pm .017\) \(4.3 \pm 0.10\)

消融实验

在固定采样样本数 \(N=20\) 下，对 Sinkhorn 熵正则化系数 \(\lambda\) 进行消融评测（基于 25 场极端暴雨事件），探究几何近似精度对归因性能的权衡。

配置 (\(\lambda\)) WGBary Sparsity \(\uparrow\) WGBary Faithfulness \(\uparrow\) WGBary \(\text{LLE}_{\cos} \downarrow (\times 10^3)\) WGBary×Grad Sparsity \(\uparrow\) WGBary×Grad Faith. \(\uparrow\) 说明
\(\lambda = 0.1\) \(0.201 \pm 0.003\) \(0.848 \pm 0.018\) \(0.06 \pm 0.01\) \(0.636 \pm 0.006\) \(0.816 \pm 0.028\) 正则化过大，过度扩散平滑
\(\lambda = 0.01\) \(0.287 \pm 0.006\) \(0.866 \pm 0.024\) \(0.09 \pm 0.01\) \(0.767 \pm 0.007\) \(0.819 \pm 0.032\) 空间集中度提升
\(\lambda = 0.001\) (默认) \(0.320 \pm 0.006\) \(0.851 \pm 0.023\) \(0.35 \pm 0.03\) \(0.829 \pm 0.006\) \(0.824 \pm 0.028\) 精度与几何保持最佳平衡点
\(\lambda = 0.0005\) \(0.329 \pm 0.006\) \(0.824 \pm 0.022\) \(0.45 \pm 0.03\) \(0.838 \pm 0.006\) \(0.798 \pm 0.035\) 趋近严格 OT，数值震荡略增

关键发现

  • 结构鲁棒性呈数量级跨越：WGBary 在衡量归因形态空间旋转与结构畸变的 \(\text{LLE}_{\cos}\) 指标上，较 SmoothGrad 实现了约 21 倍的误差下降（\(0.1 \times 10^{-3}\) vs \(2.1 \times 10^{-3}\)）。在 5 步自回归长时预报下，各基线方法的空间扭曲急剧发散，而 WGBary 仍稳固保持在 \(0.2 \times 10^{-3}\)，展现出卓越的几何抗噪稳定性。
  • 解耦的预报误差与归因漂移：最优传输质量通量分析揭示，在输入噪声比例 \(\sigma \le 0.4\) 时，神经网络本身的相对预报误差不足 1.5%，但基线梯度的归因质心却平均产生了 10 至 15 km 的位移，极值点位移在 \(t+5\) 步滚动预测中更被放大至 115 km。这实证了梯度错位是经典平滑方法的归因病态缺陷，而非物理系统本身的预测失效。
  • 双变体权衡明显：WGBary 以较低的 Gini 稀疏性完整保留了空间几何不确定性，在忠实度指标 ROAD 上始终拔得头筹；WGBary×Grad 则成功挽救了特征稀疏度，在 \(t+5\) 时 Gini 指数达到 0.750，同时维持了仅次于前者的极高忠实度。

亮点与洞察

  • 将 XAI 归因聚合重新形式化为最优传输空间均值：突破了传统计算机视觉领域默认采用欧氏逐点平滑的惯性思维，首次为连续物理场引入带几何距离约束的 Wasserstein 测地线聚合，从根源上治愈了相误差带来的虚假平滑。
  • 自回归反向传播中几何误差累积的形式化机理剖析：借助 Transformer 的均场粒子交互流与 CNN 最大池化不稳定性理论，清晰论证了高斯扰动在前向激活中如何必然转变为梯度的空间位移，为科学机器学习（Scientific ML）的可解释性提供了严密的理论基座。
  • 高效实用的卷积 Sinkhorn 工程实现：借助二维规则网格上的高斯热核卷积分解，将 Wasserstein 重心求解控制在极低的算力开销（14.6 秒/样本，仅为常规 SmoothGrad 的 1.96 倍），兼具实用性与优雅性。

局限与展望

  • 各向同性白噪声缺乏物理动力学约束：采样过程中注入的各向同性高斯噪声未考虑地转平衡等物理守恒律，生成的扰动样本在气象学上缺乏物理合理性，未来可结合集合卡尔曼滤波或物理先验构建保物理扰动协方差。
  • 单通道独立归因忽略了多变量动力耦合：当前流水线按单变量输入通道逐一提取测度，未能直接建模风场、温度场与气压场之间的三维斜压耦合传输，扩展至多通道联合 Wasserstein 重心是潜在方向。

相关工作与启发

  • vs SmoothGrad / VarGrad：SmoothGrad 与 VarGrad 均建立在欧氏空间独立同分布平稳噪声假设上，对扰动梯度做标量点对点算术平均，在连续介质场上会造成严重的空间模糊与错位；本文通过 2-Wasserstein 距离度量空间几何位移，提取测度支撑集的几何平均。
  • vs Integrated Gradients：积分梯度依赖一条从参考基线到输入的连续积分路径，然而在高度非线性的流体动力学体系中，全零或气候态均值构成的线性插值切片大多是非物理的虚拟状态；本文在真实输入邻域进行扰动，并在测度空间中解决几何一致性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从最优传输几何视角审视物理场归因的相错误，形式化了注意力机制与池化诱导的空间位移机制。
  • 实验充分度: ⭐⭐⭐⭐⭐ 基于法国气象局 TITAN 高分辨率真实预报基准，不仅开展了单步与自回归滚动评测，还设计了质心位移与 OT 质量通量等独特的几何定量指标。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨流畅，均场流体力学视角与气象物理背景自然交融，结构高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为 AI 数值天气预报与动态物理仿真系统进驻业务级决策中枢提供了关键的信任支柱。

相关论文