Résumé Technique : Dropout-GRPO pour le raisonnement latent continu

1. Énoncé du Problème

L'article traite d'une incompatibilité structurelle fondamentale entre l'optimisation de la politique relative au groupe (Group Relative Policy Optimization - GRPO) et les modèles de raisonnement latent continus (spécifiquement l'architecture COCONUT).

• L'exigence de GRPO : GRPO repose sur la diversité des $K$ déploiements (rollouts) au sein d'un groupe pour calculer un avantage moyen de groupe non nul ($A^{(k)} = r^{(k)} - \mu_r$). Si tous les déploiements produisent des trajectoires identiques, l'avantage s'effondre vers zéro, stoppant l'optimisation.
• Le défi du raisonnement latent : Contrairement aux modèles classiques de Chaîne de Pensée (Chain-of-Thought - CoT) qui tirent leur diversité de l'échantillonnage stochastique des jetons (tokens), les modèles de raisonnement latent continu (comme COCONUT) opèrent via une récurrence d'états cachés déterministe. Étant donné un prompt et des paramètres fixes, plusieurs déploiements produisent des trajectoires latentes identiques et des réponses identiques.
• La conséquence : L'application de la méthode GRPO standard à ces modèles entraîne une variance nulle ($\sigma_r = 0$), provoquant la disparition du gradient de la politique et empêchant toute génération de signal d'apprentissage.

2. Méthodologie : Dropout-GRPO

Les auteurs proposent Dropout-GRPO, une méthode qui injecte la stochasticité nécessaire dans le processus de raisonnement latent par un dropout structuré, tout en maintenant les garanties théoriques de l'inférence variationnelle.

Mécanisme Central : Dropout Variationnel à Masque Partagé

Au lieu d'un échantillonnage standard au niveau du jeton, la méthode introduit la stochasticité au niveau des paramètres :

1. Génération du Masque : Pour chaque déploiement $k$, un masque de Bernoulli unique $\xi^{(k)}$ est tiré.
2. Cohérence : Ce masque est maintenu constant à travers les $T$ étapes de récurrence latente de ce déploiement spécifique.
3. Perturbation : Le masque induit une perturbation structurée $\tilde{\theta}(\xi)$ aux poids du réseau. La récurrence latente devient $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Rejeu du Masque (Mask Replay) : Crucialement, la graine (seed) du générateur de nombres pseudo-aléatoires (RNG) pour le masque est stockée pendant la phase de déploiement. Lors de la mise à jour de la politique, le même masque est régénéré et appliqué. Cela garantit que la trajectoire évaluée pour la récompense est bit-identique à la trajectoire utilisée pour le calcul du gradient lorsque $\theta = \theta_{old}$.

Interprétation Théorique

• Moyenne de Modèle Bayésien : Suivant Gal et Ghahramani, le masque partagé traite chaque déploiement comme un échantillon postérieur unique d'une distribution structurée sur les paramètres. La politique marginale $\bar{\pi}_\theta$ est interprétée comme une moyenne de modèle bayésien.
• Impartialité du Gradient : Les auteurs prouvent que le gradient de substitution utilisant l'avantage moyen uniquement ($A^{(k)} = r^{(k)} - \mu_r$) est un estimateur impartial du gradient de l'espérance de la récompense de la politique marginale, mis à l'échelle par $(1 - 1/K)$.
• Réduction de la Variance : En réutilisant le même masque (Nombres Aléatoires Communs) lors de la mise à jour que celui utilisé lors du déploiement, la méthode réduit la variance par rapport à un ressemblage d'un nouveau masque au moment de la mise à jour.

Raffinements d'Implémentation

Quatre raffinements spécifiques sont introduits pour stabiliser l'entraînement :

