Résumé Technique : Apprendre à raisonner efficacement avec l'apprentissage par renforcement à remise

Énoncé du problème

Les Modèles de Raisonnement à Grande Échelle (LRM) résolvent des tâches complexes (par exemple, mathématiques, code) en émettant des tokens de raisonnement intermédiaires avant une réponse finale. Bien que l'apprentissage par renforcement (RL) en post-entraînement améliore la précision, il conduit souvent à des chaînes de pensée excessivement longues, augmentant les coûts de calcul, l'utilisation de la mémoire (via la croissance des caches KV) et la latence d'inférence.

L'hypothèse prédominante dans le domaine est que des réponses plus longues améliorent intrinsèquement la précision, suggérant un compromis entre la longueur de la réponse et la justesse. Cet article remet en cause cette hypothèse. Il postule que, jusqu'à un régime déterminé par la classe de modèle et l'instance du problème, il n'existe pas de compromis intrinsèque : on peut réduire la longueur de la réponse à un seuil dépendant de l'instance spécifique sans sacrifier la précision. L'objectif est d'entraîner des LRM qui atteignent des réponses correctes de manière efficace, analogue à la recherche de la trajectoire réussie la plus courte dans un problème de chemin stochastique le plus court (SSP).

Méthodologie

Cadre théorique : Optimalité de Blackwell

Les auteurs modélisent le raisonnement basé sur un vérificateur comme un processus de décision de Markov (MDP) à horizon fini avec une récompense terminale binaire (1 pour correct, 0 pour incorrect). Pour formaliser l'objectif de maximiser la précision tout en minimisant la longueur du chemin, ils utilisent l'optimalité de Blackwell.

• Définition : Une politique est optimale au sens de Blackwell si elle est optimale pour tous les facteurs d'actualisation $\gamma$ suffisamment proches de 1.
• Mécanisme : Dans les MDP orientés vers un objectif avec des récompenses terminales binaires, maximiser le retour actualisé $J_\gamma(\pi)$ lorsque $\gamma \to 1$ maximise d'abord la probabilité de succès (précision). Parmi les politiques qui atteignent la probabilité de succès maximale, le mécanisme d'actualisation brise les égalités en privilégiant celles ayant des longueurs de trajectoire attendues plus courtes.
• Insight clé : Pour une classe de politique finie $\Pi$, il existe un facteur d'actualisation critique $\gamma_{bw} < 1$ tel que pour tout $\gamma \in (\gamma_{bw}, 1)$, l'ensemble des politiques $\gamma$-optimales coïncide avec l'ensemble des politiques optimales au sens de Blackwell. Ces politiques maximisent simultanément la précision et minimisent la longueur de trajectoire attendue parmi les maximisateurs de précision.

Recette d'entraînement pratique

Guidés par les résultats théoriques, les auteurs proposent une méthodologie d'entraînement spécifique utilisant l'optimisation de politique relative par groupes (GRPO) :

1. Récompense environnementale actualisée uniquement : Le facteur d'actualisation $\gamma \in (0, 1)$ est appliqué exclusivement à la récompense environnementale extrinsèque (justesse). Les récompenses intrinsèques (mise en forme, structuration) restent non actualisées pour garantir que le modèle produise des sorties bien structurées.
2. Masquage des tokens de raisonnement : L'actualisation est appliquée uniquement aux tokens identifiés comme faisant partie de l'« intervalle de raisonnement » (par exemple, les tokens entre les balises <reasoning> et </reasoning>). Les tokens requis pour l'adhésion au prompt, la mise en forme et la réponse finale ne sont pas actualisés.
3. Régularisation KL : Une pénalité KL est appliquée par rapport à une politique de référence mobile (mise à jour périodiquement) pour empêcher l'effondrement de la politique (où le modèle arrête de raisonner trop tôt) et stabiliser l'apprentissage.
4. Budgets de tokens comparables : Pour assurer une comparaison équitable, le nombre de déroulements est ajusté afin que le nombre total de tokens traités pendant l'entraînement soit comparable entre les méthodes actualisées et non actualisées.

