Set Diffusion: Interpolating Token Orderings between Autoregression and Diffusion for Fast and Flexible Decoding

Discrete diffusion models have steadily improved in quality relative to autoregressive (AR) models. However, these models are normally constrained to fixed-length generation and do not support key-value (KV) caching. Block diffusion partially bridges diffusion and AR by unmasking token blocks left-to-right, but it is still limited to generate fixed-size blocks sequentially. Here, we present a new class of language models, set diffusion, comprised of (i) a likelihood parameterization that factorizes over flexible-position, flexible-length token sets and (ii) a set-causal diffusion architecture that supports KV cache updates after every inference step. By factorizing over token sets instead of fixed-size blocks, tokens can be decoded in arbitrarily-ordered sets, including sliding-window sets, enabling faster inference and support for any-order decoding. Set diffusion achieves better speed-quality tradeoffs on mathematical reasoning, summarization, and unconditional generation compared to prior diffusion language models while offering stronger infilling performance than block diffusion. We provide the code, along with the model weights and blog post on the project page: https://m-arriola.com/setdlms/

Large language models write text one word at a time from left to right. This works well, but it can be slow and makes some tasks awkward, such as filling in missing text in the middle of a paragraph or editing text at arbitrary positions. Recent diffusion language models can generate and edit text at many positions in parallel, but they are usually designed to produce text of a fixed length and have difficulty reusing previous computations during generation. This paper introduces a new way for language models to generate text by choosing flexible groups of positions to fill in at each step. The model can behave more like a standard left-to-right model when quality is most important, or more like a parallel generator when speed and flexibility are more important. It can also reuse previous computations as it writes, helping generation run faster. We show that this approach improves the balance between speed and quality across tasks including math reasoning, summarization, and open-ended text generation. It is especially useful for filling in missing text inside an existing passage.