📊 statistics

Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

Cet article introduit la Pass-at-k Policy Optimization (PKPO), un nouveau cadre d'apprentissage par renforcement qui dérive des estimateurs non biaisés pour optimiser directement le succès collectif d'ensembles d'échantillons (pass@k) plutôt que des tentatives isolées, améliorant ainsi l'exploration et la résolution de problèmes plus difficiles tout en maintenant ou en améliorant la performance pass@1 grâce à un recuit de k (k-annealing).

Auteurs originaux : Christian Walder, Deep Karkhanis

Publié 2026-06-11

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Christian Walder, Deep Karkhanis

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous êtes un enseignant essayant d'aider un élève à apprendre à résoudre des problèmes mathématiques difficiles.

L'ancienne méthode : Le piège du « premier essai »
Traditionnellement, lorsqu'on entraîne des modèles d'IA (comme ceux qui écrivent du code ou résolvent des mathématiques), l'ordinateur essaie de résoudre un problème, obtient un score, puis ajuste son cerveau en fonction de cette tentative unique. Si le premier essai échoue, l'ordinateur n'apprend rien des autres tentatives qu'il aurait pu faire en arrière-plan. C'est comme un élève qui passe un examen, répond mal à une question, et abandonne immédiatement, ignorant le fait qu'il aurait pu la résoudre correctement lors de sa deuxième ou troisième tentative s'il avait simplement continué.

Cette méthode optimise le Pass@1 : « Est-ce que la toute première réponse a fonctionné ? » Cela force l'IA à être prudente et conservatrice, évitant souvent les conjectures risquées et créatives nécessaires pour résoudre des problèmes vraiment difficiles.

La nouvelle idée : L'approche du « meilleur du lot »
Les auteurs de cet article proposent une nouvelle stratégie appelée Optimisation de la politique Pass@K (PKPO).

Au lieu de ne se soucier que de la première réponse, cette méthode dit : « Générons K tentatives différentes (disons 8 ou 16) pour chaque problème. Peu importe si les 7 premières sont fausses ; ce qui compte, c'est qu'au moins une d'entre elles soit correcte. »

Imaginez cela comme un filet de pêche.

Ancienne méthode : Vous lancez une seule ligne. Si vous manquez le poisson, vous la remontez et n'apprenez rien.
Méthode PKPO : Vous lancez un filet avec 16 lignes. Si même une seule ligne attrape un poisson, tout le filet est un succès. L'IA est récompensée pour la meilleure prise du filet, pas pour la moyenne de toutes les lignes.

Le tour de magie : La fiche de score
La partie difficile est de comprendre comment enseigner cela à l'IA. Si vous dites simplement à l'IA « Tu as eu un poisson sur la ligne n°4 », elle pourrait ignorer les lignes n°1, n°2 et n°3. Mais si vous lui dites « Tu as eu un poisson, donc tu as bien travaillé », elle pourrait ne pas réaliser quelle ligne a été l'héroïne.

Les auteurs ont inventé une « fiche de score » spéciale (un estimateur) qui agit comme un arbitre intelligent.

Il examine les 16 tentatives.
Il calcule un score qui récompense l'IA pour avoir n'importe quelle réponse correcte dans le groupe.
Crucialement, il accorde aussi un peu de crédit aux réponses « fausses », car elles faisaient partie du groupe qui a fini par produire le gagnant. Cela encourage l'IA à continuer d'explorer et de tenter des idées audacieuses et risquées, sachant même qu'une « mauvaise » supposition contribue au succès de l'équipe si une « bonne » supposition apparaît plus tard.

Pourquoi cela importe
L'article montre que cette méthode fonctionne comme un superpouvoir pour les tâches difficiles :

Elle débloque les problèmes difficiles : Sur des défis mathématiques et de codage très complexes où l'ancienne méthode du « premier essai » reste bloquée, cette nouvelle méthode continue d'apprendre et finit par résoudre les problèmes.
Elle est flexible : Vous pouvez dire à l'IA : « Pour la première moitié de l'entraînement, sois une prise de risque et vise le meilleur de 8 essais. Pour la seconde moitié, concentre-toi sur la réussite du premier essai. » Ce processus de « recuit » (annealing - changer progressivement les règles) aide l'IA à apprendre d'abord à explorer, puis à affiner ses compétences.
Elle fonctionne avec des modèles réels : Ils ont testé cette méthode sur des modèles open-source populaires (GEMMA2 et LLAMA3.1) et ont constaté qu'elle améliorait considérablement leur capacité à résoudre des problèmes mathématiques et à écrire du code par rapport aux méthodes précédentes.

