Estabilización del Aprendizaje por Refuerzo para Modelos de Lenguaje de Difusión
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
Resumen
La Optimización de Políticas Relativas a Grupos (GRPO) es altamente efectiva para modelos de lenguaje autorregresivos (AR) posteriores al entrenamiento, sin embargo, su aplicación directa a modelos de lenguaje grandes de difusión (dLLMs) a menudo desencadena un colapso de la recompensa. Identificamos dos fuentes de incompatibilidad. En primer lugar, GRPO se basa en ratios de importancia definidos por probabilidades de secuencia, que son intratables en dLLMs y deben estimarse (por ejemplo, mediante aproximaciones de verosimilitud basadas en ELBO o de campo medio), lo que produce ratios inherentemente ruidosos. En segundo lugar, la formulación estándar de GRPO no está diseñada para ratios estimados: su recorte condicional puede ser anómalamente evitado por el ruido de estimación independiente del modelo, produciendo picos de gradiente, mientras que su normalización de tamaño de grupo fijo amplifica las fluctuaciones de magnitud del gradiente bajo estimaciones de ratio de alta varianza. Demostramos que estos efectos forman un bucle de inestabilidad auto-reforzante que impulsa la deriva de la política y aumenta aún más la varianza del ratio. Para romper este bucle, proponemos StableDRL, una reformulación de GRPO adaptada para dLLMs que utiliza (i) recorte incondicional para suprimir los picos inducidos por valores atípicos y (ii) auto-normalización para restringir las actualizaciones dentro del casco convexo de los gradientes por muestra. Además, extendemos StableDRL a modelos de difusión por bloques mediante un mecanismo de atención escalonada.
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.