[AMS 2022] Explaining Complex 3D Atmospheric CNNs Using SHAP-Based Channel-wise XAI Techniques

Abstract In this research, we demonstrate the application of explainable artificial intelligence (XAI) techniques to investigate strategies learned by Convolutional Neural Networks (CNN)-based prediction models whose inputs are rasters composed of atmospheric and oceanic feature maps. The goal is to generate 3D explanations to reveal spatial-channel-wise features that the model is using. For example, to investigate the relative importance of near shore sea surface temperatures for the prediction of coastal fog. We present an extension of the now broadly used SHAP-based XAI technique, PartitionShap, that explains image data in terms of spatial superpixels. The proposed new version of SHAP, Channel-wise PartitionShap, evaluates the importance of superpixels within each of the channels separately. That is, a 4-channel model will have 4 separate intra-channel explanations. When gauging the relative importance of models with a large number inputs, here hundreds of feature maps, visualizing the results becomes a challenge. To gauge the relative importance of input volumes, for example if the model is looking at a gradient across channels, we are developing an interactive 3D visualization tool to investigate the explanations. Experiments are performed using a 384-channel fog prediction CNN (FogNet) by Kamangir et al. [1] and a tornado prediction CNN by Lagerquist et al. [2] with 4 feature maps. PartitionShap is an XAI technique based on the SHAP approach to calculating feature importance by approximating game-theoretic Shapley values. Each feature’s importance is based on their contribution to the output prediction. PartitionShap divides the input image into a grid of superpixels where each is considered a feature to explain. Since we are interested in intra-channel superpixels, our Channel-wise PartitionShap extension partitions first along the raster channels and then into superpixels within. We also apply SHAP using a group-based partitioning where each of five physics-based channel groups are considered features. All pixels within are permuted to test the contribution of the entire group. Most XAI techniques compatible with raster data were originally developed to explain image classification models. Many techniques permute entire pixels so that there is a single 2D explanation since the channel-wise effect was not considered. When channel-wise values are calculated, the explanation may be displayed as a set of 2D plots, one for each channel. However, we are interested in 3D explanations that may involve inter-channel regions. These are harder to visualize as separate 2D plots, so we developed a tool that displays the Channel-wise PartitionShap values as a 3D grid. The user can interact with the visualization to investigate the 3D explanation. We performed experiments using Channel-wise PartitionShap to explain two recent multi-channel models [1,2] and compare to other XAI methods that produce intra-channel explanations including Permutation Feature Importance, Saliency Maps, and Class Activation Maps. We show that the choice in grouping raster elements into features has a substantial effect on the explanation. Two methods, such as the group-based SHAP and Channel-wise PartitionShap may yield very different explanations. However, by understanding how the choice of feature partitioning changes the question posed to the model, domain expertise is used to interpret the seemingly disparate results. We find that XAI reveals strategies used by models that take advantage of 3D raster regions, such as the atmospheric vertical profile, for decision-making that agrees with forecaster knowledge. [1] Kamangir, H., Collins, W., Tissot, P., King, S. A., Dinh, H. T. H., Durham, N., & Rizzo, J. (2021). FogNet: A multiscale 3D CNN with double-branch dense block and attention mechanism for fog prediction. Machine Learning with Applications, 5, 100038. [2] Lagerquist, R., McGovern, A., Homeyer, C. R., Gagne II, D. J., & Smith, T. (2020). Deep learning on three-dimensional multiscale data for next-hour tornado prediction. Monthly Weather Review, 148(7), 2837-2861.