🤖 machine learning

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

This paper introduces Set Diffusion, a novel class of language models that combines flexible, arbitrary-order token set generation with KV cache support to achieve faster inference and superior speed-quality tradeoffs compared to existing autoregressive and block diffusion approaches.

Original authors: Marianne Arriola, Volodymyr Kuleshov

Published 2026-07-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Marianne Arriola, Volodymyr Kuleshov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Order" Problem

Imagine you are trying to write a story or solve a math problem. You have two main ways to do it:

The Strict Writer (Autoregression): You write one word at a time, strictly from left to right. You can't write the ending until you finish the beginning. This is very high-quality and accurate, but it's slow because you can't do two things at once.
The Chaos Artist (Standard Diffusion): You start with a page full of gibberish (noise) and try to fix it all at once. You can guess many words simultaneously, which is fast, but it's hard to keep the story logical. Also, you can't easily "remember" what you wrote earlier to help you write the next part, so you have to re-read the whole page every time you make a change.

The Problem: Previous attempts to mix these two (called "Block Diffusion") were like writing in rigid chunks. You could write a sentence at a time instead of a word, but you still had to finish that whole sentence before moving to the next. If you wanted to fix a word in the middle of the story, you had to wait for the whole block to finish.

The Solution: Set Diffusion

The authors introduce Set Diffusion. Think of this not as writing in a line, but as filling in a puzzle with flexible pieces.

Instead of being forced to write word-by-word (too slow) or block-by-block (too rigid), Set Diffusion lets you pick any group of words to generate next, as long as you follow a specific set of rules.

The "Sliding Window" Analogy

Imagine you are painting a long mural.

Old Way (Block Diffusion): You paint a 4-foot section, wait for it to dry completely, then move to the next 4-foot section. You can't touch the first section again until the whole block is done.
New Way (Set Diffusion): You have a "sliding window" of paint. You can paint the first 4 feet, but then you can slide that window over and paint the 2nd, 3rd, and 4th feet simultaneously with the 5th, 6th, and 7th feet. You can even jump back to fix a spot in the middle of the window while you are painting the edge.

Key Features Explained Simply

1. Flexible "Token Sets" (The Puzzle Pieces)
In this model, a "token" is just a word or a piece of code.

The Innovation: The model doesn't just guess the next word. It guesses a set of words.
The Magic: You can tell the model, "Guess the next 3 words," or "Guess the word at position 5 and position 10." It can handle groups of different sizes and in different orders. This allows it to be fast (guessing many things at once) but still smart (keeping track of the order).

2. The "Memory Bank" (KV Caching)
In AI, "KV Caching" is like a scratchpad where the model writes down the context of what it has already generated so it doesn't have to re-calculate it every time.

The Old Problem: In standard diffusion, the model had to re-read the entire text every single time it made a guess. It was like reading a whole book to remember the name of the main character.
The New Fix: Set Diffusion updates its "scratchpad" after every single step. As soon as it guesses a set of words, it saves them to memory. This makes it incredibly fast, similar to the strict writer, but with the speed of the chaos artist.

3. The "Sliding Window" Strategy
The paper introduces a specific way to choose which words to guess next, called a Sliding-Window approach.

Imagine a spotlight moving across a stage. The spotlight can cover one actor, or it can cover three actors at once.
The model uses this spotlight to decide: "I will generate the words inside this spotlight right now."
By adjusting the size of the spotlight, you can control the balance between speed (big spotlight, guess many words) and accuracy (small spotlight, guess fewer words to be more careful).

What Did They Prove?

The authors tested this new method on three main tasks:

Math Reasoning: Solving word problems (like the GSM8K dataset).
Summarization: Condensing long articles into short summaries.
Infilling: Filling in missing parts of a story (like a "Mad Libs" game).

The Results:

Speed vs. Quality: Set Diffusion found a "sweet spot" that previous models missed. It was faster than the rigid block models and more accurate than the chaotic diffusion models.
Infilling: It was significantly better at filling in missing parts of a story than the previous "Block Diffusion" method.
Flexibility: It can generate text of any length and in any order, which is crucial for tasks like editing a document or fixing code in the middle of a file.

Summary Metaphor

