Climate Grey-Box Flow Matching for Robust Climate and Weather Prediction

Deep generative models such as flow matching and diffusion models have great potential for learning complex dynamical systems, but they typically act as black boxes, neglecting underlying physical structure. In contrast, physics-based models governed by ODEs and PDEs provide interpretability and physical consistency, yet are often incomplete due to unresolved processes, missing source terms, or uncertain parameterisations. Bridging these two paradigms is a central challenge in data-driven weather and climate modelling.

We propose a Climate Grey-Box Dynamics Matching framework designed for weather and climate systems, that explicitly combines existing physical models with data-driven learning to capture unresolved dynamics where known physical operators are directly embedded into the learned dynamics. Our framework learns from observational trajectories alone and operates in a simulation-free manner inspired by gradient matching and flow matching methods. By avoiding numerical solvers, it eliminates the memory overhead, computational cost, and numerical instability associated with Neural ODE–based approaches.

To capture temporal dependencies in our simulation-free method, we introduce a lightweight attention-based temporal encoder that aggregates short-term history in a physically consistent manner. This design enables the model to represent unresolved dynamics without increasing computational complexity, making it well-suited for high-dimensional spatiotemporal climate systems. We apply this framework to weather and climate forecasting and demonstrate its effectiveness against ClimODE, a state-of-the-art solver-based grey-box model. Reformulating ClimODE as a simulation-free grey-box model reduces training complexity from Ο(L) to Ο(1), where L denotes the number of solver steps. Beyond computational gains, the simulation-free formulation yields substantial memory efficiency: training is possible on a single RTX 3060 (12 GB), whereas ClimODE requires at least 25 GB of GPU memory with a small batch size. This enables efficient training on commodity hardware and improves accessibility for large-scale climate modelling.

Experiments on weather and climate benchmarks show that the proposed method achieves improved forecast accuracy and faster convergence compared to simulation-based and fully data-driven baselines. The method demonstrates particular robustness to long horizons, as performance gains become more pronounced with extended forecast times—indicating enhanced temporal stability and resistance to error accumulation, an essential property for reliable long-range climate prediction.