Résumé Technique : Best-of-Better-N (BoBN)

Énoncé du Problème

Les méthodes d'alignement au moment de l'inférence (Inference-time alignment ou ITA), telles que le Best-of-N (BoN), visent à aligner les grands modèles de langage (LLM) avec les préférences humaines ou des exigences de tâches spécifiques sans modifier les poids du modèle. Ces méthodes reposent sur un cadre de « échantillonnage et évaluation » : elles échantillonnent $N$ réponses à partir d'une politique de référence $\pi_{ref}$ et sélectionnent celle qui présente la récompense la plus élevée selon un modèle de récompense $r$.

L'efficacité de l'ITA est fondamentalement contrainte par deux hypothèses :

1. Couverture (A1) : La politique de référence $\pi_{ref}$ doit échantillonner au moins une réponse de haute qualité (haute récompense).
2. Sélection (A2) : Le critère de sélection doit identifier avec succès la réponse de haute qualité parmi les échantillons.

La recherche actuelle en ITA se concentre largement sur l'amélioration de (A2) en affinant les critères de sélection pour atténuer le détournement de récompense (reward hacking). Cependant, un goulot d'étranglement critique subsiste dans (A1). Si $\pi_{ref}$ n'est pas alignée avec la tâche en aval (par exemple, un modèle non entraîné à la sécurité tentant de refuser des requêtes toxiques), elle assigne une probabilité négligeable aux réponses à haute récompense. Dans ces cas, aucune stratégie de sélection ne peut réussir car la région de haute récompense de l'espace de sortie est effectivement inaccessible au modèle de référence. Les solutions existantes nécessitent souvent un ajustement fin (qui est coûteux et nécessite l'accès aux poids) ou utilisent des exemples contextuels statiques qui peuvent manquer de pertinence par rapport à la tâche spécifique.

Méthodologie : Best-of-Better-N (BoBN)

Les auteurs proposent Best-of-Better-N (BoBN), un cadre d'inférence conçu pour déplacer la distribution d'échantillonnage de $\pi_{ref}$ vers la région de haute récompense sans ajustement fin ni accès aux poids. BoBN exploite l'apprentissage en contexte (In-Context Learning ou ICL) combiné à un pipeline de récupération et de restylage.

Le flux de travail de BoBN se compose de trois étapes :

1. Récupération (Retrieval) : Étant donné un prompt d'entrée $x$, le système récupère les $K$ paires prompt-réponse $(x_i, y_i)$ les plus similaires d'un ensemble de données $D$ en utilisant un modèle d'embedding. Ces exemples servent de démonstrations de comportements à haute récompense pertinents pour la requête spécifique.
2. Restylage (Restyling) : Pour garantir que les exemples récupérés s'alignent sur les contraintes, le format et la persona de la tâche cible, le modèle de référence $\pi_{ref}$ est utilisé pour « restyler » les réponses récupérées $y_i$. Plus précisément, $\pi_{ref}$ réécrit $y_i$ en $\hat{y}_i$ conditionné par un prompt tel que « Réécrivez $y_i$ pour correspondre au style $T$ ». Cette étape assure une cohérence de distribution au niveau des tokens entre le contexte et le modèle de référence, tout en préservant la logique d'alignement des données récupérées.
3. Génération et Sélection : Les exemples restylés $C_K = \{(x_i, \hat{y}_i)\}$ sont ajoutés au prompt d'entrée $x$ comme exemples en contexte. Le modèle de référence échantillonne ensuite $N$ réponses à partir de la distribution conditionnée $\pi_{ref}(\cdot | C_K, x)$. Enfin, un critère de sélection (par exemple, BoN) choisit la meilleure réponse basée sur le modèle de récompense.

En se conditionnant sur $C_K$, BoBN déplace efficacement la masse de probabilité de $\pi_{ref}$ vers la région de haute récompense, augmentant ainsi la probabilité que des candidats de haute qualité soient générés dès le départ.

Contributions Clés

1. Cadre Novateur : La proposition de BoBN, une approche basée sur l'ICL qui améliore la couverture de $\pi_{ref}$ sur les réponses de haute qualité en utilisant des exemples récupérés et restylés.
2. Analyse Théorique : Une caractérisation théorique de la manière dont la qualité de la récupération contrôle la couverture et la performance en aval. Les auteurs démontrent que la récupération pertinente réduit de manière prouvable la couverture (définie comme la divergence chi-deux entre la politique dorée et la politique de référence) par rapport à un ICL statique ou sans ICL. Ils montrent que cette réduction de la couverture se combine avec la sélection BoN pour améliorer la performance de la tâche.
3. Validation Empirique : Expérimentations approfondies sur les benchmarks d'alignement de sécurité (StrongReject, SORRY-Bench) et de raisonnement mathématique (GSM8k, MATH500). Les résultats montrent que BoBN améliore les performances sur plusieurs modèles de référence (zephyr-7b, tulu3-8b), tant lorsque $N=1$ (qualité de la réponse individuelle) que lorsque $N>1$ (combiné avec la sélection).

Résultats Expérimentaux

Les auteurs ont évalué BoBN par rapport au BoN standard (utilisant directement $\pi_{ref}$) et à URIAL (ICL classique avec des exemples statiques).

• Alignement de Sécurité : Sur le benchmark StrongReject, BoBN a nettement surpassé les bases de référence. Pour le modèle zephyr-7b, BoBN avec $N=16$ a atteint un score de 0,033 (proche de l'optimal), comparable aux modèles de pointe comme Claude Sonnet-3.7, alors que $\pi_{ref}$ seul a obtenu un score de 0,425. BoBN a également démontré une couverture supérieure ( $\log C$ plus faible) par rapport à URIAL et $\pi_{ref}$.
• Raisonnement Mathématique : Sur GSM8k et MATH500, BoBN a produit des améliorations constantes de la précision. Pour zephyr-7b sur GSM8k, BoBN a atteint une précision de 58,7 % avec $N=16$, un gain significatif par rapport aux 43,9 % obtenus par $\pi_{ref}$ avec BoN. Notamment, BoBN avec $N=4$ a souvent égalé ou dépassé les performances des bases de référence avec $N=16$, indiquant une plus grande efficacité d'échantillonnage.
• Coût et Latence : Bien que BoBN incurre une latence supplémentaire due à l'étape de restylage (nécessitant $K$ appels à $\pi_{ref}$ avant la génération), il offre une efficacité de tokens significative. BoBN avec $N=4$ a économisé jusqu'à 50 % de l'utilisation de tokens par rapport à URIAL avec $N=16$ tout en atteignant des performances similaires ou supérieures.
• Études d'Ablation :
  • Restylage : Le retrait de l'étape de restylage a entraîné un alignement de sécurité plus faible et une couverture plus élevée (moins bonne performance), confirmant la nécessité d'aligner les exemples récupérés sur le style de la tâche cible.
  • Nombre de Récupérations ($K$) : La performance s'est améliorée à mesure que $K$ passait de 1 à 8, mais a décliné lorsque $K=16$, suggérant que la récupération de trop d'exemples introduit des données moins pertinentes qui diluent le signal en contexte.

Signification et Revendications

L'article affirme que BoBN répond à une limitation fondamentale de l'alignement au moment de l'inférence : l'incapacité des modèles de référence non alignés à échantillonner des réponses à haute récompense. En déplaçant la distribution d'échantillonnage via des exemples en contexte récupérés et restylés, BoBN permet :

• Une Haute Qualité avec un Budget Fixe : Une meilleure performance de la tâche pour un nombre fixe d'échantillons $N$.
• Efficacité : Atteindre la performance cible avec un $N$ plus petit, réduisant ainsi les coûts d'inférence.
• Alignement Sans Entraînement : Améliorer l'alignement des modèles qui n'ont pas été explicitement affinés pour la tâche en aval, sans nécessiter d'accès aux poids ou aux probabilités du modèle.

Les auteurs soulignent que leur analyse théorique, basée sur les transformateurs linéaires, fournit un lien prouvable entre la qualité de la récupération, la réduction de la couverture et la performance de la tâche en aval, offrant ainsi une approche structurée pour améliorer l'ITA au-delà des simples heuristiques de sélection.