Stabilisation de l'Apprentissage par Renforcement pour les Modèles de Langage à Diffusion
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
Résumé
L'optimisation de politique relative au groupe (GRPO) est très efficace pour les modèles de langage autorégressifs (AR) après leur entraînement, mais son application directe aux grands modèles de langage par diffusion (dLLM) déclenche souvent un effondrement de la récompense. Nous identifions deux sources d'incompatibilité. Premièrement, GRPO repose sur des ratios d'importance définis par les probabilités des séquences, qui sont intraçables dans les dLLM et doivent être estimés (par exemple, via des proxys de vraisemblance basés sur l'ELBO ou par champ moyen), produisant des ratios intrinsèquement bruités. Deuxièmement, la formulation standard de GRPO n'est pas conçue pour des ratios estimés : son seuillage conditionnel peut être anormalement contourné par un bruit d'estimation indépendant du modèle, générant des pics de gradient, tandis que sa normalisation à taille de groupe fixe amplifie les fluctuations d'amplitude du gradient sous des estimations de ratio à haute variance. Nous montrons que ces effets forment une boucle d'instabilité auto-renforçante qui entraîne une dérive de la politique et augmente encore la variance des ratios. Pour briser cette boucle, nous proposons StableDRL, une reformulation de GRPO adaptée aux dLLM qui utilise (i) un seuillage inconditionnel pour supprimer les pics induits par les valeurs aberrantes et (ii) une auto-normalisation pour contraindre les mises à jour dans l'enveloppe convexe des gradients par échantillon. Nous étendons en outre StableDRL aux modèles de diffusion par blocs via un mécanisme d'attention en escalier.
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.