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

Background Discrete diffusion language models have steadily closed the quality gap with autoregressive (AR) models while offering parallel generation, but two structural handicaps remain: they are normally constrained to fixed-length generation, and their bidirectional attention prevents key-value (KV) caching. Block diffusion (BD3LM, ICLR 2025) partially bridged the two paradigms by running diffusion inside contiguous left-to-right blocks, recovering variable-length generation and per-block KV caching. This paper is the direct successor from the same Cornell group behind MDLM and BD3LM.

배경 이산 diffusion 언어모델은 병렬 생성을 제공하며 autoregressive(AR) 모델과의 품질 격차를 꾸준히 좁혀왔지만, 구조적 약점이 두 가지 남아 있다: 통상 고정 길이 생성에 묶여 있고, bidirectional attention 때문에 key-value(KV) 캐싱이 불가능하다. Block diffusion(BD3LM, ICLR 2025)은 연속된 좌→우 block 내부에서 diffusion을 돌리는 방식으로 두 패러다임을 부분적으로 이어, 가변 길이 생성과 block 단위 KV 캐싱을 복원했다. 이 논문은 MDLM과 BD3LM을 낸 Cornell 그룹의 직계 후속작이다.

Problem Block diffusion's rigid structure leaves three costs on the table: fixed left-to-right blocks preclude any-order tasks like infilling and arbitrary-position insertion; each block must finish before the next begins, capping parallelism; and because within-block denoising relies on bidirectional context, KV cache updates must wait until a block completes. It is like a toll gate that only admits a full busload at a time — even a token that is already finalized cannot enter the cache until the rest of its block is done.

문제 Block diffusion의 경직된 구조는 세 가지 비용을 남긴다: 고정된 좌→우 block은 infilling·임의 위치 삽입 같은 any-order 과제를 막고, 각 block이 끝나야 다음으로 넘어가므로 병렬성이 제한되며, block 내부 denoising이 bidirectional context에 의존하기 때문에 KV 캐시 갱신도 block이 완결될 때까지 미뤄진다. 버스 한 대가 다 차야만 통과시키는 톨게이트처럼, 이미 확정된 token조차 같은 block의 나머지가 끝날 때까지 캐시에 반영되지 못하는 셈이다.

Prior limitations Each neighboring approach gives something up: FlexMDM supports any-order flexible-length diffusion but needs a separate insertion network and a bidirectional transformer, precluding KV caching, while Esoteric LMs get caching from a causal architecture but supervise only randomly masked positions (sparser training signal) and do not handle variable-length generation. WeDLM adopts causal attention for caching but offers no tunable AR-diffusion interpolation, and classical any-order AR models (AO-ARMs) train over uniform orders with high-variance single-step estimators.

기존 한계 인접 연구들은 각자 무언가를 포기한다: FlexMDM은 any-order 가변 길이 diffusion을 지원하지만 별도 insertion network와 bidirectional transformer가 필요해 KV 캐싱이 불가능하고, Esoteric LM은 causal 구조로 캐싱은 되지만 무작위로 마스킹된 위치에만 supervision을 줘 학습 신호가 희소하며 가변 길이 생성을 다루지 않는다. WeDLM은 causal attention을 채택하나 AR-diffusion 간 조절 가능한 interpolation이 없고, 고전 any-order AR(AO-ARM) 계열은 uniform한 순서 위에서 고분산 estimator로 학습된다.

Goal Reframe the AR-diffusion interpolation axis itself: instead of varying a fixed block size, vary the *order* in which token sets are generated. Concretely, design a diffusion language model whose likelihood is autoregressive over flexible-position, flexible-length token sets, and whose architecture supports exact KV cache updates after every inference step plus arbitrary-order decoding, including sliding windows.

목표 AR-diffusion interpolation의 축 자체를 재정의한다: 고정 block 크기를 바꾸는 대신, token 집합이 생성되는 *순서*를 바꾼다. 구체적으로는 유연한 위치·길이의 token 집합에 대해 autoregressive한 우도를 갖고, 매 추론 step마다 정확한 KV 캐시 갱신과 sliding-window를 포함한 임의 순서 디코딩을 지원하는 diffusion 언어모델을 설계하는 것이 목표다.

