Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

Résumé Technique : Échantillonnage Adaptatif Guidé par l'Apprentissage par Renforcement pour le Scaling au Temps d'Inférence

Énoncé du Problème
Le passage à l'échelle au temps d'inférence (test-time scaling) est devenu une technique critique pour améliorer les capacités de raisonnement des grands modèles de langage (LLM) sans entraînement supplémentaire. Des méthodes telles que la cohérence de soi (Self-Consistency - SC), l'arbre de pensée (Tree-of-Thoughts) et l'échantillonnage Best-of-N améliorent la qualité des réponses en allouant plus de calcul lors de l'inférence. Cependant, ces gains entraînent des coûts substantiels tant en calcul total qu'en latence. Les méthodes d'échantillonnage adaptatif existantes, telles que l'Adaptive Self-Consistency (ASC) et l'Early-Stopping Self-Consistency (ESC), tentent d'atténuer ces coûts en décidant dynamiquement quand arrêter l'échantillonnage. Pourtant, ces approches souffrent de limitations significatives : elles reposent souvent sur des heuristiques conçues par l'humain ou des hypothèses distributionnelles, nécessitent des signaux auxiliaires (par exemple, la confiance interne du modèle, les états cachés ou la difficulté de la question) qui peuvent être indisponibles, et peuvent être invasives pour les pipelines d'inférence standards. Il existe un besoin pour une méthode légère et fondée sur des principes, capable d'équilibrer la justesse des réponses, la latence et le coût de calcul sans nécessiter de caractéristiques auxiliaires ou d'intervention sur le modèle.

Méthodologie
Les auteurs proposent l'Échantillonnage Guidé par l'Apprentissage par Renforcement (RL-Guided Sampling), un cadre qui formule l'échantillonnage adaptatif comme un processus de décision de Markov (MDP) à horizon fini.

