Accepted at NeurIPS 2026.

Explanation of Dynamic Physical Field Predictions using WassersteinGrad

Application to autoregressive weather forecasting.

  • Younes Essafouri1
  • Laure Raynaud2
  • Luciano Drozda3
  • Laurent Risser1
  • 1Univ. Toulouse, INSA Toulouse, CNRS UMR 5219, IMT, Toulouse, France
  • 2Météo-France, CNRS, Univ. Toulouse, CNRM, Toulouse, France
  • 3Cerfacs, CNRS/Cerfacs/IRD, CECI, Toulouse, France
arXiv Code py4cast TITAN dataset BibTeX
Which upstream winds drove the rain over Paris?

BaseGrad

Scattered, high-frequency noise.

BaseGrad attribution map over France: sparse, speckled red and blue pixels with no coherent structure.

SmoothGrad

Less variance, still spatially fragmented.

SmoothGrad attribution map over France: reduced speckle but still fragmented into disconnected patches.

WGBary Ours

One coherent region of influence.

WassersteinGrad barycenter attribution map: a single smooth, geographically coherent region of influence over northern France.

WGBary×Grad Ours

Sparse and sign-aware.

WassersteinGrad barycenter used as a mask on the base gradient: a spatially concentrated, signed attribution map.
−1+1 Normalised attribution of surface precipitation over the Paris ROI (black box) with respect to zonal wind at 250 hPa.

Figure 2Four ways of explaining the same forecast, on the same atmospheric event. Only the transport-based aggregation produces a map a meteorologist can read as a region.

01 — The problem

Neural forecasts are replacing physical solvers. Nobody can read them.

Data-driven models now match or beat numerical weather prediction at a fraction of the inference cost. But when one predicts heavy rain over Paris tomorrow, a forecaster has no way to ask what in today's atmospheric state made it say that.

Why it matters

No explanation, no operational trust

Extreme-weather warnings are life-safety decisions. A model that is right for the wrong reasons is a liability, and without attribution you cannot tell the two apart.

Why it's hard

XAI was built for static images

Gradient attribution methods assume a static image on a pixel grid. The atmosphere is a continuous, advective, physically constrained field.

Why it compounds

Forecasting is autoregressive

Modern models predict t+1 from t, then feed the prediction back in. Gradients backpropagate through every one of those T compositions, and whatever distortion each step adds accumulates.

Diagram of autoregressive prediction: a stack of 21 atmospheric channels at time t is passed through the forecasting model f repeatedly, producing states at t+1 through t+T.
Figure 1aAutoregressive use of a neural forecasting model f, producing a sequence of physical fields, here regional atmospheric states with C = 21 channels at successive lead times.

What we actually compute

Every attribution on this page is one number's sensitivity to one field: pick an input channel, pick an output channel, reduce the forecast over a region to a scalar, differentiate.

Local prediction explanation pipeline: input zonal wind at 250 hPa at time t, predicted surface precipitation at t+5 with a region of interest box over Paris, and the resulting gradient attribution map.
Figure 1bLeft to right: input zonal wind at 250 hPa at time t; predicted surface precipitation at t+5; the explanation of rain over the Paris ROI. The scalar target is the spatial mean over the box ℬ, and the attribution is G = ∇xin 𝒴target.

02 — The failure mode

Displacement, not noise.

A deep network is locally very rough, so a single gradient says more about that roughness than about the structure we want to see. Smoothing methods fix this by sampling a neighbourhood and averaging.

Three panels. Left: a smooth signal with high-frequency roughness. Middle: its gradient, where the roughness dominates and a single evaluation is essentially arbitrary. Right: averaging over a neighbourhood of N samples recovers the true underlying trend.
The shattered gradient problemA single gradient evaluation (middle, red point) is essentially arbitrary. SmoothGrad's fix is to sample N noisy neighbours and average: ASG(x) = (1/N) Σ ∇f(x + εi). On 1-D signals this recovers the trend. The question is what it recovers on a 2-D advective field.

The averaging step assumes something false

SmoothGrad averages the N maps pointwise. That step assumes the maps are noisy estimates of one attribution. On atmospheric fields, they are not.

The assumption

Noise perturbs amplitude

Each perturbed map is the true attribution plus independent, spatially stationary noise. Averaging cancels it. This is why SmoothGrad works on natural images.

What actually happens

Noise perturbs position

The maps are geometrically displaced copies of each other. Averaging spatially misaligned features does not denoise them, it blurs them, and smears mass across locations where nothing happened.

The evidence: mass migrates

Treat the clean and perturbed attribution maps as probability measures and solve for the optimal transport plan between them. Masking the diagonal isolates only the mass that physically moved to a new coordinate.