En résumé
L'article apprend à l'IA à ne plus s'inquiéter d'être parfaite dès la première tentative. Au lieu de cela, il lui apprend à générer un ensemble diversifié d'idées, à récompenser le groupe pour avoir produit n'importe quel gagnant, et à utiliser ce succès collectif pour apprendre à résoudre les énigmes les plus difficiles. Il s'agit de valoriser l'effort d'équipe de plusieurs conjectures plutôt que la performance individuelle d'une seule conjecture.

Résumé technique : Optimisation de politique Pass@K (PKPO)

Énoncé du problème

Les algorithmes d'apprentissage par renforcement (RL) appliqués aux grands modèles de langage (LLM) échantillonnent souvent plusieurs tentatives de solution ( $n > 1$ ) pour un seul problème, mais optimisent en fonction de la performance de l'échantillon individuel (pass@1). Cette approche privilégie l'utilité d'un échantillon unique à haute récompense au détriment de l'utilité collective de l'ensemble, ce qui sous-utilise efficacement la capacité d'échantillonnage. Par conséquent, les méthodes classiques de gradient de politique peuvent échouer à explorer suffisamment l'espace des solutions, particulièrement sur des tâches difficiles où la probabilité d'un échantillon correct unique est faible. Bien qu'il existe des méthodes de recherche au moment de l'inférence, le couplage entre l'inférence et les mises à jour du modèle suggère qu'optimiser naïvement l'espérance de la récompense d'un échantillon unique est sous-optimal pour les tâches où la correction peut être vérifiée mais où les solutions ne sont pas connues a priori.

Méthodologie

Les auteurs proposent la Pass-at-k Policy Optimization (PKPO), un cadre qui optimise directement la métrique pass@k (ou son analogue continu, maxg@k), définie comme l'espérance du maximum de la récompense sur $k$ échantillons indépendants tirés de la politique.

Contributions théoriques centrales

La principale contribution est la dérivation de nouveaux estimateurs à faible variance et non biaisés pour l'objectif pass@k et son gradient. Ces estimateurs fonctionnent en faisant la moyenne de simples estimateurs sur toutes les $\binom{n}{k}$ sous-populations de taille $k$ issues d'un lot de $n$ échantillons.

Récompenses binaires (Section 2) :
- Pour les étiquettes de correction binaires ( $f \in \{0, 1\}$ ), les auteurs dérivent un estimateur non biaisé $\rho(n, c, k)$ pour le pass@k, où $c$ est le nombre d'échantillons corrects.
- Ils prouvent qu'un estimateur non biaisé n'existe que si $n \ge k$ .
- Un estimateur de gradient non biaisé correspondant est dérivé, attribuant des poids aux échantillons en fonction de leur correction et de la probabilité combinatoire de leur inclusion dans un sous-ensemble réussi.
Récompenses continues (Section 3) :
- Le cadre se généralise aux récompenses continues $g(x)$ via l'objectif maxg@k : $E[\max_{i=1}^k g(x_i)]$ .
- En supposant que les récompenses sont triées ( $g_1 \le g_2 \le \dots \le g_n$ ), les auteurs dérivent un estimateur non biaisé $\rho^{(g)}(n, c, k)$ qui pondère les récompenses triées à l'aide de coefficients binomiaux.
- Un estimateur de gradient non biaisé est construit où le poids $s_i$ pour un échantillon est une combinaison linéaire de récompenses, calculée efficacement en $O(k + n \log n)$ temps.
Réduction de la variance (Section 4) :
- Pour réduire la variance des estimations de gradient sans introduire de biais, les auteurs proposent des lignes de base Leave-One-Out (LOO).
- Ils introduisent une ligne de base spécifique $s^{(loo-1)}_i$ qui soustrait l'espérance de maxg@k des sous-ensembles de taille $k-1$ (excluant l'échantillon actuel) du maxg@k des sous-ensembles de taille $k$ . Cela généralise les travaux précédents (notamment [TZSM25]) qui étaient limités au cas $n=k$ .
- Cette formulation permet une optimisation robuste même lorsque $n > k$ , un scénario courant dans l'entraînement des LLM où de grands lots sont échantillonnés.

