Accurate Precipitation Forecast by Efficiently Learning from Massive Atmospheric Variables and Unbalanced Distribution

Abstract: Short-term (0-24 hours) precipitation forecasting is highly valuable to socioeconomic activities and public safety. However, the highly complex evolution patterns of precipitation events, the extreme imbalance between precipitation and non-precipitation samples, and the inability of existing models to efficiently and effectively utilize large volumes of multi-source atmospheric observation data hinder improvements in precipitation forecasting accuracy and computational efficiency. To address the above challenges, this study developed a novel forecasting model capable of effectively and efficiently utilizing massive atmospheric observations by automatically extracting and iteratively predicting the latent features strongly associated with precipitation evolution. Furthermore, this study introduces a 'WMCE' loss function, designed to accurately discriminate extremely scarce precipitation events while precisely predicting their intensity values. Extensive experiments on two datasets demonstrate that our proposed model substantially and consistently outperforms all prevalent baselines in both accuracy and efficiency. Moreover, the proposed forecasting model substantially lowers the computational cost required to obtain valuable predictions compared to existing approaches, thereby positioning it as a milestone for efficient and practical precipitation forecasting.

• The paper introduces a latent-space iterative forecasting framework and hyperparameter-free WMCE loss that improves rare-event precipitation detection across USA and Hubei datasets.
• The model achieves USA POD/CSI/HSS scores of 0.471/0.265/0.355 and Hubei scores of 0.410/0.216/0.325, outperforming leading baselines, especially for intense precipitation.
• The approach forecasts in a 16-channel latent space and infers in 0.385 seconds, but remains limited by systematic overprediction and coarse ERA5 input resolution.

Motivation and problem setting

Short-term (0–24 hour) precipitation forecasting occupies an awkward position between radar-based nowcasting (0–2 hours) and medium-range NWP. The paper identifies three structural deficiencies in existing deep learning approaches to this task: (1) most models target ultra-short-term convection nowcasting from single-source radar data, which is insufficient once lead times extend beyond a few hours; (2) multi-source models such as the MetNet family predict either a single lead time or all lead times simultaneously, entailing high task complexity and heavy computational cost; and (3) iterative prediction schemes that do operate over multiple steps typically act in the full physical variable space, forcing the model to learn evolution dynamics for every input variable even though only a subset is strongly associated with precipitation.

A further obstacle is the extreme class imbalance of precipitation data. In the datasets used here, samples above 0.2 mm/h constitute only 6.69% of the USA dataset, and reflectivity above 20 dBZ only 1.72% of the Hubei dataset. Pixel-wise MSE/MAE losses therefore optimize almost exclusively on non-precipitation pixels, producing systematic under-detection. Notably, the authors report that importance sampling—a remedy used by MetNet-2—actually degrades performance in their experiments, because non-precipitation sequences carry precursor information relevant at longer lead times. This is a mildly contrarian claim relative to prior practice and motivates a loss-based rather than sampling-based solution.

Model architecture

The proposed framework follows an encode–iteratively-predict–project design operating in a low-dimensional latent space of dimension 16×64×64.

Encoder: each modality (MRMS QPE plus radar, ERA5 reanalysis with 243 channels spanning 37 pressure levels and surface variables, and FY4B satellite imagery where available) passes through four modality-specific multi-scale residual blocks with parallel 3×3 and 5×5 convolutions, followed by spatial alignment, channel concatenation, and a final convolution. Features are not shared across modalities.

Latent dynamics predictor (LPM): a ViT-based module takes latent features from two consecutive time steps together with time embeddings and a constant embedding, and predicts the next-step latent feature. To curb error accumulation over 24 iterative steps, the authors adopt and modify the HTA scheme (from Pangu-Weather): three LPM units predict at intervals of 1, 2, and 4 steps, with the modification that two time steps at differing intervals are supplied during training and five historical steps during inference, so that short-interval LPMs handle early lead times and the interval-4 LPM handles later ones.

Projector and reconstructor: a shared projector maps predicted latent features at all lead times in parallel to precipitation rate or near-surface reflectivity. A lightweight per-modality reconstructor reconstructs the original physical variables from the latent features; this prevents the encoder from collapsing onto features that are merely easy to predict over the short training horizon (only two unrolled iterations during training).

The WMCE loss

The core methodological contribution is the WMCE loss, which combines three terms. For pixels labeled as precipitation ($Y \ge \tau$, with $\tau$ = 0.2 mm/h or 20 dBZ), both a weighted MAE term and a weighted cross-entropy term are applied, where the weight $\log(e + y^o)$ grows with the raw intensity, ensuring heavier penalties for missed heavy precipitation. For non-precipitation pixels, only an unweighted binary cross-entropy term is used; no regression loss is applied there. Two design points deserve emphasis:

