# TRIE: An Evaluation Framework for Stochastic PDE Surrogates

arXiv: https://arxiv.org/abs/2607.00196 · HTML: https://mlsanity.com/paper/2607.00196/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 42 references that are in the map.
- First posted on arXiv 30 Jun 2026.

## Builds on (most cited first, 23 of 23)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks](https://mlsanity.com/paper/1506.03099/) (2015)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://mlsanity.com/paper/2312.15796/) (2023)
- [DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting](https://mlsanity.com/paper/2306.01984/) (2023)
- [Stochastic interpolants with data-dependent couplings](https://mlsanity.com/paper/2310.03725/) (2023)
- [Probabilistic Forecasting with Stochastic Interpolants and F"ollmer Processes](https://mlsanity.com/paper/2403.13724/) (2024)
- [When Physics Meets Machine Learning: A Survey of Physics-Informed Machine Learning](https://mlsanity.com/paper/2203.16797/) (2022)
- [DiffCast: A Unified Framework via Residual Diffusion for Precipitation Nowcasting](https://mlsanity.com/paper/2312.06734/) (2023)
- [Stabilized Neural Ordinary Differential Equations for Long-Time Forecasting of Dynamical Systems](https://mlsanity.com/paper/2203.15706/) (2022)
- [Learning to correct spectral methods for simulating turbulent flows](https://mlsanity.com/paper/2207.00556/) (2022)
- [Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics](https://mlsanity.com/paper/2110.10249/) (2021)
- [Implicit Rank-Minimizing Autoencoder](https://mlsanity.com/paper/2010.00679/) (2020)
- [Is MC Dropout Bayesian?](https://mlsanity.com/paper/2110.04286/) (2021)
- [PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations](https://mlsanity.com/paper/2505.24717/) (2025)
- [Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos](https://mlsanity.com/paper/2503.18731/) (2025)
- [Lost in Latent Space: An Empirical Study of Latent Diffusion Models for Physics Emulation](https://mlsanity.com/paper/2507.02608/) (2025)
- [Autoencoders for discovering manifold dimension and coordinates in data from complex dynamical systems](https://mlsanity.com/paper/2305.01090/) (2023)
- [Nonlinear dimensionality reduction then and now: AIMs for dissipative PDEs in the ML era](https://mlsanity.com/paper/2310.15816/) (2023)
- [Unreliable Uncertainty Estimates with Monte Carlo Dropout](https://mlsanity.com/paper/2512.14851/) (2025)

## Built on by (most cited first, 2 of 2)

- [StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics](https://mlsanity.com/paper/2608.22309/) (2026)
- [Do Neural PDE Solvers Learn the Right Dynamics?](https://mlsanity.com/paper/2610.06952/) (2026)

## Abstract (arXiv)

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×. We release our code and data at https://github.com/scailab/TRIE-SPDE-Bench to support reproducible evaluation of stochastic PDE forecasting models.