Toutes les étapes de débruitage ne se valent pas : planification des modèles pour accélérer les modèles de langage à diffusion masquée
Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models
April 11, 2026
Auteurs: Ivan Sedykh, Nikita Sorokin, Valentin Malykh
cs.AI
Résumé
Les récents progrès des modèles de langage à diffusion masquée (MDLM) réduisent l'écart de qualité avec les modèles autoregressifs, mais leur échantillonnage reste coûteux car la génération nécessite de nombreuses passes de dé-bruitage de séquence complète avec un grand Transformer et, contrairement au décodage autoregressif, ne peut pas bénéficier de la mise en cache KV. Dans ce travail, nous exploitons la flexibilité du cadre de diffusion et étudions l'ordonnancement de modèles, où un MDLM plus petit remplace le modèle complet à un sous-ensemble des étapes de dé-bruitage. Sur des modèles entraînés sur OpenWebText et LM1B, nous montrons que les premières et dernières étapes de dé-bruitage sont nettement plus robustes à un tel remplacement que les étapes intermédiaires, permettant une réduction allant jusqu'à 17% des FLOPs avec seulement une dégradation modeste de la perplexité générative en génération non conditionnelle et conditionnée par un préfixe, tout en préservant la diversité des échantillons. Nous étayons ces résultats par une analyse d'importance des étapes basée sur la perte et la divergence KL entre petits et grands modèles à travers les pas de temps, ainsi qu'une recherche exhaustive sur des segments d'étapes grossiers, qui identifient tous deux le milieu de la trajectoire de diffusion comme étant le plus sensible, de manière cohérente entre les jeux de données. Nos résultats suggèrent que des règles d'ordonnancement simples et indépendantes de l'architecture peuvent accélérer significativement l'échantillonnage des MDLM tout en préservant largement la qualité de la génération.
English
Recent advances in masked diffusion language models (MDLMs) narrow the quality gap to autoregressive LMs, but their sampling remains expensive because generation requires many full-sequence denoising passes with a large Transformer and, unlike autoregressive decoding, cannot benefit from KV caching. In this work, we exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps. Across models trained on OpenWebText and LM1B, we show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs with only modest degradation in generative perplexity under both unconditional and prefix-conditional generation, while preserving sample diversity. We support these findings with a step-importance analysis based on loss and KL divergence between small and large models across timesteps, as well as an exhaustive search over coarse step segments, both of which identify the middle of the diffusion trajectory as most sensitive consistently across datasets. Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.