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

arXiv: https://arxiv.org/abs/2603.21879 · HTML: https://mlsanity.com/paper/2603.21879/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 30 references that are in the map.
- First posted on arXiv 23 Mar 2026.

## Builds on (most cited first, 25 of 25)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting](https://mlsanity.com/paper/1506.04214/) (2015)
- [VideoGPT: Video Generation using VQ-VAE and Transformers](https://mlsanity.com/paper/2104.10157/) (2021)
- [GhostNet: More Features from Cheap Operations](https://mlsanity.com/paper/1911.11907/) (2019)
- [Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model](https://mlsanity.com/paper/1706.03458/) (2017)
- [MetNet: A Neural Weather Model for Precipitation Forecasting](https://mlsanity.com/paper/2003.12140/) (2020)
- [PredRNN++: Towards A Resolution of the Deep-in-Time Dilemma in Spatiotemporal Predictive Learning](https://mlsanity.com/paper/1804.06300/) (2018)
- [MixConv: Mixed Depthwise Convolutional Kernels](https://mlsanity.com/paper/1907.09595/) (2019)
- [SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture](https://mlsanity.com/paper/2007.04417/) (2020)
- [Broad-UNet: Multi-scale feature learning for nowcasting tasks](https://mlsanity.com/paper/2102.06442/) (2021)
- [AA-TransUNet: Attention Augmented TransUNet For Nowcasting Tasks](https://mlsanity.com/paper/2202.04996/) (2022)
- [Deep Graph Convolutional Networks for Wind Speed Prediction](https://mlsanity.com/paper/2101.10041/) (2021)
- [Wind speed prediction using multidimensional convolutional neural networks](https://mlsanity.com/paper/2007.12567/) (2020)
- [Symbolic regression for scientific discovery: an application to wind speed forecasting](https://mlsanity.com/paper/2102.10570/) (2021)
- [Multistream Graph Attention Networks for Wind Speed Forecasting](https://mlsanity.com/paper/2108.07063/) (2021)
- [GA-SmaAt-GNet: Generative Adversarial Small Attention GNet for Extreme Precipitation Nowcasting](https://mlsanity.com/paper/2401.09881/) (2024)
- [Vector Quantisation for Robust Segmentation](https://mlsanity.com/paper/2207.01919/) (2022)
- [GPTCast: a weather language model for precipitation nowcasting](https://mlsanity.com/paper/2407.02089/) (2024)
- [Graph Dual-stream Convolutional Attention Fusion for Precipitation Nowcasting](https://mlsanity.com/paper/2401.07958/) (2024)
- [SAR-UNet: Small Attention Residual UNet for Explainable Nowcasting Tasks](https://mlsanity.com/paper/2303.06663/) (2023)
- [VQUNet: Vector Quantization U-Net for Defending Adversarial Atacks by Regularizing Unwanted Noise](https://mlsanity.com/paper/2406.03117/) (2024)

## Abstract (arXiv)

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.