GenARM : Génération guidée par récompense avec modèle de récompense autorégressif pour l'alignement au moment du test

GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment

October 10, 2024

Auteurs: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh

cs.AI

Résumé

Les grands modèles de langage (LLMs) présentent des capacités impressionnantes mais nécessitent un alignement soigneux avec les préférences humaines. Les méthodes traditionnelles d'entraînement finetunent les LLMs en utilisant des ensembles de données de préférences humaines, mais entraînent des coûts importants et nécessitent un entraînement répété pour gérer des préférences utilisateur diverses. Les méthodes d'alignement au moment du test abordent cette problématique en utilisant des modèles de récompense (RMs) pour guider les LLMs figés sans re-entraînement. Cependant, les approches existantes au moment du test reposent sur des RMs au niveau de la trajectoire qui sont conçus pour évaluer des réponses complètes, les rendant inadaptés à la génération de texte autorégressive qui nécessite le calcul de récompenses de jeton suivant à partir de réponses partielles. Pour résoudre cela, nous introduisons GenARM, une approche d'alignement au moment du test qui exploite le Modèle de Récompense Autorégressif - une nouvelle paramétrisation de récompense conçue pour prédire les récompenses de jeton suivant de manière efficace et efficiente pour la génération autorégressive. Théoriquement, nous démontrons que cette paramétrisation peut guider de manière prouvable les LLMs figés vers toute distribution réalisable par des RMs traditionnels dans le cadre de l'apprentissage par renforcement régularisé par KL. Les résultats expérimentaux montrent que GenARM surpasse significativement les références d'alignement au moment du test précédentes et égale les performances des méthodes d'entraînement. De plus, GenARM permet un guidage efficace de faible à fort, alignant des LLMs plus grands avec des RMs plus petits sans les coûts élevés de l'entraînement de modèles plus grands. De plus, GenARM prend en charge l'alignement multi-objectif, permettant des compromis en temps réel entre les dimensions de préférence et répondant à des préférences utilisateur diverses sans re-entraînement.

English

Large Language Models (LLMs) exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining.