Diffusion Forcing: Next-Token Prediction Meets Full-Sequence Diffusion

A number of ideas seem notable to me here; first, they are merging the idea of sequence masking (the key training idea for LLMs) with diffusion models; they do this by keeping track of an ‘uncertainty’ level per pixel. This ‘uncertainty’ level is treated as the ‘noise’ level for the diffusion model, (a model which denoises controlled by some sort of embedding).

There are a bunch of neat things you can do with this: in particular, you can firm up parts of the image earlier than others, and thus use it for, say maze solving. They even show it controlling a robot arm moving fruit around, which is pretty wild.

In a way the title undersells the idea - this is a way to do fractional masking, since the masking level is a float - and I think is really a pretty profound and interesting idea.

However, there’s a lot not talked about in this paper; I’d be very curious to see their codebase. How exactly do you set up a maze-following task vs a video extension task? How do you hook up a robot arm to this model, and tell the model what you want done? The architecture itself deserves a significant number of papers / explication.

Anyone know of research or tools for using an existing text generating LLM with diffusion like techniques with no new pre-training, or at most, a bit of fine-tuning, such that it works with a small GPT / Phi 3 / Gwen model, for example? I know about Tree of Thoughts with MCTS etc, that are somewhat similar (though often with a different reward learned goal) but I'm interested in something closer to token level generation. Is this possible?

I work in the field and the work is presented in an extremely obtuse manner.

What is the problem you're trying to solve? Are you proposing a new generative model?

Russ is doing diffusion now? Must be very applicable to robotics.

Am I missing something about training time? Does adding per token noise cause training to slow significantly? Cool paper though!

Cool work! Curious if this can be applied back in LLMs as a discrete diffusion model with partial masking.

Very cool, but why is it called diffusion forcing?