Résumé Technique : REAR : Réalignement des Préférences au Moment de l'Inférence par Décomposition de la Récompense

Énoncé du Problème

Aligner les grands modèles de langage (LLM) avec les préférences diverses et personnalisées des utilisateurs est un défi critique. Bien que les méthodes de post-entraînement comme l'apprentissage par renforcement à partir de la rétroaction humaine (RLHF) soient efficaces, elles nécessitent une curation de données coûteuse et un entraînement supplémentaire, ce qui les rend peu flexibles face aux besoins dynamiques ou spécifiques des utilisateurs. Le passage à l'échelle au moment de l'inférence (Test-Time Scaling - TTS) offre une alternative sans entraînement en exploitant un calcul supplémentaire lors de l'inférence pour améliorer la qualité de la sortie. Cependant, les méthodes de TTS existantes ont largement été limitées aux domaines vérifiables (ex. mathématiques, codage) où la correction est facilement jugée. Appliquer le TTS à l'alignement des préférences est difficile car la qualité d'une réponse est holistique et ne peut être réduite à une simple réponse vérifiable. Les approches actuelles d'alignement au moment de l'inférence imposent souvent des surcoûts de calcul fixes ou nécessitent l'entraînement de têtes de valeur, ce qui limite leur extensibilité et leur efficacité.

Méthodologie

Le papier introduit REAR (REAlignment Reward), un nouveau cadre qui définit l'alignement des préférences au moment de l'inférence comme un problème de réalignement. L'idée centrale est que l'objectif d'entraînement d'un modèle de base peut ne pas satisfaire de manière optimale la préférence spécifique d'un utilisateur lors de l'inférence. REAR traite cela en décomposant la fonction de récompense sous-jacente en deux composantes et en redimensionnant dynamiquement leur importance.

1. Décomposition de la Récompense

Les auteurs formalisent la structure de la récompense d'un LLM pré-entraîné dans le cadre de l'apprentissage par renforcement à entropie maximale (MaxEnt RL). Ils postulent que la récompense $r(s \oplus x_p, a)$ pour un état $s$ (question) et une préférence $x_p$ peut être décomposée en :

• Récompense liée à la question ($r_0$) : La récompense implicite que le modèle de base maximise pour répondre correctement à la question.
• Récompense liée à la préférence ($r_p$) : La composante de récompense capturant la mesure dans laquelle la réponse satisfait la préférence $x_p$.

La récompense composite est modélisée comme $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$, où $\alpha$ est un coefficient d'équilibrage inhérent.

2. Récompense de Réalignement (REAR)

Pour réaligner le modèle pour un utilisateur spécifique, REAR introduit un coefficient contrôlable $\hat{\alpha}$ (ou $\lambda = \hat{\alpha}/\alpha$) pour redimensionner la composante de préférence pendant l'inférence :
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

Crucialement, les auteurs dérivent une formulation efficace en termes de calcul pour cette récompense. En exploitant la relation entre les politiques optimales et les fonctions de valeur dans le cadre MaxEnt RL, ils montrent que $r_{REAR}$ peut être exprimée comme une combinaison linéaire des log-probabilités de jetons (tokens) de la politique :
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
Cette fonction de score dépend uniquement des log-probabilités du modèle de base sous deux contextes : le prompt contenant uniquement la question et le prompt complet incluant les préférences. Cela ne nécessite ni l'entraînement d'un modèle de récompense distinct, ni la génération de jetons supplémentaires pour le scoring.

3. Intégration avec le Passage à l'Échelle au Moment de l'Inférence (TTS)

Le score REAR est intégré dans deux algorithmes de TTS :

• Échantillonnage Best-of-N (BoN) : Générer $N$ réponses et sélectionner celle ayant le score $S_{REAR}$ le plus élevé.
• Recherche en Arbre de Vérificateur Diversifié (DVTS) : Un algorithme de recherche en arbre qui utilise les scores REAR pour guider l'expansion et la sélection des branches, favorisant la diversité tout en optimisant l'alignement des préférences.

Contributions Clés

