DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Le papier propose DyJR, un cadre de régularisation innovant qui préserve la diversité dans l'apprentissage par renforcement pour les grands modèles de langage en utilisant une distribution de référence dynamique et une divergence Jensen-Shannon, surpassant ainsi les méthodes existantes comme GRPO tout en maintenant une efficacité computationnelle élevée.

L'essentiel

🧠 Le Problème : L'élève qui oublie ses meilleures idées

Imaginez que vous apprenez à résoudre des énigmes mathématiques complexes avec un robot très intelligent (un modèle de langage).
Aujourd'hui, la méthode standard (appelée GRPO) fonctionne comme un étudiant qui ne regarde que la leçon du jour.

• Il essaie une solution.
• S'il a raison, il est content.
• S'il se trompe, il recommence.
• Le problème : Une fois la leçon finie, il jette tout ce qu'il a écrit sur ses brouillons précédents. Il oublie immédiatement ses anciennes tentatives, même celles qui étaient brillantes mais qui n'ont pas tout à fait réussi.

C'est comme si un chef cuisinier goûtait un plat, le trouvait bon, mais jetait la recette dans la poubelle avant de cuisiner le plat suivant. De plus, comme le robot ne se souvient que de la dernière chose qu'il a faite, il finit par devenir "paresseux" : il répète toujours la même solution facile qu'il a trouvée par hasard, au lieu d'explorer d'autres façons ingénieuses de résoudre le problème. C'est ce qu'on appelle la perte de diversité (ou "effondrement de mode").

💡 La Solution : DyJR, le "Journal de Bord Intelligent"

Les auteurs de cet article proposent une nouvelle méthode appelée DyJR. Imaginez que nous donnons au robot un journal de bord spécial pour apprendre de son passé, mais avec deux règles très intelligentes pour ne pas se perdre.

1. Le Journal qui s'auto-nettoie (Le Tampon Dynamique)

La plupart des méthodes d'apprentissage essaient de garder toutes les anciennes recettes dans un immense coffre-fort. C'est lourd, lent et encombrant.
DyJR, lui, utilise une règle simple : "Seules les idées récentes comptent."

• L'analogie : Imaginez un tamis qui ne garde que les poissons les plus frais. Si un poisson (une solution) a été pêché il y a trop longtemps, il commence à pourrir (il ne correspond plus à ce que le robot sait faire aujourd'hui) et est jeté.
• L'astuce : Au tout début de l'apprentissage, quand le robot explore beaucoup, le tamis est grand pour garder plein d'idées variées. Dès que le robot commence à bien comprendre, le tamis rétrécit pour ne garder que les solutions les plus proches de sa façon actuelle de penser. Cela économise énormément de mémoire et de temps.

2. Le "Compas de Diversité" (La Régularisation JS)

C'est la partie la plus géniale. Au lieu de dire au robot : "Regarde cette vieille solution parfaite, fais exactement pareil !", DyJR lui dit : "Regarde cette vieille solution, assure-toi juste de ne pas t'éloigner trop de son esprit."

• L'analogie :
  • L'ancienne méthode (Mauvaise) : C'est comme un professeur qui force l'élève à copier mot pour mot la solution d'un camarade. L'élève finit par ne plus réfléchir par lui-même et perd sa créativité.
  • La méthode DyJR (Bonne) : C'est comme un professeur qui dit : "Reste dans le même quartier que les bonnes solutions, mais tu as le droit de prendre des rues différentes."
  • Techniquement, ils utilisent une mesure mathématique (la Divergence de Jensen-Shannon) qui agit comme un compas. Elle empêche le robot de s'écarter trop de la diversité qu'il a déjà trouvée, sans l'obliger à copier bêtement.

🚀 Les Résultats : Pourquoi c'est génial ?

Grâce à cette approche, le robot devient :

