Apprentissage de politiques de démasquage pour les modèles de langage par diffusion

Learning Unmasking Policies for Diffusion Language Models

December 9, 2025

papers.authors: Metod Jazbec, Theo X. Olausson, Louis Béthune, Pierre Ablin, Michael Kirchhof, Joao Monterio, Victor Turrisi, Jason Ramapuram, Marco Cuturi

cs.AI

papers.abstract

Les modèles de langage à diffusion (dLLMs) égalent désormais les performances en aval de leurs homologues autorégressifs sur de nombreuses tâches, tout en promettant une inférence plus efficace. Une variante particulièrement réussie est la diffusion discrète masquée, où un tampon rempli de jetons de masquage spéciaux est progressivement remplacé par des jetons échantillonnés depuis le vocabulaire du modèle. L'efficacité peut être améliorée en démasquant plusieurs jetons en parallèle, mais en traiter trop simultanément risque de dégrader la qualité de la génération. Ainsi, un aspect crucial de la conception des dLLMs est la procédure d'échantillonnage qui sélectionne, à chaque étape du processus de diffusion, quels jetons remplacer. De récents travaux ont effectivement montré que des stratégies heuristiques comme le seuillage de confiance conduisent à une qualité et un débit de jetons supérieurs par rapport au démasquage aléatoire. Cependant, ces heuristiques présentent des inconvénients : elles nécessitent un réglage manuel, et nous observons que leurs performances se dégradent avec des tailles de tampon plus importantes. Dans ce travail, nous proposons plutôt d'entraîner les procédures d'échantillonnage à l'aide de l'apprentissage par renforcement. Plus précisément, nous formalisons l'échantillonnage par diffusion masquée comme un processus de décision markovien où le dLLM sert d'environnement, et nous proposons une architecture de politique légère basée sur un transformeur monocouche qui mappe les confiances des jetons du dLLM vers des décisions de démasquage. Nos expériences montrent que ces politiques entraînées égalent les performances des heuristiques de pointe lorsqu'elles sont combinées à une génération semi-autorégressive, tout en les surpassant dans le cadre complet de la diffusion. Nous examinons également la transférabilité de ces politiques, constatant qu'elles peuvent généraliser à de nouveaux dLLMs sous-jacents et à des longueurs de séquence plus importantes. Cependant, nous observons aussi que leurs performances se dégradent lorsqu'elles sont appliquées à des données hors domaine, et qu'un réglage fin du compromis précision-efficacité peut s'avérer difficile avec notre approche.

English

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.