KLASS: Inferencia Rápida Guiada por KL en Modelos de Difusión Enmascarados

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

November 7, 2025

Autores: Seo Hyun Kim, Sunwoo Hong, Hojung Jung, Youngrok Park, Se-Young Yun

cs.AI

Resumen

Los modelos de difusión enmascarada han demostrado resultados competitivos en diversas tareas, incluida la generación de lenguaje. Sin embargo, debido a su proceso de refinamiento iterativo, la inferencia suele verse limitada por una velocidad de muestreo lenta y estática. Para superar este problema, presentamos `KL-Adaptive Stability Sampling' (KLASS), un método de muestreo rápido pero eficaz que aprovecha la divergencia KL a nivel de token para identificar predicciones estables y de alta confianza. Al desenmascarar múltiples tokens en cada iteración sin necesidad de entrenamiento adicional del modelo, nuestro enfoque acelera significativamente la generación manteniendo la calidad de la muestra. En benchmarks de razonamiento, KLASS logra aceleraciones de hasta 2.78x en tiempo de ejecución, mejorando el rendimiento respecto a la decodificación codiciosa estándar y alcanzando resultados de vanguardia entre los muestreadores basados en difusión. Validamos además KLASS en diversos dominios, incluyendo generación de texto, imágenes y moléculas, demostrando su eficacia como un muestreador ampliamente aplicable en diferentes modelos.

English

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.78times 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.