Probabilistic Retrofitting of Learned Simulators

Dominant approaches for modelling Partial Differential Equations (PDEs) rely on deterministic predictions, yet many physical systems of interest are inherently chaotic and uncertain. While training probabilistic models from scratch is possible, it is computationally expensive and fails to leverage the significant resources already invested in high-performing deterministic backbones. In this work, we adopt a training-efficient strategy to transform pre-trained deterministic models into probabilistic ones via retrofitting with a proper scoring rule: the Continuous Ranked Probability Score (CRPS). Crucially, this approach is architecture-agnostic: it applies the same adaptation mechanism across distinct model backbones with minimal code modifications. The method proves highly effective across different scales of pre-training: for models trained on single dynamical systems, we achieve $20–54\%$ reductions in rollout CRPS and up to $30\%$ improvements in variance-normalised RMSE (VRMSE) relative to compute-matched deterministic fine-tuning. We further validate our approach on a PDE foundation model, trained on multiple systems and retrofitted on the dataset of interest, to show that our probabilistic adaptation yields an improvement of up to $40\%$ in CRPS and up to $15\%$ in VRMSE compared to deterministic fine-tuning. Validated across diverse architectures and dynamics, our results show that probabilistic PDE modelling need not require retraining from scratch, but can be unlocked from existing deterministic backbones with modest additional training cost.

Many of the physical systems we most want to predict — turbulent flows, weather, stellar dynamics — are chaotic: tiny differences in the input produce wildly different futures. For systems like these, the right kind of forecast is not a single best guess but a distribution of plausible futures, in the spirit of ensemble weather forecasts that report "70% chance of rain" rather than a single deterministic outcome. Most machine learning models for physical simulation, however, are trained to produce a single prediction. Training probabilistic alternatives from scratch — typically based on diffusion or flow-matching — is possible but very expensive, both to train and to run at inference, and it throws away the substantial resources the community has already invested in high-quality deterministic models. We show that this trade-off is unnecessary. With a small architectural addition and a short fine-tuning step using a probabilistic loss (the Continuous Ranked Probability Score), an existing deterministic model can be retrofitted into one that produces ensembles of plausible futures. The retrofitted models are relatively cheap to train, cheap to run, and consistently outperform simply continuing to train the original deterministic model — by 20–54% on probabilistic skill and up to 30% on standard deterministic metrics. The approach works across three very different architectures and five physical systems, and extends to foundation models for partial differential equations (PDEs) trained on many systems at once. Probabilistic forecasting for chaotic physics, in other words, does not have to involve expensive training from scratch — it can be unlocked from the deterministic checkpoints we already have.