SmaAT-QMix-UNet: A Parameter-Efficient Vector-Quantized UNet for Precipitation Nowcasting

TL;DR

This paper introduces SmaAT-QMix-UNet, a compact, efficient deep learning model for precipitation nowcasting that incorporates vector quantization and mixed kernel convolutions, achieving improved performance and interpretability.

Contribution

It presents a novel UNet variant with vector quantization and mixed convolutions, reducing size and enhancing nowcasting accuracy over previous models.

Findings

Model reduces size while maintaining performance.

Full SmaAT-QMix-UNet outperforms configurations with only VQ or MixConv.

Saliency maps and UMAP embeddings provide interpretability insights.

Abstract

Weather forecasting supports critical socioeconomic activities and complements environmental protection, yet operational Numerical Weather Prediction (NWP) systems remain computationally intensive, thus being inefficient for certain applications. Meanwhile, recent advances in deep data-driven models have demonstrated promising results in nowcasting tasks. This paper presents SmaAT-QMix-UNet, an enhanced variant of SmaAT-UNet that introduces two key innovations: a vector quantization (VQ) bottleneck at the encoder-decoder bridge, and mixed kernel depth-wise convolutions (MixConv) replacing selected encoder and decoder blocks. These enhancements both reduce the model's size and improve its nowcasting performance. We train and evaluate SmaAT-QMix-UNet on a Dutch radar precipitation dataset (2016-2019), predicting precipitation 30 minutes ahead. Three configurations are benchmarked: using only VQ, only MixConv, and the full SmaAT-QMix-UNet. Grad-CAM saliency maps highlight the regions influencing each nowcast, while a UMAP embedding of the codewords illustrates how the VQ layer clusters encoder outputs. The source code for SmaAT-QMix-UNet is publicly available on GitHub: https://github.com/nstavr04/MasterThesisSnellius.