01 Baseline gradient attribution map over northern France, showing speckled red and blue values concentrated near the Paris region of interest.

Clean gradient

The attribution on the unperturbed input.

02 Perturbed gradient attribution map, visibly similar in structure but shifted relative to the clean map.

Perturbed gradient

Same event, one noisy input draw. Prediction error changes by about 1%.

03 Optimal transport mass flux map showing a structured dipole: green regions where attribution mass was sent and red regions where it was received, separated by roughly 200 kilometres.

Transport flux

Mass excavated from true locations, deposited nearby: a structured dipole, not random amplitude noise.

Received mass Sent mass

Figure 1cThe dashed reference scale is 200 km. If the perturbation were stationary amplitude noise, this panel would be structureless.

The dissociation: the forecast is fine, the explanation is not

Track the spatial centroid of the attribution map against the model's own relative prediction error, sweeping the noise level. If the instability belonged to the forecasting model, the two curves would rise together. They dissociate sharply.

10–15 km

Average centroid advection of the attribution map below σ = 0.4 while model performance degrades by less than 1%.

~115 km

Peak displacement at t+5. Each autoregressive step adds its own geometric distortion, and they compound.

< 1.5%

Relative prediction error across the whole noise sweep, at both lead times. The model is stable; its explanation is not.

Two plots at lead time t+1. Left: peak displacement in kilometres rising from about 24 to 32 km as noise level increases. Right: centroid displacement rising from about 5 to 12.5 km while relative prediction error stays within 2.5%.
Figure 4 · t+1One-step lead time. Left: peak displacement, the location of arg max |G|. Right: centroid displacement against relative prediction error on the secondary axis. Error bands are ±1 SEM across n = 100 precipitation-filtered events.
The same two plots at lead time t+5. Peak displacement now rises from about 70 to 115 km, three to four times larger than at t+1, while relative prediction error stays within 1.5%.
Figure 4 · t+5Autoregressive lead time. Peak displacement amplifies by 3–4× relative to t+1 while the prediction error stays flat. Geometric distortion accumulates across the rollout.

This is phase error: the right thing in the wrong place — a failure mode meteorological verification has known for decades.

Hoffman et al. 1995 · Keil & Craig 2007 · Ebert 2008

03 — Why it happens

The architecture makes displacement the expected outcome.

This is not incidental to weather data. Two components that appear in essentially every modern forecasting backbone convert additive input noise into spatial deformation of the backpropagated gradient.

CNNs · convolution + pooling

Max pooling routes gradients to the argmax

Pooling sends gradient strictly to the argmax coordinate within each local window. If a noise spike in a background region exceeds the true signal, the argmax operator shifts — and the backpropagated gradient is physically relocated to the noise spike's coordinate.

Separately, under Gaussian noise of standard deviation s, the expected ReLU output is z Φ(z/s) + s φ(z/s) and that second term is strictly positive even when z < 0. Neurons that should be inactive fire anyway. These phantom activations unblock gradient paths through regions irrelevant to the original feature geometry.

Transformers · self-attention

Attention routing is metastable

The forward-pass attention matrix A appears as a direct multiplicative factor in the backward pass, so the spatial structure of the attribution is anchored to the token routing A induces. Any perturbation that alters A alters the map.

Modelling tokens as an interacting particle system, self-attention is non-Lipschitz and its interaction energy admits multiple local maxima. Different noise draws can therefore drive the system into distinct basins of attraction qualitatively different clusterings, not perturbed versions of one.

Consequence: the N maps {Gi} that SmoothGrad averages are not noisy estimates of a fixed attribution. They are geometrically displaced versions of it.

04 — The key idea

Average where the mass is, not what each pixel says.

The pointwise average of two shifted bumps is two bumps. Their Wasserstein barycenter is one bump, in the consensus location. That single sentence is the method.

N perturbed measures μi

Eight ring-shaped probability measures, each shifted and deformed slightly relative to the others.
Pointwise mean of the eight rings: a smeared, broken annulus with a washed-out interior. Pointwise mean

Structure destroyed. The ring is smeared into a blur that matches none of the inputs.

Wasserstein barycenter of the same eight rings: a single sharp, well-defined annulus. Wasserstein barycenter

Structure preserved. One sharp ring — the geometric consensus of the inputs.

Both panels aggregate the same eight measures. The only thing that changes is whether the average is taken in pixel space or in Wasserstein space.

The move

Map each attribution map to a spatial probability measure μi ∝ |Gi, then ℓ1 normalisation, which preserves the magnitude of both positive and negative evidence while satisfying the positivity and mass-conservation constraints of optimal transport. Then replace the pointwise mean with the Fréchet mean in Wasserstein space: μ* = arg minμ Σi W22(μ, μi).

