Стабилизация обучения с подкреплением для диффузионных языковых моделей

Stabilizing Reinforcement Learning for Diffusion Language Models

March 6, 2026

Авторы: Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu

cs.AI

Аннотация

Метод оптимизации политики относительно групп (GRPO) демонстрирует высокую эффективность для пост-обучения авторегрессионных (AR) языковых моделей, однако его прямое применение к диффузионным большим языковым моделям (dLLM) часто приводит к коллапсу вознаграждения. Мы выявили два источника несовместимости. Во-первых, GRPO опирается на важностные отношения, определяемые вероятностями последовательностей, которые в случае dLLM являются интрактивными и должны оцениваться (например, через прокси-вероятности на основе ELBO или среднего поля), что приводит к inherently зашумленным отношениям. Во-вторых, формулировка стандартного GRPO не предназначена для оцененных отношений: её условное ограничение (clipping) может аномально обходиться из-за агностического к модели шума оценок, порождая всплески градиента, в то время как её нормировка с фиксированным размером группы усиливает флуктуации величины градиента при оценках отношений с высокой дисперсией. Мы показываем, что эти эффекты формируют самоподкрепляющийся цикл нестабильности, который провоцирует дрейф политики и дальнейший рост дисперсии отношений. Чтобы разорвать этот цикл, мы предлагаем StableDRL — реформулировку GRPO, адаптированную для dLLM, которая использует (i) безусловное ограничение для подавления всплесков, вызванных выбросами, и (ii) самонормировку для ограничения обновлений внутри выпуклой оболочки градиентов для каждого отдельного примера. Мы также расширяем StableDRL для блочных диффузионных моделей с помощью механизма ступенчатого внимания (staircase attention).

English

Group Relative Policy Optimization (GRPO) is highly effective for post-training autoregressive (AR) language models, yet its direct application to diffusion large language models (dLLMs) often triggers reward collapse. We identify two sources of incompatibility. First, GRPO relies on importance ratios defined by sequence probabilities, which are intractable in dLLMs and must be estimated (e.g., via ELBO-based or mean-field likelihood proxies), yielding inherently noisy ratios. Second, standard GRPO's formulation is not designed for estimated ratios: its conditional clipping can be anomalously bypassed by model-agnostic estimation noise, producing gradient spikes, while its fixed group-size normalization amplifies gradient-magnitude fluctuations under high-variance ratio estimates. We show these effects form a self-reinforcing instability loop that drives policy drift and further increases ratio variance. To break this loop, we propose StableDRL, a reformulation of GRPO tailored for dLLMs that uses (i) unconditional clipping to suppress outlier-induced spikes and (ii) self-normalization to constrain updates within the convex hull of per-sample gradients. We further extend StableDRL to block-wise diffusion models via a staircase attention mechanism.