中文

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