💬 NLP

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

Cet article propose la méthode SPG (Sandwiched Policy Gradient), qui améliore l'alignement des modèles de langage par diffusion avec les préférences humaines en utilisant des bornes supérieure et inférieure pour surmonter les biais des approximations unilatérales, démontrant ainsi des performances supérieures sur plusieurs tâches de raisonnement mathématique et logique.

Auteurs originaux : Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

Publié 2026-04-16

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

🍞 Le Sandwich qui sauve les modèles de diffusion : L'histoire de SPG

Imaginez que vous essayez d'enseigner à un robot comment résoudre des énigmes mathématiques ou des Sudoku. Ce robot est un modèle de langage par diffusion (dLLM). Contrairement aux robots classiques qui écrivent mot par mot (comme un humain qui écrit une lettre), ce robot fonctionne comme un peintre qui efface et repeint.

Il commence avec un tableau totalement blanc (ou rempli de masques), puis il "nettoie" progressivement le tableau pour révéler la réponse finale. C'est rapide et parallèle, mais il y a un gros problème : on ne sait pas exactement comment le robot "pense" à chaque étape.

🚧 Le Problème : La boussole cassée

Pour améliorer ce robot, on utilise souvent l'apprentissage par renforcement (RL), un peu comme un entraîneur qui donne des félicitations (récompenses) ou des remarques négatives (punitions).

Si le robot donne une bonne réponse, l'entraîneur dit : "Bravo, fais plus souvent ça !"
Si la réponse est mauvaise, il dit : "Non, ne fais plus jamais ça !"

Le problème, c'est que pour dire "Ne fais plus jamais ça", il faut calculer la probabilité exacte que le robot ait produit cette mauvaise réponse. Mais avec les modèles de diffusion, ce calcul est mathématiquement impossible (trop complexe).

Les chercheurs précédents utilisaient une "approximation" (une estimation approximative), comme regarder la boussole à travers un brouillard.

Pour les bonnes réponses, ça marchait bien.
Mais pour les mauvaises réponses, cette approximation était faussée. C'était comme si l'entraîneur disait : "Ne fais pas ça", mais en réalité, il ne savait pas combien c'était mauvais. Résultat : le robot apprenait mal et restait bloqué.

🥪 La Solution : Le Sandwich (SPG)

Les auteurs de cet article, Chenyu Wang et son équipe, ont proposé une solution géniale appelée SPG (Sandwiched Policy Gradient).

Imaginez que vous voulez mesurer la taille d'un objet caché dans une boîte, mais vous ne pouvez pas l'ouvrir.

La borne inférieure (Le pain du bas) : Vous savez que l'objet fait au moins 10 cm. C'est une estimation sûre, mais peut-être un peu basse.
La borne supérieure (Le pain du haut) : Vous savez aussi qu'il fait au plus 20 cm. C'est une limite haute.

L'objet réel est quelque part entre les deux.

L'idée du Sandwich :
Au lieu d'utiliser une seule estimation floue, SPG utilise les deux bornes en même temps, selon la situation :

Si le robot a une bonne réponse (Récompense positive) : On utilise la borne inférieure (le pain du bas). On dit : "Bravo, tu es au moins aussi bon que ça, donc améliore-toi !" C'est sûr et efficace.
Si le robot a une mauvaise réponse (Récompense négative) : On utilise la borne supérieure (le pain du haut). On dit : "Oh là là, tu es au moins aussi mauvais que cette limite haute, donc arrête-toi !" Cela permet de punir sévèrement les erreurs, ce que les anciennes méthodes ne pouvaient pas faire correctement.

En gardant la vraie valeur "coincée" entre ces deux estimations (comme un sandwich), le robot reçoit un signal d'apprentissage beaucoup plus précis et équilibré.

🧱 L'astuce supplémentaire : Le Masquage par Blocs

Pour que ce sandwich soit bien cuit, il faut aussi bien mélanger les ingrédients. Les chercheurs ont remarqué que si on mélangeait le robot de manière totalement aléatoire pendant l'entraînement, il se perdait.

