📊 statistics

Explanation of Dynamic Physical Field Predictions using WassersteinGrad: Application to Autoregressive Weather Forecasting

To address the issue of blurred feature attributions caused by geometric displacement in dynamic physical fields, this paper introduces **WassersteinGrad**, a method that uses entropic Wasserstein barycenters to aggregate perturbed attribution maps, providing more accurate and robust explanations for autoregressive weather forecasting models.

Original authors: Younes Essafouri, Laure Raynaud, Luciano Drozda, Laurent Risser

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Younes Essafouri, Laure Raynaud, Luciano Drozda, Laurent Risser

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Blurry Map" of AI Reasoning

Imagine you are a meteorologist looking at a high-tech AI weather model. The AI predicts, "A massive storm is going to hit Paris in five hours." You want to know why. You ask the AI, "Which part of the current wind patterns caused this prediction?"

The AI gives you a "heat map" (an explanation) showing the important wind areas. But there’s a problem: the map looks like a messy, vibrating smudge.

To try and fix this, scientists usually use a trick called "Smoothing." They add a little bit of random "static" (noise) to the weather data and ask the AI the same question multiple times. They then take all those different maps and average them together to get a cleaner picture.

Here is the catch: In the world of weather, "averaging" is a disaster.

Think of it like this: Imagine you are trying to find the center of a moving crowd. One noisy version of the data says the crowd is at the North Gate. Another says they are at the South Gate. If you "average" their positions, you get a result that says the crowd is in the middle of a giant, empty field. The average tells you where the crowd isn't, rather than where they actually are.

In weather forecasting, this is called a "phase error." The AI is looking at the right thing (the storm), but because of the noise, the explanation "drifts" to the wrong location.

The Discovery: Why the AI "Drifts"

The researchers discovered that this isn't just a random mistake; it's built into how AI "brains" work. They found two main culprits:

The "Winner-Takes-All" Problem (CNNs): Many AIs use a process called "pooling" to simplify data. It’s like a scout looking at a field and saying, "The tallest person is over there!" If you add a little noise, a different person might suddenly look slightly taller. The scout jumps to a completely different spot. When you average these "jumps," the explanation blurs into nothingness.
The "Social Butterfly" Problem (Transformers): Modern AIs use "Attention," where different pieces of data "talk" to each other. This is like a group of people at a party. If you nudge one person (add noise), they might suddenly start a conversation with a completely different group. This changes the entire social structure of the data. Because these "social groups" can shift wildly with tiny changes, the AI's reasoning "migrates" spatially.

The Solution: WassersteinGrad (The "Smart Moving" Average)

Instead of using a "dumb" average (which just blurs things), the researchers created WassersteinGrad.

Instead of a standard average, they use something called a "Wasserstein Barycenter."

The Analogy: The Pile of Sand

Standard Averaging is like taking ten different piles of sand and melting them into a liquid. You get a flat, thin puddle that covers everything but has no shape.
WassersteinGrad is like taking ten different piles of sand and sliding them together until they form one single, perfect, sharp pile.

Instead of blurring the maps, WassersteinGrad recognizes that the "importance" is a physical mass that has just moved slightly. It "transports" the importance from all the noisy versions back to a single, mathematically consistent location.

The Result: Clearer Skies for AI

The researchers tested this on real weather data (the TITAN dataset) and found that:

It’s much sharper: While other methods produced "fuzzy" or "shattered" maps, WassersteinGrad produced clear, localized areas of importance.
It’s more reliable over time: Weather forecasting is "autoregressive," meaning the AI predicts one hour, then uses that prediction to predict the next. In standard methods, the errors compound like a snowball rolling downhill, making the explanation useless after a few hours. WassersteinGrad stays steady, keeping the explanation "on track" even as the forecast looks further into the future.

In short: They moved from "blurring the truth" to "finding the consensus," making AI weather models much more trustworthy for the people who actually have to prepare for the storm.

Technical Summary: Explanation of Dynamic Physical Field Predictions using WassersteinGrad

1. Problem Statement

