Gumbel Distillation for Parallel Text Generation

Chi Zhang; Xixi Hu; Bo Liu; Qiang Liu

cs.CLMar 23, 2026

Gumbel Distillation for Parallel Text Generation

Authors: Chi Zhang, Xixi Hu, Bo Liu, Qiang Liu

Organizations: Department of Computer Science The University of Texas at Austin

Abstract

The slow, sequential nature of autoregressive (AR) language models has driven the adoption of parallel decoding methods. However, these non-AR models often sacrifice generation quality as they struggle to model the complex joint distribution of token sequences. To narrow this performance gap, we introduce Gumbel Distillation, a novel distillation technique that enables parallel decoders to learn this distribution effectively. Our method leverages the Gumbel-Max trick to create a deterministic mapping from a latent Gumbel noise space to the output tokens of a high-performing AR teacher. As a model-agnostic technique, Gumbel Distillation seamlessly integrates with diverse parallel decoding architectures, including MDLM and BD3-LM. Experiments on LM1B and OpenWebText show that Gumbel Distillation substantially improves the generation quality of parallel language models, achieving a 30.0% improvement in MAUVE score and 10.5% in generative perplexity over MDLM trained on OpenWebText dataset. Code available at https://github.com/hxixixh/gumbel-distill.

Figures & tables

Figure 1: A conceptual overview of Gumbel Distillation. The distillation process consists of two steps: (1) Data Generation: An autoregressive teacher model’s sampling process is combined with Gumbel noise to deterministically generate pairs of token sequences and their corresponding Gumbel noise. In practice, one could alternatively extract $\{\xi\}$ in parallel based on an offline training corpus (Section 4.1 ); (2) Student Training: A parallel student model is trained to predict a target subset of tokens, conditioned on both the context and the Gumbel noise from the teacher.

Figure 2: A toy maze problem illustrating how Gumbel Distillation helps a student model learn a structured task. The goal is to generate a valid path from a start ( <bos> ) to a target ( <eos> ) using a simple vocabulary ( up , down , left , right ). The figure visualizes 100 generated paths from each model, slightly jittered for visualization. The baseline MDLM frequently fails, producing incoherent paths, indicating its difficulty in modeling the joint distribution of the sequence. In contrast, our Gumbel-conditioned model successfully finds valid paths, closely matching the performance of the AR teacher but at a fraction of the generation cost (NFE = 3 steps vs. 10 steps).

	LM1B		OpenWebText
	MAUVE $\uparrow$	GenPPL $\downarrow$	MAUVE $\uparrow$	GenPPL $\downarrow$
AR (student-size)	0.465	36.42	0.691	14.10
MDLM	0.179	78.74	0.217	38.34
MDLM + Gumbel Distillation	0.264	67.64	0.282	34.33
BD3-LM ( $L^{\prime}=4$ )	0.193	56.98	0.251	26.40
BD3-LM ( $L^{\prime}=4$ ) + Gumbel Distillation	0.291	46.06	0.304	24.37

Table 1: Main results for unconditional text generation. AR denotes an autoregressive model with the same parameter as the student backbones (used only as a quality reference); the distillation teacher is GPT-2-Large. Models are trained for 1M steps on LM1B (length 128) and OpenWebText (length 1024). Best non-AR scores are in bold.

Model	Clarity		Grammaticality		Factuality		Style		Creativity
MDLM	2.44		2.22		2.70		2.32		2.22
MDLM + Gumbel Distillation	2.86	(+17.2%)	2.57	(+15.8%)	3.31	(+22.6%)	2.57	(+10.8%)	2.36	(+6.3%)
BD3-LM	2.89		2.95		3.21		3.34		2.75
BD3-LM + Gumbel Distillation	3.41	(+18.0%)	3.22	(+9.2%)	3.78	(+17.7%)	3.35	(+0.2%)	2.68	(-2.5%)

Table 2: Qualitative evaluation using Gemini-2.5-pro as an LLM judge on OpenWebText. Scores are on a 1-10 scale (higher is better). Gumbel Distillation shows clear improvements across most dimensions, with the most significant and consistent gains observed in Clarity and Factuality .

Figure 3: Generative Perplexity and MAUVE Score vs. Number of Function Evaluations (NFE) of MDLM on LM1B and OpenWebText . Our method ( Gumbel Distillation ) consistently outperforms the baselines, achieving lower perplexity for a given number of evaluations.

Figure 4: Zero-shot performance on common-sense reasoning and question answering benchmarks. Our Gumbel Distillation consistently improves accuracy over the baseline models across a suite of eight tasks. All scores are percentages (%).

	Baseline	+ Gumbel Distillation	$\Delta_{\mathrm{rel}}\%$
