Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Axial-UNet: A Neural Weather Model for Precipitation Nowcasting

Sumit Mamtani, Maitreya Sonawane

TL;DR

This work tackles short-term precipitation nowcasting with a lightweight Axial-UNet that fuses a UNet encoder–decoder with axial-attention to capture long-range row/column interactions while conditioning on multiple past radar frames. The model delivers fixed-lead-time predictions with modest compute and outperforms ConvLSTM, cGANs, and a plain UNet on a cleaned HKO-7 subset, achieving state-of-the-art pixel-fidelity metrics (PSNR up to 47.7, SSIM up to 0.994). The authors discuss limitations of PSNR/SSIM as sole metrics and outline avenues for meteorology-centric evaluation, probabilistic outputs, and higher-resolution data. Overall, Axial-UNet offers a scalable, efficient approach for resource-constrained, real-time precipitation nowcasting, with clear paths toward richer, physics-informed assessments.

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.

Axial-UNet: A Neural Weather Model for Precipitation Nowcasting

TL;DR

This work tackles short-term precipitation nowcasting with a lightweight Axial-UNet that fuses a UNet encoder–decoder with axial-attention to capture long-range row/column interactions while conditioning on multiple past radar frames. The model delivers fixed-lead-time predictions with modest compute and outperforms ConvLSTM, cGANs, and a plain UNet on a cleaned HKO-7 subset, achieving state-of-the-art pixel-fidelity metrics (PSNR up to 47.7, SSIM up to 0.994). The authors discuss limitations of PSNR/SSIM as sole metrics and outline avenues for meteorology-centric evaluation, probabilistic outputs, and higher-resolution data. Overall, Axial-UNet offers a scalable, efficient approach for resource-constrained, real-time precipitation nowcasting, with clear paths toward richer, physics-informed assessments.

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.
Paper Structure (17 sections, 8 equations, 2 figures, 2 tables, 2 algorithms)

This paper contains 17 sections, 8 equations, 2 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Summary and gap.
  4. Method: UNet with Axial Attention (Axial-UNet)
  5. Backbone encoder--decoder
  6. Axial attention blocks
  7. Baselines
  8. ConvLSTM
  9. cGANs
  10. Generator
  11. Discriminator
  12. Experimental Results
  13. Datasets
  14. Implementation Details
  15. Results and Comparisons
  16. ...and 2 more sections

Figures (2)

  • Figure 1: Architecture of Axial-UNet (UNet with axial attention). Part (1) is our UNet model that encompasses an encoder and decoder using Conv2D. Part (2) is the transformer metnet that attends every row and column for each location (pixel), hence an axial attention ho2019axial. Part (3) is the distribution we get as our final output, of which we take the mean to make a final prediction for T$^{th}$ minute
  • Figure 2: Training and validation loss curves for (a,b) cGANs, (c,d) UNet, and (e,f) Axial-UNet.