Résumé

Les avancées récentes dans l’alignement des grands modèles linguistiques par apprentissage par renforcement ont permis des gains remarquables dans la résolution de problèmes complexes de raisonnement, mais au prix de coûts élevés liés aux échantillonnages en politique (on-policy) et d’une exploration limitée des différentes voies de raisonnement. Dans ce travail, nous introduisons TreePO, un algorithme d’échantillonnage auto-guidé qui modélise la génération de séquences comme un processus de recherche structuré en arbre. Composé d’une politique d’échantillonnage dynamique sur arbre et d’une décodage par segments de longueur fixe, TreePO exploite l’incertitude locale pour justifier l’ajout de nouvelles branches. En amortissant le calcul sur les préfixes communs et en élaguant précocement les chemins à faible valeur, TreePO réduit essentiellement la charge de calcul par mise à jour tout en préservant ou en améliorant la diversité de l’exploration. Les contributions principales sont les suivantes : (1) un algorithme d’échantillonnage par segments qui allège la charge de la mémoire tampon des paires clé-valeur (KV cache) grâce à des segments contigus et qui permet la création de nouvelles branches en combinant un mécanisme d’arrêt anticipé ; (2) une estimation d’avantage segmentaire fondée sur une structure d’arbre, prenant en compte à la fois l’optimisation de politique proche globale et locale (proximal policy optimization) ; (3) une analyse de l’efficacité de la divergence dynamique pilotée par la probabilité et la qualité, ainsi que de la stratégie de retour (fallback). Nous validons empiriquement les gains de performance de TreePO sur plusieurs benchmarks de raisonnement, et observons une réduction de 22 % à 43 % des heures GPU nécessaires pour le processus d’échantillonnage des modèles entraînés, tout en montrant une réduction pouvant atteindre 40 % au niveau des trajectoires et 35 % au niveau des tokens pour les modèles existants. En offrant une amélioration gratuite de l’efficacité d’inférence, TreePO ouvre une voie concrète pour échelonner l’entraînement postérieur basé sur le renforcement avec moins d’échantillons et moins de calcul. La page d’accueil est disponible à l’adresse suivante : https://m-a-p.ai/TreePO.