用向量量化和混合卷积压缩模型，提升降水短时预报精度与效率
SmaAT-QMix-UNet: A Parameter-Efficient Vector-Quantized UNet for Precipitation Nowcasting
- 在编码器-解码器间加入向量量化瓶颈，减少参数量
- 用混合深度卷积替代部分模块，提升预测准确率至RMSE 0.85
- 代码开源，适合做气象预测轻量化模型的研究者参考
天气预报对经济社会活动和环境保护至关重要，但现有数值天气预报（NWP）系统计算成本高，难以满足特定应用需求。近年来数据驱动的深度学习模型在短时预报任务中表现优异。本文提出SmaAT-QMix-UNet，是在SmaAT-UNet基础上的改进版本，引入两个关键创新：在编码器-解码器桥接处加入向量量化（VQ）瓶颈，以及用混合核深度卷积（MixConv）替换部分编码器和解码器模块。这些改进有效降低模型规模并提升预报性能。我们在荷兰雷达降水数据集（2016–2019年）上进行训练与评估，预测未来30分钟降水。对比三种配置：仅使用VQ、仅使用MixConv、以及完整模型。通过Grad-CAM热力图分析影响预报的关键区域，利用UMAP可视化量化码本的聚类结构。SmaAT-QMix-UNet源代码已公开于GitHub：https://github.com/nstavr04/MasterThesisSnellius。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。