1. KL $k_3$ Huberisé : Remplace le plafonnement standard dans l'estimateur de la divergence KL par une extension de type Huber afin d'éviter que les gradients ne s'annulent pour les jetons ayant de grands log-ratios, évitant ainsi les jetons "échappés" qui dérivent sans régularisation.
2. Recul de la Zone de Confiance (Trust-Region Annealing) : Le coefficient KL $\beta$ est réduit proportionnellement au taux d'apprentissage $\eta_t$ pour éviter que la politique ne soit bloquée à un équilibre où la traction KL annule le gradient de la politique.
3. Filtre de Précision au Niveau du Groupe : Les groupes dont la précision empirique $\mu_r$ se situe en dehors d'une fenêtre configurée (ex: $[0,20, 0,95]$) sont abandonnés. Cela agit comme un curriculum émergent, concentrant le calcul sur les prompts présentant une variance non triviale.
4. Saut Coordonné DDP : Assure la cohérence de l'entraînement distribué lors de l'abandon de groupes.

3. Principales Contributions

1. Nouveau Paradigme pour le RL Latent : L'article identifie la nature déterministe de la récurrence latente comme la cause de l'échec de GRPO et propose le dropout variationnel à masque partagé comme une solution minimale et théoriquement propre. C'est la première méthode appliquant avec succès le RL relatif au groupe aux modèles de raisonnement latent continu.
2. Justification Théorique :
   • Prouve la définition correcte de la jacobienne latente $\partial h_T / \partial \theta$ sous les régularités standards des Transformers.
   • Prouve l'impartialité du gradient de substitution pour la politique marginale.
   • Démontre que le rejeu du masque fournit une réduction de la variance par les nombres aléatoires communs.
3. Validation Empirique : Applique avec succès GRPO au modèle COCONUT sur le benchmark GSM8K, un régime où le GRPO déterministe échoue complètement.
4. Implémentation de Référence : Publie le code qui réalise un rejeu de masque bit-identique sur l'ensemble de la récurrence latente.

4. Résultats Expérimentaux

Les expériences ont été menées en utilisant Qwen2.5-1.5B comme modèle de base avec une profondeur latente de $T=6$.

• Amélioration des Performances : Dropout-GRPO a amélioré la précision pass@1 de GSM8K de la ligne de base COCONUT de 27,29 % à 29,01 %.
• Récupération des Capacités : La méthode a inversé une régression de performance de 1,33 point sur l'ensemble de données SVAMP causée par l'entraînement par curriculum, ramenant la performance à 44,00 % (correspondant à la variance du modèle de base).
• Études d'Ablation :
  • GRPO Déterministe : A échoué complètement, produisant une trajectoire d'entraînement plate identique à l'initialisation SFT en raison de l'avantage nul.
  • REINFORCE + Dropout (EMA) : A échoué en raison d'une variance élevée ; le gradient d'un seul déploiement était trop bruyant pour stabiliser l'initialisation SFT.
  • Dropout-GRPO : La seule configuration ayant produit un signal d'apprentissage viable, confirmant que les bases relatives au groupe sont essentielles pour la réduction de la variance dans ce contexte.

5. Signification et Revendications

L'article positionne ce travail comme une approche pratique et théoriquement fondée pour le post-entraînement des modèles de langage (LLM) à raisonnement latent.

• Faisabilité : Il démontre que le post-entraînement par RL est possible pour les modèles raisonnant directement dans les états cachés, une capacité auparavant bloquée par l'absence de diversité des déploiements.
• Mécanisme : Il reformule le dropout non pas comme un régularisateur (souvent désactivé dans les LLM modernes) mais comme une source de stochasticité exogène structurée requise pour l'optimisation variationnelle de la politique.
• Limites : Les auteurs sont modestes quant à l'ampleur de l'amélioration (un gain de 1,72 point sur GSM8K), l'attribuant à la petite échelle du squelette (1,5B) et à la rareté des ensembles de données de raisonnement concis appropriés. Ils notent que la politique déployée (inférence sans dropout) diffère de la politique marginale optimisée, introduisant un biais dépendant des constantes de Lipschitz.

En résumé, Dropout-GRPO comble le fossé entre l'apprentissage par renforcement relatif au groupe et le raisonnement latent continu, permettant l'optimisation de modèles qui ne pouvaient auparavant pas utiliser GRPO en raison de leurs dynamiques internes déterministes.