TORNADO lets PDE ensemble forecasts vary their noise by state, cutting calibration error roughly 2x to over 5x versus representative baselines at five network evaluations
Synopsis
Probabilistic PDE surrogates often turn overconfident in few-step inference, producing predictive spreads narrower than the real uncertainty. TORNADO recasts next-step prediction as stochastic transport between the variational-autoencoder posteriors of consecutive frames, and lets a neural network set the diffusion amplitude from the latent state instead of a fixed global noise schedule. On four synthetic PDE datasets, at a fixed budget of five network evaluations, it reaches the lowest calibration error, roughly 2x to over 5x lower than representative baselines, while keeping comparable predictive accuracy.
Interpretation
A probabilistic PDE surrogate can now adjust the width of its predictive distribution to the local dynamical state instead of sharing one noise schedule across a whole trajectory. Earlier diffusion, flow-matching and stochastic-interpolant baselines wrote the injected randomness as a global diffusion coefficient depending only on auxiliary time, so noise strength was independent of the current state. On the four synthetic datasets Kolmogorov flow, KS, Burgers and Gray-Scott, at a fixed five network evaluations, TORNADO has the lowest calibration error, roughly 2x to over 5x lower than baselines on representative cases such as Kolmogorov and Burgers, with comparable MSE and CRPS.
The mechanism learns an Ito stochastic differential equation in latent space, treating the next frame's latent variable as a sample obtained by integrating from the current frame's posterior. Both the drift and the diffusion amplitude are parameterized by neural networks, with the amplitude coming from a separate diffusion head branching off the shared Transformer trunk, so noise magnitude varies with state. Training avoids backpropagating through the SDE solver: a Gaussian interpolant over auxiliary time yields closed-form conditional endpoint means and covariances, and an energy-distance term regularizes sample-level distributional alignment.
Replacing the learned diffusion amplitude with a constant or a fixed schedule worsens calibration, indicating the gain comes from state-dependent diffusion itself. In an inference ablation sharing one trained drift at the same five evaluations, deterministic Euler integration gives overconfident narrow ensembles, and constant amplitude and fixed schedule each show their own mismatch. On the Kolmogorov flow test set, TORNADO's calibration error is 0.039, versus 0.067 for the fixed schedule, 0.085 for constant amplitude and 0.286 for deterministic Euler.
Perspective
Evaluation covers four 2D synthetic datasets generated with APEBench (Kolmogorov flow, Kuramoto-Sivashinsky, Burgers and Gray-Scott) plus two real datasets, RealPDEBench combustion flame imaging and WeatherBench ERA5 reanalysis; main results come from autoregressive rollout with a fixed ensemble size and five network evaluations in the few-step regime.
The learned diffusion amplitude is tied to VAE posterior quality, and posteriors that are too narrow or too smooth both weaken calibration; supervision constrains only local endpoint moments rather than explicitly enforcing consistency with the interpolant's full continuous-time marginal evolution, and how local supervision and numerical discretization jointly shape the endpoint distribution remains open.