Source-linked AI summary
Round-Trip Consistency: Bidirectional Diffusion Models Can Predict Their Own Rollout Errors
Alexander Scheinker
TL;DR
Deployment lacks ground truth for judging how far autoregressive dynamical-system rollouts remain trustworthy. This paper uses bidirectional diffusion and round-trip consistency to estimate rollout error, predicting held-out MHD error within 1.14× at 68% coverage.
Problem
At deployment, the true future state is unavailable, preventing direct measurement of autoregressive rollout error and limiting trust assessment.
Method
A single direction-conditioned latent diffusion model predicts both temporal directions, using forward-then-backward round trips to estimate rollout error without ground truth.
Results
The round-trip signal predicts held-out MHD rollout error within 1.14× at 68% coverage and supports 15% lower incurred error at 80% coverage.
Takeaways & Limitations
Round-trip consistency provides a practical, training-free trust signal that travels with a bidirectional generative dynamics model at the cost of one extra rollout.
Takeaways & Limitations
The check is necessary but not sufficient because forward and backward errors can cancel, weakening its certification of small rollout error.
Abstract
from arXiv · showhide
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 $\mathcal{C}_i$ 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, $\mathcal{C}_i$ ranks rollout error (Spearman $0.91$-$0.98$ at fixed depth; $0.69 \pm 0.16$ within trajectories), and a simple calibrator fit on training rollouts predicts its magnitude to within $1.14\times$ ($68\%$) and $1.29\times$ ($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\times$ 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.
1 Introduction
The paper addresses compounding rollout errors by using one bidirectional conditional latent diffusion model to estimate reliability without ground truth. Its round-trip discrepancy serves as a self-supervised error proxy for calibration, selective prediction, and early stopping.
- Motivation: Compounding one-step errors cause input-distribution drift and unpredictable accuracy degradation across initial conditions during autoregressive deployment.These failures arise when models feed their own predictions back as inputs, including in plasma control, accelerator diagnostics, and weather prediction.
- Validation: 1.14 (68%) is the factor within which a training-rollout calibrator predicts held-out MHD rollout-error magnitude, one nat better than a depth-only predictor.The gain transfers to all six decoded physical fields and supports selective prediction.
- Validation: 15% lower incurred error at 80% coverage is achieved by selective prediction, three times the depth-onlydeferral baseline.The paper validates the approach on compressible MHD, an astrophysical turbulent radiative mixing layer, and natural face videos.
- Method: A single conditional latent diffusion model uses a direction flag to roll systems forward as a surrogate solver or backward as an inverse solver.The bidirectional model uses one set of weights for both temporal directions.
- Method: Round-trip discrepancy C_i provides a test-time error proxy without ground truth, ensembles, or architectural changes, using one extra rollout.Cycling a rollout forward and backward yields a trust or early-stop criterion for autoregressive generation.
2 Related Work
Prior work spans learned dynamical surrogates, diffusion-based forecasting, rollout uncertainty quantification, cycle-consistency diagnostics, and bidirectional diffusion bridges. This paper distinguishes its test-time, multi-step round-trip signal from methods requiring auxiliary heads, external observations, analytic operators, or construction-guaranteed reversibility.
- Surrogates and diffusion models for dynamics: Learned simulators include graph networks, message-passing solvers, neural operators, solver-in-the-loop hybrids, and diffusion models for probabilistic dynamics forecasting.The cited diffusion applications include turbulence benchmarking, forecasting frameworks, and operational ensemble weather forecasting.
- Uncertainty quantification for rollouts: Deep rollout uncertainty quantification commonly uses ensembles, MC dropout, and conformal calibration, while distribution shift can erode their reliability.LE-PDE-UQ evolves latent uncertainty alongside the state and reports gains over ensembles for PDE surrogates.
- Cycle consistency and nearest neighbors: Cycle consistency has diagnosed unreliable optical-flow correspondences and supervised unpaired translation and temporal correspondence, while related forecasting methods use reversal as training supervision.Examples include time-reversal regularization, an auxiliary backcasting head, and reversing predicted trajectories as a training-time loss.
- Cycle consistency and nearest neighbors: Unlike these approaches, the proposed signal needs no auxiliary head and applies at test time.The distinction is stated against prior methods using time-reversal regularization, backcasting, or reversed predicted trajectories as training procedures.
- Bidirectional diffusion bridges: Bidirectional diffusion bridges share the one-network, direction-flag design, but their forward-backward composition has no measurable defect by construction, whereas this method’s unanchored backward leg makes C_i informative.The backward leg starts from the model’s own predicted terminal pair and must return to the initial state.
3 Background
The paper formulates autoregressive surrogate rollouts for multi-channel dynamical-system fields and measures their unobservable latent-space error at depth i. It learns dynamics in a compact latent representation produced by independently encoding each physical field with a β-VAE.
- Problem setup: Autoregressive deployment feeds predicted states back as inputs to reach rollout depth i, where the true error is Ei = MSE(zt+i, ˆzt+i) in latent space.Decoded-field metrics are reported in the supplementary material.
- Latent encoding and diffusion: Each physical field is independently encoded by a β-VAE from a 512 × 512 field into a 16 × 16 × 4 latent, providing 256× compression.The weak KL penalty is β = 10−3, and the representation is designed to remain compact without collapsing detail.
- Latent encoding and diffusion: Dynamics are learned in latent space using standard denoising diffusion with network ϵθ.The latent generative-modeling setup follows Rombach et al. (2022), while the diffusion formulation follows Ho, Jain, and Abbeel (2020) and Song et al. (2021).
4 Bidirectional Latent Diffusion Dynamics
Section 4 introduces a single conditional latent diffusion model that learns bidirectional transition densities with shared weights. Uniformly training on future and past directions enables autoregressive forward rollouts and inverse solving, using a diffusion transformer with two context frames by default.
- Model and training: A single denoiser learns the bidirectional transition density.The model is trained with a conditional noise-prediction loss.
- Model and training: Uniformly sampling both direction signs forces shared weights to denoise consistently toward the future and the past.The future direction serves as a surrogate solver, while the past direction serves as an inverse solver.
- Rollouts and architecture: Rollouts compose autoregressively, with decoded fields obtained from the latent state.The decoded fields are defined as ˆx_t = D_θ(ˆz_t).
- Rollouts and architecture: The denoiser is a diffusion transformer using n=2 context frames unless noted.Physical time, direction, and denoising step condition the model, with z^(k) denoting the correspondingly noised target.
5 Round-Trip Consistency
Round-trip consistency uses a forward rollout followed by a backward rollout to produce a test-time discrepancy from measured anchor data and model-generated returns. The discrepancy enables early stopping and can certify rollout error when the learned backward map is accurate and well conditioned, but costs one extra rollout and is not sufficient in general.
- Definition and use: Round-trip discrepancy C_i compares an i-step forward rollout with its i-step backward return using only the measured anchor pair and model-generated outputs.An error-free model returns to its starting pair, while all quantities are available at test time without ground-truth rollout error E_i.
- Definition and use: A tolerance τ supports early stopping by trusting a rollout to the largest depth satisfying C_i ≤τ.The criterion selects the maximum trusted autoregressive depth using only the round-trip signal.
- Properties and cost: 2× inference overhead is required because C_i costs one backward rollout per checked depth, with no change to training.Deterministic DDIM sampling makes the cycle well defined; averaging stochastic cycles is treated as a dispersion baseline.
- Theory: Small C_i certifies small true error when the learned backward map is well conditioned and accurate on clean data, with both properties estimable offline.The relevant conditions are a not-too-small µ and a small residual δ_i; no test-time ground truth is required.
- Theory: The certificate weakens as µ →0 because a collapsing inverse can hide errors, and the round-trip check is necessary but not sufficient.The theoretical condition applies on rollout-visited pairs and does not assume invertibility of the underlying physical dynamics.
6 Experiments
Experiments show that round-trip consistency predicts rollout error, supports calibration and selective deferral, and transfers across physical fields and video. Bidirectional training also improves both temporal directions while enabling inverse reconstruction and competitive uncertainty quantification.
- Calibration: 1.14× (68%) and 1.29× (95%) are the held-out error factors achieved by the C_i-based calibrator, with 0.013 miscalibration area.It improves log-likelihood by 1.0 nat over the depth-only baseline and halves interval width.
- Decoded-field transfer: +0.32 to +0.73 nats is the positive calibration gain over depth-only baselines across all six decoded physical fields.Decoded field errors are predicted within factors of 1.17–1.30.
- Selective prediction: 15% lower mean incurred latent error results from deferring the 20% highest-error predictions, versus 5% for depth-only early stopping.At 90%/70% coverage, reductions are 7%/23% versus 2%/9% for the depth-only baseline.
- Natural face videos: Spearman 0.79–0.81 latent and 0.73–0.80 decoded-pixel correlations rank errors on 150 held-out face-video clips at rollout depths 3–21.The relationship does not decay with depth, and the same calibration recipe transfers to disjoint clips.
- Navier–Stokes benchmark: 0.244 block relative L2 is achieved by one bidirectional model on LE-PDE-UQ’s untouched test split, within 1.29× of the ten-model ensemble’s 0.190.The result is also within 1.16× of the best single model with uncertainty, which scores 0.211.
- Bidirectional capability: 0.0630 ± 0.0006 forward and 0.0621 ± 0.0006 backward losses show the bidirectional model outperforming matched-direction specialists.The direction flag provides free augmentation, while backward rollouts reconstruct preceding plasma history as a fast inverse solver.
7 Discussion and Conclusion
Round-trip consistency converts one bidirectional network’s structural reversibility into a practical, measurement-free trust signal. It requires one extra rollout and predicts held-out MHD rollout error while approaching the accuracy of a ten-model ensemble.
- Discussion and Conclusion: Round-trip consistency turns one network representing both temporal directions into an error meter that travels with the model.The signal needs no ground truth, ensembles, or governing equations.
- Discussion and Conclusion: 1.14 factor accuracy at 68%: the meter predicts held-out MHD rollout error within a factor of 1.14 with near-nominal coverage.This is the reported quantitative evidence for its calibration performance.
- Discussion and Conclusion: One extra rollout is the stated cost of obtaining the practical trust signal.The method is described as measurement-free and does not require ground truth, ensembles, or governing equations.
- Discussion and Conclusion: A single bidirectional model approaches the accuracy of a ten-model ensemble.The passage presents this as additional evidence supporting the method’s practical value.
Supplementary Materials: Round-Trip Consistency: Bidirectional … A.3 2D Compressible MHD
The supplementary materials describe a shared latent-space diffusion-transformer design and the 2D compressible MHD representation used for experiments. Systems encode multichannel physical data into conditional latents and condition denoising on time, context, anchor, and direction.
- Supplementary Materials: Round-Trip Consistency: Bidirectional: The supplementary section establishes the implementation details for the bidirectional diffusion experiments.
- A.1 Embedding into Latent Space: All datasets are first compressed into lower-dimensional latent spaces using VAEs or conditional VAEs.CelebV-HQ uses a pretrained Stable Diffusion VAE, while other datasets use field-specific conditional encoders.
- A.1 Embedding into Latent Space: Field-specific conditional VAEs use FiLM modulation, residual blocks, and lowest-resolution spatial self-attention.The six MHD fields and four turbulent-radiative-layer fields are encoded independently with one-hot field conditioning.
- A.2 Denoiser Architecture and Conditioning: The denoiser is a shared diffusion transformer whose architecture remains fixed while latent geometry varies across systems.Inputs concatenate the noised target, nearest context frames, and an anchor, with learned positional embeddings and joint self-attention.
- A.2 Denoiser Architecture and Conditioning: Denoising step, simulation-time index, and direction flag form the scalar conditioning used by adaLN-Zero transformer blocks.The direction flag takes values +1 or −1, while zero initialization makes each block the identity at training start.
- A Experimental Configurations: The standard diffusion-transformer configuration uses width 384, depth 12, 6 attention heads, MLP ratio 4, conditioning width 256, and patch size 2.Face-video runs additionally use qk-normalization to prevent attention-logit growth and late-training divergence.
- A.3 2D Compressible MHD: The MHD setup follows a conservation-law formulation for modeling two-dimensional compressible magnetohydrodynamics.The represented fields are ρ, P, vx, vy, Bx, and By.
- A.3 2D Compressible MHD: 2D compressible MHD states contain six 512×512 fields, split into single-channel inputs and encoded with one-hot FiLM conditioning.Each channel is mapped to a 16×16×4 latent with β = 10−3.
A.4 External 2D Benchmark: Turbulent Radiative Layer (the Well) · A.5 Navier-Stokes · A.6 CelebV-HQ Face Videos
The external turbulent radiative-layer benchmark evaluates bidirectional rollouts on rectangular latent grids and reveals strong compression and prediction difficulty in pressure. CelebV-HQ evaluation uses identity-disjoint clips to compare latent rollout error with round-trip consistency, while variance analysis explains why per-trajectory signals differ across systems.
- A.4 External 2D Benchmark: Turbulent Radiative Layer (the Well): 90 trajectories of 101 steps at 128 × 384 cover four physical fields under varied cooling times, using official train/valid/test splits and standard VRMSE.Each field has its own VAE, producing rectangular latent grids without square-shape assumptions.
- A.4 External 2D Benchmark: Turbulent Radiative Layer (the Well): 0.30 VRMSE is the cVAE reconstruction floor on this split, rising to 0.67 for pressure and exceeding the one-step error of the two strongest baselines.The floor comes from a 64×-compressed latent bottleneck and applies to ground-truth reconstruction alone.
- A.4 External 2D Benchmark: Turbulent Radiative Layer (the Well): 2× the reconstruction floor is the approximate one-step error in every field, with pressure error concentrated because pressure is intrinsically difficult to compress and predict.The comparison uses a single stochastic sample against deterministic regressors.
- A.4 External 2D Benchmark: Turbulent Radiative Layer (the Well): Mean-reverting predictions can be favored by the long-horizon metric, so part of the observed gap may reflect metric preference rather than sample fidelity alone.This limitation arises even in the ideal reconstruction case.
- A.5 Navier-Stokes: ρ, v, P, E, and t_cool define the radiative-layer equations through density, 2D velocity, pressure, total energy, and cooling time.The benchmark is a 2D turbulent radiative layer governed by these physical quantities.
- A.6 CelebV-HQ Face Videos: CelebV-HQ evaluates latent-space MSE on identity-disjoint test clips, with decoded LPIPS as a secondary check and deterministic-sampler round-trip consistency.Autoregressive rollout examples and their C_i values appear in Figures 12 and 13.
- A.6 CelebV-HQ Face Videos: ±28% versus ±11% multiplicative variation shows that both systems have per-trajectory difficulty factors, but their magnitudes and single-cycle metric noise differ.The variance decomposition separates between-trajectory factors from within-trajectory fluctuations after removing depth trends.
- A.6 CelebV-HQ Face Videos: 0.11 is the radiative-layer fixed-depth correlation ceiling, while the MHD calibrator reduces residuals from 0.34 to 0.13 log-units, indicating near-perfect coupling of trajectory factors.The radiative-layer result is averaged over probed depths at n=80; the MHD reduction explains approximately 85% of residual variance.
C Detailed Derivations for the Sandwich Bound
This section derives the sandwich bound linking round-trip discrepancy to forward pair error under a bi-Lipschitz backward map. It also identifies the backward model’s clean-input residual as a noise floor and explains when the certificate is tight or loose.
- Sandwich-bound derivation: Under the bi-Lipschitz assumption, backward separation is bounded by repeated factors µ and L along the visited pair trajectories.The condition is required only on the set D containing the two backward trajectories, not globally.
- Certificate interpretation: The residual δ_i measures backward-model inaccuracy on clean inputs and therefore sets the detector’s noise floor even when the forward rollout is perfect.The pair-level discrepancy uses a 1/(2n) normalization so it interconverts with frame-level MSE quantities without stray constants.
- Sandwich-bound derivation: The round-trip discrepancy equals the returned seed’s distance from the true seed, enabling triangle-inequality bounds through the intermediate backward rollout.The proof combines separation propagation with the backward residual δ_i to obtain the sandwich bound and its clipped lower branch.
- Practical estimation: All constants are estimable offline from held-out backward rollouts and model-only Jacobian–vector products or finite-difference probes.The calibrator is the practical instrument when the worst-case bound becomes loose, while the lower branch explains failure through cancellation or µ approaching zero.
D Why Bidirectional Training Can Beat Specialist Training: A Linear Analysis
Bidirectional training can outperform direction specialists because, under whitened linear dynamics, forward and reverse prediction share one parameter matrix. The shared estimator uses both regression views, reducing per-direction estimation error, especially in low-data settings.
- Empirical result: At matched compute, one bidirectional model outperforms forward-only and backward-only specialists on each specialist’s own task.The ablation uses 80 training trajectories, fixed noise, 10 seeds, and direction-matched evaluation.
- Shared parameterization: For whitened latents, the reverse dynamics matrix is A−= A⊤, so forward and backward tasks share one correctly tied parameter object.Whitening corresponds to Σ = I, which approximately follows per-channel standardization.
- Estimation benefit: At matched per-direction data n, the tied estimator’s error is σ2d2/(2n): half the specialist’s σ2d2/n.The bidirectional estimator combines forward and backward observations to estimate the same underlying matrix.
- Data regime: The predicted gain is variance reduction scaling as 1/n, largest in low-data regimes, and matches delayed overfitting around epoch ∼250 versus ∼180.The ablation measures the effect with 80 trajectories.
- Directional coverage: A forward specialist’s objective does not constrain backward queries, consistent with the measured 5-7× off-direction degradation; nonlinear tying instead shares representations.The nonlinear interpretation is framed through multitask representation complexity reduction rather than one shared matrix.
E Additional Results and Protocol Details
Additional results detail safeguards and augmentations for calibration, show stronger transfer for hydrodynamic decoded fields, and characterize LE-PDE-UQ’s held-out calibration and latent-space behavior. The analyses also identify trajectory-mean calibration as useful for completed rollouts while revealing residual deep-depth drift.
- MHD calibration: +0.67–+0.73 nats are the largest per-field calibration gains for decoded hydrodynamic quantities ρ, vx, vy, and P.Magnetic components Bx and By show smaller gains of +0.32–+0.36.
- MHD calibration: +0.26 nats comes from trajectory-mean augmentation, reducing mid-depth miscalibration area from 0.23 to 0.09 at i=40.The augmentation leaves shallow-depth calibration intact; calibrator means are verified monotone and inputs are clamped to training support.
- MHD calibration: +0.01 ± 0.02 is the gain from adding temporal-shape features, indicating that per-depth C_i values function as exchangeable measurements of trajectory-level quality.Trajectory means suit post-hoc assessment of completed rollouts, while prefix means provide a causal analogue; deepest-depth residual drift is 0.27 →0.22 at i=80.
- LE-PDE-UQ protocol: 1–3 unrolled no-gradient steps provide model-generated contexts in the disclosed pushforward-style context augmentation for LE-PDE-UQ.The additions also include Gaussian noise on context and anchor frames and a short low-learning-rate finetune, while targets remain true next frames.
- LE-PDE-UQ results: 0.88 ± 0.17 Spearman correlation links C_i and E_i within test trajectories in LE-PDE-UQ’s latent space.Across trajectories at fixed depth, the correlation is 0.32-0.59; the latent calibrator improves pooled test NLL by +0.04 nats over depth-only, with gains concentrated at shallow depths.
F Inverse Rollouts
The bidirectional model supports deterministic inverse rollouts without additional training, using backward-to-forward cycling as a mirror-image round-trip error. Backward reconstructions remain coherent at shallow-to-moderate depths, while inverse performance can differ from forward performance across datasets.
- F Inverse Rollouts: The single bidirectional model performs deterministic backward rollouts with no additional training or fine-tuning, and each backward step costs the same as a forward step.The inverse direction is selected with c_d = −1.
- F Inverse Rollouts: Backward-to-forward cycling defines the mirror-image round-trip error C−_i, which tracks backward rollout error when the model is used as an inverse solver.The cycle exchanges the roles of Φ+ and Φ− in the defining round-trip formula.
- F Inverse Rollouts: After approximately 50 steps, backward prediction is easier than forward prediction on The Well turbulent-flow test trajectories because initial conditions are similar while final conditions vary widely.This asymmetry is visible across all 9 test trajectories.
- F Inverse Rollouts: Backward reconstructions of a held-out MHD trajectory remain sharp and physically coherent at shallow-to-moderate depths, then degrade gracefully as error accumulates.The rollout is seeded only with the true terminal pair.
- F Inverse Rollouts: The backward round-trip cycle is evaluated against corresponding true rollout error on 50 held-out MHD validation trajectories, with the backward-rollout axis reversed.The same comparison is shown for forward rollouts cycled backward and backward rollouts cycled forward.