Ils ont donc inventé une stratégie de "Masquage par blocs" :
Au lieu de cacher des mots au hasard dans une phrase, on cache des blocs entiers (par exemple, 3 mots d'affilée), en laissant le début de la phrase intact. C'est comme si on demandait au robot de compléter un paragraphe entier plutôt que de deviner un mot isolé au milieu d'un chaos. Cela rend l'entraînement beaucoup plus stable et logique.

🏆 Les Résultats : Un robot qui devient un génie

Grâce à cette méthode "Sandwich", le robot a fait des progrès incroyables sur des tests difficiles :

Mathématiques (GSM8K, MATH500) : Il a gagné plusieurs points de précision par rapport aux meilleurs modèles actuels.
Logique (Sudoku, Countdown) : C'est là que la magie opère le plus ! Sur les Sudoku, le modèle a bondi de 27 % de précision. C'est énorme !

En résumé, SPG est comme un entraîneur qui a enfin une boussole fiable. Il sait exactement quand féliciter et, surtout, quand punir sévèrement les erreurs, permettant au robot de devenir beaucoup plus intelligent et rapide pour résoudre des problèmes complexes.

En résumé très simple :

Le problème : Les modèles de diffusion sont difficiles à entraîner car on ne peut pas mesurer précisément leurs erreurs.
La solution (SPG) : Utiliser deux estimations (une basse et une haute) pour "encadrer" la vérité. On utilise la borne basse pour les succès et la borne haute pour les échecs.
Le résultat : Un modèle qui apprend beaucoup plus vite et qui résout des énigmes logiques bien mieux que les précédents.

1. Problématique et Contexte

Les Modèles de Langage par Diffusion (dLLMs) émergent comme une alternative efficace aux modèles autoregressifs (AR) grâce à leur capacité à décoder plusieurs tokens en parallèle, réduisant ainsi la latence d'inférence. Cependant, l'alignement de ces modèles avec les préférences humaines ou des récompenses spécifiques à une tâche via l'Apprentissage par Renforcement (RL) reste un défi majeur.

Le problème central réside dans l'intractabilité du log-vraisemblance (log-likelihood) des dLLMs. Contrairement aux modèles AR où la probabilité d'une séquence est le produit des probabilités conditionnelles des tokens, les dLLMs (notamment les modèles à masquage comme MDLM) n'ont pas de log-vraisemblance calculable directement.

Limitation des méthodes existantes : Les travaux antérieurs contournent ce problème en utilisant des approximations comme la borne inférieure de la vraisemblance (ELBO) ou des estimations en une seule étape.
Biais de gradient : L'ELBO n'est qu'une borne inférieure ( $ELBO \le \log \pi_\theta$ ). Maximiser l'ELBO pour les échantillons à forte récompense est valide, mais minimiser l'ELBO pour les échantillons à faible récompense ne garantit pas de minimiser le vrai log-vraisemblance. Cela introduit un biais significatif dans le gradient de la politique, empêchant le modèle d'apprendre efficacement des retours négatifs (pénalités) et limitant les performances des algorithmes de RL avancés utilisant des récompenses relatives ou négatives.

2. Méthodologie : SPG (Sandwiched Policy Gradient)

Les auteurs proposent SPG, un nouvel algorithme de gradient de politique conçu pour résoudre ce biais en "sandwichant" le log-vraisemblance intractable entre deux bornes calculables.

A. Objectif "Sandwiché"

L'idée centrale est d'utiliser des bornes différentes selon le signe de l'avantage (la récompense relative) d'une séquence générée :

Pour les séquences à avantage positif ( $A_j \ge 0$ ) : On maximise la borne inférieure (ELBO). Cela encourage la génération de bonnes réponses.
Pour les séquences à avantage négatif ( $A_j < 0$ ) : On minimise une borne supérieure (EUBO - Evidence Upper Bound). Minimiser une borne supérieure garantit mathématiquement de réduire le vrai log-vraisemblance, pénalisant ainsi efficacement les mauvaises réponses.

L'objectif global $J_{SPG}(\theta)$ est défini comme suit :
$J_{SPG}(\theta) = \mathbb{E} \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \ge 0} \cdot A_j \mathcal{L}_{ELBO} + \mathbb{1}_{A_j < 0} \cdot A_j \mathcal{L}_{EUBO} \right) \right]$
Cette approche crée une borne inférieure valide de l'objectif original, permettant une optimisation robuste sans biais de gradient.

B. Estimation de la Borne Supérieure (EUBO)

Les auteurs dérivent une borne supérieure de la vraisemblance basée sur la divergence de Rényi. Pour un paramètre $\beta \ge 1$ , la borne est donnée par :
$\mathcal{L}_{EUBO} \approx \frac{1}{\beta} \sum_{i=1}^n \log \mathbb{E}_{t, z_t} \left[ w(t) \cdot \mathbb{1}(z_{t,i} = m) \cdot \pi_\theta^\beta(x_i | z_t) \right]$
Contrairement à l'ELBO où le logarithme est à l'intérieur de l'espérance, ici il est à l'extérieur, ce qui permet d'obtenir une borne supérieure.