05 — The method

WassersteinGrad

A drop-in replacement for SmoothGrad's aggregation step. Everything up to the averaging is unchanged; only the last line differs.

  1. Perturb one channel

    Draw x̃in,i = xin + εi with εi ~ 𝒩(0, σ2I). Only the selected input channel cin is replaced; the other 20 channels stay clean.

  2. Roll out and reduce

    Run the model autoregressively for T steps, take the output channel cout, and reduce it over the region of interest with a differentiable operator φ — here the spatial mean over the box.

  3. Differentiate

    Backpropagate through all T compositions to get the raw gradient Gi = ∇x̃in,i 𝒴target,i.

  4. Convert to a measure

    Normalise μi = |Gi| / Σ|Gi| so the map becomes a spatial probability density that optimal transport can act on.

  5. Take the barycenter

    Compute the entropic Wasserstein barycenter μ*λ with Sinkhorn regularisation λ = 0.001. Exact barycenters are prohibitive on a 512 × 640 grid; the entropic form is solved efficiently with the convolutional Sinkhorn algorithm.

Cost

1.96× the wall-time of SmoothGrad at λ = 0.001 (14.6 s vs 7.4 s per sample), on a single 16 GB laptop GPU. At λ = 0.01 the overhead falls to 1.13×.

Two variants

WGBary GW = μ*λ

The barycenter is the attribution map. Maximally robust geometric consensus, at the cost of losing sign information.

WGBary×Grad GW = μ*λ × ∇xinφ

The barycenter as a spatial mask on the base gradient. The barycenter supplies localisation; the gradient restores sign and amplitude contrast.

06 — Experimental setting

Kilometre-scale forecasting over France, on a meteorologist-validated model.

Dataset
TITAN derived from Météo-France's operational AROME limited-area model.
Resolution
0.025° ≈ 2.5 km, regular lat–lon grid, 1-hour analyses.
Domain
France subdomain, 512 × 640 grid points, 21 channels (5 surface + 16 upper-air across 250/500/700/850 hPa).
Model
UNetRPP, a hybrid convolutional–attention U-Net trained with Py4cast, weights validated by meteorologists. Fixed-weight regime throughout.
Splits
Train 2020–2022 (1 h step); test 2023 (3 h step).
Explained pair
Total surface precipitation (aro_tp_0m) over a Paris ROI, with respect to zonal wind at 250 hPa (aro_u_250hpa): a canonical jet-stream / precipitation link.
Sampling
N = 20 perturbed samples, σ = 0.2 × (max xt − min xt), Sinkhorn λ = 0.001, seed 42.
Evaluation
n = 100 high-intensity precipitation events from the 2023 test split, at t+1 and t+5.
Map of Western Europe showing the full AROME operational domain as a blue dashed rectangle and the smaller training and evaluation subdomain over France as a red solid rectangle.
Figure 3Blue dashed: the full AROME operational domain. Red solid: the France subdomain used for training and evaluation.

Metrics

There is no ground truth for an attribution map, so quality is assessed on three independent axes.

Faithfulness ↑

ROAD

Mask the top-p% salient pixels by linear imputation and check the prediction degrades more than under random masks of the same size. Reported as the AUC of that binary curve.

Robustness ↓

Local Lipschitz estimate

LLEℓ2 measures amplitude sensitivity to bounded input perturbation. LLEcos normalises the maps first, so it measures structural and spatial pattern change instead.

Complexity ↑

Gini index

Spatial concentration. A map spread across the whole domain is uninterpretable to a domain expert, however faithful it may be.

07 — Results

The spatial pattern stops moving.

LLEcos is the metric that directly measures the failure mode: how much the structure of an explanation rotates under perturbation. This is where the gap is not incremental.

Structural robustness under perturbation

LLEcos ↓  (×10³)  · lower is better · t + 1

At t+1, WGBary scores 0.1 against SmoothGrad's 2.1 — a ~21× reduction. Under autoregressive rollout it degrades gracefully to 0.2 while every baseline gets worse. WGBary×Grad scores worst on this axis by construction: multiplying the barycenter back onto the raw gradient reintroduces the gradient's own instability, which is the price of recovering sign and sparsity.

~21×

Reduction in LLEcos over SmoothGrad at t+1. WGBary is the lowest on both robustness metrics simultaneously.

0.815

Highest faithfulness (ROAD) at t+1, and again at t+5 with 0.762 surpassing all baselines at both lead times.

0.750

Best sparsity (Gini) at t+5 for WGBary×Grad, second-best at t+1.