La fonction objectif optimisée est :
$$\mathbb{E}_{S_1 \sim \mu, \tau \sim \pi(S_1)} [\gamma^{K(\tau)} r_e(\tau) + r_f(\tau)] - \beta \text{KL}(\pi || \pi')$$
Où $K(\tau)$ est le nombre de tokens de raisonnement, $r_e$ est la récompense environnementale, et $r_f$ est la récompense de mise en forme.

Contributions clés

1. Garantie théorique de l'absence de compromis : L'article prouve que dans toute classe de politique fixe (éventuellement restreinte), les politiques optimales au sens de Blackwell maximisent la probabilité de succès non actualisée et, parmi ces maximisateurs, minimisent la longueur de trajectoire attendue. Cela établit que la précision et la longueur du chemin ne sont pas intrinsèquement en compromis jusqu'à un seuil dépendant de l'instance.
2. Existence du facteur de Blackwell : Pour les classes de politiques finies, les auteurs prouvent l'existence d'un facteur d'actualisation de Blackwell $\gamma_{bw} < 1$. Ils fournissent des bornes sur la proximité de $\gamma$ par rapport à 1 nécessaire pour maintenir la précision tout en raccourcissant les réponses, clarifiant ainsi la sélection des hyperparamètres pour les classes de déploiement restreintes.
3. Validation empirique : En utilisant GRPO avec l'objectif actualisé, les auteurs démontrent que les modèles peuvent réduire considérablement la longueur moyenne des réponses sur des benchmarks mathématiques (GSM8K, MATH, AMC, AIME, etc.) tout en égalant la précision Pass@1 des bases de référence non actualisées.

Résultats expérimentaux

Les auteurs ont évalué quatre modèles ajustés par instruction (Qwen2.5 7B/14B, Llama 3 8B, Phi-4) sur divers ensembles de données mathématiques.

• GSM8K & MATH : Les modèles actualisés ont égalé ou légèrement dépassé la précision Pass@1 des bases de référence non actualisées tout en réduisant significativement le nombre de tokens.
  • Qwen2.5 7B (GSM8K) : La précision est restée stable (91,06 % contre 91,07 %), tandis que la longueur moyenne de la réponse est passée de 217,60 à 170,08 tokens (réduction d'environ 22 %).
  • Llama 3 8B (GSM8K) : La précision a légèrement augmenté (80,87 % contre 81,07 %), avec une longueur passant de 125,43 à 108,67 tokens (réduction d'environ 13 %).
• Benchmarks plus difficiles (AMC, AIME, MINERVA, OLYMPIAD) : La tendance s'est maintenue pour les modèles plus grands (Qwen2.5 14B, Phi-4) entraînés sur DeepScaleR. Les modèles actualisés ont atteint une précision comparable ou supérieure avec des traces nettement plus courtes (par exemple, Phi-4 sur AMC 2023 : 51,00 % $\to$ 61,00 % de précision, 1134 $\to$ 716 tokens).
• Sensibilité au facteur d'actualisation : Les expériences variant $\gamma$ ont confirmé le compromis théorique : à mesure que $\gamma$ s'éloigne de 1, les longueurs de réponse se raccourcissent de manière fiable, mais la précision finit par baisser si l'actualisation est trop agressive.

Importance et affirmations

L'article prétend fournir un mécanisme principiel et conscient de l'instance pour un raisonnement efficace qui ne repose pas sur des données curatées ou l'ingénierie de prompts. En tirant parti de l'optimalité de Blackwell, l'approche reformule le problème du « raisonnement efficace » comme un objectif RL standard avec une stratégie d'actualisation spécifique.

Les auteurs affirment que leurs résultats :

• Réfutent le compromis intrinsèque longueur-précision : Ils montrent qu'un raisonnement plus court est possible sans perte de précision jusqu'à un seuil spécifique déterminé par la capacité du modèle.
• Unifient les pénalités de longueur : Le travail suggère que les méthodes utilisant de petites récompenses négatives par étape sont équivalentes à l'actualisation dans ce contexte, impliquant que plusieurs méthodes existantes de pénalité de longueur peuvent retrouver le même ordre précision-longueur lorsqu'elles sont correctement implémentées.
• Offrent une recette pratique : La méthode proposée ne nécessite aucun changement algorithmique par rapport à l'optimisation standard de politique (comme GRPO), seulement une modification du calcul de la récompense (actualisation des tokens de raisonnement) et l'utilisation d'une politique de référence mobile.

L'article conclut en notant que, bien que des traces plus courtes améliorent l'efficacité, le potentiel du raisonnement compressé pour améliorer la généralisation reste une question ouverte pour les travaux futurs.