Lernende Entschleierungspolitiken für Diffusionssprachmodelle

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

Diffusions-(Groß-)Sprachmodelle (dLLMs) erreichen inzwischen bei vielen Aufgaben die nachgelagerte Leistung ihrer autoregressiven Gegenstücke, wobei sie die Aussicht auf höhere Effizienz während des Inferenzvorgangs bieten. Eine besonders erfolgreiche Variante ist die maskierte diskrete Diffusion, bei der ein mit speziellen Maskentokens gefüllter Puffer schrittweise durch Tokens aus dem Vokabular des Modells ersetzt wird. Effizienzgewinne lassen sich durch das parallele Entmaskieren mehrerer Tokens erzielen, doch birgt ein gleichzeitiges Entmaskern zu vieler Tokens das Risiko, die Generierungsqualität zu beeinträchtigen. Daher ist ein entscheidender Gestaltungsaspekt von dLLMs das Sampling-Verfahren, das in jedem Schritt des Diffusionsprozesses auswählt, welche Tokens ersetzt werden sollen. Tatsächlich hat die jüngste Forschung gezeigt, dass heuristische Strategien wie Confidence-Thresholding im Vergleich zum zufälligen Entmaskieren sowohl zu einer höheren Qualität als auch zu einem höheren Token-Durchsatz führen. Solche Heuristiken haben jedoch Nachteile: Sie erfordern eine manuelle Abstimmung, und wir beobachten, dass ihre Leistung mit zunehmender Puffergröße abnimmt. In dieser Arbeit schlagen wir stattdessen vor, Sampling-Verfahren mit bestärkendem Lernen zu trainieren. Konkret formalisieren wir das Sampling bei maskierter Diffusion als Markov-Entscheidungsprozess, bei dem das dLLM als Umgebung dient, und schlagen eine leichtgewichtige Policy-Architektur auf Basis eines Single-Layer-Transformers vor, die die Token-Konfidenzen des dLLM auf Entmaskierungsentscheidungen abbildet. Unsere Experimente zeigen, dass diese trainierten Policies in Kombination mit semi-autoregressiver Generierung die Leistung modernster Heuristiken erreichen und sie im vollständigen Diffusionssetting übertreffen. Wir untersuchen auch die Übertragbarkeit dieser Policies und stellen fest, dass sie auf neue zugrundeliegende dLLMs und längere Sequenzen verallgemeinern können. Allerdings beobachten wir auch, dass ihre Leistung bei der Anwendung auf datenferne Domänen abnimmt und dass eine feinabgestimmte Optimierung des Trade-offs zwischen Genauigkeit und Effizienz mit unserem Ansatz herausfordernd sein kann.

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.