Probabilistic Retrofitting of Learned Simulators

Cristiana Diaconu, Miles Cranmer, Richard E Turner, Tanya Marwah, Payel Mukhopadhyay
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:24654-24701, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-diaconu26a, title = {Probabilistic Retrofitting of Learned Simulators}, author = {Diaconu, Cristiana and Cranmer, Miles and Turner, Richard E and Marwah, Tanya and Mukhopadhyay, Payel}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {24654--24701}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/diaconu26a/diaconu26a.pdf}, url = {https://proceedings.mlr.press/v306/diaconu26a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Probabilistic Retrofitting of Learned Simulators %A Cristiana Diaconu %A Miles Cranmer %A Richard E Turner %A Tanya Marwah %A Payel Mukhopadhyay %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-diaconu26a %I PMLR %P 24654--24701 %U https://proceedings.mlr.press/v306/diaconu26a.html %V 306 %X 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.
APA
Diaconu, C., Cranmer, M., Turner, R.E., Marwah, T. & Mukhopadhyay, P.. (2026). Probabilistic Retrofitting of Learned Simulators. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:24654-24701 Available from https://proceedings.mlr.press/v306/diaconu26a.html.

Related Material