Domesticación de Modelos de Lenguaje de Difusión Enmascarada mediante Refuerzo de la Trayectoria de Consistencia con Menos Pasos de Decodificación

Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step

September 28, 2025

Autores: Jingyi Yang, Guanxu Chen, Xuhao Hu, Jing Shao

cs.AI

Resumen

Los modelos de lenguaje de difusión enmascarada (MDLMs, por sus siglas en inglés) han surgido recientemente como una alternativa prometedora a los modelos de lenguaje autorregresivos (AR), ofreciendo propiedades como decodificación paralela, órdenes de generación flexibles y el potencial de requerir menos pasos de inferencia. A pesar de estas ventajas, las estrategias de decodificación y los algoritmos de aprendizaje por refuerzo (RL, por sus siglas en inglés) adaptados para MDLMs siguen siendo poco explorados. Un enfoque ingenuo es transferir directamente técnicas bien establecidas para modelos AR a los MDLMs. Sin embargo, esto plantea una pregunta inmediata: ¿Es realmente óptima esta transferencia ingenua? Por ejemplo, 1) Las estrategias de decodificación por bloques y semi-AR no se emplean durante el entrenamiento de los MDLMs, entonces, ¿por qué superan a la decodificación de estilo de difusión completa durante la inferencia? 2) La aplicación directa de algoritmos de RL diseñados para modelos AR a los MDLMs exhibe una inconsistencia entre el entrenamiento y la inferencia, ya que la decodificación de los MDLMs no es causal (es paralela). Esto resulta en inconsistencias entre la trayectoria de despliegue y la trayectoria de optimización. Para abordar estos desafíos, proponemos el Rechazo Temprano de EOS (EOSER, por sus siglas en inglés) y el Planificador de Decodificación de Tamaño de Paso Ascendente (ASS, por sus siglas en inglés), que desbloquean el potencial de los MDLMs para realizar una decodificación de estilo de difusión completa, logrando un rendimiento competitivo con menos pasos de decodificación. Además, introducimos la Optimización de Política Relativa de Grupo de Trayectoria de Consistencia (CJ-GRPO, por sus siglas en inglés) para domar los MDLMs, que enfatiza la consistencia entre la trayectoria de despliegue y la trayectoria de optimización, y reduce los errores de optimización causados por la optimización de pasos omitidos. Realizamos experimentos extensos en tareas de razonamiento, como benchmarks matemáticos y de planificación, utilizando LLaDA-8B-Instruct. Los resultados demuestran que los mecanismos propuestos EOSER y ASS, junto con CJ-GRPO, tienen un potencial significativo para domar eficaz y eficientemente los MDLMs. Código: https://github.com/yjyddq/EOSER-ASS-RL.

English

Masked diffusion language models (MDLMs) have recently emerged as a promising alternative to autoregressive (AR) language models, offering properties such as parallel decoding, flexible generation orders, and the potential for fewer inference steps. Despite these advantages, decoding strategies and reinforcement learning (RL) algorithms tailored for MDLMs remain underexplored. A naive approach is to directly transfer techniques well-established for AR models to MDLMs. However, this raises an immediate question: Is such a naive transfer truly optimal? For example, 1) Block-wise and semi-AR decoding strategies are not employed during the training of MDLMs, so why do they outperform full diffusion-style decoding during inference? 2) Applying RL algorithms designed for AR models directly to MDLMs exhibits a training-inference inconsistency, since MDLM decoding are non-causal (parallel). This results in inconsistencies between the rollout trajectory and the optimization trajectory. To address these challenges, we propose EOS Early Rejection (EOSER) and Ascending Step-Size (ASS) decoding scheduler, which unlock the potential of MDLMs to perform full diffusion-style decoding, achieving competitive performance with fewer decoding steps. Additionally, we introduce Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) for taming MDLMs, which emphasizes the consistency between rollout trajectory and optimization trajectory, and reduces the optimization errors caused by skip-step optimization. We conduct extensive experiments on reasoning tasks, such as mathematical and planning benchmarks, using LLaDA-8B-Instruct. The results demonstrate that the proposed EOSER and ASS mechanisms, together with CJ-GRPO, hold significant promise for effectively and efficiently taming MDLMs. Code: https://github.com/yjyddq/EOSER-ASS-RL.