Résumé technique : POETS (Ensembles de politiques pour l'échantillonnage de Thompson)

Énoncé du problème

L'équilibre entre exploration et exploitation demeure un défi fondamental dans la prise de décision séquentielle et l'optimisation de boîtes noires, en particulier lorsqu'il est appliqué aux grands modèles de langage (LLM). Une exploration efficace nécessite une quantification rigoureuse de l'incertitude. Bien que l'assemblage de modèles entraînés indépendamment constitue une méthode éprouvée et agnostique au modèle pour estimer l'incertitude épistémique, l'application directe d'un assemblage naïf aux LLM est prohibitivement coûteuse en raison des contraintes de mémoire et de calcul. De plus, les méthodes d'exploration existantes pour les LLM reposent souvent sur des modèles de récompense explicites (introduisant des complexités d'optimisation imbriquées) ou dépendent fortement des capacités de suivi des instructions du modèle.

Méthodologie

L'article présente POETS (Policy Ensembles for Thompson Sampling), un cadre qui fait le lien entre la quantification de l'incertitude et l'optimisation de politiques sans nécessiter de modèle de récompense probabiliste explicite.

Insight central : Encodage implicite de la récompense

POETS s'appuie sur la dualité entre les fonctions de récompense et les politiques optimales dans le cadre d'objectifs régularisés. Plus précisément, les politiques entraînées avec une régularisation de Kullback-Leibler (KL) encodent implicitement une fonction de récompense sous-jacente. Au lieu d'apprendre un modèle de récompense puis d'ajuster une politique à celui-ci (un processus en deux étapes), POETS correspond directement les fonctions de récompense implicites encodées par un ensemble de politiques aux données en ligne observées.

L'algorithme

1. Ensemble de politiques : POETS maintient un ensemble de $n$ politiques indépendantes, $\Pi = \{\pi_1, \dots, \pi_n\}$.
2. Bootstrapping de Poisson : Pour capturer l'incertitude épistémique, les membres de l'ensemble sont diversifiés par un bootstrapping de Poisson. Pour chaque lot de données entrant, chaque membre de l'ensemble se voit attribuer un poids entier indépendant tiré d'une distribution de Poisson(1). Cela permet à l'ensemble d'apprendre une distribution a posteriori bayésienne approximative sur les fonctions de récompense sans nécessiter de jeu de données statique.
3. Fonction de perte : La méthode optimise une fonction de perte qui aligne la récompense implicite d'une politique unique sur la vraie récompense. La dérivation théorique montre que l'optimisation de cette perte est mathématiquement équivalente à une mise à jour de gradient de politique (spécifiquement, une variante de GRPO) mise à l'échelle par les poids de Poisson.
4. Échantillonnage de Thompson : En ajustant un ensemble de récompenses implicites aux vraies récompenses à l'aide de bootstraps indépendants, POETS apprend une distribution a posteriori sur les fonctions de récompense. L'exécution d'un échantillonnage de Thompson régularisé par KL est obtenue simplement en échantillonnant uniformément une politique de l'ensemble $\Pi$ et en agissant conformément à celle-ci.

Architecture efficace : « Tronc et Branches »

Pour surmonter le coût prohibitif de la maintenance de $n$ copies complètes de LLM, POETS utilise une architecture économe en paramètres :

• Tronc partagé : Un socle pré-entraîné (la majorité des poids) est partagé entre tous les membres de l'ensemble et évalué une seule fois par passage avant/arrière.
• Branches indépendantes : La diversité de l'ensemble est entièrement encapsulée dans des modules d'adaptation de faible rang (LoRA) indépendants agissant comme les dernières couches de lecture.
• Efficacité : Cette conception garantit que la génération de réponses et l'accumulation de gradients pour $n$ politiques prennent pratiquement le même temps réel qu'une seule politique, avec une surcharge mémoire ne dépendant que de la taille des branches LoRA.

Contributions clés

