Aprendizaje de Políticas de Desenmascaramiento para Modelos de Lenguaje Basados en Difusión

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

Resumen

Los Modelos de Lenguaje de Difusión (a Gran Escala) (dLLMs) ahora igualan el rendimiento en tareas posteriores de sus contrapartes autorregresivas en muchas tareas, al mismo tiempo que prometen ser más eficientes durante la inferencia. Una variante particularmente exitosa es la difusión discreta enmascarada, en la que un búfer lleno de tokens de enmascaramiento especiales es reemplazado progresivamente por tokens muestreados del vocabulario del modelo. Se puede ganar eficiencia desenmascarando varios tokens en paralelo, pero hacer demasiados a la vez corre el riesgo de degradar la calidad de la generación. Por lo tanto, un aspecto de diseño crítico de los dLLMs es el procedimiento de muestreo que selecciona, en cada paso del proceso de difusión, qué tokens reemplazar. De hecho, trabajos recientes han encontrado que las estrategias heurísticas, como el umbral de confianza, conducen a una mayor calidad y rendimiento de tokens en comparación con el desenmascaramiento aleatorio. Sin embargo, tales heurísticas tienen desventajas: requieren ajuste manual, y observamos que su rendimiento se degrada con tamaños de búfer más grandes. En este trabajo, proponemos en cambio entrenar procedimientos de muestreo utilizando aprendizaje por refuerzo. Específicamente, formalizamos el muestreo de difusión enmascarada como un proceso de decisión de Markov en el que el dLLM sirve como entorno, y proponemos una arquitectura de política ligera basada en un transformador de una sola capa que mapea las confianzas de los tokens del dLLM a decisiones de desenmascaramiento. Nuestros experimentos muestran que estas políticas entrenadas igualan el rendimiento de las heurísticas más avanzadas cuando se combinan con generación semi-autorregresiva, al mismo tiempo que las superan en el escenario de difusión completo. También examinamos la transferibilidad de estas políticas, encontrando que pueden generalizarse a nuevos dLLMs subyacentes y longitudes de secuencia más largas. Sin embargo, también observamos que su rendimiento se degrada cuando se aplican a datos fuera de dominio, y que el ajuste fino del equilibrio entre precisión y eficiencia puede ser un desafío con nuestro enfoque.

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.