Explorative modeling: Train on the best of K guesses

I was just reading this great breakdown of how diffusion Gemma works: https://newsletter.maartengrootendorst.com/p/a-visual-guide-...

In reference to the difficulties with applying this to autoregressive LLMs - I wonder if these type of hybrids might be a good candidate for this approach.

1. K-1 extra forward passes during training

2. Inaccurate sampling behavior (will sample all K modes with equal likelihood, rather than sampling them proportionally)

However, I think both of these downsides can be mitigated by adjusting the implementation a bit more (you can have the model predict K modes jointly in one forward pass, along with probabilities of each being the min-loss mode, which you can then use for properly-weighted mode sampling at inference).

That said, I'm not entirely sure if this strategy is as generally applicable as the authors are hoping. In particular:

1. For highly-conditional image generation (like modern commercial diffusion pipelines, which use a big LLM preprocessor), most of the low-frequency color/layout decisions are already made for you by the conditioning signal. The diffusion process mostly needs to generate high-frequency details, for which there are a huge number of equally-valid modes.

2. For LLMs themselves, the sequence-generation process is already factored into a discrete classification problem, and there's no mode smearing issue to fix.

Minibatch OT in flow matching also has a very similar mechanism, where samples from a noise distribution are matched to the closest data point.

There is a lot of prior work here that the authors neglect to discuss, which portrays this work as more novel than it actually is.