PuYun-LDM: A Latent Diffusion Model for High-Resolution Ensemble Weather Forecasts

TL;DR

PuYun-LDM introduces a novel latent diffusion approach for high-resolution ensemble weather forecasting, addressing spectral heterogeneity and improving short-term forecast accuracy efficiently.

Contribution

The paper presents PuYun-LDM, a new latent diffusion model with a 3D-MAE encoder and Variable-Aware Masked Frequency Modeling, tailored for multivariate meteorological data.

Findings

Outperforms traditional ensemble methods at short lead times

Generates 15-day global forecasts in five minutes on a single GPU

Maintains comparable accuracy to ENS at longer forecast horizons

Abstract

Latent diffusion models (LDMs) suffer from limited diffusability in high-resolution (<=0.25{\deg}) ensemble weather forecasting, where diffusability characterizes how easily a latent data distribution can be modeled by a diffusion process. Unlike natural image fields, meteorological fields lack task-agnostic foundation models and explicit semantic structures, making VFM-based regularization inapplicable. Moreover, existing frequency-based approaches impose identical spectral regularization across channels under a homogeneity assumption, which leads to uneven regularization strength under the inter-variable spectral heterogeneity in multivariate meteorological data. To address these challenges, we propose a 3D Masked AutoEncoder (3D-MAE) that encodes weather-state evolution features as an additional conditioning for the diffusion model, together with a Variable-Aware Masked Frequency Modeling (VA-MFM) strategy that adaptively selects thresholds based on the spectral energy distribution of each variable. Together, we propose PuYun-LDM, which enhances latent diffusability and achieves superior performance to ENS at short lead times while remaining comparable to ENS at longer horizons. PuYun-LDM generates a 15-day global forecast with a 6-hour temporal resolution in five minutes on a single NVIDIA H200 GPU, while ensemble forecasts can be efficiently produced in parallel.