If Autoregression is a relay race (one runner passes the baton to the next, strictly in order) and Standard Diffusion is a free-for-all (everyone runs at once, but it's chaotic), then Set Diffusion is a well-coordinated dance troupe.

The dancers (tokens) can move in groups (sets). They can move left-to-right, or they can jump around to fill gaps, but they always know exactly who is next to them and what the previous moves were, thanks to their shared memory (KV cache). This allows them to perform complex routines (math, stories) much faster and more flexibly than the relay team, without the chaos of the free-for-all.

Technical Summary: Set Diffusion

Problem Statement

Discrete diffusion models have emerged as a competitive alternative to autoregressive (AR) language models, offering faster generation speeds and improved controllability. However, existing diffusion approaches face significant limitations:

Fixed-Length and Bidirectional Constraints: Most discrete diffusion models rely on architectures that require full bidirectional context, preventing the use of Key-Value (KV) caching and restricting generation to fixed lengths.
Rigidity of Block Diffusion: While Block Diffusion (BD3LM) partially bridges the gap by generating fixed-size sequential blocks left-to-right, it suffers from inflexibility. It cannot support arbitrary-position insertions (infilling) or variable-length generation within a block, and KV cache updates are delayed until an entire block is decoded.
Trade-off Limitations: Current methods interpolate between AR and diffusion primarily by varying a fixed block size, which restricts decoding flexibility and parallelism.

The core challenge is to develop a language model that interpolates between autoregression and diffusion not just by block size, but by varying the order in which token sets are generated, thereby enabling flexible-length generation, KV caching, and arbitrary-order decoding while maintaining high sample quality.

Methodology

The authors introduce Set Diffusion, a framework that factorizes the generation process over flexible-position, flexible-length token sets rather than fixed blocks.

1. Probabilistic Model

Set diffusion defines an autoregressive probability distribution over a sequence of token sets $\sigma = (\sigma_1, \dots, \sigma_N)$ , where each $\sigma_n$ is a subset of token positions. The likelihood is factorized as:
$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$
Here, $\pi$ is an order policy selecting the next set of positions, and $p_\theta$ models the conditional likelihood of tokens in that set using diffusion. The model optimizes a Negative Evidence Lower Bound (NELBO) that sums the diffusion losses over these sets.

2. Set-Causal Architecture and KV Caching

A key architectural innovation is the set-causal transformer. Unlike standard bidirectional diffusion models, this architecture supports KV caching after every inference step.

Mechanism: When a token set $x_{\sigma_n}$ is generated, its key and value vectors are cached and reused to condition subsequent predictions.
Training Efficiency: The authors propose a vectorized training algorithm that processes clean context and noisy queries simultaneously in a single forward pass using a specialized attention mask. This mask ensures that noisy tokens in a set attend only to other noisy tokens in the same set and all clean tokens in preceding sets.

3. Interpolating Token Orderings

To control the trade-off between sample quality (tightness of the bound) and inference efficiency (parallelism), the authors introduce position-dependent reveal-time schedules.

Position-Offset Schedules: Each token position $\ell$ is assigned a random reveal time $R_\ell$ drawn from a distribution defined by a cumulative distribution function (CDF) $\alpha^\ell_\tau$ .
Sliding-Window Instantiation (SW-SetDLM): The authors propose a specific instantiation where tokens have an active generation interval $[a_\ell, a_\ell + w]$ $[a_{ℓ}, a_{ℓ} + w]$ . The window size $w$ $w$ controls the "left-to-right" bias:
- Small $w$ (approaching $1/L$ ) concentrates the order on left-to-right generation, recovering AR behavior and tightening the ELBO.
- Large $w$ (approaching 1) allows for order-agnostic generation, recovering standard masked diffusion.
- Intermediate values enable sliding-window decoding, allowing the model to predict flexible-length sets of tokens in arbitrary positions, supporting tasks like infilling.

4. Inference

During inference, the model samples candidate positions based on the current ordering time and the position-offset schedule. It then denoises these positions in parallel using the KV cache. Accepted tokens are committed, and the cache is updated immediately, allowing for variable-length generation and efficient decoding.

Key Contributions

Set Diffusion Framework: A new class of language models that factorizes generation over flexible-position, flexible-length token sets, generalizing autoregression, order-agnostic diffusion, and block diffusion.
Set-Causal Architecture: A transformer architecture that supports KV cache updates after every inference step, enabling efficient variable-length generation.
Order Interpolation: A principled method to interpolate between AR and diffusion by tuning position-dependent ordering schedules (specifically sliding-window schedules) rather than fixed block sizes.
State-of-the-Art Performance: Demonstration that set diffusion achieves superior speed-quality trade-offs across mathematical reasoning, summarization, and unconditional generation compared to prior diffusion models, and significantly outperforms block diffusion on infilling tasks.

Experimental Results

The authors evaluated Set Diffusion (specifically SW-SetDLM) on several benchmarks:

Mathematical Reasoning (GSM8K): Set diffusion achieves a strictly better speed-accuracy Pareto frontier than Block Diffusion. For instance, with a decoding window size $S \le 8$ , it achieves 66.41% pass@1 accuracy with a throughput of $\ge 60.42$ tokens/sec, outperforming Block Diffusion ( $S=8$ ) which achieves 63.53% accuracy at $\ge 55.39$ tokens/sec.
Infilling (ROCStories): Set diffusion substantially outperforms Block Diffusion in infilling tasks (filling 1 or 3 middle sentences), achieving higher ROUGE scores with up to ~25% faster decoding. Block diffusion struggles here due to its rigid left-to-right block structure.
Summarization (CNN/DailyMail): Set diffusion achieves competitive ROUGE scores compared to diffusion baselines and up to 10% faster decoding than Block Diffusion. It also outperforms an AR baseline, which the authors note tends to overfit on this dataset.
Unconditional Generation & Likelihood: On OpenWebText and LM1B, Set diffusion improves MAUVE scores over Block Diffusion while decoding faster. On LM1B, it achieves the best perplexity among diffusion models (27.83 PPL for $S \le 8$ ) while maintaining higher throughput than prior diffusion models.

Significance and Claims

The paper claims that Set Diffusion represents a significant advancement by decoupling the interpolation between autoregression and diffusion from rigid block structures. By interpolating via token ordering, the framework offers:

Greater Flexibility: Support for arbitrary-position insertions and variable-length generation, which are critical for tasks like infilling.
Efficiency: The ability to update KV caches after every step, a feature previously unavailable in diffusion models, leading to faster inference.
Performance: A superior trade-off between sample quality and inference speed compared to existing diffusion language models.

The authors position Set Diffusion as a unifying framework that recovers AR, order-agnostic diffusion, and block diffusion as special cases, providing a tunable mechanism to balance parallelism and sample quality based on the specific requirements of the inference task.