Method Set diffusion factorizes the sequence likelihood over token sets drawn from an ordering distribution π, marginalizing over generation orders with a variational bound; each set-conditional is modeled by masked diffusion, and a set-causal transformer keeps committed KV valid regardless of which set comes next. Orders are induced by position-offset reveal-time schedules: each position ℓ draws a random reveal time from its own interval [aℓ, aℓ+w], and sorting the sampled times yields the generation order — like a lottery where every token position is issued a ticket window for when it may be revealed, so narrow staggered windows produce near left-to-right order while one fully shared window produces fully random order. The practical instantiation, SW-SetDLM, trains on singleton sets so the sequence can be permuted into generation order and processed under one standard causal mask in a single forward pass (clean and noisy views concatenated), yielding a token-efficient, low-variance estimator.

방법 Set diffusion은 시퀀스 우도를 ordering 분포 π에서 뽑힌 token 집합들로 인수분해하고 생성 순서에 대해 variational bound로 marginalize하며, 각 집합-조건부는 masked diffusion으로, 전체는 확정된 KV가 다음에 어떤 집합이 오든 유효하게 유지되는 set-causal transformer로 모델링한다. 순서는 position-offset reveal-time schedule로 유도된다: 위치 ℓ마다 자기 구간 [aℓ, aℓ+w]에서 무작위 공개 시각을 뽑고 이를 정렬하면 생성 순서가 된다 — 각 token 위치가 자신이 공개될 수 있는 시간대 티켓을 추첨받는 복권과 같아서, 좁게 엇갈린 티켓 창은 좌→우 순서를, 전원이 공유하는 하나의 창은 완전 무작위 순서를 만든다. 실용 버전인 SW-SetDLM은 singleton 집합으로 학습해 시퀀스를 생성 순서로 permute하면 표준 causal mask 하나로 한 번의 forward pass(clean·noisy view를 concat)에서 모든 조건부 우도를 계산하는 token-efficient·저분산 estimator를 얻는다.

Key idea Causality is defined over the sampled generation order rather than over contiguous positions, so committing one token set never invalidates the KV entries of anything already committed — the cache behaves like an append-only ledger, growing exactly (no approximation) after every single inference step and never being recomputed. On top of this, a single window-width parameter w interpolates the whole model family: w=1 recovers order-agnostic MDLM, and as w→1/L the ELBO tightens to the exact AR objective (Prop. 4.1), with block diffusion and any-order AR also arising as special cases of the same likelihood.

핵심 아이디어 인과성이 연속된 위치가 아니라 샘플링된 생성 순서 위에 정의되므로, 한 token 집합을 확정해도 이미 확정된 것들의 KV가 절대 무효화되지 않는다 — 캐시는 append-only 장부처럼 매 추론 step마다 정확하게(근사 없이) 자라기만 하고 다시 계산되지 않는다. 그 위에서 window 폭 w 하나가 모델 패밀리 전체를 interpolate한다: w=1이면 order-agnostic MDLM이 되고, w→1/L이면 ELBO가 정확한 AR objective로 수렴하며(Prop. 4.1), block diffusion과 any-order AR도 같은 우도의 특수 사례로 복원된다.

Validation Four task families, all trained by the authors at small scale: GSM8K math reasoning with a student initialized from Qwen3-1.7B-Base and distilled on Qwen3-32B-Base teacher responses; CNN/DailyMail summarization (80M params); ROCStories infilling over 1,871 five-sentence stories with OWT-trained 110M models; and likelihood plus unconditional generation on OpenWebText and LM1B (perplexity, MAUVE). Throughput is measured at batch size 1 on an H100 80GB (four RTX A6000 48GB for unconditional generation), and comparisons are matched on inference prediction budget (Def. 4.2) so SW-SetDLM window sizes S≤{8,16,32} align with BD3LM block sizes S={4,8,16}.

