7 comments

[ 2.8 ms ] story [ 33.2 ms ] thread
Incredibly cool work, and a great primer on diffusion
Really interesting to see the diffusion model solve the puzzles in an iterative way, which feels more similar to how I (and probably most humans) solve them.

Outwardly, it seems to be limited by unmasking too few tokens per round, even when the heatmap shows many more high-confidence guesses available. On some of the larger puzzles it looks like it's wasting many rounds filling in the 'obvious' shapes, and then gets the interesting bit in the last round. It also doesn't seem to have learned the idea of "the background is blue with shapes drawn on top," where background is often 50% of the solution in these puzzles.

It is kind of wild that most coding tasks are editing tasks, and we humans care a lot about code editing tools, but automated tools use code generation for editing where a valid block must be generated top-to-bottom in one go.

Fixing a mistake requires re-generating the file or block of code. Or, if something generated later has implications for earlier code--a new import or function parameter's required, something like that--the only option is to go back and re-generate a big chunk. That'd be inefficient for humans, not implausible it's wrong for other code generators too.

I don't know if diffusion specifically will be the approach. (Maybe there's something to generating edit sequences?) This post's note that diffusion kills KV caching is something I hadn't even considered. It does seem right to experiment with things other than strict start-to-end generation.

Regarding the typewriter approach, I've wondered for a while if anyone has explored simple backtracking with LLMs? Like, have the LLM be able to generate a backspace/delete token that lets it "undo" previously generated tokens in an append-only fashion. Not sure how this would work with teacher forcing but seems feasible with RL.
So, not exactly the same thing at all, but the ARChitects do a really cool thing I didn't have time to talk about in this post, which is a kind of depth first search with a cumulative "minimum probability" threshold for backing out of a path. This does let the model kind of reason ahead a few tokens, and then back out if it doens't look like it's going well and try the next most likely token. https://github.com/da-fr/arc-prize-2024/blob/main/training_c...

You can image something like that for any autoregressive llm, but it probably needs some heavy heuristics. Here there are like 11 valid tokens (end of line, 1-9, or end of sequence), and other use cases are going to have way more options making this more intractable.

How would it be different from regular beam search?
Yes, I thought of the analogy with beam search but here I am proposing to add the backspace tokens to the context, not to actually rewind the context.