SSA-UNet: Advanced Precipitation Nowcasting via Channel Shuffling

TL;DR

Addresses the need for accurate yet resource-efficient precipitation nowcasting. Proposes SSA-UNet, a Shuffle Attention- and channel-shuffling-enhanced UNet variant, with shuffled depthwise separable convolutions to reduce parameters. Evaluations on KNMI Netherlands precipitation maps and the French cloud cover radar dataset show SSA-UNet and SSA-UNet(reduced) achieving comparable or improved $MSE$ and related metrics across 1-, 6-, and 12-step forecasts, with parameter counts around $3.8$M and $3.1$M respectively. Grad-CAM explainability provides insights into regions driving predictions, and the approach enables faster inference suitable for mobile/edge deployment while reducing environmental impact.

Abstract

Weather forecasting is essential for facilitating diverse socio-economic activity and environmental conservation initiatives. Deep learning techniques are increasingly being explored as complementary approaches to Numerical Weather Prediction (NWP) models, offering potential benefits such as reduced complexity and enhanced adaptability in specific applications. This work presents a novel design, Small Shuffled Attention UNet (SSA-UNet), which enhances SmaAt-UNet's architecture by including a shuffle channeling mechanism to optimize performance and diminish complexity. To assess its efficacy, this architecture and its reduced variant are examined and trained on two datasets: a Dutch precipitation dataset from 2016 to 2019, and a French cloud cover dataset containing radar images from 2017 to 2018. Three output configurations of the proposed architecture are evaluated, yielding outputs of 1, 6, and 12 precipitation maps, respectively. To better understand how this model operates and produces its predictions, a gradient-based approach called Grad-CAM is used to analyze the outputs generated. The analysis of heatmaps generated by Grad-CAM facilitated the identification of regions within the input maps that the model considers most informative for generating its predictions. The implementation of SSA-UNet can be found on our Github\footnote{\href{https://github.com/MarcoTurzi/SSA-UNet}{https://github.com/MarcoTurzi/SSA-UNet}}

Figures (8)

• Figure 1: (a) Visual representation of SSA-UNet. The rectangles indicate the data dimensions, with height representing the spatial dimension and width representing the channel dimension. Arrows and their color specify the type of operation applied to the input data. (b) Shows the composition of the shuffled double convolution block. In this block, the input passes through a sequence consisting of a shuffled depthwise separable convolution, followed by batch normalization, ReLU activation, a depthwise separable convolution, and another round of batch normalization and ReLU activation. (c) Illustrates the sequence of operations that are performed by the depthwise separable convolution and the shuffled depthwise separable convolution.
• Figure 2: Illustration of the impact of a grouped pointwise convolution followed by channel shuffling on the input channels. The grouped convolution processes channel groups independently, while channel shuffling redistributes channels to enable cross-group information exchange.
• Figure 3: Comparison of predictions generated by SmaAt-UNet, SSA-UNet and SSA-UNet(reduced) with the ground truth. The predictions demonstrate the differences in performance between the architectures, where the SSA-UNet model shows better alignment with the true image.
• Figure 4: Visualization of the hyperparameter tuning results for SSA-UNet, showing the MSE values (in logarithmic scale) achieved by each tested configuration. Each configuration is defined as a list (G) of Shuffle Attention's group size values corresponding to each encoder level.
• Figure 5: Performance metrics over time for SmaAt-UNet, SSA-UNet and SSA-UNet(reduced) in the 12-output configuration. The figure shows four plots representing the evolution of the MSE (in logarithmic scale), Accuracy, Precision, and F1 Score at 5-minute intervals over a 60-minute period.