# Axial-UNet: A Neural Weather Model for Precipitation Nowcasting

arXiv: https://arxiv.org/abs/2504.19408 · HTML: https://mlsanity.com/paper/2504.19408/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 27 references that are in the map.
- First posted on arXiv 28 Apr 2025.

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Conditional Generative Adversarial Nets](https://mlsanity.com/paper/1411.1784/) (2014)
- [Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting](https://mlsanity.com/paper/1506.04214/) (2015)
- [Deep Learning in Neural Networks: An Overview](https://mlsanity.com/paper/1404.7828/) (2014)
- [Unsupervised Learning of Video Representations using LSTMs](https://mlsanity.com/paper/1502.04681/) (2015)
- [Efficient Processing of Deep Neural Networks: A Tutorial and Survey](https://mlsanity.com/paper/1703.09039/) (2017)
- [Skillful Precipitation Nowcasting using Deep Generative Models of Radar](https://mlsanity.com/paper/2104.00954/) (2021)
- [MetNet: A Neural Weather Model for Precipitation Forecasting](https://mlsanity.com/paper/2003.12140/) (2020)
- [Transformer visualization via dictionary learning: contextualized embedding as a linear superposition of transformer factors](https://mlsanity.com/paper/2103.15949/) (2021)
- [An Introduction to Deep Generative Modeling](https://mlsanity.com/paper/2103.05180/) (2021)
- [Understanding LSTM -- a tutorial into Long Short-Term Memory Recurrent Neural Networks](https://mlsanity.com/paper/1909.09586/) (2019)
- [SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture](https://mlsanity.com/paper/2007.04417/) (2020)
- [Deep learning for precipitation nowcasting: A survey from the perspective of time series forecasting](https://mlsanity.com/paper/2406.04867/) (2024)
- [Traffic map prediction using UNet based deep convolutional neural network](https://mlsanity.com/paper/1912.05288/) (2019)
- [Distributed Deep Learning for Precipitation Nowcasting](https://mlsanity.com/paper/1908.10964/) (2019)
- [Encoder-Decoder Neural Architecture Optimization for Keyword Spotting](https://mlsanity.com/paper/2106.02738/) (2021)
- [Precipitation Nowcasting: Leveraging bidirectional LSTM and 1D CNN](https://mlsanity.com/paper/1810.10485/) (2018)

## Abstract (arXiv)

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.