SmaAT-QMix-UNet:用于降雨预报的参数高效向量量化 UNet
机器学习
2026-03-31 v2 人工智能
摘要
天气预报支持关键的社会经济活动,并补充环境保护,但操作性数值天气预报(NWP)系统仍计算密集,因而不够高效。与此同时,近期基于深度学习的数据驱动模型在 nowcasting 任务中显示出前景。本文提出了 SmaAT-QMix-UNet,这是 SmaAT-UNet 的一种增强版本,引入了两个关键创新:在编码器-解码器桥接处添加向量量化(VQ)瓶颈,以及用混合深度可分离卷积(MixConv)取代选定的编码器和解码器块。这些增强既减少了模型的大小,又提高了其 nowcasting 性能。我们在 2016-2019 年的荷兰雷达降雨数据集上训练和评估 SmaAT-QMix-UNet,预测降雨 30 分钟后。对比三种配置:仅使用 VQ、仅使用 MixConv、以及完整的 SmaAT-QMix-UNet。Grad-CAM 热力图突出显示了每个 nowcast 影响的区域,而 UMAP 嵌入图展示了 VQ 层如何对编码器输出进行聚类。SmaAT-QMix-UNet 的源代码已公开于 GitHub: https://github.com/nstavr04/MasterThesisSnellius。
关键词
引用
@article{arxiv.2603.21879,
title = {SmaAT-QMix-UNet: A Parameter-Efficient Vector-Quantized UNet for Precipitation Nowcasting},
author = {Nikolas Stavrou and Siamak Mehrkanoon},
journal= {arXiv preprint arXiv:2603.21879},
year = {2026}
}
备注
6 pages, 5 figures