One-step (t + 1) Autoregressive (t + 5)
Method Spar. ↑Faith. ↑LLEℓ2 ↓LLEcos ↓ Spar. ↑Faith. ↑LLEℓ2 ↓LLEcos ↓
BaseGrad 0.547±.0010.772±.0100.068±.0013.5±0.07 0.531±.0020.714±.0190.128±.0064.2±0.10
IntegratedGrad 0.628±.0040.810±.0110.080±.0012.9±0.10 0.632±.0040.739±.0180.087±.0066.3±0.10
SmoothGrad 0.539±.0020.809±.0120.046±.0012.1±0.05 0.526±.0010.757±.0170.083±.0042.6±0.06
VarGrad 0.871±.0040.809±.0120.024±.0012.3±0.06 0.710±.0060.740±.0130.058±.0043.0±0.10
WGBary 0.340±.0020.815±.0110.023±.0010.1±0.01 0.330±.0020.762±.0160.047±.0020.2±0.01
WGBary×Grad 0.849±.0030.800±.0100.045±.0013.6±0.08 0.750±.0040.760±.0170.075±.0044.3±0.10

Bold = best Underline = second best Mean ± SEM over n = 100 events LLEcos reported ×10³ Sparsity: Gini · Faithfulness: ROAD · Robustness: LLE

Side by side

Drag to wipe between the pointwise average and the geometric consensus, on the same event and the same region of interest.

SmoothGrad attribution at t+5: incoherent, fragmented patches spread across western France. WassersteinGrad barycenter attribution at t+5: a single broad, coherent region of influence near the Atlantic coastline. SmoothGrad WGBary

Lead time t+5, autoregressive. Use the arrow keys when focused.

What changes at longer lead times

At t+5 the differences sharpen. BaseGrad and SmoothGrad remain dominated by incoherent patterns with no stable spatial organisation.

WGBary keeps producing a geographically coherent attribution. Its support does broaden relative to t+1.

WGBary×Grad instead yields a sparser attribution concentrated near the Atlantic coastline, highlighting the dominant contributing regions.

08 — Limitations & what's next

What this work does not yet settle.

Limitations

  • One model, one dataset. Evaluation focuses on UNetRPP with pretrained, meteorologist-validated parameters. The mechanistic analysis applies to any model containing convolution/pooling or Transformer layers, but generalisation to other state-of-the-art architectures and spatial domains remains to be shown.
  • The noise is physically uninformed. Standard SmoothGrad and our method both perturb inputs with white Gaussian noise, which creates atmospheric states that cannot physically exist.
  • WGBary is sign-free. The barycenter represents the geometric uncertainty of displaced attributions, so it is less sparse than the alternatives and discards the direction of influence. WGBary×Grad trades geometric interpretability back for sparsity and sign.
  • Quality is measured by proxy. The metrics are proxies, not evidence that an operational decision actually improves.

Future work

  • Physically admissible perturbations. Replace isotropic Gaussian noise with perturbations constrained by climatology, so the perturbed samples remain physically realisable atmospheric states.
  • A metric for physical plausibility. The open problem: how do you quantify that an attribution map is physically right, rather than merely stable?
  • Meteorological case studies. Move from benchmark events to real storms and extremes.
  • Global, graph-based models. Ongoing work extends the framework to ECMWF's AIFS through an Anemoi plugin, running against an operational checkpoint with attributions through the multi-step rollout.

Early signal on a global model

Preliminary, on a different architecture, resolution and lead time, a graph encoder–processor–decoder at ~0.25° with 6-hour steps, attributing 2 m temperature over Paris to 10 m zonal wind at 12 h.

SmoothGrad attribution on the AIFS global model: fragmented, sign-alternating red and blue speckle across western Europe with no stable spatial organisation.
SmoothGrad — fragmented, sign-alternating speckle, no stable spatial organisation.
WassersteinGrad attribution on the AIFS global model: a single bright coherent region of influence stretching across northern France and the Channel.
WGBary — a single coherent region of influence.

Preliminary result over 10 events at 12 h lead time: faithfulness 0.796 vs 0.716, LLEcos 0.00009 vs 0.00071 for SmoothGrad.

09 — Citation

BibTeX

@article{essafouri2026wassersteingrad,
  title   = {Explanation of Dynamic Physical Field Predictions using
             WassersteinGrad: Application to Autoregressive Weather Forecasting},
  author  = {Essafouri, Younes and Raynaud, Laure and
             Drozda, Luciano and Risser, Laurent},
  journal = {arXiv preprint arXiv:2604.22580},
  year    = {2026}
}
Acknowledgements

This work has benefited from the AI Interdisciplinary Institute ANITI. ANITI is funded by the France 2030 program under the Grant agreement n°ANR-23-IACL-0002.