KLASS: KL-Guided Fast Inference in Masked Diffusion Models

2025/11/07 by S. Kim, Kim, Seo Hyun, Hong, Sunwoo +6 · 3 citations
Computer Science · Social Sciences · #Generative Adversarial Networks and Image Synthesis #Topic Modeling #Computational and Text Analysis Methods

paper · pdf · doi:10.48550/arxiv.2511.05664

Abstract

Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its iterative refinement process, the inference is often bottlenecked by slow and static sampling speed. To overcome this problem, we introduce `KL-Adaptive Stability Sampling' (KLASS), a fast yet effective sampling method that exploits token-level KL divergence to identify stable, high-confidence predictions. By unmasking multiple tokens in each iteration without any additional model training, our approach speeds up generation significantly while maintaining sample quality. On reasoning benchmarks, KLASS achieves up to 2.78× wall-clock speedups while improving performance over standard greedy decoding, attaining state-of-the-art results among diffusion-based samplers. We further validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness as a broadly applicable sampler across different models.

Citations

• Diffusion Language Models Know the Answer Before Decoding
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
• Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
• Dimple: Discrete Diffusion Multimodal Large Language Model with Parallel Decoding
• Remasking Discrete Diffusion Models with Inference-Time Scaling
• Implicit Search via Discrete Diffusion: A Study on Chess
• Large Language Diffusion Models
• Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
• Path Planning for Masked Diffusion Model Sampling
• Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms
• Simple Guidance Mechanisms for Discrete Diffusion Models
• Beyond Autoregression: Fast LLMs via Self-Distillation Through Time
• Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning
• Distillation of Discrete Diffusion through Dimensional Correlations
• Think While You Generate: Discrete Diffusion with Planned Denoising
• How Discrete and Continuous Diffusion Meet: Comprehensive Analysis of Discrete Diffusion Models via a Stochastic Integral Framework
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling
• Informed Correctors for Discrete Diffusion Models
• Simple and Effective Masked Diffusion Language Models
• Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data
• Simplified and Generalized Masked Diffusion for Discrete Data
• MACM: Utilizing a Multi-Agent System for Condition Mining in Solving Complex Mathematical Problems
• Convergence Analysis of Discrete Diffusion Model: Exact Implementation through Uniformization
• Executable Code Actions Elicit Better LLM Agents
• One-step Diffusion with Distribution Matching Distillation
• Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
• Reflexion: Language Agents with Verbal Reinforcement Learning
• Concrete Score Matching: Generalized Score Matching for Discrete Data
• Elucidating the Design Space of Diffusion-Based Generative Models
• A Continuous Time Framework for Discrete Denoising Models
• MaskGIT: Masked Generative Image Transformer
• Training Verifiers to Solve Math Word Problems
• Program Synthesis with Large Language Models
• Structured Denoising Diffusion Models in Discrete State-Spaces
• Evaluating Large Language Models Trained on Code
• Measuring Mathematical Problem Solving With the MATH Dataset
• Score-Based Generative Modeling through Stochastic Differential Equations
• Denoising Diffusion Probabilistic Models
• The Curious Case of Neural Text Degeneration
• Improved Techniques for Training GANs
• Deep Unsupervised Learning using Nonequilibrium Thermodynamics
• A general method for numerically simulating the stochastic time evolution of coupled chemical reactions

Cited by

• Rethinking the Generation Order of Block Diffusion Language Models
• Learning Unmasking Policies for Diffusion Language Models
• Where and When to Commit: Candidate-Aware Decoding for Diffusion Language Models

Related

• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
• FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion
• Remasking Discrete Diffusion Models with Inference-Time Scaling
• WeFT: Weighted Entropy-driven Fine-Tuning for dLLMs
• No Compute Left Behind: Rethinking Reasoning and Sampling with Masked Diffusion Models