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

2026-03-23

AI Summary

This work proposes SmaAT-QMix-UNet to address the high computational cost of traditional numerical weather prediction and its inefficiency in supporting short-term precipitation forecasting. The model introduces a vector quantization (VQ) bottleneck at the U-Net encoder-decoder bridge and replaces selected encoder-decoder modules with mixed-depthwise convolution (MixConv), augmented with attention mechanisms. This design significantly reduces parameter count while improving 30-minute precipitation forecast accuracy. Experiments on Dutch radar rainfall data from 2016–2019 demonstrate that the synergistic integration of VQ and MixConv enhances model performance beyond existing baselines. Furthermore, interpretability is provided through Grad-CAM and UMAP analyses, offering insights into the model’s decision-making process.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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 \footnote{\href{https://github.com/nstavr04/MasterThesisSnellius}{https://github.com/nstavr04/MasterThesisSnellius}}.

Problem

• precipitation nowcasting
• parameter efficiency
• computational intensity
• deep data-driven models
• weather forecasting

Innovation

• Vector Quantization
• MixConv
• Parameter-Efficient
• Precipitation Nowcasting
• UNet