Stabilisierung von Reinforcement Learning für Diffusionssprachmodelle
Stabilizing Reinforcement Learning for Diffusion Language Models
March 6, 2026
Autoren: Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu
cs.AI
Zusammenfassung
Group Relative Policy Optimization (GRPO) ist äußerst effektiv für autoregressive (AR) Sprachmodelle nach dem Training, doch seine direkte Anwendung auf Diffusions-Großsprachmodelle (dLLMs) löst häufig einen Reward Collapse aus. Wir identifizieren zwei Ursachen für diese Inkompatibilität. Erstens basiert GRPO auf Importance Ratios, die durch Sequenzwahrscheinlichkeiten definiert sind. Diese sind in dLLMs nicht handhabbar und müssen geschätzt werden (z.B. über ELBO-basierte oder Mean-Field-Likelihood-Proxys), was inhärent verrauschte Ratios liefert. Zweitens ist die Formulierung des Standard-GRPO nicht für geschätzte Ratios ausgelegt: Ihr bedingtes Clipping kann durch modellagnostisches Schätzrauschen anomal umgangen werden, was Gradient Spikes erzeugt, während ihre Normalisierung mit fester Gruppengröße Schwankungen der Gradientenstärke unter hochvarianen Ratioschätzungen verstärkt. Wir zeigen, dass diese Effekte eine sich selbst verstärkende Instabilitätsschleife bilden, die Policy Drift antreibt und die Ratio-Varianz weiter erhöht. Um diese Schleife zu durchbrechen, schlagen wir StableDRL vor, eine Neuformulierung von GRPO für dLLMs, die (i) unbedingtes Clipping zur Unterdrückung von Ausreißer-induzierten Spikes und (ii) Selbstnormalisierung verwendet, um Updates innerhalb der konvexen Hülle der Gradienten pro Stichprobe zu beschränken. Wir erweitern StableDRL außerdem auf blockweise Diffusionsmodelle über einen Stufen-Attentionsmechanismus.
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.