Round-Trip Consistency: Bidirectional Diffusion Models Can Predict Their Own Rollout Errors
Alexander Scheinker
stat.ML, cs.LG, physics.comp-ph, physics.plasm-ph
2026-08-01
A bidirectional latent diffusion model estimates its own rollout error at test time: the forward-then-back round-trip gap is a ground-truth-free proxy, flagging OOD inputs at AUROC 0.98 on MHD.
Learned surrogate models now stand in for numerical solvers across scientific computing, and diffusion models serve as probabilistic simulators for everything from turbulence to operational weather forecasting. All of them roll out autoregressively: each step's prediction is fed back as the next input, so small one-step errors compound, the input distribution drifts from training data, and accuracy degrades at a depth that varies unpredictably across initial conditions. The catch is that at deployment, in plasma control, accelerator diagnostics, or weather prediction, the true state is unknown, so a model cannot say how far into the future it can still be trusted.
Standard uncertainty tools (deep ensembles, MC dropout, learned variance, conformal intervals, the spread of repeated samples) all measure dispersion: how much the model disagrees with itself. None of them measures whether the learned dynamics are being applied accurately to the trajectory at hand, and dispersion-based signals degrade under exactly the distribution shift that autoregressive rollouts induce.
The authors train a single conditional latent diffusion model with a direction flag cd ∈ {+1, -1}, so one set of weights represents both the forward evolution (the surrogate solver) and the time reverse (an inverse solver). The denoiser conditions on the two most recent frames. A β-VAE compresses each physical field into a latent space (for MHD, 512×512 to 16×16×4, a 256× compression), a diffusion transformer (DiT) does the denoising, sampling is deterministic DDIM, and the scalar conditions (diffusion step, time index, direction flag) enter via adaLN-Zero.
The core idea is plain: reversibility becomes a checkable invariant. An accurate model composed with its own inverse is the identity, so the discrepancy left after rolling forward i steps and back i steps, the round-trip consistency error Ci, is a self-supervised proxy for the true rollout error Ei. Every quantity in Ci is available at test time (the anchor pair is encoded measured data, the returned pair is model-generated), at the cost of one extra rollout (2× inference), with no ground truth, no ensemble, and no governing equations.
A theorem (Proposition 1) sandwiches Ci between two affine functions of the true error under a bi-Lipschitz assumption on the backward map (both an upper bound L and a lower bound μ). That lower bound μ is the anti-cancellation condition: it forbids the backward map from collapsing distinct terminal states onto the same returned seed, which would let forward and backward errors cancel and fool the check. So the check is necessary but not sufficient: small Ci certifies small error only when the backward map is well-conditioned (μ not too small) and accurate on clean data (small residual δ), and both are measurable offline.
Validated across five systems: compressible magnetohydrodynamics (MHD, 512×512) with the Orszag-Tang vortex as a deliberate out-of-distribution stress test, an astrophysical turbulent radiative mixing layer (the Well benchmark), natural face videos (CelebV-HQ), and turbulent Navier-Stokes (the LE-PDE-UQ benchmark).
| Metric | Result |
| Spearman of Ci with true error (fixed depth, across trajectories) | 0.91-0.98 (0.97 at i=20) |
| Within-trajectory Spearman | 0.69 ± 0.16 |
| Calibrator error factor | 1.14× (68%), 1.29× (95%), near-nominal coverage |
| Log-likelihood gain over a depth-only predictor | +1.0 nat, positive for all six decoded fields |
| OOD detection AUROC (Orszag-Tang vortex) | 0.98 (1.0 by depth 10) |
| Error reduction at 80% coverage (selective prediction) | 15% (depth-only early stopping manages 5%, three times less) |
A few points deserve attention:
One nuance matters: the per-trajectory signal is valuable in proportion to how sharply trajectories differ relative to the metric's noise. It is large on MHD (±28% factor), marginal on the radiative layer (±11%, where a depth-only calibrator is already near-optimal), and maximal for natural video, where clip-to-clip heterogeneity dominates.
For anyone building scientific ML, PDE surrogates, physics simulators, or weather and climate models, this is a deployable artifact: an error meter that travels with the model and needs no ground truth to say how far an autoregressive prediction can still be trusted. It fills exactly the hole dispersion signals cannot. When the model is systematically wrong (given a context, it goes wrong in the same way every time), self-disagreement sees nothing, but round-trip drift does.
More broadly, treating reversibility as a trust signal is a transferable idea: any generative model that can learn a time reverse (including, possibly, time-reversed language models) could self-check via round-trip consistency, provided its backward map is well-conditioned.