PA-Net: Precipitation-Adaptive Mixture-of-Experts for Long-Tail Rainfall Nowcasting
Authors: Xinyu Xiao, Sen Lei, Eryun Liu, Shiming Xiang, Hao Li, Cheng Yuan, Yuan Qi, Qizhao Jin
Organizations: Fudan University · the School of Information Science and Technology, Southwest Jiaotong University · Zhejiang University · the Institute of Automation, Chinese Academy of Sciences · the Public Meteorological Service Center, China Meteorological Administration
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.
Figures & tables
Intensity
RL
LR
MR
HR
RS
Precipitation(mm)
[0, 0.1)
[0.1, 4.0)
[4.0, 13.0)
[13.0, 25.0)
[25.0, ∞ )
ERA5
80.76%
18.22%
0.96%
0.06%
0.006%
WeatherBench
83.59%
15.94%
0.44%
0.02%
0.003%
TABLE I: Empirical frequency of precipitation intensity categories across the ERA5 and WeatherBench archives. Categories are ordered by ascending severity: Rainless (RL), Light Rain (LR), Moderate Rain (MR), Heavy Rain (HR), and Rainstorm (RS).
Fig. 1: Architecture of the Precipitation-Adaptive Network (PA-Net): a Transformer backbone equipped with Dual-Axis Compressed Latent Attention (DACLA) for scalable spatiotemporal modeling and Precipitation-Adaptive MoE (PA-MoE) for rainfall-intensity-driven expert allocation, jointly enabling enhanced nowcasting skill from light drizzle through torrential downpours.
Fig. 2: Overview of the Precipitation-Adaptive MoE (PA-MoE): tokens associated with intense rainfall activate a larger expert ensemble, while rainless and light rain are routed to fewer experts, thereby concentrating representational capacity on the rare yet high-impact tail of the precipitation distribution.
Model
Params(M)
MACs(G)
TS ↑
IoU ↑
ConvLSTM
14.79
463.74
20.68
37.83
ConvGRU
17.19
604.8
20.41
38.03
TrajGRU
12.72
720.64
19.50
36.09
PredRNN
25.84
1103.64
17.95
36.58
PFST
40.65
472.97
18.95
38.05
SCCN
48.1
2990.48
19.77
38.01
TABLE II: Category-averaged performance on the ERA5 dataset. Bold entries mark the highest scores.
Model
RL
LR
MR
HR
RS
PFST
88.13
50.12
22.88
17.27
12.04
SCCN
87.13
51.42
20.16
17.5
13.85
SimVP
87.28
50.97
25.89
21.02
14.04
SimVP+gSTA
87.23
51.71
25.13
19.36
13.88
TAU
88.41
52.30
26.33
19.31
12.58
PA-Net
88.89
57.78
32.45
24.37
23.41
TABLE III: Per-category IoU (%) on the ERA5 dataset stratified by rainfall intensity.
Metric
Method
Prediction time (hours)
0 ∼ 1
1 ∼ 2
2 ∼ 3
3 ∼ 4
4 ∼ 5
5 ∼ 6
IoU ↑
Pers.
31.10
28.64
26.98
25.78
24.89
24.22
Clim.
13.36
13.36
13.36
13.36
13.36
13.36
W-Clim.
16.68
16.68
16.68
16.68
16.68
16.68
ConvLSTM
40.52
40.10
38.84
37.39
35.85
34.29
ConvGRU
40.01
39.71
38.83
37.75
36.56
35.31
TABLE IV: Hourly IoU (%) and TS (%) on the ERA5 dataset across six one-hour forecast windows.
Fig. 3: Six-hour precipitation forecast visualization on the ERA5 dataset. Columns from up to down: PA-Net and ground truth. Orange colors indicate heavier rainfall.
PA-MoE
DACLA
IATP
IoU ↑
TS ↑
40.93
21.04
✓
43.08
22.79
✓
✓
✓
45.38
24.62
TABLE V: Component ablation on the ERA5 dataset. PA-MoE: Precipitation-Adaptive MoE; DACLA: Dual-Axis Compressed Latent Attention; IATP: Intensity-Aware Training Protocol.
School of Automation, Southeast University · Nanjing XinDa Institute of Meteorological Science and Technology · Beijing Leninainfo Technology Co., Ltd. +5