KLASS: Inferenza Rapida Guidata da KL nei Modelli di Diffusione con Mascheramento

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

November 7, 2025

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

cs.AI

Abstract

I modelli di diffusione mascherata hanno dimostrato risultati competitivi in varie attività, inclusa la generazione linguistica. Tuttavia, a causa del loro processo di raffinamento iterativo, l'inferenza è spesso limitata da una velocità di campionamento lenta e statica. Per superare questo problema, introduciamo il "KL-Adaptive Stability Sampling" (KLASS), un metodo di campionamento rapido ma efficace che sfrutta la divergenza KL a livello di token per identificare previsioni stabili e ad alta confidenza. Smascherando più token in ogni iterazione senza alcun addestramento aggiuntivo del modello, il nostro approccio accelera significativamente la generazione mantenendo la qualità del campione. Su benchmark di ragionamento, KLASS raggiunge accelerazioni temporali reali fino a 2,78 volte, migliorando le prestazioni rispetto alla decodifica greedy standard e ottenendo risultati allo stato dell'arte tra i campionatori basati su diffusione. Convalidiamo ulteriormente KLASS in diversi domini, inclusi la generazione di testo, immagini e molecole, dimostrandone l'efficacia come campionatore ampiamente applicabile su diversi modelli.

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.