Aprendendo Políticas de Desmascaramento para Modelos de Linguagem de Difusão

Learning Unmasking Policies for Diffusion Language Models

December 9, 2025

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

cs.AI

Resumo

Os Modelos de Linguagem de Difusão (Grandes) (dLLMs) agora igualam o desempenho em tarefas específicas (downstream) de suas contrapartes autoregressivas em muitas tarefas, ao mesmo tempo que prometem ser mais eficientes durante a inferência. Uma variante particularmente bem-sucedida é a difusão discreta mascarada, na qual um buffer preenchido com tokens especiais de máscara é progressivamente substituído por tokens amostrados do vocabulário do modelo. A eficiência pode ser obtida ao desmascarar vários tokens em paralelo, mas fazer muitos de uma vez corre o risco de degradar a qualidade da geração. Assim, um aspecto crítico de projeto dos dLLMs é o procedimento de amostragem que seleciona, a cada passo do processo de difusão, quais tokens substituir. De fato, trabalhos recentes descobriram que estratégias heurísticas, como o limiar de confiança (confidence thresholding), levam a uma maior qualidade e taxa de processamento de tokens (throughput) em comparação com o desmascaramento aleatório. No entanto, tais heurísticas têm desvantagens: elas exigem ajuste manual, e observamos que seu desempenho se degrada com tamanhos de buffer maiores. Neste trabalho, propomos, em vez disso, treinar procedimentos de amostragem usando aprendizado por reforço. Especificamente, formalizamos a amostragem por difusão mascarada como um processo de decisão de Markov no qual o dLLM atua como o ambiente, e propomos uma arquitetura de política leve (lightweight policy) baseada em um transformador de camada única que mapeia as confianças dos tokens do dLLM para decisões de desmascaramento. Nossos experimentos mostram que essas políticas treinadas igualam o desempenho das heurísticas state-of-the-art quando combinadas com geração semi-autoregressiva, ao mesmo tempo que as superam no cenário de difusão completa. Também examinamos a transferibilidade dessas políticas, constatando que elas podem generalizar para novos dLLMs subjacentes e comprimentos de sequência maiores. No entanto, também observamos que seu desempenho se degrada quando aplicadas a dados fora do domínio (out-of-domain), e que o ajuste fino do compromisso (trade-off) entre precisão e eficiência pode ser desafiador com nossa abordagem.

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.