1. Formalisation du Réalignement : Le papier formalise l'alignement des préférences au moment de l'inférence comme un problème de réalignement, proposant REAR comme une récompense efficace calculée à partir d'un objectif décomposé.
2. Scoring sans Entraînement : REAR est dérivé uniquement des probabilités internes du modèle de base, ne nécessitant aucun entraînement supplémentaire, aucun entraînement de modèle de récompense, ni étapes de génération spéculative.
3. Algorithmes Scalables : Les auteurs développent deux algorithmes de test-time scalables (BoN et DVTS) guidés par REAR qui peuvent améliorer la performance de génération avec des budgets d'échantillonnage variables.
4. Généralisabilité : La méthode montre qu'elle se généralise au-delà de l'alignement explicite des préférences vers le raisonnement mathématique et les tâches visuelles en injectant du texte de préférence spécifique à la tâche.

Résultats Expérimentaux

Les auteurs ont évalué REAR sur plusieurs benchmarks en utilisant les modèles Qwen et Llama :

• Benchmarks d'Alignement de Préférences : Sur PrefEval (couvrant les préférences explicites/implicites et les choix), Multifaceted Bench, et Ping-Pong (jeu de rôle), les méthodes guidées par REAR (tant BoN que DVTS) ont systématiquement surpassé les baselines, incluant le décodage glouton (greedy decoding), Amulet, le Linear Alignment (LA), et le Best-of-N avec des Modèles de Récompense Génératifs (GenRM). Par exemple, DVTS avec REAR a amélioré le décodage glouton de 11,6 % en moyenne sur les métriques.
• Robustesse au Contexte Long : REAR a démontré une robustesse supérieure dans les scénarios de contexte long (jusqu'à 16k tokens) par rapport à GenRM, qui se dégrade considérablement à mesure que la longueur du contexte augmente.
• Raisonnement Mathématique : Appliqué aux tâches mathématiques (MATH500, AIME, AMC) avec un texte de préférence approprié (ex. "Agissez en tant que mathématicien rigoureux"), la sélection guidée par REAR a surpassé le vote majoritaire, montrant des gains de mise à l'échelle constants.
• Tâches Visuelles : Sur MMHal-Bench pour la détection d'hallucinations visuelles, REAR a amélioré les scores globaux et réduit les taux d'hallucination par rapport au décodage gluton et à GenRM.
• Transfert Inter-Modèles : La méthode s'est généralisée à Llama-3.1-8B-Instruct sans réglage par modèle, confirmant que la décomposition de la récompense n'est pas spécifique à la famille Qwen.
• Efficacité : Les méthodes guidées par REAR étaient nettement plus rapides que les baselines nécessitant des modèles de récompense externes ou une optimisation itérative (comme Amulet), car elles évitent de charger des modèles séparés et reposent sur un seul passage direct (forward pass) pour le scoring.

Signification et Revendications

Le papier affirme que REAR fournit une solution scalable et générale pour aligner les LLM avec les préférences des utilisateurs au moment de l'inférence sans le surcoût de l'entraînement ou de modèles de récompense externes. En décomposant la récompense en composantes de question et de préférence, REAR permet un contrôle dynamique du compromis entre répondre à la question et respecter les préférences.

Les auteurs soulignent que REAR permet au passage à l'échelle au moment de l'inférence de s'étendre aux domaines subjectifs et ouverts (comme l'alignement des préférences et le jeu de rôle) où les vérificateurs traditionnels échouent. Ils notent que bien qu'un $\lambda = 20$ par défaut fonctionne de manière robuste sur de nombreuses tâches, la méthode permet de régler la force du réalignement. Ce travail se positionne comme une étape vers la rendre l'alignement au moment de l'inférence plus accessible et efficace, comblant le fossé entre le succès du TTS dans les domaines vérifiables et la complexité de l'alignement des préférences subjectives.

Limitations Reconnues : Les auteurs notent que REAR expose un hyperparamètre réglable ($\lambda$) à l'utilisateur, ce qui peut nécessiter de petites analyses de validation pour des tâches ayant des distributions de préférences significativement différentes. De plus, bien que le TTS soit plus léger que le fine-tuning, identifier le budget de calcul optimal reste une direction pour les travaux futurs.