Estabilizando o Aprendizado por Reforço para Modelos de Linguagem de Difusão
Stabilizing Reinforcement Learning for Diffusion Language Models
March 6, 2026
Autores: Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu
cs.AI
Resumo
O Group Relative Policy Optimization (GRPO) é altamente eficaz para modelos de linguagem autorregressivos (AR) após o treinamento, mas sua aplicação direta a modelos de linguagem grandes de difusão (dLLMs) frequentemente desencadeia um colapso da recompensa. Identificamos duas fontes de incompatibilidade. Primeiro, o GRPO depende de razões de importância definidas por probabilidades de sequência, que são intratáveis em dLLMs e devem ser estimadas (por exemplo, via *proxies* de verossimilhança baseados em ELBO ou de campo médio), produzindo razões inerentemente ruidosas. Segundo, a formulação padrão do GRPO não foi projetada para razões estimadas: seu recorte condicional pode ser anormalmente contornado por ruídos de estimação independentes do modelo, produzindo picos de gradiente, enquanto sua normalização de tamanho de grupo fixo amplifica flutuações na magnitude do gradiente sob estimativas de razão de alta variância. Mostramos que esses efeitos formam um ciclo de instabilidade auto-reforçado que conduz ao desvio da política e aumenta ainda mais a variância da razão. Para quebrar esse ciclo, propomos o StableDRL, uma reformulação do GRPO adaptada para dLLMs que utiliza (i) recorte incondicional para suprimir picos induzidos por *outliers* e (ii) auto-normalização para restringir as atualizações dentro do casco convexo dos gradientes por amostra. Estendemos ainda o StableDRL para modelos de difusão baseados em blocos através de um mecanismo de atenção em escada.
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.