Het aanleren van ontmaskeringsbeleid voor diffusie-taalmodellen

Learning Unmasking Policies for Diffusion Language Models

December 9, 2025

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

cs.AI

Samenvatting

Diffusie (Large) Language Models (dLLM's) evenaren nu de downstream-prestaties van hun autoregressieve tegenhangers op veel taken, met de belofte van efficiëntere inferentie. Een bijzonder succesvolle variant is gemaskeerde discrete diffusie, waarbij een buffer gevuld met speciale maskertokens geleidelijk wordt vervangen door tokens die uit de vocabulaire van het model worden bemonsterd. Efficiëntiewinst kan worden behaald door meerdere tokens parallel te demaskeren, maar te veel tokens tegelijk vervangen riskeert een verslechtering van de generatiekwaliteit. Een cruciaal ontwerpaspect van dLLM's is daarom de samplingprocedure die in elke stap van het diffusieproces selecteert welke tokens moeten worden vervangen. Recent onderzoek heeft inderdaad aangetoond dat heuristische strategieën, zoals drempelwaardebepaling op basis van confidentie, leiden tot zowel hogere kwaliteit als een hogere tokenverwerkingssnelheid vergeleken met willekeurig demaskeren. Dergelijke heuristieken hebben echter nadelen: ze vereisen handmatige afstemming, en wij observeren dat hun prestaties afnemen bij grotere buffergroottes. In dit werk stellen wij voor om samplingprocedures te trainen met reinforcement learning. Concreet formaliseren we gemaskeerde diffusie-sampling als een Markov-beslissingsproces waarin de dLLM als omgeving fungeert, en introduceren we een lichtgewicht policy-architectuur gebaseerd op een single-layer transformer die de tokenconfidenties van de dLLM afbeeldt op demaskeringsbeslissingen. Onze experimenten tonen aan dat deze getrainde policies de prestaties van state-of-the-art heuristieken evenaren in combinatie met semi-autoregressieve generatie, en deze overtreffen in de volledige diffusie-instelling. We onderzoeken ook de overdraagbaarheid van deze policies en constateren dat ze kunnen generaliseren naar nieuwe onderliggende dLLM's en langere sequentielengtes. We observeren echter ook dat hun prestaties afnemen wanneer ze worden toegepast op out-of-domain data, en dat een fijnmazige afstemming van de nauwkeurigheid-efficiëntie-afweging uitdagend kan zijn met onze aanpak.

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.