TRIE: An Evaluation Framework for Stochastic PDE Surrogates
2026-06-30 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how to predict complex systems that have a lot of randomness, like certain chaotic physical processes described by stochastic partial differential equations (SPDEs). They point out that simple neural networks usually fail to give good uncertainty estimates and long-term statistical predictions. To fix this, they created TRIE, a way to test if models can reliably capture the overall statistical behavior and uncertainty of these systems. They found that generative models work better than traditional methods and that special latent models can make predictions faster while still being accurate. The authors also shared their code and data to help others test these forecasting methods.
stochastic partial differential equationsneural surrogatesuncertainty quantificationgenerative modelsinvariant measuresMonte Carlo dropoutheteroscedastic Gaussian likelihoodKuramoto-Sivashinsky equationKolmogorov flowCRPS (Continuous Ranked Probability Score)
Authors
Bharat Srikishan, Javier E. Santos, Nikhil Muralidhar, Charles D. Young
Abstract
Many scientific systems exhibit uncertainty from stochastic forcing, unresolved degrees of freedom, or imperfect observations, making reliable surrogate forecasting fundamentally distributional rather than pointwise. For such systems, deterministic neural surrogates fail to capture statistical measures and forecast uncertainty. We introduce TRIE, an evaluation framework for stochastic PDE surrogates that asks whether models reproduce invariant measures, provide trustworthy predictive uncertainty, and scale to efficient probabilistic generation. We demonstrate TRIE on two stationary chaotic spatially extended SPDEs, stochastic Kuramoto--Sivashinsky and stochastic Kolmogorov flow, across 11 parameter values. Our evaluation shows that standard pointwise-trained neural surrogates can produce plausible short rollouts while failing to match long-time statistical structure. Approximate uncertainty methods such as Monte Carlo dropout and heteroscedastic Gaussian likelihoods produce stochastic forecasts, but are often miscalibrated and overconfident under temporal and spatial uncertainty diagnostics. Across these criteria, generative models provide the most consistent performance, accurately capturing invariant measure statistics and achieving the lowest CRPS in all reported probabilistic settings. Finally, we show that latent generative models with automatic dimension discovery retain much of this statistical fidelity while reducing Kolmogorov inference time by roughly $12\times$. We release our code and data at https://github.com/scailab/TRIE-SPDE-Bench to support reproducible evaluation of stochastic PDE forecasting models.