검증 네 과제군을 저자들이 소규모로 직접 학습해 평가했다: Qwen3-1.7B-Base에서 초기화해 Qwen3-32B-Base teacher 응답으로 distill한 GSM8K 수학 추론, 80M 파라미터 CNN/DailyMail 요약, OWT로 학습한 110M 모델의 ROCStories infilling(5문장 스토리 1,871개), 그리고 OpenWebText·LM1B 우도 및 무조건부 생성(perplexity, MAUVE). 처리량은 H100 80GB에서 batch size 1로 측정하고(무조건부 생성은 RTX A6000 48GB 4장), inference prediction budget을 맞춰(Def. 4.2) SW-SetDLM window S≤{8,16,32}를 BD3LM block S={4,8,16}과 정렬해 비교한다.

Results On GSM8K, SW-SetDLM (S≤8) reaches 66.41% zero-shot pass@1 at ≥60.4 tok/s versus block diffusion's 63.53% at ≥55.4 tok/s — a strictly better speed-accuracy Pareto frontier, though the AR baseline still leads raw accuracy at 75.74%. On ROCStories infilling it beats BD3LM on ROUGE (R-1 11.6 vs 9.2 on infill 1/5) while decoding ~25% faster; on CNN/DailyMail it decodes up to 10% faster at competitive ROUGE (and beats the overfitting AR baseline); on unconditional OWT generation MAUVE improves 64.91 → 71.98 over BD3LM at higher throughput. On LM1B it posts the best diffusion test perplexity (≤27.83, S≤8) together with the highest diffusion throughput (≥98.5 tok/s), and the token-efficient objective cuts gradient variance by ~3×.

결과 GSM8K에서 SW-SetDLM(S≤8)은 66.41% zero-shot pass@1을 ≥60.4 tok/s로 달성해 block diffusion의 63.53%·≥55.4 tok/s를 속도-정확도 양면에서 이기며 Pareto frontier가 전 구간에서 우위다(단 AR 베이스라인은 75.74%로 정확도 자체는 여전히 앞선다). ROCStories infilling에서는 BD3LM 대비 ROUGE-1 11.6 vs 9.2(infill 1/5)에 ~25% 빠른 디코딩, CNN/DailyMail에서는 최대 10% 빠른 디코딩에 대등한 ROUGE(overfitting하는 AR 베이스라인은 능가), 무조건부 OWT 생성에서는 MAUVE가 BD3LM 대비 64.91→71.98로 오르면서 처리량도 높다. LM1B에서는 diffusion 계열 최고 test perplexity(≤27.83, S≤8)를 최고 처리량(≥98.5 tok/s)과 함께 기록했고, token-efficient objective는 gradient variance를 ~3× 줄인다.

Comparison Baselines span the retraining lane broadly: AR transformers, masked diffusion (MDLM, SEDD, D3PM, UDLM, Duo), DiffuGPT and ASSD on infilling, and BD3LM as the main hybrid rival; Eso-LM, FlexMDM, WeDLM, and AO-ARMs are distinguished at the mechanism level rather than in head-to-head tables. Training-free dLLM cache approximations (dKV-Cache, Fast-dLLM) appear only as motivation and as a component — Fast-dLLM's confidence-aware parallel decoding is reused as the sampler — not as baselines (reviewer judgment: a fair boundary, since set diffusion changes the training recipe rather than adding a post-hoc inference trick).

비교 베이스라인은 재학습 계열을 폭넓게 덮는다: AR transformer, MDLM·SEDD·D3PM·UDLM·Duo 등 masked diffusion, infilling에서는 DiffuGPT·ASSD, 그리고 핵심 라이벌인 BD3LM. Eso-LM·FlexMDM·WeDLM·AO-ARM과는 수치 표가 아니라 메커니즘 수준에서 구분 짓는다. training-free dLLM 캐시 근사(dKV-Cache, Fast-dLLM 등)는 베이스라인이 아니라 motivation과 부품으로만 등장한다 — Fast-dLLM의 confidence-aware parallel decoding은 sampler로 재사용된다 (리뷰어 판단: set diffusion은 사후 추론 트릭이 아니라 학습 레시피를 바꾸는 계열이므로 타당한 경계 설정이다).