Implémentation

La méthode est conçue comme un "remplacement direct" pour les transformations de récompense standard dans n'importe quel algorithme de gradient de politique (par exemple, PPO). Étant donné un lot de $n$ échantillons et leurs récompenses brutes, les récompenses sont transformées à l'aide des estimateurs dérivés (spécifiquement la fonction sloo_minus_one dans le code fourni) avant d'être injectées dans l'étape de mise à jour de la politique.

Résultats clés

Les auteurs valident la PKPO via des expériences de type "toy" et des benchmarks du monde réel utilisant GEMMA2 (2B et 9B paramètres) et LLAMA3.1 (8B paramètres).

Expériences de type "toy" :
- Dans une expérience de politique gaussienne 1D, la PKPO parvient à déplacer la moyenne de la politique optimale pour être plus "tolérante au risque" pour un $k$ plus élevé, permettant au modèle d'explorer des régions avec une probabilité plus faible mais un potentiel de récompense plus élevé.
- L'analyse de la variance confirme que l'estimateur sloo_minus_one proposé réduit considérablement la variance du gradient par rapport au partitionnement naïf ou aux estimateurs sans lignes de base.
Benchmarks du monde réel (MATH, Coding, ARC-AGI-1) :
- Optimisation sélective : Fixer la cible d'optimisation $k_{opt}$ égale à la métrique d'évaluation $k_{eval}$ produit les meilleurs résultats. Par exemple, optimiser pour le pass@8 améliore significativement les scores d'évaluation pass@8 par rapport à une optimisation pour le pass@1.
- Résolution de problèmes plus difficiles : Sur le benchmark exigeant ARC-AGI-1, l'optimisation conventionnelle du pass@1 stagne, échouant à améliorer le taux de résolution cumulé. En revanche, la PKPO avec $k_{opt} > 1$ "débloque" l'apprentissage, atteignant des taux de résolution et des scores pass@k plus élevés pour tous les $k$ .
- Stratégie d'annealing (recuit) : Les auteurs démontrent qu'un recuit de $k$ pendant l'entraînement (en commençant par un $k_{opt}$ élevé pour encourager l'exploration, puis en réduisant à $k_{opt}=1$ pour consolider) permet d'obtenir de fortes performances en pass@1 parallèlement à des gains significatifs en pass@k, évitant le compromis typique où l'amélioration du pass@k dégrade le pass@1.
- Comparaison : La PKPO surpasse le cas particulier de $n=k$ (comme dans [TZSM25]) et les bases de régularisation d'entropie, particulièrement sur des tâches difficiles où l'exploration est critique.

Signification et affirmations

L'article affirme que la PKPO fournit une méthode robuste, efficace et théoriquement fondée pour optimiser l'utilité collective d'un ensemble d'échantillons plutôt que simplement la moyenne ou le meilleur échantillon individuel.

Déblocage de l'exploration : Les auteurs soutiennent qu'en privilégiant l'utilité conjointe (la probabilité qu'au moins un des $k$ échantillons réussisse) plutôt que l'utilité individuelle de l'échantillon, la méthode améliore l'exploration. Ceci est crucial pour résoudre des problèmes plus difficiles où le "signal" (une solution correcte) est rare dans l'espace de recherche.
Flexibilité : Contrairement aux méthodes précédentes qui étaient restreintes à $k=n$ (optimisation pour le meilleur de l'ensemble du lot) ou $k=1$ , la PKPO permet une optimisation robuste pour n'importe quel $k$ arbitraire $\le n$ .
Impact pratique : La méthode permet aux praticiens de régler le compromis exploration-exploitation via le paramètre $k_{opt}$ , permettant aux modèles de résoudre des tâches de raisonnement et de codage plus complexes qui sont autrement inaccessibles à l'optimisation standard du pass@1.

Ce travail se positionne comme une généralisation et une réduction de la variance des méthodes existantes de gradient de politique basées sur les statistiques d'ordre, offrant une voie pratique pour passer à l'échelle l'entraînement par RL pour les LLM sur des tâches nécessitant une exploration profonde.