1. Plus créatif : Il n'abandonne pas ses idées originales. Il continue d'explorer plusieurs chemins pour résoudre un problème, au lieu de se bloquer sur le premier qui marche.
2. Plus rapide et léger : Il n'a pas besoin de stocker des montagnes de données. Il garde juste l'essentiel, le "frais".
3. Plus performant : Sur des tests de mathématiques et de génération de code (SQL), DyJR bat largement les méthodes actuelles. Il trouve plus de bonnes solutions, même pour des problèmes très difficiles.

🎯 En résumé

DyJR, c'est comme donner à un apprenti génie un journal de bord intelligent qui :

• Garde les idées fraîches et jette les vieilles.
• Lui rappelle de rester curieux et diversifié, sans le forcer à copier aveuglément le passé.

Résultat : Le robot apprend plus vite, utilise moins d'énergie et devient beaucoup plus brillant pour résoudre des énigmes complexes !

Résumé technique

Titre : DyJR : Préservation de la diversité en Apprentissage par Renforcement avec Récompenses Vérifiables via une Replay Dynamique de Jensen-Shannon

1. Problématique

L'apprentissage par renforcement (RL) a considérablement amélioré les capacités de raisonnement des grands modèles de langage (LLM), notamment via des algorithmes on-policy comme GRPO (Group Relative Policy Optimization). Cependant, ces méthodes souffrent d'une inefficacité d'échantillonnage majeure : elles rejettent les données de déploiement (rollouts) après une seule mise à jour.

Les méthodes existantes d'Experience Replay (rejeu d'expérience) tentent de résoudre ce problème en réutilisant des échantillons historiques pour des mises à jour directes du gradient. Toutefois, l'article identifie deux limitations critiques de cette approche :

1. Effondrement de mode (Mode Collapse) : La maximisation directe de la vraisemblance des trajectoires historiques force le modèle à s'adapter excessivement à des chemins de solution spécifiques, réduisant ainsi son potentiel d'exploration et sa diversité.
2. Coût computationnel et mémoire : Les méthodes actuelles (comme RLEP) nécessitent le stockage massif de l'historique complet des trajectoires, entraînant une surcharge mémoire GPU importante. De plus, les données historiques ne sont pas uniformément utiles ; leur valeur diminue rapidement à mesure que le modèle évolue.

L'hypothèse centrale de l'article est que la valeur principale des données historiques ne réside pas dans le renforcement de la précision, mais dans le maintien de la diversité des stratégies de raisonnement, en particulier durant les phases initiales de l'entraînement où l'entropie du modèle est élevée.

2. Méthodologie : DyJR

Les auteurs proposent DyJR (Dynamic Jensen-Shannon Replay), un cadre de régularisation simple mais efficace qui redéfinit l'utilisation du replay d'expérience. L'approche repose sur deux innovations majeures :

A. Construction Dynamique des Données (Time-Sensitive Dynamic Buffer)
Au lieu de stocker massivement toutes les données, DyJR utilise une stratégie de tampon adaptatif :

• FIFO et Âge Maximal (Max Age) : Le tampon ne conserve que les échantillons générés récemment (dans une fenêtre temporelle $M$). Les données trop anciennes sont évacuées (FIFO - First-In-First-Out) pour éviter le décalage de distribution (distribution shift) par rapport à la politique actuelle.
• Sélection Adaptative : Une stratégie de sélection par niveau de confiance (confidence-stratified) est utilisée. Elle priorise les trajectoires parfaites ($r=1$) mais ajuste dynamiquement le taux de remplissage du tampon.
• Phase de Warm-up : Durant les premières étapes (ex: 20 itérations), le taux de remplissage est temporairement augmenté pour capturer un spectre large de motifs de raisonnement à haute entropie avant que le modèle ne converge.

B. Régularisation par Divergence de Jensen-Shannon (JS)
Au lieu d'effectuer des mises à jour de gradient directes sur les données de replay (ce qui cause le surapprentissage), DyJR introduit une contrainte de distribution :

