Explanation of Dynamic Physical Field Predictions using WassersteinGrad: Application to Autoregressive Weather Forecasting
arXiv:2604.22580v1 [stat.ML] 24 Apr 2026
Younes Essafouri1 [email protected]
Laure Raynaud2 [email protected]
Luciano Drozda3 [email protected]
Laurent Risser1 [email protected] 1
Univ. Toulouse, INSA Toulouse, CNRS UMR 5219, IMT, Toulouse, France 2 Météo-France, CNRS, Univ. Toulouse, CNRM, Toulouse, France 3 Cerfacs, CNRS/Cerfacs/IRD, CECI, Toulouse, France
Abstract As the demand to integrate Artificial Intelligence into high-stakes environments continues to grow, explaining the reasoning behind neural-network predictions has shifted from a theoretical curiosity to a strict operational requirement. Our work is motivated by the explanations of autoregressive neural predictions on dynamic physical fields, as in weather forecasting. Gradient-based feature attribution methods are widely used to explain the predictions on such data, in particular due to their scalability to high-dimensional inputs. It is also interesting to remark that gradient-based techniques such as SmoothGrad are now standard on images to robustify the explanations using pointwise averages of the attribution maps obtained from several noised inputs. Our goal is to efficiently adapt this aggregation strategy to dynamic physical fields. To do so, our first contribution is to identify a fundamental failure mode when averaging perturbed attribution maps on dynamic physical fields: stochastic input perturbations do not induce stationary amplitude noise in attribution maps, but instead cause a geometric displacement of the attributions. Consequently, pointwise averaging blurs these spatially misaligned features. To tackle this issue, we introduce WassersteinGrad, which extracts a geometric consensus of perturbed attribution maps by computing their entropic Wasserstein barycenter. The results, obtained on regional weather data and a meteorologist-validated neural model, demonstrate promising explainability properties of WassersteinGrad over gradient-based baselines across both single-step and autoregressive forecasting settings.
1
Introduction
Deep learning models have achieved remarkable complexity and performance across a wide array of tasks. However, their inherent opaque nature often remains a critical barrier to operational trust in various domains [9, 30]. This opacity is especially problematic when deploying deep neural models in high-stakes applications, where erroneous predictions can have severe physical or societal consequences. A striking example that motivated our work is weather forecasting, where deep learning is increasingly competing with traditional physical solvers [33, 28, 3, 10]. As these systems approach deployment in high-stakes decision pipelines, a fundamental question that arises to trust their predictions is: Which atmospheric features drive given forecasts? Explainable AI (XAI) aims Preprint. Under review.
to address this transparency gap by systematically revealing the decision-making processes behind model predictions [6]. Among existing XAI strategies, gradient-based feature attribution methods have become standard for highlighting salient input regions related to a prediction, when the inputs are sampled on a regular and high dimensional spatial domain [50, 32]. However, a fundamental limitation of utilizing raw input gradients to explain deep-learning-based predictions is that they are often visually noisy and locally shattered due to the highly non-linear nature of deep networks [8]. To reduce this noise, widely adopted smoothing techniques like SmoothGrad [45] first sample several noisy neighbors of the input observation with a spatially independent Gaussian noise, and then average the gradients explaining the predictions obtained using the perturbed inputs. Although these strategies have proven to be particularly effective on image data, we argue that the reliance on pointwise averaging leads to poorly localized explanations, when used to forecast the state of dynamic physical phenomena. In our driving application, meteorological phenomena are characterized by their continuous structural geometry and advective movement, as formalized by the governing fluid dynamics (see e.g. [25], pages 56–60). When explaining the predictions of noisy meteorological inputs, it appeared to us that the attribution masses tend to migrate from their true spatial location to nearby and physically incorrect locations. This is formally recognized as phase error in meteorological verification [24, 19, 26], predicting the right thing but in the wrong location. We showcase an empirical evidence for this phenomenon in Figure 1(c). This failure is not incidental. We formally demonstrate in our paper that this is compounded by the intrinsic structure of the layers that are commonly used in deep neural forecasting models. In addition, we empirically show that the geometric displacement of the explanations is reinforced when the prediction model is used in an autoregressive manner, as illustrated in Figure 1(a), which has become standard in weather forecasting [28, 33]. Each autoregressive step likely introduces an independent geometric distortion. When backpropagating gradient from final lead times to explain noisy inputs, these spatial misalignments accumulate. This makes it particularly interesting to develop novel XAI solutions that take into account the intrinsic properties of modern autoregressive forecast models. This is what we do in our paper, with the introduction of the transport-based aggregation framework WassersteinGrad.
(a) Auto-regressive predictions for dynamic physical data
(c) impact of input noise on gradient locations
... Baseline grad
(b) Local prediction explanation
Perturbed grad
200km
ROI for Received mass of information
aro_u_250hpa channel
Sent mass of information
Figure 1: (a) Autoregressive use of a neural-based forecasting model f , producing a sequence of physical fields (here regional atmospheric states with C = 21 channels) at successive lead times t + 1, . . . , t + T . (b - from left to right) input zonal wind at 250 hPa at time t; predicted surface precipitation at time t + 5; explanation of rain prediction in Paris area using the gradient attribution method of [7]. (c - top) Gradients obtained on baseline and perturbed input data for T = 1. (c bottom) impact of the noise on the spatial motion of gradient information (see Section 3.3). Note that although the injected noise has a perciptible effect on the gradient-based attribution map, it only induced about 1% of relative prediction error. 2
Our main contributions are as follows: 1. Mechanistic analysis of gradient displacement: We provide an in-depth discussion of why input perturbations induce geometric rather than stationary noise in attribution maps. By tracing noise propagation through two ubiquitous architectural components: (i) pooling argmax instability in CNNs, which physically relocates backpropagated gradients, and (ii) attention switching in Transformers, formalized via mean-field interacting particle dynamics [22, 39]. We identify structural inductive biases that make geometric displacement of backpropagated gradients architecturally expected. 2. Transport based Aggregation Framework: We propose WassersteinGrad, a transportbased aggregation framework. By mapping raw gradients into the space of spatial probability measures, we utilize the entropic Wasserstein barycenter [2, 16] to extract a geometric consensus of perturbed explanations. 3. Empirical Validation on Neural Weather Forecasting: We provide the first empirical characterization of geometric displacement in gradient attributions for standard and autoregressive neural weather forecasting models. We first use an optimal transport mass flux analysis to show that the deformation of attribution maps is coherent and spatially structured. We then show that this manifests in practice as a spatial displacement of the attribution centroid, with high stochastic variance. Finally, we compare WassersteinGrad with widely used gradient-based baselines across faithfulness, robustness, and sparsity metrics, demonstrating consistent improvements.
2
Related Works
2.1
Gradient-Based Methods and Smoothing
The gradient of a model’s prediction reflects its sensitivity to infinitesimal changes in the input space. Methods such as Saliency [7, 44] and Input×Gradient [5, 43] assign importance scores to input features by computing the partial derivatives of the output with respect to the input. However, the resulting gradient maps are often visually noisy and locally discontinuous [23, 45]. As highlighted by [8], this noise emerges due to the high local variation of gradients in deep networks with non-linear activations (e.g., ReLU), giving rise to the Shattered Gradients Problem (SGP) [8]. Finally, Integrated Gradients [47] computes attributions by integrating gradients along weighted averages between the input and a baseline input, assumed to be neutral. It, however, requires defining such a baseline, and the resulting intermediate inputs may be physically meaningless in non-linear physical settings. SmoothGrad. To mitigate the SGP, [45] approximate a Gaussian-smoothed gradient via Monte Carlo sampling over N noisy inputs x̃i ∼ N (x, σ 2 ): N
1 X ∇f (x̃i ) GSG (x) = N i=1 Other popular methods extend SmoothGrad along different axes: NoiseGrad and FusionGrad [12] introduce stochasticity into model weights rather than inputs, and VarGrad [1] replaces the mean with the variance of perturbed gradients. None of these methods, however, specifically address the geometric displacement we identify. 2.2
The Fragility of Explanations and Spatial Uncertainty
Recent literature has exposed fundamental limitations in gradient-based smoothing. For instance, some argue that SmoothGrad acts as a band-pass filter, suppressing both high and low-frequency components [32]. This can lead to spectral bias and inconsistencies, where the resulting explanations change drastically with hyperparameter choices (e.g., the noise variance σ 2 ). Due to noise sampling, [52] shows that SmoothGrad suffers from out of distribution error. While the empirical fragility of attribution maps is well-documented, the theoretical mechanics by which input noise propagates through specific architectures remain under-explored in XAI. Existing theoretical works primarily analyze forward-pass noise propagation in the context of adversarial robustness [14] or study gradient magnitudes (e.g., exploding/vanishing gradients) [37]. Critically, all these works 3
conceptualize gradient noise as a stationary phenomenon: noise affects attribution amplitude at each spatial location independently. On the architectural side, [27] demonstrated that standard attention layers lack Lipschitz continuity, making their gradients inherently unstable. To formally understand these dynamics, recent studies have elegantly modeled Transformer attention mechanisms [48] as interacting particle systems [39, 22, 11]. In this work, we build upon this interacting-particle framework to analyze how architectural inductive biases transform additive input noise into geometric spatial deformation in the attention layer. 2.3
Optimal Transport and Wasserstein Barycenters
The Wasserstein metric measures the minimal transport cost between probability measures, incorporating the underlying geometry of the space [49]. Within XAI, OT has primarily been used as a diagnostic tool: [38] leverage it to compare saliency maps, and [36] use it to identify features driving dataset shifts. To aggregate distributions, [2] introduced the Wasserstein barycenter which generalizes the concept of an average by finding the Fréchet mean in Wasserstein space. Unlike pointwise averaging, the Wasserstein barycenter preserves geometric structure [16]. Made tractable via entropic regularisation and Sinkhorn iterations [16], Wasserstein barycenter has been applied in graphics [46],clustering [51], and model ensembling [17].
3
Methodology
3.1
Pipeline for Explainability and Attribution Setup
Let xt ∈ RH×W ×C denote the physical field state at time t, defined on a spatial grid of H × W points with C input channels (e.g. wind speed, temperature, pressure). Let f : RH×W ×C → RH×W ×C denote the neural forecasting model, applied autoregressively over T steps: x̂t+T = f T (xt ) = f ◦ f ◦ · · · ◦ f (xt ). | {z }
(1)
T
As illustrated in Figure 1(b), we use the following four-step pipeline to generate localised attributions: 1. Input channel selection. We select a single input channel cin ∈ {1, . . . , C} on which spatial attributions will be represented (e.g. zonal wind at 250 hPa). We denote the corresponding spatial slice as xin = [xt ]cin ∈ RH×W . 2. Output channel selection. We select a single output channel cout ∈ {1, . . . , C} of the predicted field to explain (e.g. 10 m surface wind speed): Yout = f T (xt ) cout ∈ RH×W . Gradients backpropagate through all T compositions of f . 3. Region of Interest (ROI). We define a spatial bounding box B over a region of interest (e.g. Paris). We then define an aggregation operator ϕ : RH×W → R that maps the forecast fieldPto a scalar target. In this work, ϕ is the spatial average over B: Ytarget = ϕ(Yout ) = 1 u∈B Yout (u). Note that ϕ can be any differentiable aggregation operator. |B| 4. Attribution computation. We compute the attribution map as the gradient of the scalar target with respect to the selected input channel: G = ∇xin Ytarget ∈ RH×W . To simplify notation in the following sections, we drop the time index t and denote the input as x. Let the perturbed spatial slice be: x̃in,i = xin + εi where εi ∼ N (0, σ 2 I), for i ∈ {1, . . . , N }. We then construct the channel-wise perturbed input, x̃i , such that the selected channel is replaced by the noisy slice, while all other channels remain unperturbed: x̃in,i if c = cin [x̃i ]c = [x]c if c ̸= cin The scalar target under this channel-specific perturbation is evaluated as: Ytarget,i = ϕ [f T (x̃i )]cout and the corresponding attribution map is computed as the gradient of the target with respect to the perturbed input slice: Gi = ∇x̃in,i Ytarget,i . 4
3.2
How do structural inductive biases displace backpropagated gradients?
This section is diagnostic rather than formal. We build on established analyses of two foundational architectural paradigms in modern deep learning: Convolutional Neural Networks (CNNs) [29] and Transformer-based attention mechanisms [48]. These two architectures exhibit fundamentally different spatial inductive biases [18]. In both cases, we analyze how noise introduced during the forward pass propagates through the backward pass, and demonstrate that this propagation fundamentally manifests as a geometric displacement of the resulting gradients. For CNNs, we formally derive the effect via an analysis of the backward pass. For attention layers, we provide a mathematical motivation by modeling the mechanism as an interacting particle system [22, 39], identifying architectural properties that make such geometric displacement expected. Empirical evidence is provided in Section 4.2. 3.2.1
Convolution layers
As detailed in Appendix I, CNNs exhibit local geometric deformation due to the interplay of nonlinearities and spatial pooling. When additive Gaussian noise ε ∼ N (0, σ 2 I) is applied to a reference input observation x, it first propagates linearly through the convolutional kernels K during the forward pass. The resulting noise at the pre-activation is itself Gaussian: η ∼ N 0, σ̃ 2 , where σ̃ 2 = σ 2 ∥K∥22 . During backpropagation, this propagated noise disrupts the backward gradient flow through two inherent phenomena: x + η), 1. For a ReLU activation y = max(0, the expected forward activation under Gaussian noise becomes E[y] = x Φ σ̃x + σ̃ ϕ σ̃x . The strictly positive term σ̃ϕ(x/σ̃) causes dead neurons to suddenly fire. During backpropagation, these phantom activations open spurious, structurally incorrect gradient paths in the attribution map. 2. Spatial displacement via Pooling: CNNs rely on spatial downsampling, such as Max Pooling, which route gradients strictly to the argmax coordinate within a local window. If a noise spike in a background region exceeds the true signal, the argmax operator spatially shifts. Consequently, the backpropagated gradient is physically displaced from the true object coordinate to the noise spike coordinate. Overall, input perturbations in CNNs inherently manifest as geometric spatial displacement in the resulting attribution maps. More technical details are provided in Appendix I. 3.2.2
Attention-based layers
Forward pass and attention matrix. Let Oin ∈ Rn×d denote the input token sequence to an attention layer and Oout ∈ Rn×dv its output. The forward pass first computes queries, keys, and values via learned projections Q = Oin WQ , K = Oin WK and V = Oin WV . The output is then obtained via scaled dot-product attention Oout = AV , where QK ⊤ A = softmax √ ∈ Rn×n , (2) dk is the attention matrix encoding pairwise token similarities and WQ , WK ∈ Rd×dk and WV ∈ Rd×dv are the learned projection weights [48]. Gradient and dependence on attention matrix. respect to the input sequence Oin takes the form:
As derived in Appendix J, the gradient with
1 1 ⊤ ∇Oin Ytarget,i = A⊤ ∆O WV⊤ + √ ∆S KWQ⊤ + √ ∆⊤ QWK . Dk Dk S
(3)
where ∆O = ∂Ytarget,i /∂Oout is the gradient propagated from subsequent layers and ∆S is the gradient through the softmax. The forward-pass attention matrix A defined in Eq. (2) appears directly as a multiplicative factor in the backward pass. Therefore, the spatial structure of the attribution map ∇Oin Ytarget,i is directly anchored to the token routing induced by A. Consequently, any perturbation that alters A directly alters the attribution map. We show hereafter that input perturbations do precisely this. 5
Attention layers as interacting particle systems. To understand why input perturbations can produce changes in A, we adopt the interacting particle framework of [22, 39], which models the n input tokens as particles {xi (s)}ni=1 evolving on the unit sphere Sd−1 under self-attention dynamics: ẋi (s) = P⊥
xi (s)
1 Zβ,i (s)
β⟨WQ xi (s),WK xj (s)⟩ e W x (s) , V j j=1
Pn
(4)
x (0) = x i i Pn β⟨WQ xi (s),WK xk (s)⟩ where β > 0 is the inverse temperature, Zβ,i (s) = is the partition k=1 e ⊥ function, and Px y = y − ⟨x, y⟩x is the orthogonal projection onto the tangent space Tx Sd−1 . Full derivation is given in Appendix K. PN In the mean-field limit (n → ∞), the empirical measure of the tokens µs = N1 i=1 δxi (s) evolves according to the continuity equation ∂s µs +∇(µs X [µ(s)]) = 0, where the vector field X [µ(s)] drives the particles. As shown in recent literature [11, 39, 22], this dynamics corresponds RR β⟨x,y⟩ to a Wasserstein 1 e dµ(x) dµ(y) [22, gradient flow aimed that maximizes the interaction energy Eβ (µ) = 2β 39] driving tokens toward semantic clusters. The energy Eβ admits multiple local maxima on P(Sd−1 ), corresponding to qualitatively distinct token clustering configurations [39, 11]. These multiple attractors are a key source of the instability we characterize below. Effect of input perturbation on attention routing. Input perturbation methods such as SmoothGrad injects additive Gaussian noise ϵi ∼ N (0, σ 2 I) to the input x. As formalized in Appendix K.4, within the dynamical system framework of Eq. (4), this is equivalent to perturbing the initial conditions: xi (0) = xi + ϵi , which induces a perturbation of the initial measure µ0 . Two properties of the dynamics lead to geometric displacement in the attribution maps: First, [27] demonstrated that the self-attention operator is non-Lipschitz: there is no global constant bounding the sensitivity of the output to perturbations of the input. Consequently, nearby initial conditions xi (0) and xi (0) + ϵi offer no guarantee of nearby trajectories under Eq. (4). Second, the presence of multiple local maxima in Eβ induces metastability [39, 11], meaning that different noise realizations ϵi can drive the system into distinct basins of attraction, corresponding to different subsets of spatial grid coordinates that the attention clusters around. Together, these two properties imply that the perturbed attention matrix A(x + ϵi ) is not a scaled version of the clean matrix A(x), but reflects a qualitatively different token routing configuration. Since the spatial structure of ∇Oin Ytarget is anchored to A through Eq. (3), this routing change propagates directly into the attribution maps. The N noisy attribution maps {Gi }N i=1 are therefore not noisy estimates of a fixed attribution, but geometrically displaced versions of it. Direct empirical evidence for this claim is provided in Section 4.2. 3.3
Attribution Aggregation via the Wasserstein Barycenter
The attribution maps Gi assign an importance score Gi (u) ∈ R to each spatial coordinate u ∈ Ω, where Ω ⊂ Z2 is the spatial grid. As these scores represent output derivatives with respect to local inputs, they can be negative and possess arbitrary magnitudes, making them incompatible with OT, which operates on non-negative measures. We therefore map each Gi to a discrete spatial probability distribution by taking absolute values followed by ℓ1 normalization: |Gi (u)| , v∈Ω |Gi (v)|
µi (u) = P
(5)
where µi (u) represents the spatial density of feature importance at location u of the input channel cin , preserving the magnitude of both positive and negative evidence, while satisfying the positivity and mass-conservation constraints of OT. Our goal is now to find a single explanation that represents the geometric consensus of the displaced attribution measures {µ1 , µ2 , . . . , µN }. To do so, we consider 2-Wasserstein distances W2 (µ, ν) that account for the underlying geometry of the grid Ω by computing the minimal cost to transport mass from µ to ν [49]. To aggregate the N samples µi , we specifically compute the Wasserstein Barycenter µ∗ , PN defined as the Fréchet mean in the Wasserstein space [2]: µ∗ = arg minµ∈P(Ω) i=1 W22 (µ, µi ). 6
Algorithm 1 WassersteinGrad Require: Model f , input x ∈ R|Ω| , number of samples N , noise variance σ 2 , Sinkhorn regularization level λ, input channel of interest cin , function ϕ for output region/channel of interest extraction Ensure: Attribution map GW 1: Initialize measure list M ← [ ] 2: for i = 1 to N do 3: Generate noisy channel slice: x̃in,i ← [x]cin + εi , where εi ∼ N (0, σ 2 I) 4: Construct perturbed tensor x̃i by replacing channel cin of x with x̃in,i 5: Compute the perturbed prediction f T (x̃i ) 6: Average the prediction in a region/channel of interest Ytarget,i ← ϕ [f T (x̃i )]cout 7: Compute raw gradients Gi ← ∇x̃in,i Ytarget,i |Gi | 8: Convert to probability measure on channel cin : µi ← P |Gi | channel=cin 9: Append µi to M 10: end for 11: Compute µ∗ λ λ , the entropic Wasserstein barycenter [46] of M with regularization ∗ T 12: return GW = µ∗ (if WG ) or G = µ × ∇ ϕ [f (x)] (if WG ) Bary W x c Bary×Grad in out λ λ
Since exact Wasserstein barycenter estimations are computationally prohibitive for high-dimensional grids [15], we smooth the transport problem using Entropic Regularization [15]: µ∗λ = PN 2 arg minµ∈P(Ω) i=1 W2,λ (µ, µi ), where λ > 0 is the regularization parameter. This formulation enables efficient computation of the barycenter via the Convolutional Sinkhorn algorithm [46]. Computing µ∗λ introduces a 1.96× wall-time overhead compared to SmoothGrad (see Appendix C.2). As summarized in Algorithm 1, two variants based on the barycenter µ∗λ are finally considered for prediction explanations on the input channel cin : • WGBary . The barycenter µ∗λ can be used as the attribution map. This yields a geometrically robust explanation that captures the consensus spatial structure of the displaced measures, at the cost of losing sign information • WGBary×Grad . To recover the sign of the gradient and improve the spatial concentration, µ∗λ can be used as a spatial mask on the base gradient, so the attribution map becomes µ∗λ × ∇xin ϕ [f T (x)]cout . The barycenter provides geometric localization and the base gradient provides amplitude contrast and sign.
4
Experiments
4.1
Experimental Settings
Dataset and Model. We evaluate on TITAN [35], a high-resolution meteorological benchmark derived from the AROME limited-area model at Météo-France, providing kilometre-scale (0.025◦ ≈ 2.5 km) analyses over Western Europe. We use a subdomain over France (512 × 640 grid points, 21 channels), with January–December 2023 as the test split. Our forecasting backbone is UNetRPP [41], a pretrained hybrid convolutional–attention U-Net trained using Py4cast [34], predicting x̂t+1 = f (xt ) at 1-hour lead time with weights validated by meteorologists. Full details in Appendices A–B. Setup and Hyperparameters. All methods are evaluated in a fixed-weight regime. We explain predictions of total surface precipitation (aro_tp_0m) over a Paris bounding box with respect to zonal wind at 250 hPa (aro_u_250hpa), a canonical pair reflecting the well-established jet stream influence on surface precipitation [25]. All stochastic methods use N =20 perturbed samples with σ=0.2 × (max(xt ) − min(xt )), WGBary and WGBary×Grad use Sinkhorn λ=0.001, selected via a faithfulness–sparsity–robustness tradeoff (Appendices D.1– D.2). 7
4.2
The empirical reality of Spatial Displacement
To visualize how attribution mass is redistributed under perturbation, we compute the 2-Wasserstein distance between the clean (µclean ) and perturbed (µpert ) gradient measures and recover the optimal transport plan T ∗ (Appendix E.1). By masking the diagonal of T ∗ , we isolate only the mass that physically migrated to a new spatial coordinate. As shown in Figure 1(c), attribution mass is systematically excavated from true feature locations (red) and deposited into adjacent coordinates (green), forming a structured dipole pattern that is consistent with geometric displacement rather than random amplitude noise. Quantitatively, we measure this geometric deformation by tracking the spatial centre of mass (cx , cy ) and the peak location of the attribution maps across noise levels σ ∈ [0.1, 1.0] (Appendix E.2). To confirm this instability is a failure mode of the attribution method rather than a degradation of the underlying forecasting model, we track this spatial displacement alongside the relative prediction error Erel (σ) = RMSEσ /RMSE0 . As shown in Figure 4, the two curves dissociate sharply: below σ=0.4, model performance degrades by less than 1%, yet attribution centroids advect by an average of 10–15 km, with peak displacements exhibit even larger spatial jumps. Furthermore, this geometric displacement compounds under autoregressive rollout, reaching ∼ 115 km at t+5. 4.3
Qualitative results
One-step forecasting. Figure 2 compares attribution maps for a representative atmospheric event. BaseGrad [7, 44] produces spatially scattered, high-frequency noise. SmoothGrad [45] reduces this variance but remains spatially fragmented. In contrast, WGBary yields a geographically coherent attribution localized over northern France.WGBary×Grad combines this geometric localization with gradient amplitude, producing a spatially concentrated map that preserves the sign of the attribution. Autoregressive forecasting (t + 5). At longer lead times, differences between methods become more pronounced. BaseGrad and SmoothGrad remain dominated by incoherent patterns with no stable spatial organization. WGBary continues to produce geographically coherent attributions, though its spatial support broadens relative to t+1 (Figure 2), reflecting increased uncertainty at longer horizons. In contrast, WGBary×Grad yields a sparser attribution, concentrated near the Atlantic coastline, highlighting dominant contributing regions. BaseGrad
t+1
(a)
t+5
(a)
BaseGrad
1.00
(b)
SmoothGrad
1.00
(c)
WGBary (Ours)
1.00
(d)
WGBary × Grad (Ours)
1.00
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.25
0.25
0.25
0.25
0.50
0.50
0.50
0.50
0.75
0.75
0.75
0.75
1.00 1.00
1.00 1.00
1.00 1.00
1.00 1.00
(b)
SmoothGrad
(c)
WGBary (Ours)
(d)
WGBary × Grad (Ours)
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.25
0.25
0.25
0.25
0.50
0.50
0.50
0.50
0.75
0.75
0.75
0.75
1.00
1.00
1.00
1.00
Figure 2: Qualitative comparison of gradient attribution methods at one-step (t + 1, top) and autoregressive (t + 5, bottom) lead times. The attribution target is surface precipitation aggregated over the Paris region of interest (black box), with respect to input zonal wind at 250 hPa.
4.4
Quantitative Results
We compare our methods against four standard gradient-based baselines: BaseGrad [44], SmoothGrad [45], VarGrad [1], and Integrated Gradients [47]. Quantitative results, reported in Table 1, are aggregated over n=100 atmospheric events from the 2023 test split characterized by high-intensity precipitation over the Paris ROI. We assess performance across both one-step (t+1) and autoregressive 8
Table 1: Quantitative evaluation of attribution quality at one-step (t+1) and autoregressive (t+5) lead times. Metrics include Sparsity: Gini↑, Faithfulness: ROAD↑, Robustness: LLEℓ2 ↓ and LLEcos ↓ (×103 ) on n=100 events. Mean ± SEM. Bold: best. Underline: second best. Method Spar.↑ BaseGrad 0.547±.001 IntegratedGrad 0.628±.004 SmoothGrad 0.539±.002 VarGrad 0.871±.004 WGBary WGBary×Grad
One-step (t + 1) Faith.↑ LLEℓ2 ↓ LLEcos ↓ 0.772±.010 0.068±.001 3.5±0.07 0.810±.011 0.080±.001 2.9±0.10 0.809±.012 0.046±.001 2.1±0.05 0.809±.012 0.024±.001 2.3±0.06
Spar.↑ 0.531±.002 0.632±.004 0.526±.001 0.710±.006
Autoregressive (t + 5) Faith.↑ LLEℓ2 ↓ LLEcos ↓ 0.714±.019 0.128±.006 4.2±0.10 0.739±.018 0.087±.006 6.3±0.10 0.757±.017 0.083±.004 2.6±0.06 0.740±.013 0.058±.004 3.0±0.10
0.340±.002 0.815±.011 0.023±.001 0.1±0.01 0.330±.002 0.762±.016 0.047±.002 0.2±0.01 0.849±.003 0.800±.010 0.045±.001 3.6±0.08 0.750±.004 0.760±.017 0.075±.004 4.3±0.10
(t+5) lead times. Full mathematical formulations for all baseline methods and evaluation metrics (Sparsity, Faithfulness, and Robustness) are provided respectively in Appendices F and G. Sparsity. WGBary represents the geometric uncertainty of displaced attributions and is therefore less sparse than the other methods. In contrast, WGBary×Grad recovers sparsity (best at t + 5, second-best at t + 1) by using the barycenter as a spatial mask on the base gradient, trading geometric interpretability for sparse, sign-aware attributions. Faithfulness. WGBary achieves the highest faithfulness at both lead times, surpassing all baselines. WGBary×Grad offers competitive faithfulness with substantially higher sparsity, making it preferable when sign-aware, spatially precise attributions are required. Robustness. LLEℓ2 measures attribution amplitude sensitivity, LLEcos measures spatial pattern rotation under perturbation. WGBary achieves the lowest score on both metrics simultaneously at both lead times: a ∼21× reduction in LLEcos over SmoothGrad at t+1 and degrades gracefully under autoregressive rollout, still outperforming all baselines at t+5. Notably, under autoregressive rollout, each additional forecasting step compounds geometric distortion in baseline methods, whereas WGBary ’s geometric consensus absorbs this accumulation.
5
Conclusion
In this work, we identified a failure mode in standard gradient-smoothing XAI techniques when applied to the prediction of dynamic physical phenomena on continuous spatial domains. Through an in-depth analysis of two common layer architectures, CNN pooling/argmax and Transformer attention as well as an empirical validation on the TITAN meteorological benchmark, we demonstrated that input perturbations do not merely induce high-frequency amplitude noise in attribution maps. Rather, they cause systematic geometric displacement of attribution mass toward physically incorrect spatial coordinates. This phenomenon is interestingly analogous to phase error in meteorological verification. We then introduced WassersteinGrad to mitigate this issue by replacing pointwise averaging in SmoothGrad with transport-based geometric consensus via the entropic Wasserstein barycenter. Our quantitative results demonstrate that WassersteinGrad outperforms SmoothGrad and other gradientbased baselines across different metrics on meteorological data. In particular, it produces attribution maps that are geometrically coherent and physically localized. Broader Impacts. As deep learning increasingly replaces traditional physical solvers in high-stakes domains such as weather forecasting, the demand for reliable, physically meaningful explanations is an operational necessity. A good alignment between where a model truly attends and where its explanation reports it attends may, in practice, lead forecasters to trust neural-based predictions. We also hope this work will serve as an invitation to the community to critically reconsider the direct transfer of XAI methodology from computer vision to the physical sciences, where spatial structure and physical causality impose fundamentally different requirements on explanation quality. Limitations and Future work. Our evaluation focused on the UNetRPP forecasting model with pretrained parameters validated by meteorologists. Although our mechanistic analysis of backpropagated gradient displacements is applicable to any prediction model containing convolution/pooling layers or Transformer layers, its generalisation to other state-of-the-art architectures and spatial domains remains to be explored. A natural extension of our work concerns the noise injection strategy itself. 9
Standard SmoothGrad and our method both perturb inputs with white Gaussian noise, which is physically uninformed. Replacing isotropic noise with physically informed perturbations could improve the physical realism of the perturbed samples. We leave this as a promising direction for future work. Acknowledgements. This work has benefited from AI Interdisciplinary Institute ANITI. ANITI is funded by the France 2030 program under the Grant agreement n°ANR-23-IACL-0002.
References [1] Julius Adebayo, Justin Gilmer, Ian Goodfellow, and Been Kim. Local explanation methods for deep neural networks lack sensitivity to parameter values. In Proceedings of ICLR Workshop, 2018. [2] Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904–924, 2011. [3] Ferran Alet, Ilan Price, Andrew El-Kadi, Dominic Masters, Stratis Markou, Tom R Andersson, Jacklynn Stott, Remi Lam, Matthew Willson, Alvaro Sanchez-Gonzalez, et al. Skillful joint probabilistic weather forecasting from marginals. arXiv preprint arXiv:2506.10772, 2025. [4] David Alvarez-Melis and Tommi S Jaakkola. On the robustness of interpretability methods. In Proceedings of ICML Workshop on Human Interpretability in Machine Learning, 2018. [5] Marco Ancona, Enea Ceolini, Cengiz Öztireli, and Markus Gross. Towards better understanding of gradient-based attribution methods for deep neural networks. In Proceedings of International Conference on Learning Representations (ICLR), 2017. [6] Alejandro Barredo Arrieta, Natalia Díaz-Rodríguez, Javier Del Ser, Adrien Bennetot, Siham Tabik, Alberto Barbado, Salvador García, Sergio Gil-López, Daniel Molina, Richard Benjamins, et al. Explainable artificial intelligence (XAI): Concepts, taxonomies, opportunities and challenges toward responsible AI. Information fusion, 58:82–115, 2020. [7] David Baehrens, Timon Schroeter, Stefan Harmeling, Motoaki Kawanabe, Katja Hansen, and Klaus-Robert Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 11:1803–1831, 2010. [8] David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In International conference on machine learning, pages 342–350. PMLR, 2017. [9] J.M. Benitez, J.L. Castro, and I. Requena. Are artificial neural networks black boxes? IEEE Transactions on Neural Networks, 8(5):1156–1164, 1997. [10] Philine Lou Bommer, Marlene Kretschmer, Anna Hedström, Dilyara Bareeva, and Marina M-C Höhne. Finding the right XAI method—a guide for the evaluation and ranking of explainable AI methods in climate science. Artificial Intelligence for the Earth Systems, 3(3):e230074, 2024. [11] Giuseppe Bruno, Federico Pasqualotto, and Andrea Agazzi. A multiscale analysis of mean-field transformers in the moderate interaction regime. In The Thirty-ninth Annual Conference on Neural Information Processing Systems (NeurIPS), 2025. [12] Kirill Bykov, Anna Hedström, Shinichi Nakajima, and Marina M-C Höhne. Noisegrad—enhancing explanations by introducing stochasticity to model weights. In Proceedings of the AAAI Conference on Artificial Intelligence, 2022. [13] Prasad Chalasani, Jiefeng Chen, Amrita Roy Chowdhury, Xi Wu, and Somesh Jha. Concise explanations of neural networks using adversarial training. In International Conference on Machine Learning (ICML), pages 1383–1391, 2020. [14] Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1310–1320, 2019. 10
[15] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems (NeurIPS), 26, 2013. [16] Marco Cuturi and Arnaud Doucet. Fast computation of Wasserstein barycenters. In Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 685–693, Bejing, China, 22–24 Jun 2014. [17] Pierre Dognin, Igor Melnyk, Youssef Mroueh, Jerret Ross, Cicero Dos Santos, and Tom Sercu. Wasserstein barycenter model ensembling. In Proceedings of International Conference on Learning Representations (ICLR), 2019. [18] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the 9th International Conference on Learning Representations (ICLR), 2020. [19] Elizabeth Ebert. Fuzzy verification of high-resolution gridded forecasts: A review and proposed framework. Meteorological Applications, 15:51 – 64, 03 2008. [20] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021. [21] Rémi Flamary, Cédric Vincent-Cuaz, Nicolas Courty, Alexandre Gramfort, Oleksii Kachaiev, Huy Quang Tran, Laurène David, Clément Bonet, Nathan Cassereau, Théo Gnassounou, Eloi Tanguy, Julie Delon, Antoine Collas, Sonia Mazelet, Laetitia Chapel, Tanguy Kerdoncuff, Xizheng Yu, Matthew Feickert, Paul Krzakala, Tianlin Liu, and Eduardo Fernandes Montesuma. Pot python optimal transport (version 0.9.5), 2024. [22] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. Bulletin of the American Mathematical Society, 62(3):427–479, 2025. [23] Adrian Hill, Neal McKee, Johannes Maeß, Stefan Bluecher, and Klaus Robert Muller. Smoothed differentiation efficiently mitigates shattered gradients in explanations. In The Thirty-ninth Annual Conference on Neural Information Processing Systems (NeurIPS), 2025. [24] Ross N. Hoffman, Zheng Liu, Jean-Francois Louis, and Christopher Grassoti. Distortion representation of forecast errors. Monthly Weather Review, 123(9):2758 – 2770, 1995. [25] Eugenia Kalnay. Atmospheric Modeling, Data Assimilation and Predictability. Cambridge University Press, 2003. [26] Christian Keil and George C. Craig. A displacement-based error measure applied in a regional ensemble forecasting system. Monthly Weather Review, 135(9):3248 – 3259, 2007. [27] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pages 5562–5571. PMLR, 2021. [28] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al. Learning skillful medium-range global weather forecasting. Science, 382(6677):1416–1421, 2023. [29] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1(4):541–551, 1989. [30] Zachary C Lipton. The mythos of model interpretability: In machine learning, the concept of interpretability is both important and slippery. Queue, 16(3):31–57, 2018. 11
[31] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of International Conference on Learning Representations (ICLR), 2019. [32] Amir Mehrpanah, Erik Englesson, and Hossein Azizpour. On spectral properties of gradientbased explanation methods. In European Conference on Computer Vision, pages 282–299. Springer, 2024. [33] Gabriel Moldovan, Ewan Pinnington, Ana Prieto Nemesio, Simon Lang, Zied Ben Bouallègue, Jesper Dramsch, Mihai Alexe, Mario Santa Cruz, Sara Hahner, Harrison Cook, et al. An update to ecmwf’s machine-learned weather forecast model AIFS. arXiv preprint arXiv:2509.18994, 2025. [34] Météo-France. Py4cast: Weather forecasting with deep learning. https://github.com/ meteofrance/py4cast. [35] Météo-France. Titan: Training inputs & targets from arome for neural networks. https: //huggingface.co/datasets/meteofrance/titan, 2024. [36] Philip Naumann, Jacob Kauffmann, and Grégoire Montavon. Wasserstein distances made explainable: Insights into dataset shifts and transport phenomena. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2026. [37] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310–1318. Pmlr, 2013. [38] Anthony Rhodes, Yali Bian, and Ilke Demir. Quantifying explainability with multi-scale gaussian mixture models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pages 8223–8228, June 2024. [39] Philippe Rigollet. The mean-field dynamics of transformers. arXiv preprint arXiv:2512.01868, 2025. [40] Yao Rong, Tobias Leemann, Vadim Borisov, Gjergji Kasneci, and Enkelejda Kasneci. A consistent and efficient evaluation strategy for attribution methods. In Proceedings of International Conference on Machine Learning (ICML), 2022. [41] Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan. Unetr++: delving into efficient and accurate 3d medical image segmentation. IEEE Transactions on Medical Imaging, 43(9):3377–3390, 2024. [42] Anna Shalova and André Schlichting. Solutions of stationary mckean–vlasov equation on a high-dimensional sphere and other riemannian manifolds. Advances in Nonlinear Analysis, 15(1):20250141, 2026. [43] Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje. Learning important features through propagating activation differences. In Proceedings of the 34th International Conference on Machine Learning - Volume 70, ICML’17, page 3145–3153, 2017. [44] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. In Workshop Track Proceedings of International Conference on Learning Representations (ICLR), 2013. [45] Daniel Smilkov, Nikhil Thorat, Been Kim, Fernanda Viégas, and Martin Wattenberg. Smoothgrad: removing noise by adding noise. In ICML workshop on Visualization for Deep Learning, 2017. [46] Justin Solomon, Fernando de Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: efficient optimal transportation on geometric domains. ACM Transactions on Graphics (TOG), 34(4), July 2015. [47] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pages 3319–3328. PMLR, 2017. 12
Table 2: Meteorological variables composing the state tensor xt . Surface diagnostics provide 5 channels; upper-air profiles at four pressure levels contribute 16 channels, for C = 21 total. Variable
Physical Description
Level
Surface and Near-Surface Variables aro_t2m 2-metre temperature aro_r2 2-metre relative humidity aro_tp_0m Total accumulated precipitation aro_u_10m 10-metre zonal wind component (u) aro_v_10m 10-metre meridional wind component (v)
2m 2m Surface 10 m 10 m
Upper-Air Profiles (250, 500, 700, 850 hPa) aro_t Air temperature aro_u Zonal wind component (u) aro_v Meridional wind component (v) aro_z Geopotential height
×4 levels ×4 levels ×4 levels ×4 levels
[48] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NeurIPS), 30, 2017. [49] Cédric Villani. Optimal transport : old and new / Cédric Villani. Grundlehren der mathematischen Wissenschaften. Springer, Berlin, 2009. [50] Yongjie Wang, Tong Zhang, Xu Guo, and Zhiqi Shen. Gradient based feature attribution in explainable ai: A technical review. arXiv preprint arXiv:2403.10415, 2024. [51] Jianbo Ye, Panruo Wu, James Z Wang, and Jia Li. Fast discrete distribution clustering using wasserstein barycenter with sparse support. IEEE Transactions on Signal Processing, 65(9):2317–2332, 2017. [52] Linjiang Zhou, Chao Ma, Zepeng Wang, Libing Wu, and XIAOCHUAN SHI. Adaptgrad: Adaptive sampling to reduce noise. In The Thirty-ninth Annual Conference on Neural Information Processing Systems (NeurIPS), 2025.
A
The TITAN Meteorological Dataset
We utilize the TITAN dataset, a high-resolution meteorological benchmark developed by MétéoFrance for deep learning-based weather forecasting [35]. TITAN is derived from the operational AROME limited-area model, enabling kilometre-scale forecasting over Western Europe. Spatial Resolution. TITAN is projected onto a regular lat–lon grid at a native resolution of 0.025◦ (≈ 2.5 km). The full AROME domain spans [−12◦ , 16◦ ] × [37.5◦ , 55.4◦ ] (1121 × 717 grid points). For all experiments we use a centred subdomain over France: [−6◦ , 9.975◦ ] × [40.125◦ , 52.9◦ ], corresponding to a spatial grid of H × W = 512 × 640 grid points. Both domains are illustrated in Figure 3. Temporal Resolution. The dataset provides analyses at a 1-hour time step. Our models are trained on lead times ∆t ∈ {+1h, . . . , +6h}, with a context window of one preceding analysis step as input. State Variables. We define the atmospheric state at timestep t as a multi-channel spatial tensor xt ∈ RH×W ×C , where C is the number of meteorological channels. Table 2 lists all variables included in the state tensor. Preprocessing and Normalization. Following the standard pipeline in the py4cast framework [34], all physical variables are standardized to zero mean and unit variance (Z-score normalization). 13
55°N
AROME domain 52°N 49°N
46°N
43°N
Training & evaluation domain
40°N
16°E
12°E
8°E
4°E
0°
4°W
8°W
AROME full domain Training & evaluation subdomain 12°W
37°N
Figure 3: Geographic extent of the TITAN/AROME dataset. The blue dashed rectangle delineates the full AROME operational domain; the red solid rectangle indicates the subdomain used for model training and evaluation.
B
Model Architecture and Training Configuration
B.1
Architecture: UNetRPP
We use UNetRPP [41], a hybrid convolutional–attention U-Net architecture that combines a hierarchical encoder–decoder structure with multi-head self-attention at each resolution stage. The encoder downsamples the input spatial tensor via strided convolutions, producing a sequence of feature maps at progressively coarser resolutions. Skip connections from each encoder stage are projected and injected into the corresponding decoder stage, following the standard U-Net architecture. The decoder uses linear upsampling followed by a projection head. Positional embeddings are learned via a small MLP and Instance Normalization is applied throughout. The full architecture hyperparameters are listed in Table 3. B.2
Training Configuration
Training strategy. The model was trained using the Py4cast framework [34], with the AdamW optimiser [31] (β1 = 0.9, β2 = 0.95), an initial learning rate of 10−3 decayed to 3 × 10−7 , and a linear warm-up of 1,000 steps. Full training details including the training strategy are provided in [34]. Data Splits. Training split is drawn from the years 2020-2022 at 1-hour observation steps, as summarised in Table 4. B.3
Hyperparameter Settings
Table 5 summarises all hyperparameters used for attribution methods and evaluation metrics. All stochastic methods use a fixed random seed of 42 for reproducibility. 14
Table 3: UNetRPP architecture and forecasting hyperparameters. Hyperparameter
Value
Architecture Hidden size Encoder attention heads Decoder attention heads Depth per stage (encoder blocks) Spatial downsampling rate Decoder projection size Encoder projection sizes Skip connections Linear upsampling Positional embedding Normalization Convolution operator
1024 16 4 [3, 3, 3, 3] 4× 64 [64, 64, 64, 32] ✓ ✓ Learned MLP Instance Norm Conv2d
Forecasting Training strategy Number of auto-regressive steps Loss function
Scaled auto-regressive 1 (per training iteration) MSE
Table 4: Dataset temporal splits. The validation and test sets use a 3-hour stride between initial conditions. Split
Start
End
Sample step
Train Test
2020-01-01 2023-01-01
2022-12-31 2023-12-31
1h 3h
C
Computational Cost and Resources
C.1
Hardware and implementation
All experiments were conducted on a single NVIDIA RTX A4500 Laptop GPU (16 GB VRAM). Implementations use PyTorch 2.5.1 with CUDA 12.1. Forecasting inference uses the Py4cast framework [34] with pretrained UNetRPP weights. Attribution methods, including WassersteinGrad, are our own implementations. For the computation of the entropic Wasserstein barycenter, we utilize the ot.bregman.convolutional_barycenter2d function from the POT library [20, 21]. C.2
Wall-time per sample
Wall-time measurements capture the full attribution pipeline for a single sample, including gradient computation, noise sampling where applicable, and aggregation. We exclude data loading and model weight initialization from all measurements. Each method was timed over samples drawn sequentially from the 2023 test split. All methods were evaluated under identical conditions: same samples, same spatial grid resolution 512 × 640, N = 20 perturbed samples for all stochastic methods. Table 6 reports mean wall-time and standard deviation across 100 timed samples for each attribution method.
D
Ablation Studies
D.1
Number of perturbed samples N
We perform an ablation study on the number of perturbed inputs N , utilizing a fixed Sinkhorn regularization parameter of λ = 0.001. Table 7 reports the impact of N on Sparsity, Faithfulness (ROAD), and Robustness (ℓ2 -LLE and cosine-LLE) evaluated over n = 25 test events. We select N = 20 for all subsequent experiments, as it strikes a favorable balance between minimizing computational overhead and maximizing attribution stability. 15
Table 5: Hyperparameter configurations. Summary of settings for the gradient-based attribution methods and parameters used for the evaluation metrics. Parameter
Value
Description
(IntegratedGrad, SmoothGrad, VarGrad, WGBary , WGBary×Grad ) N 20 Number of perturbed samples σ 0.2 × (max(xin ) − min(xin )) Input perturbation noise std WassersteinGrad only λ 0.001 Variant 1 WGBary Variant 2
Sinkhorn entropic regularisation Barycenter used directly as attribution map Barycenter masked on base gradient
WGBary×Grad
Faithfulness (ROAD) pmax 15% p range {1%, . . . , 15%} Krand 5 Imputation Linear + Gaussian noise σimp 0.1 × (max(xin ) − min(xin ))
Maximum masking percentage Masking percentages evaluated Random mask repetitions per percentage Feature removal strategy [40] Imputation noise std
Robustness (LLEℓ2 , LLEcosine ) KLLE 7 σLLE 0.1 × (max(xin ) − min(xin )) Distance ℓ2
Perturbation iterations Perturbation noise std ∥G(xin ) − G(xin + ε)∥2
Sparsity (Gini Index) Sorting Ascending
Required by Gini formula
Table 6: Wall-time per sample. Method BaseGrad SmoothGrad IntegratedGrad VarGrad WGBary (λ = 0.01) WGBary (λ = 0.001) WGBary×Grad (λ = 0.001)
D.2
Mean (s)
Std (s)
Overhead vs. SmoothGrad
0.37 7.42 8.78 7.70 8.40 14.60 14.90
0.03 0.12 0.13 0.10 0.11 0.18 0.2
(no sampling) 1× (reference) 1.14× 1.03× 1.13× 1.96× 1.98×
Sinkhorn Regularisation λ
We ablate the entropic regularisation parameter λ of the Sinkhorn barycenter for both WassersteinGrad variants. Table 8 reports Sparsity, Faithfulness (ROAD), and Robustness (ℓ2 -LLE and cosine-LLE) on n = 25 events. We select λ = 0.001 as the default based on the best faithfulness score and the sparsity-robustness tradeoff for WGBary and WGBary×Grad , but also by the requirement that lower values of λ provide a more accurate approximation of the true 2-Wasserstein barycenter.
E
Experiment Details
E.1
Optimal Transport reveals spatial displacement
We model the clean and perturbed saliency maps as discrete probability measures over the spatial domain Ω ⊂ R2 : |Gclean (u)| , v∈Ω |Gclean (v)|
µclean (u) = P
µpert (u) = P 16
|Gpert (u)| , v∈Ω |Gpert (v)|
(6)
Table 7: Ablation on Number of perturbed samples N . Top: WGBary (barycenter alone). Bottom: WGBary×Grad (barycenter masked on base gradient). All metrics on n = 25 events with highest mean precipitation in the region of interest, mean ± SEM. ↑ higher is better. N
Sparsity↑
Faithfulness↑
LLEℓ2 ↓
LLEcosine ↓(×103 )
Wall-time
0.815±0.030 0.825±0.027 0.851±0.023 0.837±0.027 0.837±0.027
0.037±0.002 0.032±0.002 0.029±0.003 0.028±0.003 0.027±0.003
0.50±0.03 0.41±0.01 0.35±0.03 0.34±0.02 0.33±0.02
3.70±0.18 7.64±0.20 14.60±0.20 22.51±0.20 32.34±0.20
0.840±0.031 0.821±0.027 0.824±0.028 0.815±0.029 0.825±0.026
0.042±0.003 0.042±0.003 0.042±0.003 0.042±0.003 0.042±0.003
3.64±0.22 3.63±0.22 3.63±0.22 3.63±0.20 3.63±0.20
3.96±0.18 7.70±0.20 14.90±0.20 22.61±0.20 32.46±0.20
WGBary 5 10 20 30 40
0.319±0.006 0.319±0.007 0.320±0.006 0.316±0.006 0.316±0.006
WGBary×Grad 5 10 20 30 40
0.827±0.007 0.828±0.006 0.829±0.006 0.828±0.006 0.828±0.006
Table 8: Ablation on Sinkhorn regularisation λ. Top: WGBary (barycenter alone). Bottom: WGBary×Grad (barycenter masked on base gradient). All metrics on n = 25 events with highest mean precipitation in the region of interest, mean ± SEM. ↑ higher is better. λ
Sparsity↑
Faithfulness↑
LLEℓ2 ↓
LLEcosine ↓(×103 )
0.201±0.003 0.287±0.006 0.320±0.006 0.329±0.006
0.848±0.018 0.866±0.024 0.851±0.023 0.824±0.022
0.028±0.002 0.027±0.002 0.029±0.003 0.031±0.004
0.06±0.01 0.09±0.01 0.35±0.03 0.45±0.03
0.816±0.028 0.819±0.032 0.824±0.028 0.798±.035
0.063±0.004 0.054±0.003 0.042±0.003 0.041±.003
3.48±0.18 3.59±0.19 3.63±0.22 3.72±0.23
WGBary 0.1 0.01 0.001 0.0005
WGBary×Grad 0.1 0.636±0.006 0.01 0.767±0.007 0.001 0.829±0.006 0.0005 0.838±.006
where Gclean = ∇xin ϕ [f T (x)]cout and Gpert = ∇x̃in ϕ [f T (x̃)]cout are the clean and perturbed attribution maps respectively. The optimal transport plan T ∗ solves: T ∗ = arg
X
min
T ∈Π(µpert , µclean )
Tuv · ∥u − v∥2 ,
(7)
u, v∈Ω
where Π(µpert , µclean ) is the set of joint distributions with marginals µpert and µclean . The marginal constraints are: X X Tuv = µpert (u), Tuv = µclean (v). (8) v∈Ω
u∈Ω
∗ To isolate displaced mass, we mask the diagonal of T ∗ : entries Tuu correspond to mass that remains at the same spatial coordinate u ∈ Ω and carry no information about displacement. The off-diagonal ∗ entries Tuv , u ̸= v, encode the spatial flow of attribution mass migrating from coordinate u to coordinate v under perturbation, and are visualised in Figure 1(c).
E.2
Displacement of the center of mass
Let xin denote the clean input slice and let G(xin ) denote the explanation map produced by the explainer. We consider additive Gaussian perturbations of the form (p,r)
xin
= xin + ϵ(p,r) ,
ϵ(p,r) ∼ N (0, σp2 I), 17
(9)
where p indexes the perturbation level and r the repetition. The noise level is defined as σp = αp (max(xin ) − min(xin )) ,
(10)
(p,r) with αp a scalar controlling the perturbation intensity. For each perturbed input xin , we compute
the corresponding explanation map:
(p,r) G(p,r) = G xin ,
(11)
G(0) = |G(xin )|.
(12)
and the clean explanation:
Centroid of the explanation. The spatial centroid of the explanation map is computed using the centre of mass: P P (p,r) (p,r) i,j i Gij i,j j Gij (p,r) (p,r) cx = P , c = . (13) P y (p,r) (p,r) i,j Gij i,j Gij (0)
(0)
Similarly, the centroid of the clean explanation is (cx , cy ). Centroid displacement. The displacement induced by the perturbation is measured using the Euclidean distance between the two centroids: r 2 2 (0) (0) (p,r) (p,r) (p,r) . (14) − cy + cy d = cx − c x Aggregation.
For each perturbation level p, we repeat the experiment R times and compute v u R R u1 X X 2 1 (p,r) d¯p = d , sp = t d(p,r) − d¯p . (15) R r=1 R r=1
Peak displacement as complementary measure. We also track the spatial location of the maximum attribution, (px , py ) = arg maxi,j Gij , and compute its displacement similarly to the centroid. Aggregation across events. For each noise level σp , we compute the per-event mean across R samples, then aggregate across n = 100 atmospheric events from the 2023 test split filtered by mean precipitation over the Paris ROI. The reported error bands are ±1 SEM across events. Autoregressive amplification. Figure 4 reports both centroid and peak displacement at singlestep (t + 1) and five-step (t + 5) autoregressive horizons. The peak displacement amplifies by approximately 3–4× between t + 1 (saturating around 32 km) and t + 5 (reaching 115 km), while the relative prediction error remains below 1.5% at both horizons. This autoregressive amplification of geometric displacement suggests that each autoregressive step introduces independent geometric distortions that compound through the rollout.
F
Gradient based methods
F.1
BaseGrad
Introduced by [44], Base Gradient computes the partial derivatives of the output with respect to the input. GVanilla (xin ) = ∇xin ϕ [f T (x)]cout . (16) F.2
IntegratedGrad
IntegratedGrad [47] addresses the local sensitivity of Base Gradient by integrating the gradient along a straight path from a baseline xbase to the input xin . In practice, the integral is approximated via a Riemann sum over N interpolation steps x̃in,i = xbase + Mi (xin − xbase ): N
xin − xbase X GIG (xin ) = ∇x̃in,i ϕ [f T (x̃i )]cout . N I=1
18
(17)
13 12
32 30 28 26
1.020
10
1.015
9 7
1.005
6
1.000
5 0.2
0.4 0.6 0.8 Noise level (fraction of input range)
1.0
0.2
Peak displacement
0.4 0.6 0.8 Noise level (fraction of input range)
Centroid displacement Relative prediction error
1.0 Mean ± SEM across n=116 events
1.0150
16
110 100 90 80 70
1.0125
Relative prediction error
Centroid displacement (km)
120
t+5
1.010
8
22
Peak displacement (km)
1.025
11
24
130
Centroid displacement Relative prediction error
Relative prediction error
Peak displacement
Centroid displacement (km)
t+1
Peak displacement (km)
34
1.0100
14
1.0075 12
1.0050 1.0025
10
1.0000
60 0.2
0.4 0.6 0.8 Noise level (fraction of input range)
1.0
0.2
0.4 0.6 0.8 Noise level (fraction of input range)
1.0 Mean ± SEM across n=116 events
Figure 4: Centroid and peak displacement of gradient attributions under input perturbations, at single-step (t + 1, top) and autoregressive (t + 5, bottom) lead times. Left: peak displacement (location of arg max |G|). Right: centroid displacement alongside relative prediction error (secondary axis). Peak displacement amplifies by 3–4× between t + 1 and t + 5, while the prediction error remains below 1.5% in both cases. Error bands: ±1 SEM across n = 100 precipitation-filtered events. We define xbase as the climatological mean of the training set, representing a neutral atmospheric state. Since our inputs are Z-standardized, this baseline corresponds to a neutral atmospheric state where xbase = 0. F.3
SmoothGrad
To mitigate local gradient shattering, SmoothGrad [45] computes the expected gradient over a local Gaussian neighbourhood, approximated via Monte Carlo sampling over N perturbed inputs x̃in,i = xin + εi , where εi ∼ N (0, σ 2 I): N
GSG (xin ) =
F.4
1 X ∇x̃in,i ϕ([f T (x̃i )]cout ). N i=1
(18)
VarGrad
VarGrad [1] replaces the mean of SmoothGrad with the variance of the perturbed gradients, highlighting regions where the gradient is highly sensitive to local noise: N
GVar (xin ) =
G
2 1 X ∇x̃in,i ϕ([f T (x̃i )]cout ) − GSG (xin ) . N i=1
(19)
Evaluation Metrics
To evaluate the quality of our attribution maps within the continuous spatial domain of weather forecasting, we adopt a suite of quantitative metrics inspired by recent XAI benchmarks [10]. We evaluate explanations across three fundamental axes: Faithfulness, Robustness, and Complexity. 19
Faithfulness (ROAD-MSE ↑). Features marked as highly salient must strictly influence the model’s prediction. We evaluate this using the RemOve And Debias (ROAD) framework [40]. Because weather forecasting is a regression task, raw RMSE degradations vary by orders of magnitude across events , making cross-event aggregation of continuous scores unstable. We therefore adapt ROAD to this setting via a binary comparison restricted to the region of interest B. For each masking percentage p ∈ {1%, 2%, . . . , pmax }, we mask the top-p% salient pixels using linear imputation [40] and compare the resulting prediction degradation against the mean of K random masks of identical size: # " K X 1 ϕ [f T (x)]cout − ϕ [f T (xrand,p,k )]cout . binary(p) = 1 ϕ [f T (x)]cout − ϕ [f T (xsal,p )]cout > K k=1 (20) where xsal,p denotes the input with top-p% salient pixels masked, xrand,p,k denotes the k-th random mask of the same size, and ϕ is the ROI aggregation operator defined in Section 3.1. The faithfulness score is the AUC of the binary curve: R pmax binary(p) dp qROAD = 1 . (21) pmax − pmin Robustness (LLEℓ2 ↓). A reliable explanation must remain stable under noise. We quantify this using the Local Lipschitz Estimate (LLE) [4], which measures the worst-case instability of the attribution map G under a bounded input perturbation: qLLEℓ2 = max
∥ε∥2 ≤δ
∥G(xin ) − G(xin + ε)∥2 , ∥ε∥2
(22)
A lower score indicates greater robustness to local perturbations. Robustness (LLEcosine ↓). To measure structural rather than amplitude change, we normalize the attribution maps before comparison: qLLEcosine = max
∥ε∥2 ≤δ
∥Ĝ(xin ) − Ĝ(xin + ε)∥2 , ∥ε∥2
(23)
where Ĝ = G/∥G∥2 denotes the unit-normalized attribution map. This metric penalises methods whose attribution structure changes under perturbation. Complexity (Gini Index ↑). A highly distributed explanation is often uninterpretable for domain experts. We quantify spatial concentration using the Gini Index [13]: Pd (2i − d − 1) |G(xin )(i) | qSPA = i=1 Pd , (24) d i=1 |G(xin )(i) | where d = H × W is the total number of spatial grid points and |G(xin )(i) | denotes the i-th smallest absolute attribution value. A higher Gini coefficient indicates a sparser, more localised explanation.
H
Mathematical Preliminaries
H.1
Optimal Transport and Wasserstein Barycenters
H.1.1
The 2-Wasserstein Distance
Let µ, ν ∈ P(Ω) be two probability measures defined on the spatial grid Ω. The 2-Wasserstein distance is defined via the Kantorovich formulation: Z W22 (µ, ν) = min ∥u − v∥22 dπ(u, v), (25) π∈Π(µ,ν)
Ω×Ω
where Π(µ, ν) denotes the set of all joint probability couplings with marginals µ and ν. 20
H.1.2
Discrete 2-Wasserstein Distance
In our setting, spatial explanations are defined on a discrete grid of n locations. We represent them as normalized non-negative vectors µ, ν ∈ ∆n , where ( ) n X n n ∆ = p ∈ R+ pi = 1 . (26) i=1 n×n
Let M ∈ R be the cost matrix, where Mij = ∥ui − uj ∥22 encodes the squared Euclidean distance between spatial coordinates ui , uj ∈ Ω. The discrete 2-Wasserstein distance is: W22 (µ, ν) = min ⟨M, π⟩ n×n
s.t.
π1 = µ,
π ⊤ 1 = ν,
(27)
π∈R+
P where π is a transport plan and ⟨M, π⟩ = i,j Mij πij . This linear program scales as O(n2 ) variables and is prohibitive for high-resolution grids. H.1.3
Entropic Regularization and Sinkhorn Iterations
To obtain a tractable formulation, we use entropic regularization [15]: X 2 πij (log πij − 1), W2,λ (µ, ν) = min⟨M, π⟩ + λ π
(28)
i,j
subject to the same marginal constraints. The entropy term makes the objective strictly convex and yields the scaling form π ⋆ = diag(a) Kλ diag(b),
Kλ = exp(−M/λ),
(29)
where the scaling vectors a and b are computed via alternating matrix-vector normalizations (Sinkhorn iterations), which are highly parallelisable on GPUs [15]. H.2 H.2.1
Mean-Field Limits and Gradient Flows The Mean-Field Regime
The mean-field regime approximates the behaviour of a large number of interacting particles (n → ∞) by tracking their statistical distribution rather than individual trajectories. Consider a discrete system of n particles (tokens) {xi (s)}ni=1 evolving according to a coupled ODE: n
ẋi (s) =
1X K(xi (s), xj (s)), n j=1
(30)
where K : RD × RD → RD is an interaction kernel (e.g. the attention mechanism) and D is the token embedding dimension. We introduce the empirical measure µn (s), a discrete probability measure describing the state of the system at time s: n
1X µn (s) = δx (s) , n i=1 i
(31)
where δx denotes the Dirac mass at position x. The Mean-Field Limit (n → ∞). As n → ∞, the empirical measure µn (s) converges weakly to a deterministic continuous probability measure µs . The evolution of this limiting measure is governed by the continuity equation: ∂s µs + ∇ · (µs vs [µs ]) = 0, (32) where the velocity field vs [µs ] at a point x depends on the entire distribution µs via convolution with the kernel: Z vs [µs ](x) = K(x, y) dµs (y). (33) This equation describes the transport of mass: the local change in density ∂s µs is fully accounted for by the divergence of the flux µs vs . 21
H.2.2
Wasserstein Gradient Flows
A gradient flow describes the evolution of a system following the path of steepest descent (or ascent) to minimise (or maximise) an energy functional. In the space of probability measures, the geometry is defined by the Wasserstein metric. The Gradient Flow Structure. We consider an energy functional E : P(RD ) → R. In the context of attention, this is the interaction energy: ZZ 1 eβ⟨x,y⟩ dµ(x) dµ(y). (34) Eβ (µ) = 2β A curve of measures (µs )s≥0 is a Wasserstein gradient flow for E if it satisfies the continuity equation with velocity field equal to the negative Wasserstein gradient of the energy: δE , (35) vs = −∇W2 E(µs ) = −∇ δµ where δE δµ denotes the first variation of the energy.
I
Noise Effect in Convolutional Architectures
In Section 3.2.1, we established that input perturbations in CNNs manifest as spatial geometric displacement in the resulting attribution maps. In this section, we formally derive this phenomenon by tracing the propagation of additive Gaussian noise through the forward pass (convolutions, nonlinearities, and pooling) and analysing its impact on the backward pass. I.1
Gradient Path
We analyse the backward pass using the path formulation of the gradient, following [23, 8]. Applying the multivariable chain rule along the computational graph, the partial derivative of the output Ytarget = f T (x) with respect to the i-th input feature xi is the sum of derivatives along all paths connecting them: X Y T ∂f T (x) ∇f (x) i = = [Jf n (αn )]ρ , (36) ∂xi n∈ρ ρ∈Pi→c
where αn denotes the input activations to function f n . Factorization of Linear and Nonlinear Contributions. Following [23], we factorize the terms in this path product into two categories: 1. Linear or affine functions (e.g. dense and convolutional layers): their Jacobian entries are constant weights, [Jf n ]ρ = Wρn . 2. Nonlinear functions (e.g. ReLU, max pooling): their Jacobian entries depend on the input activation αn , [Jf n (αn )]ρ = gρn (αn ). Separating the constant linear weights from the input-dependent nonlinear derivatives yields: ! ! X Y Y ∂f T (x) u v v = Wρ gρ (α ) . ∂xi u∈ρ v∈ρ ρ∈Pi→c
We collapse these products for notational clarity: Y Wρ = Wρu (path-weight), u∈ρ
Γρ (x) =
Y
gρv (αv )
(path-indicator function),
v∈ρ
22
(37)
giving: X ∂f T (x) Wρ Γρ (x). = ∂xi
(38)
ρ∈Pi→c
Previous works view the fragility of Γρ (x) under input perturbations through the lens of the Shattered Gradient Problem [23, 8]. I.2
Linear Propagation in Convolutions
Consider a standard 2D convolutional layer l. Let ol ∈ RH×W denote the input feature map and l+1 Kl ∈ R(2V +1)×(2U +1) denote the convolutional kernel. The clean pre-activation zh,w at spatial location (h, w) is: V U X X l+1 zh,w = Klm,n olh+m, w+n . (39) m=−V n=−U
We assume the input to layer l is corrupted by i.i.d. additive Gaussian noise: õlh,w = olh,w + εlh,w ,
εlh,w ∼ N (0, σ 2 ).
(40)
Substituting into the convolution yields: X X l+1 Klm,n olh+m, w+n + z̃h,w = Klm,n εlh+m, w+n . m,n
|
(41)
m,n
{z
l+1 clean signal zh,w
}
|
{z
}
l+1 noise term ηh,w
Because the noise variables are independent Gaussians and the convolution is linear, the noise term l+1 ηh,w remains Gaussian with variance scaled by the squared ℓ2 norm of the kernel: l+1 ηh,w ∼ N 0, σ 2 ∥Kl ∥22 ,
∥Kl ∥22 =
V X
U X
Klm,n
2
.
(42)
m=−V n=−U l+1 l+1 l+1 Thus z̃h,w = zh,w +ηh,w . The noise propagates linearly until it interacts with the non-linear activation φ.
I.3
Impact on Smooth Activations (General Case)
For smooth, twice-differentiable activation functions (e.g. GELU, SiLU), we expand via a secondl+1 order Taylor series around the clean pre-activation zh,w : l+1 l+1 õl+1 h,w = φ zh,w + ηh,w 2 l+1 l+1 l+1 l+1 l+1 l+1 3 = φ(zh,w ) + φ′ (zh,w ) ηh,w + 21 φ′′ (zh,w ) ηh,w + O (ηh,w ) . (43) Taking the expectation and using E[η] = 0 and E[η 2 ] = σ 2 ∥Kl ∥22 : 1 ′′ l+1 2 l 2 l+1 E õl+1 φ (zh,w ) σ ∥K ∥2 + O(σ 3 ). (44) h,w = φ(zh,w ) + 2 The expected activation is shifted from the clean activation proportionally to the propagated noise variance and the local curvature φ′′ . I.4
Analysis for Piecewise Linear Activations (ReLU)
For piecewise linear functions such as ReLU, the second derivative is zero almost everywhere, rendering the Taylor expansion uninformative. Let z denote the clean pre-activation and η ∼ N (0, σ̃ 2 ) the propagated noise. We compute the expected value: E[y] = E[ReLu(z + η)] = E[max(0, z + η)]. 23
Using the Gaussian density ϕσ̃ (ε) = √
ε2 1 exp − 2 , 2σ̃ 2πσ̃
we integrate: Z ∞ max(0, z + η) ϕσ̃ (η) dη
E[y] = −∞ Z ∞
=
(z + η) ϕσ̃ (η) dη z z = zΦ + σ̃ ϕ σ̃ σ̃ −z
where Z u
2 1 √ e−t /2 dt (Gaussian CDF) 2π −∞ 1 −u2 /2 ϕ(u) = √ e (Gaussian PDF) 2π
Φ(u) =
Creation of Phantom Gradients: Even when the clean pre-activation is negative (z < 0) and the neuron should be strictly inactive, the term σ̃ϕ(z/σ̃) remains strictly positive. Consequently, the injected noise creates phantom activations: neurons that would be strictly inactive in the clean pass now exhibit a positive expected output. During the backward pass of attribution generation, these phantom activations prematurely unblock gradient flow through the ReLU gates, opening spurious attribution paths Γρ (x) in (38) that highlight regions entirely irrelevant to the original feature geometry [23].
J
Backpropagation and Noise Effect in Attention Layers
Setup in the XAI context. In our attribution framework, the scalar target Ytarget is the regional forecast defined in Section 3.1. The attribution map ∇xin Ytarget measures the sensitivity of this scalar forecast to the input channel xin . Within a Transformer attention layer, the input token sequence is Oin ∈ Rn×D (the feature map at that stage, with n spatial positions and embedding dimension D), and Oout ∈ Rn×Dv is its output. We derive ∇Oin Ytarget by backpropagating through the attention mechanism. Forward pass. Queries, keys, and values are computed via learned projections WQ , WK ∈ RD×Dk , WV ∈ RD×Dv : Q = Oin WQ , Scaled dot-product attention yields: QK ⊤ S= √ , Dk
K = Oin WK ,
V = Oin WV .
A = softmax(S) ∈ Rn×n ,
where softmax is applied row-wise: Aij = exp(Sij )/
P
Oout = AV,
(45)
(46)
j ′ exp(Sij ). ′
Gradients with respect to projections. Let ∆O = ∂Ytarget /∂Oout ∈ Rn×Dv denote the upstream gradient propagated from subsequent layers. By the chain rule: ∂Ytarget = A⊤ ∆O , ∂V ∂Ytarget ∆A = = ∆O V ⊤ . ∂A 24
(47) (48)
Defining ∆S = ∂Ytarget /∂S: ∂Ytarget 1 = √ ∆S K, ∂Q Dk
(49)
1 ∂Ytarget = √ ∆⊤ Q. ∂K Dk S
(50)
The row-wise softmax Jacobian gives, for row i: ! X (∆S )ij = Aij (∆A )ij − Aik (∆A )ik .
Gradient through the softmax.
(51)
k
In matrix form via the Hadamard product: ∆S = A ⊙ ∆A − (∆A ⊙ A)11⊤ ,
(52)
where 1 ∈ Rn is a vector of ones, making (∆A ⊙A)11⊤ a matrix whose columns equal the row-sums of ∆A ⊙ A. Total gradient with respect to the input. Applying the chain rule backward through Eqs. (45)– (50): 1 1 ⊤ ∇Oin Ytarget = A⊤ ∆O WV⊤ + √ ∆S KWQ⊤ + √ ∆⊤ QWK . (53) Dk Dk S
K
The Mean-Field Dynamics of Transformers
The Transformer architecture has become the dominant framework in modern deep learning. While many variants exist to stabilize the architecture, the core mathematical interest lies in understanding how the attention mechanism itself regulates token dynamics [39]. K.1
Discrete to Continuous Time
We denote autoregressive forecasting steps as t, and the continuous-time depth within a single model pass as s. Transformers process data through a sequence of discrete layers k = 1, . . . , L. This recursive structure naturally suggests a discrete-time dynamical system. As the depth increases, we view the layer index as a continuous variable s, transitioning from discrete updates to a continuous-time nonlinear flow between tokens [39, 22]: ẋ(s) = Fs (x(s)).
(54)
The attention mechanism defines a nonlocal velocity field, coupling each token to all others through a kernel that depends on their pairwise similarities. K.2
The Dynamics of Self-Attention
We model the n tokens as particles xi (s) evolving on the unit sphere SD−1 . K.2.1
Standard Self-Attention (SA)
The dynamics of the standard self-attention mechanism, incorporating the geometric constraint of layer normalization (keeping particles on the sphere), are given by [39, 22]: n X 1 ẋi (s) = P⊥ eβ⟨WQ xi (s), WK xj (s)⟩ WV xj (s) , (55) xi (s) Zβ,i (s) j=1 where: 25
• β > 0 is the inverse temperature parameter, Pn • Zβ,i (s) = ℓ=1 eβ⟨xi (s),xℓ (s)⟩ is the partition function, D−1 • P⊥ . x y = y − ⟨x, y⟩x is the orthogonal projection onto the tangent space Tx S
For simplicity we take WV = WQ = WK = Id. This equation describes n particles interacting through the kernel K(x, y) = eβ⟨x,y⟩ . K.2.2
Unnormalized Self-Attention (USA)
To simplify analysis, we consider a variant that omits the normalization denominator Z. The unnormalized self-attention dynamics are: 1 Pn ẋ (s) = P⊥ β⟨xi (s),xj (s)⟩ e x (s) , j i xi (s) n j=1 (56) xi (0) = xi . In practice, the behaviour of USA closely mirrors that of SA, making it a valuable proxy for theoretical study [39, 22]. K.3 K.3.1
Gradient Flows and Transport Mean-Field Limit
AsPthe number of tokens n → ∞, we describe the system using the empirical measure µs = n 1 i=1 δxi (s) . The evolution of this density is governed by the continuity equation: n ∂s µs + ∇ · (µs X [µs ]) = 0, where the vector field is X [µ](x) = P⊥ x
Z e
β⟨x,y⟩
(57)
y dµ(y)
(58)
for the unnormalized case. K.3.2
Maximizing Interaction Energy
Both SA and USA dynamics are driven by the objective to maximize the interaction energy [39, 22]: ZZ 1 Eβ (µ) = eβ⟨x,y⟩ dµ(x) dµ(y). (59) 2β While the objective is the same, the cost of movement differs between the two models: 1. USA as Wasserstein gradient flow: The unnormalized dynamics correspond to a gradient flow with respect to the standard Wasserstein metric W2 . Attraction can accelerate exponentially (v ∝ eβ ) [39, 22]. 2. SA as weighted gradient flow: Standard self-attention corresponds to a gradient flow with respect to a weighted Wasserstein metric. The normalization term Z cancels exponential growth, keeping the velocity field bounded and regulating clustering speed [39, 22]. K.4 K.4.1
The Noisy Regime Continuous Noise (SDE)
When noise is injected at every time step, the dynamics of a token Xi (s) become a stochastic differential equation (SDE) on the sphere [39]: n X √ 1 dxi (s) = P⊥ eβ⟨xi (s),xj (s)⟩ Xj (s) ds + 2κ−1 dWi (s), (60) xi (s) | {z } n j=1 diffusion (noise) | {z } drift (clustering force)
where κ > 0 controls the noise magnitude and Wi (s) denotes Brownian motion on the sphere, pushing particles apart. 26
Fokker-Planck dynamics and free energy. In the mean-field limit, the probability distribution µs evolves according to the Fokker-Planck equation, combining advective and diffusive transport [42, 39]: ∂s µs = ∇ · (µs ∇V[µs ]) + κ−1 ∆µs , (61) where V[µs ] is the potential field. The driving energy changes from pure interaction energy to free energy: Free energy = Interaction energy + Entropy term. • Low noise (high κ): The clustering force dominates; tokens group semantically. • High noise (low κ): The diffusive force dominates; the distribution becomes uniform and model output can become incoherent. K.4.2
Input Perturbation
Alternatively, noise can be added solely to the input as an initial condition, leaving the dynamics deterministic (ODE) for s > 0: • The initial condition is random: x0 ∼ µ0 . • Due to the non-Lipschitz nature of attention and the complex energy landscape (metastability), different realisations of the initial noise can lead the system into different basins of attraction. • Result: Potentially different final clusterings or attention patterns for the same semantic input.
27