1. Nouveau cadre : Introduction de POETS, qui capture l'incertitude épistémique via des ensembles de politiques et réalise implicitement un échantillonnage de Thompson régularisé par KL sans modèle de récompense explicite.
2. Garanties théoriques : Preuve que POETS hérite des fortes bornes de regret cumulatif doux de $O(\sqrt{T\gamma_T})$ associées à l'échantillonnage de Thompson régularisé par KL, où $\gamma_T$ est le gain d'information maximal.
3. Efficacité computationnelle : Démonstration qu'une architecture LoRA « Tronc et Branches » permet un assemblage profond avec une surcharge minimale (par exemple, une augmentation du temps d'exécution d'environ 7,7 % pour $n=16$ sur un GPU H200).
4. Performance empirique : Efficacité d'échantillonnage de pointe sur divers domaines d'optimisation de boîtes noires combinatoires (affinement de FAQ, recherche de protéines, conception de circuits quantiques) et amélioration de la robustesse dans les contextes d'apprentissage par renforcement, en particulier dans les scénarios hors politique avec rejouage d'expérience.

Résultats expérimentaux

Les auteurs ont évalué POETS par rapport à des bases de référence incluant GRPO, TOSFIT (qui utilise un modèle de récompense explicite), FIBO (échantillonnage de Thompson en contexte) et EVOSEARCH (recherche évolutionnaire).

• Optimisation de boîte noire : POETS a atteint une efficacité d'échantillonnage de pointe dans la recherche de protéines, la conception de circuits quantiques et l'affinement de FAQ. Notamment, il a surpassé les bases de référence même sans tampon de rejouage, alors que le GRPO standard échouait à maintenir la diversité et surajustait aux candidats précoces.
• Apprentissage par renforcement (RL) : Dans le cadre du jeu de données AIME (bandit contextuel avec récompenses vérifiables), POETS a amélioré les trajectoires d'optimisation pour les contextes en politique et hors politique (rejouage d'expérience). Il a réussi à atténuer les problèmes de surajustement observés dans GRPO lorsqu'il était exposé à des contextes répétés.
• Diversité : L'ensemble a maintenu la diversité de génération (mesurée par la divergence de Jensen-Shannon) tout au long de l'entraînement, empêchant l'effondrement de l'entropie courant dans les méthodes de gradient de politique standard.
• Ablations : Les expériences ont confirmé que le bootstrapping de Poisson est essentiel pour la diversité, que de petits ensembles ($n=8$) suffisent pour des performances élevées, et que l'architecture « Tronc et Branches » capture efficacement l'incertitude avec un coût computationnel minimal.

Importance et affirmations

L'article affirme que POETS représente une avancée significative en permettant une exploration fondée sur des principes et consciente de l'incertitude pour les LLM dans la découverte scientifique et le RL, sans le fardeau computationnel de la modélisation explicite de récompense ni la fragilité du suivi des instructions en contexte.

• Théorique : Il établit que les ensembles de politiques avec bootstrapping suffisent à mettre en œuvre un échantillonnage de Thompson régularisé par KL, héritant de bornes de regret optimales.
• Pratique : Il fournit une solution évolutive qui rend l'assemblage profond réalisable pour les LLM, permettant une utilisation efficace du rejouage d'expérience et une exploration robuste dans des régimes de forte incertitude.
• Portée : Ce travail est présenté comme une contribution algorithmique fondamentale pour le post-entraînement des LLM et la découverte scientifique, ciblant spécifiquement les problèmes d'optimisation combinatoire où les évaluations sont coûteuses.

Les auteurs notent que, bien que POETS améliore l'efficacité d'échantillonnage et la robustesse, il ne traite pas de la mauvaise spécification des récompenses, de l'équité ou de l'alignement, qui restent des responsabilités des concepteurs de systèmes. Des travaux futurs sont suggérés pour étendre le cadre des bandits contextuels aux contextes multi-étapes avec attribution de crédit fine.