A Framework for Explainable AI in Weather Forecasting: Diagnosing Deep Learning Models via Gradient-Based Attributions

Each day, potentially critical decisions made by governments and organizations depend on accurate weather forecasts, determining whether to evacuate for a storm or simply to carry an umbrella. In this context, Deep Learning (DL) models are becoming a popular and computationally efficient alternative to traditional Numerical Weather Prediction (NWP) models, offering the potential to capture complex data patterns which may be missed using physical explicit equations (Lam et al., 2023). However, their opaque (black-box) nature remains a barrier to operational trust.

Explainable AI (XAI) aim to address this opacity by revealing the decision process behind predictions. Indeed, classical XAI techniques reveal when DL models rely on spurious correlations rather than causal physical mechanisms to deliver predictions (Geirhos et al., 2020). However, their direct application to meteorological data often yields attribution maps that are noisy (Kim et al., 2019) and difficult to interpret due to their high dimensionality. It additionally remains unclear whether these tools can consistently identify the complex physical drivers inherent in NWP (Bommer et al., 2024).

Based on previous works (Bommer et al., 2024; Kim et al., 2023; Yang et al., 2024), we establish a framework to generate compact and interpretable explanations of local weather forecast predictions obtained using deep neural networks. These explanations build on the output of gradient-based methods like VanillaGrad and SmoothGrad (Smilkov et al., 2017), which are scalable to high-dimensional data. More specifically, our framework first allows for targeted analysis by selecting a region of interest (e.g., Paris area) and a target variable (e.g., accumulated precipitation). It therefore answers the question: "Why did the neural network predict this feature at this location?" To do so, it first computes dense attribution maps with respect to all input variables (e.g., wind components at varying altitudes). Traditionally, bounding boxes are used to define the region of importance in these maps (Kim et al., 2023). However, they are unable to provide detailed directional information. We therefore propose in our framework to determine regions of importance using "confidence ellipses" that summarize the center, main directions, and importance of the most concentrated regions. Unlike bounding boxes, the representation of these ellipses, with the raw attribution maps as a background, provides rich and easily interpretable information regarding the directionality and spatial spread of the model's focus.

Preliminary results on the hybrid transformer-convolutional-based model UNETR++ (Shaker et al., 2024) trained and tested on the TITAN dataset from Météo-France (comprising hourly surface and vertical profiles of wind, temperature, and geopotential over metropolitan France) demonstrate our framework's pertinence for explaining predictions from deep neural networks. We were able to verify that different trained models successfully capture the vertical hierarchy of atmospheric variables, evidenced by an effective receptive field that expands with increasing altitude. More interestingly, our framework allowed us to identify systematic biases learned during training that correlate with known physical occurrences. These findings serve as a foundational step for future work on developing novel explainability methods to detect whether trained models capture complex physical mechanisms.