Résumé Technique : Aligner les politiques de recherche arborescente avec des budgets de jetons fixes dans le passage à l'échelle au moment de l'inférence des LLM

Énoncé du Problème

Les grands modèles de langage (LLM) peuvent améliorer la qualité de leurs réponses grâce au passage à l'échelle au moment de l'inférence (Test-Time Scaling), qui alloue un calcul supplémentaire lors de l'inférence sans mettre à jour les paramètres du modèle. Bien que le décodage par recherche arborescente (par exemple, MCTS - Monte Carlo Tree Search) soit une forme efficace de passage à l'échelle, les déploiements réels imposent souvent un budget de jetons fixe par requête.

Les politiques de recherche arborescente existantes sont largement indépendantes du budget (budget-agnostic). Elles traitent généralement le budget de jetons simplement comme une condition de terminaison plutôt que comme une contrainte dynamique guidant le comportement de recherche. Ce décalage entraîne deux modes de défaillance principaux :

1. Sur-élargissement tardif : La recherche continue d'étendre des nœuds peu profonds tard dans le processus, épuisant le budget avant de pouvoir affiner ou achever les candidats prometteurs.
2. Terminaison prématurée : La recherche s'arrête prématurément, laissant une partie du budget alloué inutilisée.

Les variantes d'arrêt précoce actuelles se concentrent sur l'économie de jetons mais manquent d'une politique pour changer dynamiquement la stratégie de recherche, passant d'une exploration large à un raffinement à mesure que le budget restant diminue.

Méthodologie : BG-MCTS (Budget-Guided MCTS)

Les auteurs proposent le BG-MCTS (Budget-Guided MCTS), un algorithme de décodage par recherche arborescente qui conditionne explicitement sa politique de recherche au budget de jetons restant. Le mécanisme central repose sur un ratio de suffisance du budget $\rho \in [0, 1]$, défini comme :
$$\rho = 1 - \frac{C_{utilisés}}{B}$$
où $C_{utilisés}$ est le nombre cumulé de jetons générés et $B$ est le budget total. À mesure que $\rho$ diminue (se rapprochant de 0), la politique passe de l'exploration vers l'exploitation et l'achèvement.

BG-MCTS introduit deux mécanismes de contrôle conditionnés par le budget :

1. Sélection Guidée par le Budget (BG-PUCT)

La sélection PUCT standard est modifiée pour atténuer le bonus d'exploration et introduire un biais d'achèvement :
$$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$$

• Atténuation de l'exploration : Le terme d'exploration est multiplié par $\rho$. À mesure que le budget s'épuise ($\rho \to 0$), l'incitation à explorer de nouvelles branches diminue.
• Biais d'achèvement : La valeur accumulée $\tilde{W}(s, \rho)$ est corrigée par un terme proportionnel à la profondeur du nœud $d(x)$ et à $(1-\rho)$. Cela oriente la sélection vers des nœuds plus profonds qui sont plus proches de générer une réponse finale, sans artificiellement booster les nœuds qui contiennent déjà une réponse complétée.

2. Élargissement d'Arbre Guidé par le Budget

Pour éviter un élargissement tardif inutile, BG-MCTS traite l'« élargissement » (générer un nouvel enfant à partir d'un parent) comme une action sélectionnable aux côtés des enfants existants. Un enfant génératif virtuel $s_{gen}(p)$ est ajouté à l'ensemble de sélection. Son score est :
$$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$$

• $\mu(p)$ et $\sigma^2(p)$ représentent la moyenne et la variance des scores des enfants existants.
• Le terme est mis à l'échelle par $\rho$. Lorsque le budget est abondant, l'élargissement est encouragé (particulièrement si les enfants divergent). À mesure que le budget diminue, l'incitation à élargir est supprimée, forçant la recherche à approfondir les trajectoires existantes prometteuses.

Algorithme Unifié

L'algorithme suit une boucle MCTS standard mais rend à la fois la décision de sélection et d'élargissement dépendantes de $\rho$. Cela induit un programme de large-à-profond (wide-to-deep schedule) : une exploration large au début de la recherche, suivie d'un raffinement ciblé et de l'achèvement des réponses à mesure que le budget est épuisé.

Contributions Clés

1. Conception de Politique Sensible au Budget : Le papier introduit le premier cadre de décodage par recherche arborescente qui aligne dynamiquement le comportement de recherche (exploration vs exploitation vs élargissement) avec le budget de jetons restant, plutôt que de traiter le budget comme un simple signal d'arrêt statique.
2. Mécanismes Nouveaux : La proposition du BG-PUCT (avec exploration conditionnée par le budget et correction de valeur biaisée par la profondeur) et des Enfants Génératifs Virtuels pour un élargissement sensible au budget.
3. Validation Empirique : Évaluation complète sur des benchmarks de raisonnement mathématique (MATH500, AIME24/25) et un benchmark de raisonnement physique (UGPhysics) utilisant des LLM en poids ouverts (Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B).

Résultats Expérimentaux

Les auteurs ont évalué BG-MCTS par rapport à des bases de référence agnostiques au budget (MCTS standard, AB-MCTS-M, LiteSearch) et des bases de référence de mise à l'échelle (Échantillonnage répété, Raffinement séquentiel) sous des budgets fixes de 10k, 20k et 30k jetons.

• Performance : BG-MCTS a systématiquement surpassé toutes les bases de référence sur tous les modèles et budgets. Par exemple, sur MATH500 avec Qwen2.5-7B, BG-MCTS a atteint 71,1 % de précision à 30k jetons, contre 65,9 % pour la deuxième meilleure base de référence de recherche arborescente (MCTS).
• Utilisation du Budget : Contrairement à LiteSearch, qui s'arrête souvent prématurément, BG-MCTS continue d'améliorer la précision à mesure qu'il approche de l'épuisement du budget.
• Comportement Qualitatif : L'analyse des arbres de recherche a révélé que BG-MCTS maintient une largeur importante au début, mais bascule vers l'approfondissement de branches spécifiques en fin de recherche. En revanche, le MCTS standard continue souvent d'élargir des nœuds peu profonds jusqu'à épuisement du budget, échouant ainsi à achever des chemins prometteurs.
• Raisonnement Physique : L'efficacité de la méthode s'est étendue à UGPhysics, démontant que le comportement "large-à-profond" aligné sur le budget est bénéfique au-delà du pur raisonnement mathématique.

Signification et Revendications

Le papier affirme que conditionner les décisions de recherche arborescente sur le budget restant est plus efficace que de traiter le budget simplement comme une condition d'arrêt.

• Mise à l'échelle Prévisible : BG-MCTS offre une interface simple pour un passage à l'échelle prévisible au moment de l'inférence, où les utilisateurs spécifient un budget de jetons, et la politique adapte automatiquement son programme d'exploration vers l'achèvement dans cette enveloppe de coût.
• Implications pour la Synthèse de Données : Les auteurs notent que BG-MCTS tend à produire moins de nœuds répondus par instance de problème, mais avec une précision (exactitude) nettement plus élevée. Ce régime de "moins d'instances, meilleures réponses" est souligné comme particulièrement précieux pour la synthèse de données, où les traces de raisonnement de haute qualité sont plus critiques que la couverture brute.
• Robustesse : La méthode reste efficace même en utilisant des modèles de récompense de processus (PRM) imparfaits qui fournissent des signaux faibles ou saturés, suggérant que les changements structurels sensibles au budget peuvent compenser les signaux d'évaluation bruités.

Le travail conclut qu'aligner les politiques de recherche avec des contraintes de jetons fixes est une étape cruciale pour le déploiement des LLM dans des environnements réels à ressources limitées.