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.
Figures & tables
Fig. 1: (a) SmaAT-QMix-UNet architecture: Rectangles represent feature maps, with height indicating spatial resolution and width the channel dimension. MixConv blocks are used in the last two encoder levels and the first decoder stage, while a VQ layer discretizes the B×18×18×512 bottleneck tensor. (b) Vector-quantization module: Latent features are flattened, each 512-D vector is assigned to its nearest codebook entry ( K=32 ), and reshaped into a quantized feature map. Training optimizes the combined codebook and β -weighted commitment losses.
Fig. 2: Mixed depthwise convolution (MixConv). The input tensor is split in two disjoint groups. First group is processed by a 3x3 depthwise convolution and the second group by a 5x5 depthwise convolution. The two outputs are then concatenated along the channel dimension.
Model
Parameters
Inference Time (ms)
MSE (px)
Precision
Recall
Accuracy
F1 Score
Persistence
–
–
0.0248
0.678
0.643
0.756
0.660
SmaAT-UNet
4 M
45
0.0122
0.730
0.850
0.829
0.786
SmaAT-Q-UNet
4 M
41
0.0119
0.748
0.820
0.832
0.782
SmaAT-Mix-UNet
2.5 M
42
0.0129
0.670
0.866
0.794
0.756
SmaAT-QMix-UNet
2.5 M
39
0.0120
0.763
0.812
0.838
0.787
TABLE I: Performance comparison at 30-min lead time on the NL-50 test set, including persistence, SmaAT-UNet, and proposed models, with model size and inference time reported. Best values are in bold
Fig. 3: Comparison of predictions generated by different models. The SmaAT-QMix-UNet model shows better alignment with the ground truth.
Fig. 4: (a) UMAP visualization of encoder feature vectors before and after vector quantization in SmaAT-QMix-UNet, where grey points denote pre-VQ representations and colored points indicate their assigned codewords. (b) Hyperparameter tuning results for the VQ module, showing validation performance across 16 combinations of codebook size and commitment cost, with K=32 and β=0.75 achieving the best performance.
Fig. 5: Heatmaps generated with Grad-CAM for SmaAT-QMix-UNet, showing activation regions across the five encoder and four decoder levels, including responses from the convolutional blocks (DoubleDSC or MixConv) and the CBAM modules in the encoder.