• Principe : Le mélange des politiques historiques dans le tampon sert d'ancrage distributionnel dynamique.
• Fonction de Perte : Au lieu d'optimiser directement la vraisemblance, le modèle minimise la Divergence de Jensen-Shannon (JS) entre la politique actuelle et le mélange des politiques historiques.
• Avantage : Contrairement à la divergence KL directe (Forward KL) qui tend à couvrir tous les modes (mode-covering) et peut lisser excessivement la distribution, la JS divergence est symétrique et bornée. Elle agit comme une régularisation flexible qui empêche le modèle de s'éloigner des chemins de succès diversifiés sans bloquer l'optimisation vers de nouvelles solutions.

L'objectif global est la somme de la perte GRPO (sur les données en ligne) et de la perte de régularisation JS (sur les données de replay) :
$$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$$

3. Contributions Clés

1. Changement de Paradigme : Passage d'une optimisation de la précision via le replay à une régularisation de la diversité.
2. Stratégie de Tampon Dynamique : Introduction d'un mécanisme non uniforme qui adapte la taille du stockage selon la phase d'entraînement (expansion initiale pour la diversité, contraction ultérieure pour l'efficacité), réduisant drastiquement les besoins en mémoire.
3. Analyse Fine : Démonstration que la divergence JS est supérieure à la divergence KL dans ce contexte de données non stationnaires, et analyse détaillée de l'évolution des probabilités des tokens (Rank-k) pour prouver la préservation de la diversité.

4. Résultats Expérimentaux

Les expériences ont été menées sur des tâches de raisonnement mathématique (benchmarks AIME, AMC, HMMT, etc.) et de génération SQL (BIRD, Spider), utilisant des modèles comme Qwen3-4B et Llama-3.1-8B.

• Performance Supérieure : DyJR surpasse significativement GRPO et d'autres méthodes de replay (RLEP, Ex-GRPO, DPH-RL).
  • Sur les benchmarks mathématiques, DyJR atteint une précision moyenne de 34,1 %, contre 29,8 % pour GRPO (gain de +4,3 %).
  • Sur les tâches SQL, DyJR améliore la précision Pass@1 de +3,3 % (BIRD) et +5,0 % (Spider) par rapport à GRPO.
• Efficacité Mémoire : Contrairement à RLEP qui nécessite de stocker des milliers de paires de questions/réponses, DyJR maintient un tampon très restreint (environ 1 600 séquences dans les expériences), réduisant la surcharge mémoire GPU à un niveau négligeable (< 1 Go).
• Analyse de la Diversité (Rank-k) :
  • GRPO : Montre un effondrement rapide de l'entropie et une probabilité de token Rank-1 dépassant 90 %, indiquant une perte d'exploration.
  • DyJR : Maintient une distribution de probabilité plus équilibrée entre les tokens Rank-1, Rank-2 et Rank-3, prouvant que le modèle explore plusieurs chemins de raisonnement tout en restant performant.
• Évolutivité (Pass@k) : DyJR démontre une meilleure scalabilité lorsque le budget d'échantillonnage ($k$) augmente, suggérant une capacité à générer des solutions variées et correctes.

5. Signification et Impact

Ce travail est significatif car il résout le compromis classique entre l'efficacité de l'échantillonnage et la diversité du modèle en RL pour les LLM.

• Efficacité : Il permet d'utiliser des données historiques sans les coûts prohibitifs de stockage ni les risques de surapprentissage.
• Robustesse : En évitant l'effondrement de mode, DyJR rend les modèles plus robustes face à la complexité des tâches de raisonnement, évitant la convergence prématurée vers des solutions sous-optimales.
• Généralité : La méthode s'applique à différents types de tâches (mathématiques, SQL) et architectures, suggérant qu'elle pourrait devenir un standard pour l'entraînement RLVR (Reinforcement Learning with Verifiable Rewards) à grande échelle.

En conclusion, DyJR démontre que la valeur des données passées en RL réside dans leur capacité à maintenir un espace d'exploration riche, et non simplement à renforcer les succès passés, offrant ainsi une voie plus stable et efficace pour l'amélioration des capacités de raisonnement des LLM.