🤖 machine learning

Probabilistic Retrofitting of Learned Simulators

This paper introduces a computationally efficient, architecture-agnostic method to retrofit pre-trained deterministic PDE simulators into probabilistic models using the Continuous Ranked Probability Score (CRPS), achieving significant performance improvements in uncertainty quantification without the need for retraining from scratch.

Original authors: Cristiana Diaconu, Miles Cranmer, Richard E. Turner, Tanya Marwah, Payel Mukhopadhyay

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Cristiana Diaconu, Miles Cranmer, Richard E. Turner, Tanya Marwah, Payel Mukhopadhyay

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Average" Prediction Trap

Imagine you are trying to predict the weather. In the past, scientists used super-complex computer models that gave you one single answer: "It will rain at 2 PM."

But the real world is messy. Sometimes it rains at 2 PM, sometimes at 3 PM, and sometimes the storm splits in two. A single prediction can't capture that chaos. If you force a computer to give just one answer for a chaotic system, it tends to "average out" the possibilities. Instead of predicting a sharp, realistic storm, it predicts a blurry, gray cloud that never quite happens. It's like trying to describe a jazz solo by playing the average note of every instrument at once—it sounds smooth, but it's not the real music.

Most current AI models for physics (like fluid flow or heat) are built this way: they are deterministic. They are trained to minimize error by giving that single, "safest" average answer.

The Old Solution: Start from Scratch (The Expensive Way)

To fix this, scientists realized they needed probabilistic models. Instead of saying "It will rain at 2 PM," the model should say, "There's a 30% chance of rain at 2 PM, a 50% chance at 3 PM, and a 20% chance it won't rain at all."

The problem? Building these probabilistic models from the ground up is incredibly expensive. It's like trying to build a new, high-tech car engine from scratch every time you want to add a turbocharger. It takes massive amounts of time, money, and computing power.

The New Solution: "Retrofitting" (The Smart, Cheap Way)

This paper proposes a clever shortcut. Instead of building a new engine, they take the existing, high-performance deterministic engines (the models we already have) and retrofit them.

Think of it like taking a reliable, standard sedan and adding a "chaos kit" to it. You don't throw the car away; you just add a few new parts that allow it to handle bumpy roads (uncertainty) much better.

How do they do it?

Injecting "Noise": They add a small, random "noise" signal into the model. Imagine you are a chef following a perfect recipe. The deterministic model is the chef who follows the recipe exactly every time. The new method tells the chef, "Okay, but this time, shake the salt shaker a tiny bit differently."
The "Fair" Scoring Rule (CRPS): This is the secret sauce. Usually, to train a model to be probabilistic, you need complex math. The authors use a specific scoring rule called CRPS (Continuous Ranked Probability Score).
- The Analogy: Imagine a dartboard.
  - Deterministic training tells the player: "Hit the bullseye." If they miss, they get penalized.
  - CRPS training tells the player: "Don't just hit the bullseye; spread your darts out so that if you throw 100 times, the pattern of darts matches the actual shape of the target." It rewards the model for being accurate (hitting near the target) AND diverse (not clustering all darts in one spot).

What They Found

The team tested this "retrofitting" method on three different types of AI models (called Walrus, Lola, and Poseidon) and five different physics simulations (like swirling fluids, shockwaves, and heat).

Here are the results, translated:

It Works Everywhere: Whether they used a simple model trained on one specific problem or a massive "foundation model" trained on many problems, the retrofitting worked.
Better Accuracy: The retrofitted models didn't just get better at guessing probabilities; they actually became more accurate at predicting the future state of the system. In some cases, they reduced errors by up to 54%.
No Blurry Pictures: The old models tended to blur out sharp details (like shockwaves in a gas explosion). The new retrofitted models kept those sharp details crisp because they weren't forced to average everything out.
The "Ensemble" Trick: Because the model now includes random noise, you can run it 16 times with slightly different "noise" inputs. This gives you 16 different possible futures. The more times you run it, the clearer the picture becomes. It's like asking 16 different experts for their opinion and taking the average of their best guesses.

Why This Matters

The paper argues that we don't need to waste billions of dollars retraining massive AI models from scratch to make them probabilistic. We can take the powerful, deterministic models we have already built, add a small "noise injection" layer, and train them with this new scoring rule.

The Bottom Line:
They took a car that was great at driving on straight highways (deterministic) and added a suspension kit that lets it handle off-road chaos (probabilistic) without needing to build a whole new vehicle. It's cheaper, faster, and the car drives better.

What They Did Not Claim

They did not claim this works for medical diagnosis or clinical uses.
They did not claim this solves climate change (though it helps model physics).
They did not claim it works for every possible AI task, only for simulating physical systems governed by equations (PDEs).

The core message is simple: Don't throw away your old, good models. Just tweak them to understand uncertainty, and they will become even better.

Technical Summary: Probabilistic Retrofitting of Learned Simulators

Problem Statement

The dominant paradigm for modeling Partial Differential Equations (PDEs) using deep learning relies on deterministic training frameworks that minimize Mean Squared Error (MSE) or Mean Absolute Error (MAE). While effective for stable, fully observed systems, this approach is fundamentally ill-suited for chaotic, noisy, or partially observable physical systems. Deterministic models operating at coarse resolutions often fail to capture the inherent uncertainty of such dynamics, leading to "blurry" predictions where the model converges to the mean of possible future states, resulting in physically inconsistent rollouts.

