Axial-UNet: A Neural Weather Model for Precipitation Nowcasting

2025/04/28 by Sumit Mamtani, Mamtani, Sumit, Maitreya Sonawane +1
Earth and Planetary Sciences · Computer Science · #Precipitation Measurement and Analysis #Meteorological Phenomena and Simulations #Image Enhancement Techniques

paper · pdf · doi:10.48550/arxiv.2504.19408

Abstract

Accurately predicting short-term precipitation is critical for weather-sensitive applications such as disaster management, aviation, and urban planning. Traditional numerical weather prediction can be computationally intensive at high resolution and short lead times. In this work, we propose a lightweight UNet-based encoder-decoder augmented with axial-attention blocks that attend along image rows and columns to capture long-range spatial interactions, while temporal context is provided by conditioning on multiple past radar frames. Our hybrid architecture captures both local and long-range spatio-temporal dependencies from radar image sequences, enabling fixed lead-time precipitation nowcasting with modest compute. Experimental results on a preprocessed subset of the HKO-7 radar dataset demonstrate that our model outperforms ConvLSTM, pix2pix-style cGANs, and a plain UNet in pixel-fidelity metrics, reaching PSNR 47.67 and SSIM 0.9943. We report PSNR/SSIM here; extending evaluation to meteorology-oriented skill measures (e.g., CSI/FSS) is left to future work. The approach is simple, scalable, and effective for resource-constrained, real-time forecasting scenarios.

Citations

• Deep learning for precipitation nowcasting: A survey from the perspective of time series forecasting
• Generative Adversarial Networks
• Encoder-Decoder Neural Architecture Optimization for Keyword Spotting
• Skilful precipitation nowcasting using deep generative models of radar
• Transformer visualization via dictionary learning: contextualized embedding as a linear superposition of transformer factors
• An Introduction to Deep Generative Modeling
• Identifying Opportunities for Skillful Weather Prediction with Interpretable Neural Networks
• SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture
• MetNet: A Neural Weather Model for Precipitation Forecasting
• Traffic map prediction using UNet based deep convolutional neural network
• Understanding LSTM -- a tutorial into Long Short-Term Memory Recurrent Neural Networks
• Precipitation Nowcasting: Leveraging bidirectional LSTM and 1D CNN
• Attention Is All You Need
• Efficient Processing of Deep Neural Networks: A Tutorial and Survey
• Efficient Processing of Deep Neural Networks: A Tutorial and Survey
• Image-to-Image Translation with Conditional Adversarial Networks
• Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting
• U-Net: Convolutional Networks for Biomedical Image Segmentation
• Unsupervised Learning of Video Representations using LSTMs
• Conditional Generative Adversarial Nets
• Deep Learning in Neural Networks: An Overview

Related

• Machine Learning for Precipitation Nowcasting from Radar Images
• SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture
• Data-driven Precipitation Nowcasting Using Satellite Imagery
• SSA-UNet: Advanced Precipitation Nowcasting via Channel Shuffling
• WF-UNet: Weather Fusion UNet for Precipitation Nowcasting