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

As deep learning models achieve remarkable complexity and performance across various tasks, their inherent opacity remains a critical barrier to operational trust in high-stakes applications. This opacity is particularly problematic in weather forecasting, where erroneous predictions can have severe physical or societal consequences. In recent years, deep learning has increasingly competed with traditional physical solvers in weather forecasting, becoming a new predictive tool. However, as these systems approach deployment in high-stakes decision pipelines, explaining their predictions becomes a pressing issue. Among existing explainable AI strategies, gradient-based feature attribution methods have become standard due to their scalability to high-dimensional spatial domains. However, these methods have a fundamental limitation: they often appear visually noisy and locally shattered due to the highly non-linear nature of deep networks. To reduce this noise, widely adopted smoothing techniques like SmoothGrad first sample several noisy neighbors of the input observation with spatially independent Gaussian noise and then average the gradients explaining the predictions obtained using the perturbed inputs. Although these strategies have proven effective on image data, we argue that reliance on pointwise averaging leads to poorly localized explanations when used to forecast the state of dynamic physical phenomena.

In predicting dynamic physical fields, traditional gradient-based feature attribution methods exhibit a fundamental failure mode: stochastic input perturbations cause geometric displacement in attribution maps rather than stationary amplitude noise. This displacement results in pointwise averaging blurring these spatially misaligned features, affecting the interpretability of predictions. Specifically, in meteorological phenomena, this geometric displacement manifests as attribution masses migrating from their true spatial location to nearby physically incorrect locations. This phenomenon is formally recognized as phase error in meteorological verification, predicting the right thing but in the wrong location. Additionally, when the prediction model is used in an autoregressive manner, this geometric displacement phenomenon is reinforced. Each autoregressive step likely introduces an independent geometric distortion, and when backpropagating gradients from final lead times to explain noisy inputs, these spatial misalignments accumulate. Therefore, developing novel explainable AI solutions that take into account the intrinsic properties of modern autoregressive forecast models is particularly interesting.

The core innovation of this paper lies in introducing a transport-based aggregation framework called WassersteinGrad to address the geometric displacement issue in dynamic physical field predictions. Specifically, WassersteinGrad extracts a geometric consensus of perturbed attribution maps by computing their entropic Wasserstein barycenter. This innovation not only enhances the transparency of model predictions but also provides a more reliable explanation tool for AI applications in high-stakes environments. Compared to existing methods like SmoothGrad, WassersteinGrad not only addresses the geometric displacement issue but also achieves significant improvements in interpretability and robustness. By mapping perturbed attribution maps into the space of spatial probability measures and utilizing the entropic regularization of the Wasserstein barycenter, WassersteinGrad extracts a geometric consensus, achieving more interpretable predictions in high-dimensional inputs. This innovation opens new pathways for the application of explainable AI in dynamic physical fields.

• �� Input Channel Selection: Select an input channel cin on which spatial attributions will be represented (e.g., select a spatial slice).

• �� Output Channel Selection: Select an output channel cout of the predicted field to explain (e.g., select surface precipitation).

• �� Region of Interest (ROI): Define a region of interest B and compute a scalar target over this region.

• �� Attribution Computation: Compute the gradient of the target with respect to the selected input channel.

• �� Perturbed Input Construction: Perturb the input with Gaussian noise to construct channel-specific perturbed inputs.

• �� Attribution Map Calculation: Compute the gradient of the target with respect to the perturbed input slice.

• �� Attribution Map Mapping: Map each attribution map to a discrete spatial probability distribution.

• �� Wasserstein Barycenter Calculation: Extract geometric consensus by computing the entropic Wasserstein barycenter of the perturbed attribution maps.

The experimental design includes using the AROME high-resolution meteorological benchmark dataset, derived from the AROME limited-area model at Météo-France, providing kilometer-scale analyses over Western Europe. A subdomain over France is used as the test split, spanning January to December 2023. The forecasting backbone is a pretrained hybrid convolutional-attention U-Net trained using Py4cast, predicting meteorological states at 1-hour lead time. In the experiments, total surface precipitation is chosen as the prediction target, and zonal wind at 250 hPa is used as the input channel. All stochastic methods use N perturbed samples with noise variance σ.2x t)x t)), WGBary and WGBary×Grad use Sinkhorn λ=0.001, selected via a faithfulness-sparsity-robustness tradeoff.

Experimental results show that WassersteinGrad outperforms gradient-based baseline methods in interpretability metrics. In both single-step and autoregressive forecasting settings, WassersteinGrad demonstrates superior interpretability, particularly in spatial consistency and physical localization. Specific data indicate that at noise levels σ<0.4, model performance degrades by less than 1%, yet attribution centroids advect by an average of 10-15 km. Additionally, in autoregressive forecasting, WassersteinGrad absorbs geometric distortion accumulation through geometric consensus, exhibiting robustness superior to other baseline methods.

WassersteinGrad has direct application potential in high-stakes environments such as weather forecasting. By enhancing the transparency of model predictions, WassersteinGrad helps build trust in AI systems, particularly in high-risk areas such as weather forecasting. Additionally, the geometric consensus mechanism offers new insights for explainable AI research in other domains. In the future, this method could also be applied to predictions in other dynamic physical fields, such as oceanography and atmospheric sciences, helping scientists better understand and explain complex natural phenomena.

WassersteinGrad's computational complexity on high-dimensional grids is relatively high, despite optimization through entropic regularization and Sinkhorn iterations. Further improvements in computational efficiency are needed. Additionally, the generalizability of this method across different predictive models and spatial domains remains to be validated, particularly in other physical fields with complex dynamics. The current noise injection strategy still uses white Gaussian noise, lacking physical information. Future work should explore more physically meaningful perturbation strategies.