Résumé technique : SAGE – Façonnage d'ancres pour une exploration guidée dans l'RLVR des LLMs

1. Énoncé du problème

Les récentes avancées en apprentissage par renforcement avec récompenses vérifiables (RLVR) ont démontré leur capacité à améliorer la précision pass@1 des grands modèles de langage (LLM) sur des tâches de raisonnement. Cependant, une limitation critique est apparue : si le RLVR augmente souvent les taux de réussite sur un seul échantillon, il échoue fréquemment à produire des gains comparables en pass@k (la probabilité de trouver une réponse correcte parmi $k$ échantillons).

Des analyses antérieures suggèrent que cette divergence provient du fait que le RLVR standard induit un effondrement de mode (ou « rétrécissement du support »). Au lieu de découvrir de nouvelles stratégies de raisonnement, la politique s'effondre sur un sous-ensemble étroit de modes de raisonnement préexistants et à haute densité présents dans le modèle de référence de base. Ce comportement est attribué aux contraintes structurelles du terme de régularisation standard Kullback-Leibler inverse (reverse-KL). Bien que le reverse-KL assure la stabilité nécessaire de l'entraînement en ancrant la politique à une distribution de référence ($\pi_{ref}$), il supprime intrinsèquement les trajectoires à faible densité, empêchant efficacement le modèle d'explorer des voies de raisonnement alternatives et valides qui se situent en dehors du support immédiat du modèle de référence.

Les tentatives pour résoudre ce problème en supprimant le terme KL ou en le remplaçant par un forward-KL se sont révélées sous-optimales. La suppression du terme KL conduit souvent à des piratages de récompense et à une instabilité, tandis que le forward-KL, bien qu'il élargisse le support, alloue fréquemment une masse de probabilité à des régions sans rapport avec la récompense, échouant ainsi à améliorer le pass@1. Le défi central consiste à conserver les avantages stabilisateurs du reverse-KL tout en permettant une expansion contrôlée du support empirique.

2. Méthodologie : Cadre SAGE

Les auteurs proposent SAGE (Shaping Anchors for Guided Exploration), un cadre qui réinterprète la régularisation reverse-KL non pas comme une contrainte statique à affaiblir, mais comme une ancre dynamique dont la distribution cible peut être remodelée pour guider l'exploration.

Mécanisme central

Le RLVR reverse-KL standard optimise :
$$J(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref})]$$
Cela conduit à une solution stationnaire où la politique est une inclinaison exponentielle de la référence : $\pi_\theta \propto \pi_{ref} \exp(r/\beta)$. Cette mise à jour multiplicative amplifie les régions à haute densité de $\pi_{ref}$ tout en supprimant celles à faible densité.

