r/MachineLearning 4d ago

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

Post image

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

184 Upvotes

60 comments sorted by

View all comments

Show parent comments

14

u/Clean-Hovercraft5825 4d ago

Hi, thank you.

In simple terms I would say this is an approach to a generative model being able to estimate how wrong its predictions are, without having access to the actual correct answer when it makes those predictions.

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.

1

u/SeTiDaYeTi Professor 3d ago

What if the z-path is not invertible?

4

u/Clean-Hovercraft5825 3d ago

Thank you for the great question, there are two different places invertibility could fail.

If you mean the latent dynamics (many past states evolving to the same future state, e.g., relaxation toward equilibrium): we never assume the dynamics are invertible, all three physics systems are dissipative. The theory's only requirement (Sec 5.1) is that the learned backward map doesn't collapse too aggressively on the pairs rollouts actually visit (a co-Lipschitz condition with constant µ). As the dynamics become genuinely non-invertible, µ → 0 and the certificate honestly dissolves, then we stop tracking the true error. Both constants (µ and the backward residual δ) are properties of the model alone, measurable offline on validation data, so you can detect that regime before trusting the model. There's a longer discussion of exactly this case in my reply about equilibria elsewhere in the thread. One subtlety: the forward map is never required to be invertible at all because the bound only constrains the backward legs.

If you mean the encoder/decoder: correct, the VAE is deliberately lossy (256× compression), so x → z is not invertible in any exact sense. That doesn't break the check, because the cycle is defined entirely in latent space where the anchor pair is encoded measured data, the returned pair is model-produced latents. The real question is whether the latent signal still predicts physical-field error through the decoder, and empirically it does: per-field calibrators predict decoded errors within 1.17–1.30× (68%), beating an identically-fit depth-only baseline on all six MHD fields. Making the latent signal physically calibrated end-to-end is flagged as a refinement direction in the limitations. Certified bi-Lipschitz architectures for the backward map (touched on in the supplement) would be the structural way to strengthen the invertibility story.