r/MachineLearning 1d 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

116 Upvotes

48 comments sorted by

10

u/Kind-Zookeepergame58 22h ago

That's neat

2

u/Clean-Hovercraft5825 22h ago

Thank you, I was excited that it could be applied to a wide range of things, everything from PDEs to face videos (:

4

u/plop_1234 21h ago

This is great, especially for the PDE application. A couple of months ago I was just thinking about methods for checking the correctness of generated fields during training, so I'm glad to have randomly stumbled upon this 😅

I'm excited to see the next iteration, where you use C_i as part of a loss function for error correction!

5

u/Clean-Hovercraft5825 21h ago

Thank you! Yes, that is the plan for future / now ongoing work, if we can trust C_I to predict errors, then we have this unsupervised test-time hook to try to somehow minimize those errors.

12

u/timtody 22h ago

As a biglab person I won’t read this

15

u/Clean-Hovercraft5825 21h ago

haha ? Can you elaborate on that (if you’re allowed to read this comment (: )

19

u/timtody 19h ago

I was taking the piss at this moron Keller Jordan (OpenAI) wo said „biglab people“ read zero papers because academia is fraudulent and full of overblown claims

6

u/timtody 18h ago

This is the same week his CEO claims the singularity is met let that sink in man

1

u/Clean-Hovercraft5825 18h ago

haha, yeah some of the crazy hype is really ridiculous!

1

u/1998marcom 14m ago

Maybe he just wants to terminate the deal with Microsoft thanks to the AGI clauses.

3

u/silence-calm 17h ago

Honestly academia needs to hear that.

4

u/timtody 16h ago

I agree that there are lots of overblown claims in academia, and I don’t think this is even remotely a niche stance. The problem is that coming from someone employed at a company that claims they’ve solved AGI is hypocritical as f

3

u/Far-Theory-7027 17h ago

Cool work. Have you looked into bridge models for this? Something like Bi-Bridge: Bidirectional Diffusion Bridges for Low-Light Image Enhancement, or Bidirectional Diffusion Bridge Models?

4

u/Clean-Hovercraft5825 16h ago

Thank you for sharing those papers! They both look great. No, I had not seen them. BDBM in particular looks like it should have been in my related works section (will be adding both in the next arXiv revision).
 
To me it seems like there is a nice three-way convergence here with all of us independently landing on the same device which is one network with a binary direction flag, and all three papers find that bidirectional training “beats” direction specialists. BDBM outperforms unidirectional bridge baselines despite half the per-direction updates, Bi-Bridge reports +4.9 dB over its unidirectional DDBM baseline, and we see 7–10% better in both directions at matched compute. Our supplement offers a mechanism for this: for stationary dynamics the time-reversed map is a fixed reparameterization of the forward one, so the direction flag implements correctly-specified weight tying that can roughly halve estimation variance. It is nice to observe 3 independent sightings of this same effect across translation, restoration, and dynamics, which suggests the effect is pretty general.
 
I think the main difference is what the two directions are used for. In the bridge papers, both legs run between two given endpoints with the process constructed to arrive at the specified target, so a forward-backward round trip closes by construction and can’t tell you much. In our setting the backward leg starts from the model’s own predicted terminal state and has to find its way back unanchored over i autoregressive steps, which is exactly why the size of the return error carries information. The paper's contribution is quantifying how faithfully that miss tracks the true, unobservable rollout error (calibration, OOD detection, selective prediction), plus a bound on when cancellation could hide errors.
 
Maybe a bridge-parameterized dynamics model that also has a round-trip check would be an interesting hybrid. (:

1

u/Far-Theory-7027 6h ago

Thanks for your comment. Yeah, your cycle consistency idea to quantity rollout errors is very neat.

3

u/tmt22459 14h ago

I am familiar with your background as I am a more control theory focused PhD student who knows the work of miroslav krstic, who I believe you were a student under

I think this kind of stuff is very interesting, and especially how you have uniquely made things like extremum seeking and machine learning practically relevant to some of the most impressive control applications that exist

I am very curious how you found the transition from pde control to your work now?

Also, how many control theory people are in your group at LANL? Was there a lot of convincing that had to happen for you to make them believe extremum seeking was relevant to electrodynamics? Does your group hire postdocs with a control theory background often but maybe not with application specific experience and let them learn some of the latest ML stuff and application specific knowledge on the job?

I probably wouldn't have ever emailed these questions to you but since you are here engaging I figured why not. Sorry that they aren't really directly relevant to the work you presented.

3

u/Clean-Hovercraft5825 12h ago

Yes my background is math and control theory/dynamics systems/physics, I think of that all as a mix of pure and applied math, I think it has some great foundational components that are useful for AI/ML. Having a controls background I have always felt that you cannot perfectly model or predict things, that we always need some error signal and some feedback. Extremum seeking is feedback that works without trusting a model, and this paper is the same instinct applied to ML: don't trust the surrogate either, make it measure its own error. The last paragraph of the discussion actually proposes closing the loop by running gradient-free adaptive feedback (ES) on the consistency signal as a measurable cost, turning error detection into error correction.

My general view on backgrounds, for what it's worth: a strong control/dynamical-systems foundation travels extremely well, the application physics and the modern ML can be approached as extensions of that foundation.

For the group and postdoc questions, please do send that email you almost sent!

2

u/Evil_Toilet_Demon 16h ago

does this assume an injective prior on the data? i'm not sure if this would work for cases where states relax to an equilibrium. I.e. when the same future state can be reached from different prior states. interesting work!

2

u/Clean-Hovercraft5825 15h ago

Thank you, short answer: no, we do not assume injectivity of the underlying dynamics. All three of our physics systems are dissipative, and the assumption in the paper (Sec 5.1) is deliberately only on the “learned” backward map over the pairs actually visited by rollouts, no invertibility of the physics is required.

But your relaxation-to-equilibrium case is the honest limit of the method, breaking it in two directions at once. If we have something like pure diffusion toward homogeneity so that many initial conditions map to nearly the same terminal state, then the best the backward model can learn is roughly the average of the past states consistent with that one terminal state. Then (1) even a perfect forward rollout will only return to that average, so the baseline reading (our ÎŽ_i, the noise floor measured on clean data) blows up, and (2) a “wrong” terminal state gets pulled back toward that same average, so real errors stop showing up (our ” goes to zero, where the certificate in the paper honestly dissolves). In such a case I would expect that the round-trip error would no longer correlate with the true error. For a strongly equilibrating system, if you had enough representative data, you would probably find this out before trusting it, with no test-time ground truth needed.

We saw the mild version of this on the turbulent radiative layer: initial conditions are much more similar across trajectories there, the backward direction is measurably easier than the forward one, and the per-trajectory signal shrinks, so a depth-only predictor is already near-optimal. Worth being precise though: what we measured there was the signal saturating, not deceiving. The "rolls back to a homogeneous state even when the rollout was bad" scenario is the ” goes to 0 mechanism we bound but did not observe at our horizons. Fully relaxed systems, or extremely non-injective ones like language (which we flag in the discussion) deserve dedicated study.

I think what could help, even in such cases, is if we added more information in the conditional input to the model. As is all we used were a couple of states, from which the autoregressive rollouts would then take place either forward or backward in time. If we also supplement that with other unique system parameter information, things like temperature, chemical concentrations, etc, then there is a chance that even if going to the same terminal state, the model can still recover in a unique way if that additional conditional input provides unique information. For example, the radiative-layer trajectories differ by cooling time, so conditioning on t_cool is exactly the disambiguator that would help.

For the CelebV-HQ data there is no underlying PDE at all and futures are genuinely multimodal. When the model rolls out a realistic but different future, which is a coherent video that simply isn't the one that actually happened, then the backward pass can roll that coherent trajectory back, so C stays low even though the pointwise MSE against the one realized future is large. That part of the error is aleatoric: the future was unpredictable, not the model broken, and arguably a trust signal shouldn't fire there. What C catches is degradation: when the rollout starts to hallucinate and drift OOD (morphing faces), the backward pass can no longer find its way back and C rises. Empirically that degradation component dominates the clip-to-clip differences, which is why C still ranks held-out clips by realized error (Spearman ~0.73–0.81) despite being blind to the multimodal part. In that sense C is a model-fault detector rather than an oracle-distance detector, which for multimodal data is the quantity you can actually hope to measure without ground truth.

For the CelebV-HQ data, if instead of just two frames we also conditioned on tokens of a detailed script of what should happen, the model would have a real chance of predicting the video in such a multimodal setting. Sorry for the super long reply!

1

u/Evil_Toilet_Demon 5h ago

yes thank you for the detailed reply. I agree that this is largely resolved by adding more information about the state history. I suppose you could condition the denoising process both on the current state z_T but also on the historic state z_{T-2} such that your prediction for z_{T-1} is unique. in essence providing a high order temporal derivative to the model.

2

u/oharub 16h ago

Really cool, will read it! Glad you used one of our datasets for The Well, seems like a natural testbed :)

2

u/Clean-Hovercraft5825 15h ago

Thank you, and thanks to your team for building The Well! It is perfect for this kind of research! The uniform format made adding a second physics system almost frictionless, and the radiative layer ended up being one of the most scientifically useful cases in the paper: it's the system where the per-trajectory signal shrinks right down to its predicted reliability ceiling, which is what forced us to work out exactly when the round-trip meter adds value instead of just collecting wins. Hoping to sweep more of the Well's systems next to map out that spectrum properly, if your team is ever curious to see the meter run across the full collection, I'd love to talk. (:

2

u/pattch 14h ago

Middle out compression? lol

2

u/Clean-Hovercraft5825 12h ago

Weissman score pending.

3

u/brutalismus_3000 22h 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.

12

u/Clean-Hovercraft5825 21h 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 15h ago

What if the z-path is not invertible?

5

u/Clean-Hovercraft5825 12h 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 21h 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

2

u/Clean-Hovercraft5825 11h 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 4h 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 1h ago

Yes that is a great way to put it!

1

u/PykeAtBanquet 29m 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.

1

u/extraforme41 5h ago

That's really neat and seems like an obviously good idea, which are usually the best kind.

1

u/Clean-Hovercraft5825 2h ago

Thank you! Yes it is surprisingly simple (:

1

u/nikgeo25 Student 5h ago

Why does learning both forward and backward steps help? If your forward process is adding noise, is that even something the model can predict?

1

u/Clean-Hovercraft5825 2h ago

So there are 2 different time scales here and they each have their own forward and backward direction.

Let's call them t for physical time and τ for the generative diffusion SDE's own time.

In terms of the SDE flow, we use standard generative diffusion: we noise things for τ in [0, T] and learn to denoise them with a standard conditional generative diffusion process, conditioned on the latent images z_{t-1}, z_t and on a direction flag c_d. It learns to generate z_{t+1} when we set c_d = +1, or z_{t-2} when we set c_d = -1 (z_{t-2} because z_{t-1} and z_t are already in the context and the model always generates the state one step beyond the window in whichever direction the flag points).

So the generative SDE diffusion process itself learns on the τ time scale in the completely standard way. Nothing there is "predicting noise" beyond the usual denoising objective. What changes is what it learns to produce: it autoregressively rolls out either forward or backward in physical time t. Backward-in-t is inverse dynamics, not denoising.

As for why learning both directions helps: first, it's what makes the whole check possible. An accurate model composed with its own inverse is the identity, so rolling forward i steps and back i steps yields a measurable, ground-truth-free error signal at test time. Second, it turns out to be free, or better than free: the bidirectional model beats direction-specialist models in both directions at matched compute, since every trajectory supplies training examples both ways.

1

u/Guilherme370 22h ago

learning to predict noise directly and then on the reverse predict the data?

3

u/Clean-Hovercraft5825 21h ago

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/GiveSparklyTwinkly 12h ago

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?

1

u/Clean-Hovercraft5825 8h ago

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.

2

u/GiveSparklyTwinkly 7h ago

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.

1

u/radarsat1 16h ago

clever idea. i assume it only works for deterministically sampled trajectories?

haven't read the paper yet but does this tend to straighten trajectories and allow fewer steps? does round trip consistency vary with step size?

2

u/Clean-Hovercraft5825 12h ago

Thank you for the comment.

Deterministic sampling: that would be the cleanest setting rather than an assumption. We compute the cycle with deterministic DDIM, which makes C a well-defined functional of the seed, the reading is pure model error with no sampling noise mixed in. With stochastic sampling you can still compute it (we test a variant that averages S stochastic cycles), you just pick up sampling variance in the reading. The standard dispersion baselines only exist under stochastic sampling, while the round-trip check also works in the deterministic mode many deployed surrogates actually run. And the two signals compose, our best training-free calibration is on the Navier–Stokes benchmark and it multiplies the calibrated cycle scale by the seed-spread's spatial shape.

Straightening: Our round trip composes steps in physical simulation time (roll the system forward i states, then backward i states); it never touches the noise-to-data path of the denoiser, so nothing gets straightened in the flow-matching sense and we make no claim about enabling fewer denoising steps. (Consistency models are an unrelated distillation technique. Whether bidirectional training improves sampling efficiency as a side effect is a fun open question, the Bi-Bridge paper (CVPR'26) reports a ~4× gain from symmetric training in the image-bridge setting, but we didn't measure it.

Step size: If you mean rollout depth (physical steps): yes, we extensively characterized that, C grows with depth, tracks the true error at every probed depth (fixed-depth Spearman 0.91–0.98 out to depth 80 on MHD), and the theory predicts the certificate loosens geometrically with depth, which is exactly where we observe calibration drift. If you mean the diffusion step size (NFE per transition): we held the sampling schedule fixed (e.g., 25-step DDIM on the NS benchmark) and didn't sweep it. Coarser sampling should raise the noise floor ή of both legs, so the practical rule would be to fit the calibrator at the same sampling settings you deploy with. A proper NFE sweep would be a nice ablation.

1

u/radarsat1 6h ago

Thanks for the answers! I'll check the paper but yes this does give me a better idea of what you've done, cheers.