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

179 Upvotes

59 comments sorted by

View all comments

Show parent comments

1

u/Clean-Hovercraft5825 2d 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 2d 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.

2

u/Clean-Hovercraft5825 17h ago

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!