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

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 LLMs, 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 LLMs 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/

• The paper introduces set diffusion, a novel framework that factorizes generation likelihood over arbitrary token sets, interpolating between autoregressive and diffusion models.
• It employs a set-causal transformer architecture with efficient KV caching and sliding-window decoding, delivering superior speed-quality tradeoffs.
• Empirical results demonstrate improved performance in summarization, infilling, and reasoning tasks, achieving faster throughput and reduced gradient variance.

Set Diffusion: Flexible Language Modeling via Token Set Factorization

Motivation and Context

The paper introduces set diffusion, a paradigm for discrete generative modeling in language tasks that interpolates between traditional autoregressive (AR) models and discrete denoising diffusion models (D3PMs). Unlike previous frameworks—block diffusion, masked diffusion, or AO-ARM—the central innovation is factorizing the generative likelihood over arbitrarily-positioned and variable-length token sets, rather than contiguous left-to-right blocks or single-token increments. This enables greater decoding flexibility, parallelism, and efficiency while permitting practical features such as variable-length output and KV caching after every inference step.

Discrete diffusion models have recently shown improved controllability and faster generation relative to AR baselines but often lack support for flexible-length generation and efficient KV cache utilization, hampering their applicability to tasks such as infilling and arbitrary-position insertion. Block diffusion partially addresses these limitations but still imposes rigidity via fixed-size blockwise decoding. Set diffusion generalizes these approaches, recovering AR, fully non-AR diffusion, block diffusion, and any-order AR models through its ordering distribution.

Methodological Framework

Probabilistic Model

Set diffusion parameterizes the likelihood over flexible token sets. The sequence is partitioned into disjoint sets, each generated according to a chosen ordering distribution $T(o)$ and modeled by a conditional diffusion process. The likelihood is marginal over generation orders:

$$P(x) = \sum_o \prod_{n=1}^N T(o_n | x_{<o_n}) \cdot p_\epsilon(x_{o_n} | x_{<o_n})$$

Each token set is sampled and denoised conditioned on the previously revealed context. Training efficiently marginalizes over orders via variational inference, typically with parameter-independent $T(o)$ for tractable optimization.

Set-Causal Architecture

The architectural backbone is a transformer modified to support arbitrary-position token sets via attention masks that enforce causality along sampled generation orders. The network outputs token probabilities and KV cache updates at each step, decoupling the need for full bidirectional context as in masked diffusion.

KV caching is enabled after every inference step, rather than only after block completion, facilitating practical acceleration for large-scale models. The architecture supports sliding-window and arbitrary set decoding, which enhances both throughput and generation flexibility.

Interpolation Between AR and Diffusive Paradigms

Set diffusion flexibly interpolates between AR and diffusion by tuning (1) the size of decoded token sets and (2) the ordering distribution $T(o)$. Strong left-to-right biases replicate AR; uniform distributions recover order-agnostic diffusion; contiguous block schedules recover block diffusion; singleton sets with uniform orders recover AO-ARM.

Position-Dependent Ordering Schedules

The authors propose structured ordering schedules, assigning reveal times to each token position. Varying the window size and reveal distributions allows fine-grained control over generation order and parallelization tradeoff. As the active interval width decreases, the NELBO tightens toward the AR objective.

Sliding-Window Set Diffusion LLMs (SW-SetDLM)

A practical instantiation, SW-SetDLM, implements set diffusion with position-offset ordering distributions and causal transformers. Training specializes to singleton sets for efficiency and lower gradient variance, attaining a token-efficient loss where every token contributes likelihood gradients per ordering.

Inference proceeds with sliding-window decoding, accepting arbitrary sets of tokens at each step and updating the KV cache, maximizing parallelism and flexibility. The expected inference prediction budget is tunable for hardware and task constraints.

Empirical Evaluation

Speed-Accuracy Tradeoffs

Set diffusion models are evaluated on GSM8K mathematical reasoning, CNN/DailyMail summarization, ROCStories infilling, unconditional OWT generation, and LM1B likelihood estimation. Across all tasks, set diffusion establishes strictly better speed-quality Pareto frontiers than block diffusion and masked diffusion baselines.

Key numerical results:

• GSM8K zero-shot pass@1: SW-SetDLM achieves up to 66.41% with 60.42 tokens/sec at S≤8, outperforming block diffusion and masked diffusion.
• Infilling (ROCStories): Higher ROUGE-1/2/L scores and ~25% faster throughput than block diffusion.
• Summarization: Comparable or higher ROUGE scores to AR and diffusion baselines, up to 10% faster decoding than block diffusion.
• OWT (unconditional generation): MAUVE score 71.98 for SW-SetDLM(S≤8) with >91 tokens/sec, exceeding block diffusion in throughput and sample quality.
• LM1B: Best diffusion test perplexity (≤27.83) for SW-SetDLM(S≤8) at 98.5 tokens/sec.

Ablations confirm the criticality of efficient training via comprehensive likelihood factorization and KV caching, with gradient variance reduced by ~3x and convergence accelerated relative to masked diffusion objectives.

Infilling and Arbitrary-Order Decoding

Set diffusion's flexible token set prediction allows for efficient infilling of one or multiple spans, outperforming block diffusion (whose rigid block order constrains insertions). MDLM is notably slower due to lack of KV caching, requiring repeated full-context denoising.

Likelihood Estimation and Policy Gradient

The objective supports one-pass likelihood computation for entire sequences, facilitating efficient policy gradient steps in reinforcement learning frameworks.

Related Work Comparison

Set diffusion differs from MDLMs and block diffusion by enabling position-dependent generation scheduling and supporting flexible-length generation with causality and KV caching. Unlike FlexMDM and AO-ARM, set diffusion does not require auxiliary insertion networks and achieves tighter NELBO via structured ordering priors.

Esoteric LLMs and WeDLM employ causal architectures for KV caching but lack variable-length generation or interpolation between AR and diffusion via order distributions.

Implications and Future Directions

Set diffusion advances practical and theoretical understanding of discrete generative modeling for language, demonstrating that interpolating between AR and diffusion via generation orders is more effective and flexible than blockwise factorization. The method supports efficient KV caching, hardware-aware parallel decoding, infilling with arbitrary spans, and flexible-length generation.

The position-offset ordering schedules can be hand-tuned for task/hardware constraints or learned during training (as in [Kim et al., 2026], [Wang et al., 2025c]), which may further optimize speed-quality tradeoffs. The token-efficient loss opens doors for scalable reinforcement-learning-based optimization.

As the accuracy-parallelization gap narrows between diffusion and AR, set diffusion is poised for deployment in high-throughput language generation, editing, and reasoning systems requiring both speed and flexibility.

Conclusion

Set diffusion generalizes blockwise and autoregressive generative paradigms by factorizing likelihoods over token sets governed by flexible ordering schedules. This enables variable-length, flexible-position generation, efficient KV cache utilization, and superior speed-quality tradeoffs across multiple language modeling tasks. The architecture and objective are broadly applicable, representing a unified framework for controllable, efficient, and flexible decoding in discrete text generation (2607.01775).