Gumbel Distillation for Parallel Text Generation

Paper: 2603.22216 Authors: Chi Zhang, Xixi Hu, Bo Liu, Qiang Liu Categories: cs.CL Venue: ICLR 2026

Abstract

The slow, sequential nature of autoregressive (AR) language models has driven the adoption of parallel decoding methods. However, non-AR models often sacrifice generation quality as they struggle to model the complex joint distribution of token sequences. Gumbel Distillation is a novel technique that enables parallel decoders to learn this distribution effectively by leveraging the Gumbel-Max trick to create a deterministic mapping from a latent Gumbel noise space to the output tokens of an AR teacher. As a model-agnostic technique, it seamlessly integrates with diverse parallel decoding architectures including MDLM and BD3-LM. Experiments on LM1B and OpenWebText show 30.0% improvement in MAUVE score and 10.5% in generative perplexity over MDLM.

Key Contributions

• Gumbel Distillation framework: Transforms the difficult joint-distribution matching problem into a supervised learning problem using the Gumbel-Max reparameterization
• Model-agnostic plug-and-play: Integrates with MDLM, BD3-LM, and Medusa with minimal architectural changes
• Significant quality gains: 30% MAUVE improvement and 10.5% generative perplexity reduction on OpenWebText
• Theoretical grounding: Provides principled approach to capturing token dependencies that parallel decoders typically miss

The Fundamental Problem

Parallel decoders generate multiple tokens simultaneously but must assume conditional independence:

$p_\theta(x_I | x_{\neg I}) = \prod_{i \in I} p_\theta(x_i | x_{\neg I})$

This assumption ignores dependencies within the target token set. For example, predicting “San Francisco” requires knowing that “Francisco” depends on “San” — but parallel models predict both independently. This leads to token repetition, incoherent text, and grammatical errors.

The key question: Can we make the learning problem easier by providing a “blueprint” from a powerful AR teacher?

The Gumbel-Max Trick

The core insight leverages the Gumbel-Max trick for categorical sampling. To sample from a distribution with logits l = (l₁, …, l_V):

1. Draw i.i.d. Gumbel noise: ξ_k ~ G(0,1)
2. Compute: Y = argmax_k (l_k + ξ_k)

The crucial observation: once the logits and noise are known, the argmax is fully deterministic. The randomness is entirely externalized into the Gumbel noise vector.

Two-Stage Framework

Stage 1: Data Generation

For each training example:

1. Run the AR teacher to generate a token sequence
2. Record both the output tokens AND the Gumbel noise vectors used in sampling
3. Store (noise, token) pairs as training data

This creates a dataset where each Gumbel noise vector serves as a latent “blueprint” — uniquely encoding the sampling decisions that produced the corresponding text.

Stage 2: Student Training

Train the parallel student model to reconstruct text conditioned on:

• The input context (as usual)
• The Gumbel noise vectors (new conditional input)

This transforms joint-distribution matching into supervised regression: given the noise “blueprint,” predict the deterministic output.

Why It Works

Traditional distillation asks the student to match the teacher’s output distribution — a hard problem when the student assumes conditional independence. Gumbel Distillation sidesteps this by:

1. Externalizing randomness: The Gumbel noise captures all stochastic decisions
2. Deterministic mapping: Given noise, the mapping from noise → tokens is deterministic
3. Supervised learning: The student learns input-output pairs, not distributions
4. Implicit dependencies: Token dependencies are encoded in the noise vectors, allowing the student to capture them without explicit joint modeling

Results

MDLM + Gumbel Distillation on OpenWebText

Metric	MDLM	MDLM + Gumbel	Improvement
MAUVE ↑	baseline	+30.0%	Significant
Gen. Perplexity ↓	baseline	-10.5%	Significant

Integration with Multiple Architectures

• MDLM: Masked diffusion language model — improvements across all metrics
• BD3-LM: Block-decomposed diffusion — consistent gains with block-wise generation
• Medusa: Multi-token prediction heads — enhanced speculative decoding quality

Takeaways

• The conditional independence assumption in parallel decoders is the primary quality bottleneck — Gumbel Distillation provides a principled solution
• By externalizing randomness via Gumbel noise, the intractable distribution-matching problem becomes tractable supervised learning
• The approach is model-agnostic: any parallel decoder can benefit from Gumbel Distillation as a plug-and-play module
• This work addresses a fundamental limitation shared by masked diffusion models, multi-token prediction, and other non-AR approaches
• Published at ICLR 2026 with code available, enabling broad adoption