GPT-2-Small backbone (OpenWebText)
Head 0 (NTP)	1.000	1.000	-
Head 1	0.445	0.465	(+4.5%)
Head 2	0.310	0.344	(+11.0%)
Head 3	0.264	0.322	(+22.0%)
Accepted Length	1.609	1.691	(+5.1%)

Table 3: Per-position acceptance rates of Medusa MTP heads. We report conditional acceptance rates under typical acceptance.

\Delta_{\mathrm{rel}}\%

denotes relative improvement over the baseline. Top: frozen GPT-2-Small trained on OpenWebText (self-distillation). Bottom: scaling to frozen Vicuna-7B trained on 60k ShareGPT samples (Medusa recipe; Head 0 uses the frozen LM head).

Method	MAUVE $\uparrow$	GenPPL $\downarrow$
MDLM	0.179	78.74
+ Token-level KD	0.166	95.88
+ Sequence-level KD	0.169	99.48
+ APD	0.203	57.61
+ Gumbel Distillation	0.264	67.64
+ Gumbel Distillation + APD	0.255	49.28

Table 4: Comparison to outcome-based distillation baselines and APD on MDLM (LM1B). Token-level and sequence-level KD are outcome-based AR to NAR distillation baselines; APD is applied only at inference time. Best MAUVE and best GenPPL are in bold.

Parallel Gumbel Extraction	Sequential Gumbel Extraction	Gumbel Noise	Gaussian Noise	Uniform Noise	MAUVE Score( $\uparrow$ )	Gen. PPL ( $\downarrow$ )
✓		✓			0.264	67.64
	✓	✓			0.189	86.38
✓			✓		0.242	81.43
✓				✓	0.097	-

Table 5: Consolidated ablation study using MDLM on LM1B. Parallel Gumbel Extraction is superior to the sequential method. The specific Gumbel distribution is critical, as replacing it with Gaussian noise degrades performance, and Uniform noise leads to mode collapse.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Model scale	Layers $L$	Hidden $H$	Vocab $V$	Added params $VH$	Overhead
GPT-2 Small	12	768	50,257	38.6M	33.0%
GPT-2 Large	36	1,280	50,257	64.3M	8.3%
LLaMA-7B	32	4,096	32,000	131.1M	1.9%
LLaMA-13B	40	5,120	32,000	163.8M	1.3%
Qwen-7B	32	4,096	152,000	622.6M	8.9%

Appendix

Table 6: Parameter overhead of the Gumbel projection matrix (

W_{\xi}\in\mathbb{R}^{H\times V}

). Overhead is

VH/\text{TotalParams}

. Large vocabularies increase overhead, while larger backbones reduce it.

Figure 5: Gumbel Distillation Integrated with MDLM. Left: original MDLM architecture, masked positions are predicted based solely on the unmasked context. Right: Gumbel-conditioned architecture: distillation pairs $(\bm{\xi},\bm{x})$ are obtained from an AR teacher; then MDLM is modified to take Gumbel as conditional input, learning to predict masked positions with corresponding Gumbel noise.

Figure 6: Gumbel Distillation integrated with Medusa. The base backbone model is frozen when training the MTP heads. We first obtain distillation pairs $(\bm{\xi},\bm{x})$ from the AR backbone as teacher, then each MTP head is modified to take Gumbel as conditional input.

Figure 7: Prompt used for LLM-based evaluation using Gemini-2.5-pro.

NFE	=1	=2	=4	=8	=16
MDLM	53.0	64.0	69.0	81.0	94.0
MDLM + Gumbel Distillation	82.0	86.0	97.0	99.0	98.0

Appendix

Table 7: Success rates vs. NFEs for MDLM and MDLM + Gumbel Distillation on the maze toy.

	PTB	Wikitext	LM1B	Lambada	AG News	Pubmed	Arxiv
MDLM	92.54	34.52	66.10	48.63	65.32	43.67	37.47
MDLM + Gumbel Distillation	35.12	15.57	26.01	15.56	20.03	20.81	16.85
Block Diffusion	96.27	30.82	61.14	49.51	63.09	42.59	39.39
Block Diffusion + Gumbel Distillation	37.54	13.86	22.69	15.78	18.19	19.78	17.02

Appendix

Table 8: Validation perplexities (

\downarrow

) of models trained on OWT.

Model	GenPPL $\downarrow$
MDLM	55.1
MDLM + Gumbel Distillation	47.8

