PA-Net: Precipitation-Adaptive Mixture-of-Experts for Long-Tail Rainfall Nowcasting

TL;DR

PA-Net introduces a rainfall-adaptive Transformer framework that dynamically allocates computational resources to improve long-term rainfall nowcasting, especially for rare heavy rainfall events, by focusing on societal impact-critical data.

Contribution

The paper presents PA-Net, a novel Transformer with a precipitation-adaptive mixture-of-experts and dual-axis attention, effectively modeling long-tail rainfall distribution with improved accuracy.

Findings

Significant performance gains in heavy-rain regimes

Effective handling of long spatiotemporal contexts

Enhanced focus on rare extreme rainfall events

Abstract

Precipitation nowcasting is vital for flood warning, agricultural management, and emergency response, yet two bottlenecks persist: the prohibitive cost of modeling million-scale spatiotemporal tokens from multi-variate atmospheric fields, and the extreme long-tailed rainfall distribution where heavy-to-torrential events -- those of greatest societal impact -- constitute fewer than 0.1% of all samples. We propose the Precipitation-Adaptive Network (PA-Net), a Transformer framework whose computational budget is explicitly governed by rainfall intensity. Its core component, Precipitation-Adaptive MoE (PA-MoE), dynamically scales the number of activated experts per token according to local precipitation magnitude, channeling richer representational capacity toward the rare yet critical heavy-rainfall tail. A Dual-Axis Compressed Latent Attention mechanism factorizes spatiotemporal attention with convolutional reduction to manage massive context lengths, while an intensity-aware training protocol progressively amplifies learning signals from extreme-rainfall samples. Experiment on ERA5 demonstrate consistent improvements over state-of-the-art baselines, with particularly significant gains in heavy-rain and rainstorm regimes.