How TORNADO works
A visual walkthrough of the method: VAE posteriors, conditional endpoint targets, and inference with state-dependent diffusion.
Latent states as distributions
Each PDE state is encoded as a diagonal Gaussian VAE posterior.
Conditional endpoint targets
Closed-form moments of a Gaussian interpolant supervise drift and diffusion.
Adaptive stochastic transport
A latent SDE with state-dependent diffusion generates the forecast ensemble.
Input state
is a snapshot of 2D Kolmogorov flow. As in other latent neural PDE surrogates, the time evolution is learned in a compressed latent space rather than on the grid.
VAE encoder posterior
The encoder maps each field to a diagonal Gaussian posterior. Each snapshot becomes a local continuous distribution, from which different latent representations of the same field can be sampled.
Posterior of the next state
Encoding a pair of consecutive frames yields two local posteriors. The second is available only during training, where it defines the target.
Latent target distribution
Since is unobserved, the only supervision for the next latent state is the encoder posterior. Averaged over likely next states it defines the target, which minimizes the forward-KL step objective (Prop. 3.1).
Gaussian interpolant
Independent endpoints and are joined by a Gaussian interpolant over an auxiliary time τ ∈ [0, 1]. Each slice shows its marginal at one value of τ.
Bridge state
Training draws an interpolation time τ and a bridge state from the corresponding marginal. The supervision targets are conditioned on the pair .
Conditional endpoint law
Because both posteriors are Gaussian, the endpoint given is Gaussian with closed-form moments (Prop. 3.2). The cone connects to the 2σ region of this conditional law on the next posterior.
Dependence on τ
For small τ the bridge state carries little information about the endpoint and : the target covers nearly the whole next posterior. As τ → 1 the conditional law concentrates, and it degenerates to a point at τ = 1.
Drift and diffusion objectives
The remaining evolution is approximated by a one-step transition from . Its mean displacement is matched to and its variance to , so training never differentiates through an SDE solver. An energy-distance term on sampled endpoints completes the objective.
Conditional target along the bridge
Set τ directly, draw it at random, or resample the bridge state. The bar reports the spread of the conditional target relative to the next posterior.
Encoding the initial condition
At inference, is encoded into and a single latent sample is drawn. It is the initial state of every ensemble member.
SDE integration
Each member is an independent realization of the learned Itô SDE, integrated with Euler–Maruyama over N steps. N is the number of function evaluations; the main experiments use N = 5.
State-dependent diffusion
Classical stochastic interpolants inject noise through a prescribed, state-agnostic schedule . In TORNADO the diffusion is learned from data, so the ensemble spread adapts to the local latent dynamics.
Decoding the ensemble
The decoder maps the K latent samples to physical space, giving the predictive mean and a 95% confidence interval. For multi-step forecasts, each member is propagated autoregressively.
Calibration
The probability integral transform ranks each ground-truth value within its predictive ensemble. A calibrated forecast gives a flat histogram; bell shapes indicate excessive spread and U shapes insufficient spread.
Uncertainty follows the error
Over a 50-step autoregressive rollout of a chaotic flow, the mean forecast drifts from the ground truth. The predicted 95% interval widens accordingly and stays on the scale of the realized error.
Across benchmarks
At a fixed budget of five function evaluations, TORNADO has the lowest calibration error on all four synthetic PDEs and both real-world datasets, with competitive MSE and CRPS.
Further details
Proofs, ablations and results on Kolmogorov flow, Kuramoto–Sivashinsky, Burgers, Gray–Scott, combustion and ERA5 are in the paper.