Apprendimento di politiche di smascheramento per modelli linguistici basati sulla diffusione

Learning Unmasking Policies for Diffusion Language Models

December 9, 2025

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

cs.AI

Abstract

I modelli linguistici di diffusione (dLLM) ora eguagliano le prestazioni downstream delle loro controparti autoregressive in molte attività, mantenendo al contempo la promessa di una maggiore efficienza durante l'inferenza. Una variante particolarmente riuscita è la diffusione discreta mascherata, in cui un buffer riempito con token mascheratura speciali viene progressivamente sostituito con token campionati dal vocabolario del modello. È possibile ottenere efficienza smascherando più token in parallelo, ma smascherarne troppi in una volta rischia di degradare la qualità della generazione. Pertanto, un aspetto progettuale critico dei dLLM è la procedura di campionamento che seleziona, ad ogni passo del processo di diffusione, quali token sostituire. Infatti, lavori recenti hanno rilevato che strategie euristiche come la soglia di confidenza portano sia a una qualità più elevata che a un throughput di token superiore rispetto allo smascheramento casuale. Tuttavia, tali euristiche presentano degli svantaggi: richiedono una regolazione manuale e osserviamo che le loro prestazioni si degradano con dimensioni del buffer maggiori. In questo lavoro, proponiamo invece di addestrare le procedure di campionamento utilizzando l'apprendimento per rinforzo. Nello specifico, formalizziamo il campionamento per diffusione mascherata come un processo decisionale di Markov in cui il dLLM funge da ambiente, e proponiamo un'architettura policy leggera basata su un transformer a singolo strato che mappa le confidenze dei token del dLLM alle decisioni di smascheramento. I nostri esperimenti mostrano che queste policy addestrate eguagliano le prestazioni delle euristiche più all'avanguardia se combinate con la generazione semi-autoregressiva, superandole nell'impostazione di diffusione completa. Esaminiamo anche la trasferibilità di queste policy, riscontrando che possono generalizzare a nuovi dLLM sottostanti e a lunghezze di sequenza maggiori. Tuttavia, osserviamo anche che le loro prestazioni si degradano quando applicate a dati fuori dominio, e che una regolazione fine del compromesso precisione-efficienza può essere impegnativa con il nostro approccio.

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.