TRIE: An Evaluation Framework for Stochastic PDE Surrogates

Reliable forecasting of scientific systems often requires modeling stochasticity rather than deterministic trajectories. While deterministic neural surrogates are computationally efficient, they fail to capture the statistical distributions and predictive uncertainty inherent in stochastic partial differential equations (SPDEs). The authors introduce TRIE, an evaluation framework designed to assess whether surrogates can reproduce long-time invariant measures, provide trustworthy predictive uncertainty, and scale to efficient probabilistic generation.

TRIE evaluates models based on three critical criteria:

• Invariance: Does the model reproduce the long-time statistical structure (invariant measure) of the system, rather than just matching short-term trajectories?
• Trustworthiness: Does the model provide calibrated and sharp uncertainty estimates, as measured by the continuous ranked probability score (CRPS)?
• Efficiency: Can the model generate probabilistic samples at a speed that makes it practical for large-scale scientific applications?

The authors test these criteria on two complex systems: the stochastic Kuramoto–Sivashinsky equation and the stochastic Kolmogorov flow. They compare various architectures, including standard ResNets, UNets, and PDE-Transformers, against probabilistic methods like Monte Carlo dropout, heteroscedastic Gaussian likelihoods, and stochastic interpolants.

Pointwise-trained deterministic models consistently fail to capture the long-time statistical structure of the systems, often converging to incorrect attractors. While approximate uncertainty methods like Monte Carlo dropout provide stochastic outputs, they are frequently miscalibrated and overconfident. Stochastic interpolants emerge as the most consistent performers, accurately capturing invariant measures and achieving the lowest CRPS across all tested regimes. Furthermore, by integrating implicit rank-minimizing autoencoders, the authors demonstrate that latent-space stochastic interpolants can significantly reduce inference time—by roughly 12x in the Kolmogorov flow experiments—while retaining high statistical fidelity.