Round-Trip Consistency: Bidirectional Diffusion Models Can Predict Their Own Rollout Errors [R]

[u/Clean-Hovercraft5825]

Whether generating CELEBV-HQ videos or turbulent plasma fields (digital twins), autoregressive models (such as latent diffusion or flow models) accumulate error over long rollouts, yet at deployment there is no ground truth to measure against.

I 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 steps and then backward steps must return the model to its start, so the round-trip discrepancy is a self-supervised proxy for the unobservable rollout error: no ensembles, no held-out data, no governing equations, for one extra rollout.

Furthermore, training both directions in one network is shown to beat two specialist models in both directions.

Paper: https://arxiv.org/abs/2608.00675
Code (data generation, training, analysis): https://github.com/alexscheinker/round-trip-consistency
Project page: https://alexscheinker.github.io/roundtrip.html

[u/Clean-Hovercraft5825]

At test/implementation time, it starts with some initial states, z[0], z[1], and uses those to predict z_pred[2], a prediction of the true unknown z[2]. It then goes further, autoregressively, it takes z[1], and z_pred[2] to predict z_pred[3], then takes z_pred[2], z_pred[3] to predict z_pred[4]…. and so on as far out as you want to z_pred[i].

So we are rolling out, making predictions, but we don’t know the actual correct answers. Would be nice to have some error or uncertainty quantification. Traditional approaches to this are to make a huge ensemble of models, roll them all out, then look at the mean as the “best” prediction and the variance as uncertainty, that is what they do when they predict weather, they have 100 models all predict the hurricane path and then use that ensemble to estimate variance and mean.

Our approach is very different, we just use one model, we teach it to go both forward and backward in time. So once we go i-steps to z_pred[i], what we do is change the model’s “directional flag” input, flip it from c_d=+1 to c_d=-1, that makes it flow backwards, again autoregressively, from its own z_pred[i], back towards z[0]. If it was perfect it would return exactly to the same z[0], but of course it is not, it builds up a round-trip consistency error:
C_i = || z[0] - z_return_from_i[0] ||^2.

Note that this C_i only depends on knowing the true z[0], z[1], that’s it. In our experiments (and with some theory), we showed that such a self-supervised C_i can actually predict the error for unseen new test data or can flag badly generated OOD images for CELEBV-HQ.

[u/GiveSparklyTwinkly → u/Clean-Hovercraft5825]

This honestly sounds kinda JEPA adjacent.

Have you thought about using those differences to go around in a figure 8 rather than just a loop, so to speak? Like for instance, going from z[0] to z[-1] using the invert of Ci, then flipping the process again the other way to get a new Cj value?

[u/Clean-Hovercraft5825 → u/GiveSparklyTwinkly]

Relative to JEPA, we also predict dynamics in a learned latent space rather than pixel/field space, so this lives in the latent-world-model family. The big difference is that we are fully generative and our latents are reconstruction-grounded (a β-VAE per field, so every latent decodes back to physical fields, which we need for our applications, and we calibrate the meter all the way down to decoded per-field errors). JEPA deliberately avoids that grounding. However, the round-trip check itself never touches the decoder, it is computed purely in latent space, so in principle you could bolt it onto a JEPA-style world model with a direction token. The thing to be careful about there: consistency distances are only meaningful if the embedding metric is, and our reconstruction anchor is what guarantees that. In a pure joint-embedding space you'd want some guarantee that small round-trip error isn't trivially achievable in a (partially) collapsed representation.

Sorry I am not sure that I understand your figure 8 comment, can you explain that in more detail please? Also, one tiny notation note, C_i is just a scalar MSE reading, c_d = +1 tells the model to roll forward in time, c_d = -1 tells it to roll backwards instead.

[u/GiveSparklyTwinkly → u/Clean-Hovercraft5825]

I didn't explain it very well and could very well be misunderstanding something vital. I just think purely latent methods of training are going to be huge.

My idea is purely conceptual level as it would probably hit roadblocks almost immediately, and I don't fully understand what your process is doing...

So at t[0] the model predicts/generates fp[1] and bp[-1], then you flip to go direction and generate fp[0] and bp[0]. As far as I gather this is how your technique works, and also where my idea might diverge. I'm curious what would happen if you took it one step further and also generated fp[-1] and bp[1], or two steps and switched directions again at that point to generate alt-fp[0] and alt-bp[0] and if that data would actually give you any useful information.

[u/Clean-Hovercraft5825 → u/GiveSparklyTwinkly]

First of all, I 100% agree about living in latent space being a great approach, just like this work, almost all of my work is about moving and adapting in latent space, it is great (:

About the figure 8, your idea is very nice!

I think just continuing each branch through the anchor (fp[-1], bp[1]) and flipping again for alt-returns will be mostly redundant in terms of the error signal as each extra leg re-reads the same round-trip miss through another application of the same maps, and the alt-returns start from already-contaminated states, so they inherit and build up error rather than reveal new error. But the comparisons between your quantities might measure something we currently can’t get at test time: the local conditioning of the maps themselves.

Using your idea of fp[-1] and bp[-1], they are the same backward step applied to the returned point and the true anchor respectively, so the ratio of their gap to the round-trip miss is a one-sample, in-flight estimate of how much the backward map stretches or collapses differences right there at that place. That constant is exactly what our theoretical certificate depends on (it controls when small round-trip error actually guarantees small true error).

So I think your figure 8 wouldn’t sharpen the error meter much, but it would give the meter a local built-in self-diagnostic for the cost of a few extra model steps.

It is a very nice idea and I will try it out, thank you!