Existing alternatives, such as diffusion or flow-based generative models, offer probabilistic capabilities but impose prohibitive computational burdens. They require specialized, resource-intensive training regimes and often suffer from slow, iterative inference processes. Furthermore, training high-fidelity probabilistic foundation models from scratch is becoming economically and environmentally unsustainable given the rapid scaling of parameter counts and training data. There is a critical gap in efficiently transforming existing, high-performing deterministic backbones into probabilistic models without discarding the significant resources already invested in them.

Methodology

The authors propose a training-efficient strategy to retrofit pre-trained deterministic models into probabilistic emulators. This approach is architecture-agnostic and requires minimal code modifications. The core methodology involves two key components:

Noise Injection via Conditional Layer Normalization (AdaLN):
The deterministic architecture is modified to include a stochastic branch. A global noise vector $\epsilon \sim \mathcal{N}(0, I)$ is sampled and passed through a Multi-Layer Perceptron (MLP) to generate noise embeddings. These embeddings are projected into modulation parameters ( $\gamma, \beta, \alpha, \delta$ ) via conditional layer normalization layers. These parameters scale and shift the normalized inputs and gate the residual connections within the backbone (e.g., in transformer attention blocks or U-Net encoder layers). This mechanism allows the model to generate an ensemble of diverse, sharp predictions from a single input history without substantially altering the core learned features.
Optimization with Continuous Ranked Probability Score (CRPS):
Instead of minimizing point-wise error, the model is fine-tuned using the Continuous Ranked Probability Score (CRPS), a strictly proper scoring rule. The empirical CRPS is computed from an ensemble of $M$ predictions $\{x_j\}$ against a ground truth $y$ :
$\text{CRPS}(\{x_j\}_{j=1}^M, y) = \frac{1}{M}\sum_{j=1}^M |x_j - y| - \frac{1}{2M^2}\sum_{j=1}^M\sum_{k=1}^M |x_j - x_k|$
To address bias in finite ensembles, the authors employ an unbiased "fair" CRPS estimator. The first term encourages individual ensemble members to be accurate (minimizing absolute error), while the second term encourages diversity (dispersing ensemble members), preventing mode collapse.

Training Strategy:
The retrofitting process leverages the compatibility between the CRPS objective and the MSE/MAE losses used in pre-training. Key strategies include:

Initializing backbone parameters from the pre-trained deterministic model.
Using differential learning rates: smaller rates for the backbone to preserve features, and larger rates for the newly introduced noise branch.
Carefully initializing the conditional normalization layers to small values to ensure the model starts near its deterministic behavior, allowing for smooth learning of stochasticity.

Key Contributions

The paper presents four primary contributions:

A Unified, Architecture-Agnostic Recipe: The authors demonstrate that the same retrofitting strategy is effective across three distinct architectures (Walrus, Lola, and Poseidon) representing both pixel-space and latent-space modeling paradigms, and across five physical systems.
Efficiency and Performance: The method achieves substantial gains over continued deterministic fine-tuning. It yields reductions of 20–54% in rollout CRPS and up to 30% improvements in Variance-Normalized Root Mean Square Error (VRMSE) for single-system models. For foundation models, it shows up to 40% CRPS and 15% VRMSE improvements over deterministic fine-tuning.
Inference-Time Flexibility: The retrofitted models exhibit test-time scaling properties. Predictive skill monotonically improves with the number of ensemble members generated, allowing users to dynamically trade off compute for accuracy without retraining.
Ablations and Insights: Extensive ablation studies provide practical recommendations for hyperparameter tuning (learning rates, ensemble size, noise dimensionality) and demonstrate the method's robustness.

Results

The approach was validated on single-system models and a multi-system foundation model (HalfWalrus) using datasets from The Well (Rayleigh-Bénard, Euler Multi-Quadrants, Shear Flow, Turbulent Radiative Layer, and Viscoelastic Instability).

Single-System Models: Across Walrus, Lola, and Poseidon architectures, CRPS retrofitting consistently outperformed deterministic fine-tuning. For instance, on the Rayleigh-Bénard dataset, the Walrus model achieved a 43.1% improvement in CRPS and a 21.5% reduction in VRMSE. Qualitative analysis showed that while deterministic baselines collapsed into over-smoothed predictions at long horizons, the CRPS ensemble preserved high-frequency turbulent structures.
Foundation Models: When applied to the HalfWalrus foundation model, retrofitting yielded 34–40% improvements in CRPS and >13% in VRMSE for chaotic regimes (Rayleigh-Bénard and Turbulent Radiative Layer). In chaotic viscoelastic regimes, the probabilistic model significantly outperformed the deterministic baseline, which suffered from over-smoothing.
Comparison to Baselines:
- Deep Ensembles: CRPS retrofitting outperformed deep ensembles (4x compute cost) in both accuracy and calibration (spread-to-skill ratio).
- Diffusion Models: In a compute-matched comparison within the Lola framework, CRPS retrofitting matched or exceeded diffusion-based retrofitting in accuracy while offering single-step inference (vs. iterative denoising). Crucially, CRPS retrofitting proved more robust to dataset heterogeneity (e.g., Shear Flow) and more effective at leveraging pre-trained deterministic representations than diffusion models.

Significance and Claims

The paper argues that probabilistic PDE modeling need not require retraining from scratch. By adopting a retrofitting strategy based on CRPS and noise injection, the community can unlock probabilistic capabilities from existing deterministic backbones with modest additional training cost.

The authors claim this approach offers a viable alternative to the prohibitive costs of training generative foundation models from scratch. It addresses the economic and environmental unsustainability of retraining massive models while providing a unified framework that bridges the gap between deterministic efficiency and probabilistic fidelity. The method is particularly significant for the emerging era of foundation models for dynamical systems, enabling the repurposing of robust deterministic checkpoints to handle the inherent uncertainty of chaotic physical systems.