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

2026-07-02

AI Summary

Existing discrete diffusion language models are constrained by fixed-length generation and lack support for KV caching, while block diffusion models suffer from limited decoding flexibility and parallelism due to their reliance on sequentially fixed blocks. This work proposes Set Diffusion Language Models, which generalize token decoding to arbitrary ordered sets for the first time. By introducing a set-factorized likelihood parameterization and a set-causal diffusion architecture, the approach enables dynamic KV cache updates and decoding in any order. The method naturally accommodates flexible strategies such as sliding windows, achieving superior speed-quality trade-offs in mathematical reasoning, summarization, and unconditional text generation, while also significantly improving performance in text infilling tasks.

Abstract

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/

Problem

• discrete diffusion
• autoregressive models
• KV caching
• flexible decoding
• token ordering

Innovation

• set diffusion
• flexible-order decoding
• KV caching
• discrete diffusion models
• arbitrary token ordering