Appendix

Table 9: Infilling on LM1B. Gumbel conditioning improves GenPPL on a non-AR infilling task, suggesting it does not overfit to AR-style sampling.

Method	MAUVE $\uparrow$	GenPPL $\downarrow$
MDLM (Small)	0.179	78.74
MDLM (Small) + Gumbel Distillation	0.264	67.64
MDLM (Medium, partial)	0.186	74.22
MDLM (Medium, partial) + Gumbel Distillation	0.253	61.34

Appendix

Table 10: Gumbel Distillation improves both MAUVE and GenPPL for different model scales.

Panel of Figure 8 — Figure 8: Sample from MDLM of length 1024 and T = 1024 steps. The generative perplexity of the sample under GPT-2-Large is 62.36 and its entropy is 5.56.

Panel of Figure 9 — Figure 9: Sample from MDLM + Gumbel Distillation of length 1024 and T = 1024 steps. The generative perplexity of the sample under GPT-2-Large is 35.19 and its entropy is 5.32.

Panel of Figure 10 — Figure 10: Sample from BD3-LM of length 1024 and T = 1024 steps. The generative perplexity of the sample under GPT-2-Large is 37.7 and its entropy is 5.21.

Panel of Figure 11 — Figure 11: Sample from BD3-LM + Gumbel Distillation of length 1024 and T = 1024 steps. The generative perplexity of the sample under GPT-2-Large is 27.2 and its entropy is 5.19.

Panel of Figure 12 — Figure 12: Sample from AR of length 1024 and T = 1024 steps. The generative perplexity of the sample under GPT-2-Large is 15.7 and its entropy is 5.29.

Non-Autoregressive Text Generation

Non-autoregressive text generation produces multiple tokens in parallel to reduce decoding time while maintaining coherent, useful outputs.

See all 31 papers

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?
Feb 26, 2026Pengxiang Li, Dilxat Muhtar, Tianlong Chen +2Parallel Decoding for Diffusion Language Models Diffusion Language Model Decoding
The Hong Kong Polytechnic University, Hong Kong, China · ELLIS Institute Tübingen, Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany +3
Parallel Language Models via Variational Decoding
Oct 19, 2026Felix Draxler, Robin C. Geyer, Daniel Geyfman +13Parallel Decoding LLM Inference Acceleration
University of California, Irvine, CA, USA · Department of Computer Science, ETH Zurich, Switzerland · Institute for Artificial Intelligence in Medicine, Charité Berlin, Germany +1
DiffusionGemma Technical Report
Jul 31, 2026DiffusionGemma Team, Adrien Ali Taïga, James Assiene +41Diffusion Language Model Decoding Blockwise Diffusion Language Models

Feb 26, 2026J/K move · Enter open · S save

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?

Pengxiang Li, Dilxat Muhtar, Tianlong Chen, Lu Yin+1

The Hong Kong Polytechnic University, Hong Kong, China · ELLIS Institute Tübingen, Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany+3

Parallel Decoding for Diffusion Language Models Diffusion Language Model Decoding Diffusion Language Models Non-Autoregressive Text Generation Parallel Reasoning in Language Models LLM Reasoning

Open paper AskPDF arXiv

Diffusion Language Models (DLMs) are often advertised as enabling parallel token generation, yet practical fast DLMs frequently converge to left-to-right, autoregressive (AR)-like decoding dynamics. In contrast, genuinely non-AR generation is promising because it removes AR's sequential bottleneck, better exploiting parallel hardware to reduce synchronization/communication overhead and improve latency scaling with output length. We argue that a primary driver of AR-like decoding is a mismatch between DLM objectives and the highly sequential structure of widely used training data, including standard pretraining corpora and long chain-of-thought (CoT) supervision. Motivated by this diagnosis, we propose NAP (Non-Autoregressive Parallel DLMs), a proof-of-concept, data-centric approach that better aligns supervision with non-AR parallel decoding. NAP curates examples as multiple independent reasoning trajectories and couples them with a parallel-forced decoding strategy that encourages multi-token parallel updates. Across math reasoning benchmarks, NAP yields stronger performance under parallel decoding than DLMs trained on standard long CoT data, with gains growing as parallelism increases. Our results suggest that revisiting data and supervision is a principled direction for mitigating AR-like behavior and moving toward genuinely non-autoregressive parallel generation in DLMs. Our code is available at https://github.com/pixeli99/NAP.

Abstract

Non-Autoregressive Text Generation

Explore similar work

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?

Parallel Language Models via Variational Decoding

DiffusionGemma Technical Report