KLASS: Inferência Rápida Guiada por KL em Modelos de Difusão com Mascaramento

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

Resumo

Os modelos de difusão mascarada demonstraram resultados competitivos em várias tarefas, incluindo geração de linguagem. No entanto, devido ao seu processo iterativo de refinamento, a inferência é frequentemente limitada por uma velocidade de amostragem lenta e estática. Para superar este problema, introduzimos o "KL-Adaptive Stability Sampling" (KLASS), um método de amostragem rápido e eficaz que explora a divergência KL a nível de token para identificar previsões estáveis e de alta confiança. Ao desmascarar múltiplos tokens em cada iteração sem qualquer treino adicional do modelo, a nossa abordagem acelera significativamente a geração, mantendo a qualidade da amostra. Em benchmarks de raciocínio, o KLASS alcança acelerações de até 2,78× em tempo real, melhorando o desempenho em relação à descodificação greedy padrão e obtendo resultados de última geração entre os amostradores baseados em difusão. Validamos ainda o KLASS em diversos domínios, incluindo geração de texto, imagem e moléculas, demonstrando a sua eficácia como um amostrador amplamente aplicável em 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.