Gumbel Distillation for Parallel Text Generation

TL;DR

Gumbel Distillation is a model-agnostic recipe for closing the quality gap between parallel (non-autoregressive) text decoders and their autoregressive teachers. The core trick is the Gumbel-Max identity: sampling a token from a categorical distribution is equivalent to taking the argmax of (logits + Gumbel noise), so a deterministic function from a latent Gumbel-noise tensor to a full token sequence captures the AR teacher’s joint distribution exactly. The student parallel decoder is then trained to imitate this deterministic noise→tokens map. The method drops into existing parallel-decoding architectures (MDLM, BD3-LM) and on LM1B / OpenWebText delivers +30.0% MAUVE and +10.5% generative perplexity over MDLM trained on OpenWebText.

Key claims

Sampling from a categorical token distribution is equivalent to argmax(logits + g) with g i.i.d. standard Gumbel — so the joint distribution p_AR(x_{1:T}) is encoded as a deterministic function f: G → x_{1:T} where G is a Gumbel-noise tensor [§ Method, Gumbel-Max trick].
A parallel (non-AR) student trained to match this deterministic Gumbel→sequence mapping inherits the AR teacher’s joint distribution, including the inter-token dependencies that vanilla parallel decoders fail to model [§ Abstract, § Method].
The recipe is model-agnostic and integrates with diverse parallel-decoding architectures, demonstrated on MDLM (masked discrete-diffusion LM) and BD3-LM (block-discrete-diffusion LM) [§ Abstract].
On OpenWebText, Gumbel Distillation improves MDLM by 30.0% MAUVE and 10.5% generative perplexity [§ Abstract, Experiments].
The technique is framed as knowledge distillation: an AR teacher provides supervision, and the parallel student is trained to replicate the teacher’s joint via the latent noise space, rather than via per-token or per-marginal targets [§ Keywords, § Abstract].

Method

The construction has two pieces. First, the Gumbel-Max trick: for a categorical with logits ℓ ∈ R^V, the random variable argmax_v (ℓ_v + g_v) with g_v ~ Gumbel(0, 1) i.i.d. is distributed as a sample from softmax(ℓ). Extending this to a full sequence under an AR factorization gives a deterministic map from a T × V Gumbel-noise tensor (plus the AR model’s logits at each step) to a length-T token sequence — the AR sampling procedure reduces to “given this noise tensor, what tokens does the teacher produce?”

Second, the distillation objective: train a parallel (non-AR) student whose input is the latent Gumbel-noise tensor and whose output is the corresponding token sequence the AR teacher would have sampled. Because the same noise tensor deterministically produces the same sequence under the teacher, the student has a clean regression target that encodes the full joint. The student architecture is left open — the paper instantiates it on MDLM and BD3-LM, both masked discrete-diffusion LMs that natively support parallel decoding.

The technique should be read as the Gumbel-noise analog of standard distillation: rather than matching teacher logits or sampled token sequences, the student matches the teacher’s deterministic latent-to-output function. This is conceptually adjacent to flow/diffusion distillation (where the student matches the teacher’s noise→sample ODE map) — see Diffusion Distillation.

Results

Reported headline: +30.0% MAUVE Score and +10.5% generative perplexity over MDLM trained on OpenWebText [§ Abstract].
Experiments span LM1B and OpenWebText with MDLM and BD3-LM as the underlying parallel-decoding architectures [§ Abstract].

(The PDF body was not retrievable at filing time; further per-table numbers, AR-teacher details, and ablations are not summarized here. The OpenReview forum page lists the paper as ICLR 2026 submission #9680.)

Why it’s interesting

Three threads on the wiki converge on this paper.

First, Diffusion Language Models has been tracking dLLMs along three production axes — throughput, AR-serving deployability, RL stability — with quality parity established at the 7-8B scale by Mercury: Ultra-Fast Language Models Based on Diffusion, Introspective Diffusion Language Models, and WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference. Gumbel Distillation is orthogonal to all three: it’s a training-objective fix that targets the same fundamental problem — that parallel decoders factor the joint into per-position marginals and lose the inter-token dependencies — without changing serving infra, RL recipe, or per-token throughput. The ParallelBench paper (ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs) made the case that the quality ceiling under parallel decoding is structural, lower-bounded by the conditional total correlation of the parallel-decoded set; Gumbel Distillation is a candidate route around that ceiling — not by reducing TC, but by giving the student a noise-conditioned deterministic target that encodes TC structure implicitly via the AR teacher’s sampling map.

Second, the recipe is structurally a Diffusion Distillation move applied to language. The image/video side has settled on distilling deterministic teacher trajectories (PF-ODE, mean-flow, transition-matching) into few-step students; Gumbel Distillation is the discrete-token analog where the “trajectory” is the Gumbel-Max deterministic map and the “few-step student” is a parallel decoder. This pairs cleanly with TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward‘s argument that deterministic trajectories enable unbiased per-step reward estimation — the same property may enable Gumbel-Distillation-trained students to be RL-finetuned with rewards computed at the endpoint and credit-assigned cleanly back through the noise tensor.

Third, Jiaming’s framing — “multi-token prediction and diffusion LLMs are the same thing” — is supported by this paper’s architectural agnosticism: the same Gumbel-noise → sequence target works on MDLM (full-sequence masked diffusion) and BD3-LM (block-diffusion, closer to MTP). If the gap-closing mechanism is the same, the engineering details (block size, mask schedule, attention pattern) may be tunable hyperparameters rather than fundamentally different approaches. ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs‘s evidence that block-size cannot simultaneously optimize local and global dependencies suggests Gumbel Distillation’s per-sample-noise-tensor formulation may be a way to sidestep the block-size dilemma.

Caveats: results are reported on LM1B / OpenWebText (academic-scale benchmarks, not chat or code), and the headline 30% MAUVE / 10.5% PPL improvement is over an MDLM baseline — there’s no head-to-head against AR teachers of matched compute, so this lands as “closes the gap” rather than “exceeds AR.” The PDF body was not retrievable at filing, so the magnitude of the remaining gap to the AR teacher is not summarized here.