Round-Trip Consistency: Bidirectional Diffusion Models Can Predict Their Own Rollout Errors
Authors: Alexander Scheinker
Organizations: Los Alamos National Laboratory Los Alamos, NM, 87544, USA
Abstract
Autoregressive models accumulate error over long rollouts, yet at deployment there is no ground truth to measure it against. We train a single conditional latent diffusion model that steps a dynamical system forward or backward in time via a direction flag, and show that this bidirectionality supplies a measurement-free test-time error signal: rolling forward i steps and then backward i steps must return the model to its start, so the round-trip discrepancy Ci is a self-supervised proxy for the unobservable rollout error: no ensembles, no held-out data, no governing equations, for one extra rollout. We validate on compressible magnetohydrodynamics (MHD), an astrophysical turbulent radiative mixing layer, and natural face videos (CelebV-HQ). On held-out MHD trajectories, Ci ranks rollout error (Spearman 0.91-0.98 at fixed depth; 0.69±0.16 within trajectories), and a simple calibrator fit on training rollouts predicts its magnitude to within 1.14× (68%) and 1.29× (95%) with near-nominal coverage - one nat beyond a depth-only predictor, transferring to all six decoded physical fields. The same signal flags the out-of-distribution Orszag-Tang vortex (AUROC 0.98; 1.0 by depth 10) exactly where sampling-dispersion baselines invert, and it cuts incurred error by 15% at 80% coverage - three times the depth-only baseline. Bidirectional training comes at negative cost, beating direction specialists in both directions, and the backward direction doubles as a fast inverse solver. On LE-PDE-UQ's turbulent Navier-Stokes benchmark, a single bidirectional model reaches accuracy within 1.3× of their ten-model ensemble at a tenth of the training cost, with the best training-free pixel-level calibration. Round-trip consistency turns reversibility into a practical trust signal for generative models.
This work presents a new bidirectional autoregressive latent diffusion approach for predicting the evolution of multiple fields (mass density, pressure, velocity, and magnetic field components) for magnetohydrodynamics. We show that this bidirectional flow can be used as a self-supervised consistency metric for uncertainty and error estimation, which enables the model to estimate test-time uncertainty and error without access to ground truth, by comparing how closely flowing forwards and backwards in time returns to the same predicted fields. We also demonstrate this methods's potential to serve as a non-invasive plasma diagnostic, and show how adaptive feedback can be used to make the model more robust based on sparse diagnostics or limited views/measurements.
A trained flow or diffusion model is usually run with only a handful of solver steps, and the integration error this leaves behind is unevenly distributed across the image. We ask where that error is injected and how it reaches the endpoint, and answer with a signed source-and-transport accounting of few-step integration error, tested to first order. A perturbation experiment on five models at 256^2 resolution shows the learned dynamics spread local disturbances widely: near the start of sampling, under 10% of the summed endpoint response remains at the source. Signed one-step truncation residuals, propagated through the model's own linearized dynamics, reconstruct much of the endpoint error's direction and regional structure (cosine 0.81-0.87), and a region's error owes more to what arrives from elsewhere than to its own injection. Structure-destroying nulls, with protocols frozen before evaluation, locate what carries the account: randomizing contribution signs halves it, and reassigning which region receives each contribution, with content, norms, and signs intact, destroys it entirely. Where the injections land is readable from the model itself. The variation of its velocity or prediction field along the trajectory, a structure that emerges during training, predicts the final per-region gap (within-image rho of 0.57-0.70 on fine trajectories, weaker from the cheap solve alone). The prediction is partial because endpoint error depends not only on injected magnitude but on its sign, timing, and transport through the learned dynamics. A training penalty on the injected variation lowers few-step error, so the structure is one a model can be trained to change.
Looped transformers scale computational depth without increasing parameter count by repeatedly applying a shared transformer block and can be used for iterative refinement, where each loop rewrites a full fixed-size prediction in parallel. On difficult problems, such as those that require search-like computation, reaching a highly structured solution starting from noise can require long refinement trajectories. Learning such trajectories is challenging when training specifies only the target solution and provides no supervision over the intermediate refinement path. Diffusion models tackle this issue by corrupting data with varying magnitudes of noise and training the model to reverse it in a \textit{single step}. However, this process misaligns training and testing behaviour. We introduce Denoising Recursion Models, a method that similarly corrupts data with noise but trains the model to reverse the corruption over \textit{multiple} recursive steps. This strategy provides a tractable curriculum of intermediate states, while better aligning training with testing and incentivizing non-greedy, forward-looking generation. Through extensive experiments, we show this approach outperforms the Tiny Recursion Model (TRM) on ARC-AGI, where it recently achieved breakthrough performance.