Stabilisatie van Reinforcement Learning voor Diffusion Language Models
Stabilizing Reinforcement Learning for Diffusion Language Models
March 6, 2026
Auteurs: Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu
cs.AI
Samenvatting
Group Relative Policy Optimization (GRPO) is zeer effectief voor getrainde autoregressieve (AR) taalmodel(len), maar de directe toepassing ervan op diffusion large language models (dLLM's) leidt vaak tot reward collapse. Wij identificeren twee bronnen van incompatibiliteit. Ten eerste vertrouwt GRPO op important ratios die zijn gedefinieerd door sequentiekansen, die in dLLM's onberekenbaar zijn en geschat moeten worden (bijvoorbeeld via ELBO-gebaseerde of mean-field likelihood proxies), wat inherent ruisrijke ratios oplevert. Ten tweede is de formulering van standaard GRPO niet ontworpen voor geschatte ratios: de conditionele clipping kan op een afwijkende manier worden omzeild door model-agnostische schattingsruis, wat gradient spikes produceert, terwijl de vaste normalisatie van de groepsgrootte fluctuaties in de gradient-grootte versterkt onder hoog-variantie ratioschattingen. Wij tonen aan dat deze effecten een zichzelf versterkende instabiliteitslus vormen die policy drift veroorzaakt en de ratio-variantie verder vergroot. Om deze lus te doorbreken, stellen wij StableDRL voor, een herformulering van GRPO die is afgestemd op dLLM's en gebruikmaakt van (i) onvoorwaardelijke clipping om door uitschieters veroorzaakte spikes te onderdrukken en (ii) zelf-normalisatie om updates te begrenzen binnen het convexe omhulsel van per-voorbeeld gradients. Wij breiden StableDRL verder uit naar bloksgewijze diffusionmodellen via een staircase attention-mechanisme.
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.