Gumbel Distillation for Parallel Text Generation

2026-03-23

AI Summary

Parallel text generation models typically underperform autoregressive counterparts due to the difficulty of modeling the complex joint distribution over token sequences. This work proposes Gumbel Distillation, a novel approach that introduces the Gumbel-Max trick into knowledge distillation for the first time, establishing a deterministic mapping from an implicit Gumbel noise space to the outputs of a high-performing autoregressive teacher model. This enables parallel decoders to effectively learn the sequence-level joint distribution without relying on any specific architectural assumptions, making it applicable to various non-autoregressive models such as MDLM and BD3-LM. Experiments on LM1B and OpenWebText demonstrate substantial improvements, including a 30.0% increase in MAUVE score and a 10.5% reduction in generation perplexity, significantly narrowing the performance gap with autoregressive models.

Technology Category

• Natural Language Processing: Generation
• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• parallel text generation
• non-autoregressive language models
• generation quality
• joint distribution modeling

Innovation

• Gumbel Distillation
• parallel text generation
• non-autoregressive modeling
• Gumbel-Max trick
• knowledge distillation