As deep learning models are increasingly deployed in high-stakes domains like weather forecasting, the need for Explainable AI (XAI) has shifted from a theoretical interest to an operational necessity. Traditional gradient-based feature attribution methods (e.g., Saliency, SmoothGrad) are used to identify which input features drive a prediction.

However, the authors identify a fundamental failure mode when applying these methods to dynamic physical fields (e.g., atmospheric states). Standard smoothing techniques like SmoothGrad rely on pointwise averaging of gradients obtained from noisy inputs. The authors demonstrate that in physical systems, stochastic input perturbations do not merely cause "amplitude noise" (random fluctuations in importance); instead, they induce geometric displacement (spatial shifts). In meteorological terms, this is analogous to a phase error: the model identifies the correct phenomenon but at the wrong spatial location. This misalignment is exacerbated in autoregressive forecasting, where geometric distortions accumulate over successive time steps.

2. Methodology

The authors propose WassersteinGrad, a transport-based aggregation framework designed to find a "geometric consensus" rather than a pointwise average.

Mechanistic Analysis: The paper provides a theoretical basis for why this displacement occurs by analyzing two ubiquitous architectures:
- CNNs: Spatial displacement is caused by the instability of the argmax operator in Max Pooling and the creation of "phantom gradients" in ReLU activations due to noise-induced firing of dead neurons.
- Transformers: Using an interacting particle system framework, the authors show that self-attention is non-Lipschitz and possesses multiple local energy maxima (metastability). Input noise can push the system into different "basins of attraction," resulting in qualitatively different token routing (attention patterns) for nearly identical inputs.
The WassersteinGrad Framework:
1. Perturbation: Generate $N$ noisy versions of the input channel.
2. Gradient Computation: Compute the raw gradient maps $G_i$ for each noisy input.
3. Measure Mapping: Convert these gradients into spatial probability distributions $\mu_i$ by taking absolute values and $\ell_1$ normalization.
4. Aggregation via Barycenter: Instead of averaging the pixels, they compute the entropic Wasserstein barycenter ( $\mu^*$ ). This finds the Fréchet mean in Wasserstein space, which preserves the geometric structure and "moves" the mass to a central consensus location.
Two Variants:
- WGBary: Uses the barycenter directly as the explanation (highly robust but loses sign information).
- WGBary $\times$ Grad: Uses the barycenter as a spatial mask applied to the original clean gradient to recover sign information and improve sparsity.

3. Key Contributions

Identification of Geometric Displacement: A formal demonstration that input noise in physical forecasting causes spatial migration of attribution mass rather than simple amplitude noise.
Architectural Analysis: A deep dive into how CNN pooling and Transformer attention mechanisms structurally facilitate this displacement.
WassersteinGrad Algorithm: A novel XAI aggregation method that utilizes Optimal Transport (OT) to achieve geometric robustness.
Empirical Validation: The first characterization of this phenomenon in the context of autoregressive neural weather forecasting using the TITAN dataset.

4. Results

The method was evaluated on a hybrid CNN-Attention U-Net (UNetRPP) predicting precipitation over France.

Qualitative: While BaseGrad and SmoothGrad produced fragmented, noisy, or spatially incoherent maps (especially at $t+5$ lead times), WassersteinGrad produced geographically coherent and localized attributions.
Quantitative:
- Faithfulness (ROAD): WassersteinGrad achieved the highest faithfulness scores, meaning the identified features were more critical to the actual model output.
- Robustness (LLE): It showed a massive reduction in spatial pattern rotation (LLE $_{cos}$ ) compared to SmoothGrad, maintaining stability even as errors accumulated in autoregressive rollouts.
- Sparsity (Gini): The WGBary $\times$ Grad variant achieved superior sparsity, providing clean, interpretable maps that focus on specific atmospheric drivers (e.g., the jet stream).

5. Significance

This work is significant because it challenges the "one-size-fits-all" approach of transferring XAI methods from Computer Vision to the Physical Sciences. It proves that for continuous, advective physical phenomena, spatial geometry matters more than pixel-wise intensity. By providing a mathematically grounded way to aggregate explanations, WassersteinGrad enables higher operational trust in AI-driven weather forecasting and other high-stakes physical modeling applications.