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

2026-04-24

AI Summary

Existing gradient-based attribution methods suffer from blurred attribution maps in dynamic physical fields—such as numerical weather prediction—due to geometric displacements, undermining their reliability for interpretation. This work proposes WassersteinGrad, an attribution aggregation framework grounded in entropy-regularized Wasserstein barycenters, which aligns attribution maps generated from multiple perturbations via optimal transport, thereby overcoming the limitations of conventional pointwise averaging. WassersteinGrad reveals, for the first time, that attribution discrepancies primarily stem from geometric shifts rather than amplitude noise, enabling the construction of geometrically consistent attribution consensus. Evaluated on regional meteorological data using an autoregressive neural weather forecasting model validated by meteorologists, WassersteinGrad consistently outperforms existing gradient-based baselines in both single-step and multi-step forecasts, yielding significantly sharper and more stable explanations.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

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.

Problem

• dynamic physical fields
• attribution maps
• geometric displacement
• feature attribution
• autoregressive forecasting

Innovation

• WassersteinGrad
• attribution map alignment
• dynamic physical fields
• autoregressive forecasting
• explainable AI