Укрощение маскированных диффузионных языковых моделей через обучение с подкреплением на основе траекторий согласованности с уменьшением числа шагов декодирования

Аннотация

Маскированные диффузионные языковые модели (MDLM) недавно появились как перспективная альтернатива авторегрессивным (AR) языковым моделям, предлагая такие свойства, как параллельное декодирование, гибкие порядки генерации и возможность сокращения количества шагов вывода. Несмотря на эти преимущества, стратегии декодирования и алгоритмы обучения с подкреплением (RL), адаптированные для MDLM, остаются недостаточно изученными. Наивный подход заключается в прямом переносе хорошо зарекомендовавших себя методов для AR-моделей на MDLM. Однако это сразу же вызывает вопрос: действительно ли такой наивный перенос оптимален? Например, 1) Блочные и полуавторегрессивные стратегии декодирования не используются при обучении MDLM, так почему же они превосходят полное диффузионное декодирование на этапе вывода? 2) Применение алгоритмов RL, разработанных для AR-моделей, непосредственно к MDLM демонстрирует несоответствие между обучением и выводом, поскольку декодирование MDLM является некаузальным (параллельным). Это приводит к несоответствиям между траекторией rollout и траекторией оптимизации. Для решения этих проблем мы предлагаем механизмы EOS Early Rejection (EOSER) и Ascending Step-Size (ASS) для планирования декодирования, которые раскрывают потенциал MDLM для выполнения полного диффузионного декодирования, достигая конкурентоспособной производительности с меньшим количеством шагов декодирования. Кроме того, мы представляем Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) для управления MDLM, который подчеркивает согласованность между траекторией rollout и траекторией оптимизации и уменьшает ошибки оптимизации, вызванные пропуском шагов. Мы проводим обширные эксперименты на задачах рассуждения, таких как математические и планировочные бенчмарки, используя LLaDA-8B-Instruct. Результаты показывают, что предложенные механизмы EOSER и ASS вместе с CJ-GRPO обладают значительным потенциалом для эффективного и результативного управления MDLM. Код: 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.

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

Аннотация

Support