• Formulation du MDP :

  • État ($s_t$) : L'état est construit uniquement à partir des statistiques du pool actuel de réponses échantillonnées ($D_t$). Il comprend les comptes des $K$ classes de réponses les plus fréquentes, le nombre total d'échantillons générés et l'entropie de ces comptes. Crucialement, l'état ne dépend pas d'informations spécifiques à la requête ou de signaux internes du modèle.
  • Action ($a_t$) : À chaque tour, le contrôleur sélectionne une action parmi un ensemble discret : soit arrêter l'échantillonnage (action 0), soit générer un nombre spécifique de réponses candidates supplémentaires en parallèle (actions $k_1, k_2, \dots$).
  • Fonction de Récompense : La récompense équilibre la qualité finale de la réponse par rapport aux coûts d'inférence.
    • Pénalité par étape : Une pénalité est infligée pour chaque tour d'échantillonnage supplémentaire ($\lambda_{lat}$) et pour chaque nouvel échantillon généré ($\lambda_{comp}$).
    • Récompense Terminale : Lors de la terminaison (soit par l'arrêt du contrôleur, soit en atteignant le budget maximum), le modèle reçoit une récompense basée sur la réponse issue du vote majoritaire. La cible de justesse ($y^\star$) est définie comme le vote majoritaire qui serait obtenu si l'échantillonnage se poursuivait jusqu'au budget maximum $N$, découplant ainsi la récompense des étiquettes de vérité terrain (ground-truth).
  • Optimisation : Le contrôleur est entraîné via l'Optimisation de Politique Proximale (PPO) pour maximiser la récompense cumulée attendue. Les auteurs notent que cet objectif peut être interprété comme la relaxation lagrangienne d'un problème d'optimisation sous contraintes où l'objectif est de maximiser la précision sous des contraintes de latence et de calcul.

• Architecture et Déploiement :
  Le contrôleur est un perceptron multicouche (MLP) léger à quatre couches. Il nécessite uniquement les réponses finales générées pour fonctionner, ce qui le rend non invasif et déployable sur CPU. Il ne nécessite pas de réglage fin (fine-tuning) du LLM sous-jacent.

Contributions Clés

1. Formulation MDP pour l'Échantillonnage Adaptatif : L'article introduit une nouvelle formulation de l'échantillonnage adaptatif sous forme de MDP, passant de règles heuristiques vers une politique apprise qui navigue explicitement dans le compromis performance-coût.
2. Conception Légère et Non Invasive : La méthode repose purement sur les statistiques du pool de réponses échantillonnées, ne nécessitant aucun signal auxiliaire (comme les scores de confiance) et aucune modification du processus de raisonnement du LLM. Le contrôleur est suffisamment petit pour être entraîné et déployé sur des CPU.
3. Interprétation Théorique : Les auteurs fournissent un lien théorique entre leur objectif de RL et l'optimisation sous contraintes, montrant que la fonction de récompense pondérée agit comme une relaxation lagrangienne pour maximiser la précision sous des contraintes de ressources strictes.
4. Généralisation : Le cadre démontre une forte transférabilité, où les politiques entraînées sur un jeu de données ou un modèle (ex: Qwen-3-0.6B) peuvent être appliquées à différents benchmarks et même à des modèles différents et plus performants (ex: GPT-4.1-nano) avec une dégradation minimale des performances.

Résultats Expérimentaux
La méthode a été évaluée sur trois benchmarks de raisonnement mathématique (AIME24, AIME25, HMMT 2025) en utilisant diverses familles d'échantillonneurs LLM (famille Qwen-3 et GPT-4.1-nano).

• Performance vs Baselines : Comparée à des bases solides comme ASC et ESC, l'échantillonnage guidé par RL offre de meilleurs compromis :
  • vs ASC : Réduit les tours d'échantillonnage d'environ 3 à 4$\times$ et les échantillons totaux d'environ 30 % tout en maintenant une précision comparable. Cela traite le goulot d'étranglement de la latence élevé de l'échantillonnage séquentiel de l'ASC.
  • vs ESC : Réduit les besoins totaux en échantillons d'environ 33 % et les tours d'échantillonnage d'environ 10 %, en atteignant une précision plus élevée en évitant un arrêt précoce trop agressif.
  • vs SC : Bien que toutes les méthodes adaptatives réduisent le nombre total d'échantillons par rapport à la SC à budget fixe, l'échantillonnage guidé par RL atteint cela avec nettement moins de tours que l'ASC et une meilleure efficacité que l'ESC.
• Comportement de Mise à l'Échelle (Scaling) : Les courbes de mise à l'échelle indiquent que l'échantillonnage guidé par RL surpasse systématiquement les bases selon divers paramètres, offrant une meilleure précision pour un nombre identique ou moindre d'échantillons et de tours.
• Allocation des Ressources : L'analyse montre que la politique apprise alloue intuitivement plus d'échantillons aux requêtes présentant une entropie de réponse plus élevée (diversité), s'alignant sur la nécessité d'avoir plus d'échantillons pour converger vers un vote majoritaire fiable.
• Études d'Ablation : La cible "Running Majority" (vote majoritaire au budget max $N$) s'est avérée supérieure à l'utilisation des étiquettes de vérité terrain ("Real Label"), laquelle dégradait les performances en introduisant du bruit et en empêchant l'apprentissage de critères d'arrêt universels.

Signification et Revendications
L'article affirme que l'échantillonnage guidé par RL offre une alternative fondée sur des principes et légère aux méthodes d'échantillonnage adaptatif existantes. En formulant le problème comme un MDP et en entraînant un petit contrôleur via le RL, la méthode équilibre efficacement la justesse des réponses, la latence et le coût de calcul sans nécessiter de modifications invasives du LLM ou l'accès aux états internes du modèle. Les auteurs soulignent que cette approche est particulièrement précieuse pour le déploiement pratique, car elle peut être entraînée sur des modèles open-source peu coûteux et transférée pour réduire les coûts d'inférence de systèmes plus puissants, propriétaires ou fermés. Ce travail représente une première étape vers l'échantillonnage adaptatif basé sur le RL, avec des directions futures suggérant l'incorporation de signaux d'état plus riches et de mesures de coûts réels.