SAGE modifie cela en introduisant une fonction de guidage $q(x, y) > 0$ qui remodelle la distribution d'ancre elle-même. Le nouvel objectif devient :
$$J'(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta \tilde{D}_{KL}(\pi_\theta \parallel q \cdot \pi_{ref})]$$
Ici, $\tilde{D}_{KL}$ est une pseudo-divergence KL définie pour des fonctions non normalisées. La solution stationnaire devient :
$$\pi_\theta(y | x) \propto q(x, y) \pi_{ref}(y | x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$
En concevant $q(x, y)$ pour qu'elle soit grande pour des modes de raisonnement sous-explorés mais potentiellement valides, SAGE amplifie leur masse de probabilité sans rompre la stabilité fournie par l'ancre $\pi_{ref}$.

Instantiation via des signaux intrinsèques

Pour rendre $q(x, y)$ pratique sans nécessiter d'informations de récompense explicites lors de la construction du guide, les auteurs proposent d'utiliser des signaux intrinsèques dérivés du processus de raisonnement lui-même :

1. Exploration aléatoire : Une ligne de base utilisant des perturbations stochastiques.
2. Exploration au niveau des tokens : Utilise la surprise ($-\log p$) pour biaiser vers des choix lexicaux rares. Les auteurs notent que cela capture des déviations de surface mais peut ne pas refléter des différences structurelles de raisonnement.
3. Exploration au niveau des branches (SAGE-Branch) : La variante principale proposée. Elle utilise l'entropie ($H_t$) comme signal de bifurcation du raisonnement. La fonction de guidage est définie comme suit :
   $$q_{Branch}(y_t | x, y_{<t}) = 1 + \gamma (H_t - \tau)_+$$
   où $H_t$ est l'entropie au niveau du token, $\tau$ est un seuil d'activation, et $\gamma$ contrôle l'amplitude. Cela amplifie sélectivement l'exploration aux points de décision à forte incertitude (nœuds de bifurcation) tout en supprimant le bruit dans les étapes symboliques déterministes.

Garanties théoriques

L'article fournit une analyse théorique (Théorème 4.3) démontrant que, contrairement au RLVR standard qui est borné par le support empirique du modèle de référence, SAGE peut théoriquement étendre le support empirique. Plus précisément, si la fonction de guidage $q$ est suffisamment grande pour une solution valide par la récompense $y^*$ qui se trouve actuellement en dehors du support empirique de la politique de référence, SAGE peut lui attribuer une probabilité non négligeable, récupérant ainsi efficacement des modes de raisonnement « cachés ».

3. Contributions clés

1. Analyse des limites du RLVR : Les auteurs identifient que l'ancre reverse-KL, bien que stabilisatrice, est la cause structurelle principale de l'effondrement de mode dans le RLVR. Ils soutiennent que supprimer ou remplacer simplement le terme KL est insuffisant, car cela perturbe le compromis efficacité-couverture.
2. Cadre SAGE : Ils proposent une méthode fondée sur des principes pour permettre une expansion contrôlée du support empirique en remodelant la distribution d'ancre KL via une fonction de guidage $q(x, y)$, plutôt que de supprimer l'ancre.
3. Intégration de signaux intrinsèques : Ils démontrent comment des signaux intrinsèques légers (spécifiquement l'entropie et la surprise) peuvent être utilisés pour construire des fonctions de guidage efficaces ne nécessitant aucun passage avant supplémentaire ni données externes.
4. Validation empirique : Des expériences extensives montrent que les variantes de SAGE, en particulier l'exploration au niveau des branches, améliorent de manière cohérente à la fois le pass@1 et le pass@k sur divers benchmarks de raisonnement mathématique.

4. Résultats expérimentaux

Les auteurs ont évalué SAGE sur les ensembles de données AIME 2024-2025, AMC23 et MATH-500 en utilisant le modèle Qwen2.5-Math-7B-Base et l'algorithme GRPO.

• Gains de performance : La variante GRPO + Branch a obtenu les performances moyennes les plus fortes, améliorant le pass@1 de 0,266 (GRPO standard) à 0,284 et le pass@256 de 0,714 à 0,729. Elle a surpassé les lignes de base sur les trois ensembles de données.
• Comparaison avec des alternatives :
  • GRPO sans KL : A montré des récompenses d'entraînement plus élevées mais des performances lors du test nettement inférieures (pass@1 et pass@256), confirmant le risque de piratage de récompense sans l'ancre KL.
  • Forward-KL : A légèrement amélioré le pass@256 mais n'a pas réussi à égaler les gains de pass@1 de SAGE, indiquant qu'il alloue une masse à des régions hors cible.
  • Autres algorithmes RL : SAGE a été intégré avec succès avec BNPO et DAPO, résolvant de manière cohérente le compromis où ces méthodes amélioraient généralement le pass@1 au détriment du pass@k.
• Analyse des motifs de raisonnement : Une analyse qualitative a révélé que si le RLVR standard renforçait des motifs dominants comme la « configuration des contraintes », SAGE a considérablement augmenté la fréquence des preuves par l'absurde (passant de 0,8 % dans GRPO à 3,3 % dans GRPO+Branch), un mode de raisonnement rare mais valide.
• Robustesse : Dans les évaluations hors distribution (benchmark Chevaliers et Valets), les modèles entraînés avec SAGE ont maintenu leurs performances sous un changement de distribution, tandis que les modèles GRPO standards ont montré une dégradation significative, suggérant que SAGE atténue l'effondrement de mode.
• Généralisation : La méthode s'est également révélée efficace sur des modèles à raisonnement long (DeepSeek-R1-Distill-Qwen-7B) et sur des modèles plus petits (Qwen2.5-1.5B/3B), démontrant une robustesse à travers différentes échelles de modèles et paradigmes de raisonnement.

5. Importance et affirmations

L'article affirme que SAGE offre une approche pratique et fondée sur des principes pour atténuer l'affinement des modes dans le RLVR reverse-KL. En traitant la régularisation KL comme un instrument contrôlable plutôt que comme une contrainte rigide, SAGE permet aux LLM de découvrir et d'utiliser des trajectoires de raisonnement diversifiées qui sont valides mais sous-représentées dans le modèle de base.

Les auteurs soulignent que SAGE est modulaire et orthogonal aux techniques existantes (telles que le façonnage d'avantage ou l'échantillonnage dynamique), car il modifie la distribution d'ancre elle-même plutôt que la structure d'optimisation. Ils positionnent SAGE comme un cadre général pour une expansion stable et compatible avec KL du support empirique.

Limites et portée :
Les auteurs reconnaissent modestement que SAGE opère dans le support de la politique de référence fixe. Il redistribue la masse de probabilité parmi les trajectoires qui ont déjà une probabilité non nulle sous $\pi_{ref}$ ; il ne peut pas générer de nouveaux motifs de raisonnement entièrement nouveaux qui seraient impossibles à produire pour le modèle de base. Ils notent que si la distillation à partir de modèles enseignants plus puissants peut élargir la frontière du raisonnement, SAGE la complète en permettant une exploration stable à l'intérieur de cette frontière. L'instanciation actuelle repose sur des signaux intrinsèques légers, et les auteurs suggèrent que l'apprentissage de fonctions de guidage plus expressives à partir de modèles plus puissants est une direction prometteuse pour les travaux futurs.