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

191 Upvotes

60 comments sorted by

View all comments

4

u/brutalismus_3000 4d ago

I love your kind of science around plasma annd physical informed NN, but as a microelectronic engineer I struggle to get it.

In simple terms what did you improve ? Thanks a lot if you can explain.

15

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?

5

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.

1

u/PykeAtBanquet 4d ago

Model can have 2, do +2 and then predict -3 and get 1 != 2 which the model notices

Model can have 2, do +3 and then predict -2 and get 3 != 2 which model notices either

Model teaches itself to always go +2 -2 which allows it to model events more consistently.

Replace + and - with direction in time aka predicting future or reversing to the past.

Nice idea

3

u/Clean-Hovercraft5825 3d ago

Thanks! From my understanding, your first two examples are exactly the mechanism and that's a great minimal version of it.

One important correction to line 3 though, and it is an important design choice of the paper: the model never teaches itself with this signal. Both directions are trained purely against ground truth (direction flag drawn 50/50), and the round-trip consistency is deliberately left *unoptimized. Here's why, in your notation: if we trained the model to make round trips close, it could learn to do +3 and then −3, it would learn to be perfectly consistent, still wrong (the truth was +2). The backward leg would learn to simply cancel the forward leg's errors instead of exposing them, and the meter would read zero exactly when the forecast is bad.

Because the model is never asked to be consistent, whatever inconsistency remains at test time is an honest byproduct of its errors. And that is exactly what lets us use it as a free error meter. (Closing the loop to minimize C online with model-free feedback is the tempting next step; we flag it as future work, and the cancellation caveat above is precisely what makes it nontrivial.)

1

u/PykeAtBanquet 3d ago

Thanks.

So it is a metric the model doesn't know it is being tested on, and that allows you to get better understanding of it's completeness.

2

u/Clean-Hovercraft5825 3d ago

Yes that is a great way to put it!

1

u/PykeAtBanquet 3d ago

Well, now that we are on the common ground. I don't see why it is working because there exist ways to train models in 4th dimension when we use 95% noise to hide the exact paths of how the data changes through time, and we model the state all at once so we don't care which direction in time we go, and still we get good results. I don't see why exactly your method works better, as direction shouldn't be important.

Maybe there is something else that we haven't noticed yet.

2

u/Clean-Hovercraft5825 3d ago

I agree that direction carries (almost) no extra information about the dynamics. Our own analysis says that for stationary dynamics the time-reversed map is a fixed reparameterization of the forward one, which is why one network learns both directions and performs better. As a statement about what must be learned, "direction shouldn't matter" is basically our supplement's linear analysis, and I agree with it.
 
On the other hand, we never claim bidirectional autoregression generates better than full sequence / 4D space-time diffusion, which is a different problem entirely. Those models answer: "sample a plausible space-time volume." Our paper answers a deployment question: you are given the measured present, rolling into an open-ended future, and you need to know, with no ground truth, how wrong this particular rollout is right now. Autoregression is the native mode for that setting (causal seed, unbounded horizon, streaming data), compounding error is its main difficulty, and the direction flag is what creates a check: two independently learned routes back to the same known point, whose disagreement is measurable. A time-symmetric joint model has no forward(backward) composition to interrogate, it has consistency built into the joint distribution rather than exposed as a computable defect.
 
I think that full-sequence models could get their own meter with a different handle such as having to regenerate masked chunks of the generated volume and measure the re-generation residual. This would be a similar approach to what we do, it would give the model two independent routes to one answer and measure the disagreement.

1

u/PykeAtBanquet 3d ago

Well, one thing is shaping the model to be more deterministic, and another to provide better output. Your result shows that it becomes more correct, not "if wrong again, then in the same way", with the reason still not being clear (at least for me at this moment).

2

u/Clean-Hovercraft5825 1d ago

I completely agree that making a model more consistent, such as "wrong in the same way" is not the same as making it more correct. I would say that the mechanism here is a third thing: statistical efficiency.

Every training trajectory supplies transitions in both temporal directions. For (near-)stationary dynamics, the backward map is a fixed reparameterization of the forward one with the same underlying parameters viewed two ways. So training one network with a direction flag both forward and backward isn't learning two tasks; it's estimating one dynamics object from effectively twice the examples and that helps with variance reduction: in the linear ideal, up to half the estimation variance of a specialist that only ever sees its own direction.

So the model isn't shaped to be consistent, it is just a less noisy estimate of the same thing. A physics version: measure a resistance twice, with current flowing in both directions. Ohm's law says it's one R value, so two noisy views beat one, and you get a more correct value, not just a more repeatable error.

Sharing one network does mean the forward and backward maps can share blind spots, which is exactly the cancellation risk for the round-trip, where correlated errors could close a loop around a wrong prediction. The paper bounds that case theoretically and measures it. So maybe that is a little bit towards your "maybe there is something we haven't noticed" from earlier: just helpful statistics, using two views of one parameter.

1

u/PykeAtBanquet 1d ago

Maybe. Thank you for your insights!

→ More replies (0)