• Applying MAE only to precipitation pixels reduces task complexity and avoids letting abundant zero-valued pixels dominate optimization. The authors argue this is only sound because their latent-space iteration decouples each lead-time prediction from earlier predicted intensities; they explicitly note the loss would be inappropriate for models iterating in physical space.
• Unlike focal loss, WMCE has no tunable hyperparameters. Ablations show that weighted MAE plus focal loss yields high POD but near-zero CSI/HSS due to persistent false alarms, and exhibits unstable training across epochs, whereas WMCE improves monotonically.

Experimental results

Experiments cover two heterogeneous datasets: a USA dataset (2015–2018, MRMS QPE at 6 km, hourly, 24-hour horizon) and a Hubei dataset (2022–2023, near-surface radar reflectivity at 2 km, 15-minute resolution, 6-hour horizon). Baselines include ConvGRU, SimVP, SimVP_HTA, MetNet-2, and EarthFarseer, evaluated with POD, CSI, and HSS averaged over 24 lead times.

On the USA dataset, the proposed model achieves POD/CSI/HSS of 0.471/0.265/0.355 at the 0.2 mm/h threshold versus 0.369/0.181/0.254 for the best baseline, corresponding to reported average improvements of roughly 60.6%, 36.6%, and 34.0% over the second-best method across thresholds. On the Hubei dataset it attains POD/CSI/HSS of 0.410/0.216/0.325 at 20 dBZ, with gains of about 58.4%, 20.4%, and 20.2%. The advantage widens at higher thresholds—for example, POD at 30 dBZ reaches 0.281 versus 0.106 for SimVP—indicating that the weighting scheme specifically benefits rare, intense events. MetNet-2 degrades sharply on the Hubei dataset, which the authors attribute to the nonlinear Z-R relationship between reflectivity and rain rate and to coarser interval partitioning; this supports the claim of robustness across target modalities, though it also reflects a baseline trained under conditions mismatched to the task.

Qualitative cases show the model reproducing initiation, eastward propagation, and dissipation of precipitation bands that baselines dissipate prematurely or blur heavily. Supplementary experiments substituting operational HRRR analyses for ERA5 confirm the approach remains viable without reanalysis latency, supporting operational applicability.

Ablation findings

Three ablations substantiate the design choices. First, replacing pixel-wise MAE with progressively stronger formulations (log-weighted MAE, CE, weighted CE) yields monotonic metric improvements, with the largest gains at the ≥8 mm/h threshold. Second, iterative prediction in physical space collapses within a few lead hours—even with ERA5 added—and latent-space iteration with ERA5 dominates all alternatives, confirming both the latent-space hypothesis and the value of multi-source inputs. Third, applying MetNet-style importance sampling reduces mid-to-late-stage POD and worsens metrics at ≥8 mm/h, consistent with the precursor-signal argument.

An Integrated Gradients attribution analysis shows that upper-air variables (geopotential, wind components, vertical velocity) dominate at longer lead times while near-surface radar attribution decays with altitude and lead time, a pattern the authors connect to mid-latitude synoptic dynamics; this lends physical plausibility to the learned representations.

Efficiency

Despite having the largest parameter count among compared models (173M vs. 103M for EarthFarseer), the model trains in 10.9 hours (comparable to SimVP's 10.2) and infers in 0.385 seconds—the fastest of all methods, roughly 3× faster than ConvGRU and over 13× faster than EarthFarseer. The efficiency gain stems directly from iterating on 16-channel latent features rather than hundreds of physical variables. This jointly addresses accuracy and cost, which the authors frame as the main practical contribution.

Limitations and open questions

The paper concedes two limitations explicitly: predictions exhibit systematic overprediction (also visible qualitatively as inflated precipitation coverage), and the 0.25° ERA5 resolution is coarse relative to the 2–6 km prediction grids, likely capping achievable detail. Additional caveats bear on generality: training uses only two unrolled iterations with reconstruction losses compensating for longer horizons, and the claim that masking MAE on non-precipitation pixels is safe rests on the latent-space architecture rather than being validated more broadly. Whether the WMCE formulation transfers to probabilistic or ensemble forecasting, and how the overprediction bias can be calibrated, remain open.

Conclusion

This paper presents a latent-space iterative forecasting framework paired with a hyperparameter-free WMCE loss tailored to severely imbalanced precipitation data. Across two spatiotemporally distinct datasets and multiple target modalities, it delivers consistent improvements in event detection and intensity skill over strong baselines while achieving the lowest inference latency, supported by ablations isolating the contributions of the loss design, latent-space iteration, multi-source inputs, and the rejection of importance sampling. Its principal unresolved issues are overprediction and dependence on coarse reanalysis fields.