Search papers, labs, and topics across Lattice.
This paper introduces set diffusion, a novel class of language models that enhances the decoding process by allowing flexible-position and flexible-length token sets, thereby overcoming the limitations of fixed-size blocks in existing block diffusion models. By employing a set-causal diffusion architecture that supports key-value caching after each inference step, set diffusion enables faster inference and arbitrary token ordering. The results demonstrate superior speed-quality tradeoffs in tasks such as mathematical reasoning, summarization, and unconditional generation, while also improving infilling performance over previous methods.
Set diffusion achieves faster and more flexible decoding by allowing arbitrary token ordering, outperforming traditional diffusion models in key tasks.
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 generating token blocks left-to-right, but its fixed-size sequential blocks limit decoding flexibility and parallelism. 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/