Significance It generalizes the "exact KV cache via retraining" lane from fixed blocks to arbitrary sets, so that AR, MDLM, block diffusion, and any-order AR all become different settings of one dial — the reveal-window width w — on a single likelihood. For efficient dLLM serving this matters because the cache is an *exact* object refreshed every step: AR-side KV machinery (paging, quantization, eviction) has a well-defined target to act on, in contrast to approximate-cache dLLMs where cached states silently go stale. Accepted at ICML 2026.

의의 "재학습으로 정확한 KV 캐시를 얻는" 계열을 고정 block에서 임의 집합으로 일반화한 것으로, AR·MDLM·block diffusion·any-order AR이 전부 하나의 다이얼(reveal-window 폭 w)의 서로 다른 눈금이 된다. 효율적 dLLM 서빙 관점에서 중요한 이유는 캐시가 매 step 갱신되는 *정확한* 객체라는 점이다: paging·quantization·eviction 같은 AR 쪽 KV 인프라가 작동할 대상이 잘 정의되며, 캐시된 상태가 조용히 낡아가는 근사-캐시 dLLM과 대비된다. ICML 2026 채택작이다.

Limitations The paper's own limitations section concedes that an accuracy-parallelism tradeoff persists in tuning the position-offset schedules, which are currently hand-designed rather than learned. Scale is modest — 110M/80M-parameter models plus one distilled 1.7B GSM8K student — and the AR baseline still wins GSM8K accuracy outright (75.74% vs 66.41%). SW-SetDLM's one-shot factorized denoising also assumes conditional independence among tokens within a set, the same approximation parallel-decoding heuristics make (reviewer judgment: the most likely failure point as set sizes grow).

한계 논문 스스로 한계로 명시하듯, position-offset schedule 튜닝에는 여전히 정확도-병렬성 tradeoff가 남아 있고 schedule은 학습되는 것이 아니라 손으로 설계된다. 규모도 작다 — 110M/80M 모델과 distill된 1.7B GSM8K student가 전부이며, GSM8K 정확도는 AR 베이스라인이 75.74% vs 66.41%로 확실히 앞선다. 또한 SW-SetDLM의 one-shot 인수분해 denoising은 집합 내 token들의 조건부 독립을 가정하는데, 이는 parallel-decoding 휴리스틱들과 같은 근사다 (리뷰어 판단: 집합 크기를 키울 때 가장 먼저 깨질 지점이다).

Future work Explicitly deferred: learned ordering distributions, which are compatible with the framework but require propagating gradients through discrete sampling — the authors leave this to future work and point to progressive-unmasking and learned-order approaches as candidates. The limitations section also suggests tuning schedules to the parallelization profile of the inference hardware, effectively treating the ordering distribution as a serving-time knob.

향후 과제 명시적으로 미룬 것은 학습된 ordering 분포다: 프레임워크와 호환되지만 이산 샘플링을 통과하는 gradient 전파가 필요해 future work로 남기며, progressive-unmasking·learned-order 계열 연구를 후보로 지목한다. Limitations 절에서는 추론 하드웨어의 병렬화 특성에 맞춰 schedule을 조정하는 방향도 제안하는데, 사실상 ordering 분포를 serving-time knob으로 취급하는 셈이다.

Released resources Code is public at kuleshov-group/setdlms on GitHub, with model weights and an explanatory blog post on the authors' project page. Supplement L documents training configurations down to optimizer, batch sizes, and tokenizers, making it a sufficient starting point for reproduction.

자원 공개 코드는 GitHub kuleshov-group/setdlms에, 모델 가중치와 해설 블로그 포스트는 저자 운영 프로젝트 페이지에 공개되어 있다. 부록 L에 optimizer·batch size·tokenizer까지 학습 설정이 상세히 기록되어 있어 재현 출발점으로 충분하다.