C. Stratégie de Masquage par Blocs (Block-wise Masking)

Pour estimer ces bornes via Monte Carlo de manière stable et efficace, les auteurs proposent une stratégie de masquage différente du masquage aléatoire pur :

La séquence est divisée en blocs.
Un bloc aléatoire est sélectionné pour être masqué partiellement.
Les blocs précédents restent propres (clean), et les blocs suivants sont entièrement masqués.
Cette stratégie aligne mieux la distribution des données lors du déroulement de la politique (rollout) et de l'optimisation, imitant le processus de décodage semi-autoregressif utilisé en inférence.

D. Mélange de Bornes (Mixture)

Pour stabiliser l'entraînement et réduire la variance du gradient (surtout pour les traces à avantage négatif), les auteurs combinent l'ELBO et l'EUBO dans un objectif mixte :
$\mathcal{L}_{Mix} = \omega \cdot \mathcal{L}_{EUBO} + (1-\omega) \cdot \mathcal{L}_{ELBO}$
Théoriquement, ce mélange permet de réduire la variance du gradient par rapport à l'utilisation d'une seule borne, en pondérant les tokens incertains et confiants de manière adaptative.

3. Contributions Clés

Algorithme SPG : Un nouvel algorithme de gradient de politique qui élimine le biais de gradient inhérent aux méthodes basées uniquement sur l'ELBO pour les modèles de diffusion, en utilisant une optimisation "sandwichée" basée sur la récompense.
Dérivation théorique de l'EUBO : Une nouvelle borne supérieure de la vraisemblance pour les modèles de diffusion masqués, basée sur la divergence de Rényi, permettant la pénalisation correcte des échantillons négatifs.
Stratégie de Masquage par Blocs : Une technique d'estimation Monte Carlo qui améliore la stabilité de l'entraînement et l'efficacité de l'optimisation en alignant les distributions de données.
Réduction de variance : Une preuve théorique et empirique montrant que le mélange optimal des bornes supérieure et inférieure réduit strictement la variance du gradient par rapport à l'utilisation d'une seule borne.

4. Résultats Expérimentaux

Les auteurs ont évalué SPG sur quatre benchmarks de raisonnement mathématique et logique, ainsi que sur des tâches de codage, en comparant avec l'état de l'art (D1, WD1, UniGRPO, LLaDA).

Performances sur le Raisonnement (GSM8K, MATH500, Countdown, Sudoku) :

GSM8K : Amélioration de +3.6% par rapport à la méthode la plus performante précédente (atteignant 86.1% de précision).
MATH500 : Amélioration de +2.6%.
Countdown : Amélioration massive de +18.4%.
Sudoku : Amélioration spectaculaire de +27.0% (atteignant 94.0% de précision).
SPG surpasse systématiquement les méthodes de base (D1, UniGRPO) sur toutes les longueurs de génération (128, 256, 512).

Tâches de Codage (HumanEval, MBPP) :

Amélioration de +1.9% sur HumanEval et +4.7% sur MBPP.

Analyse et Robustesse :

Convergence : SPG montre une convergence plus rapide et plus stable des récompenses durant l'entraînement (Figure 3).
Généralisation : L'algorithme fonctionne bien avec différentes stratégies d'inférence (décodage semi-autoregressif, masquage aléatoire, etc.), démontrant une forte robustesse.
Efficacité : Le coût computationnel par mise à jour de gradient est comparable à celui des méthodes de base (D1, UniGRPO), malgré l'utilisation de l'estimation Monte Carlo.

5. Signification et Impact

Ce travail est significatif car il résout un problème fondamental théorique et pratique dans l'application du RL aux modèles de langage par diffusion.

Correction du biais : Il démontre que l'utilisation exclusive de l'ELBO pour le RL est sous-optimale et biaisée, surtout pour les récompenses négatives.
Nouvel état de l'art : SPG établit un nouveau standard de performance pour les dLLMs, prouvant qu'ils peuvent rivaliser, voire surpasser, les modèles autoregressifs sur des tâches complexes de raisonnement une fois correctement alignés.
Fondation pour le futur : La méthode proposée ouvre la voie à l'utilisation de dLLMs dans des scénarios nécessitant un apprentissage par renforcement complexe, en fournissant un cadre mathématiquement solide pour l'estimation des gradients de politique.

En résumé, SPG transforme les modèles de diffusion de langage en candidats viables et performants pour l'alignement par RL, en surmontant l'obstacle majeur de l'intractabilité du log-vraisemblance grâce à une approche de bornes sandwichées ingénieuse.