Échantillonner, ne pas chercher : Repenser l'alignement en phase de test pour les modèles de langage

Sample, Don't Search: Rethinking Test-Time Alignment for Language Models

April 4, 2025

Auteurs: Gonçalo Faria, Noah A. Smith

cs.AI

Résumé

L'augmentation du calcul au moment du test s'est imposée comme une direction prometteuse pour améliorer les performances des modèles de langage, en particulier dans les scénarios où le fine-tuning du modèle est impraticable ou impossible en raison de contraintes computationnelles ou de poids de modèle privés. Cependant, les méthodes existantes de recherche au moment du test utilisant un modèle de récompense (RM) voient souvent leur qualité se dégrader à mesure que le calcul augmente, en raison de la sur-optimisation de ce qui sont intrinsèquement des proxies de récompense imparfaits. Nous introduisons QAlign, une nouvelle approche d'alignement au moment du test. À mesure que nous augmentons le calcul au moment du test, QAlign converge vers un échantillonnage à partir de la distribution alignée optimale pour chaque prompt individuel. En adoptant les avancées récentes en chaînes de Markov Monte Carlo pour la génération de texte, notre méthode permet d'obtenir des sorties mieux alignées sans modifier le modèle sous-jacent ni même nécessiter un accès aux logits. Nous démontrons l'efficacité de QAlign sur des benchmarks de raisonnement mathématique (GSM8K et GSM-Symbolic) en utilisant un RM spécifique à la tâche, montrant des améliorations constantes par rapport aux méthodes existantes de calcul au moment du test comme best-of-n et le vote majoritaire. De plus, lorsqu'il est appliqué avec des RM plus réalistes entraînés sur le jeu de données de préférences Tulu 3, QAlign surpasse l'optimisation directe des préférences (DPO), best-of-n, le vote majoritaire et le vote majoritaire pondéré sur une gamme variée de jeux de données (GSM8K, MATH500, IFEval, MMLU-Redux et TruthfulQA). Une solution pratique pour aligner les modèles de langage au moment du test en utilisant un calcul supplémentaire sans dégradation, notre approche repousse les limites des capacités que l'on peut obtenir à partir de modèles de langage prêts à l'emploi sans formation supplémentaire.

English

Increasing test-time computation has emerged as a promising direction for improving language model performance, particularly in scenarios where model finetuning is impractical or impossible due to computational constraints or private model weights. However, existing test-time search methods using a reward model (RM) often degrade in quality as compute scales, due to the over-optimization of what are inherently imperfect reward proxies. We introduce QAlign, a new test-time alignment approach. As we scale test-time compute, QAlign converges to sampling from the optimal aligned distribution for each individual prompt. By adopting recent advances in Markov chain Monte Carlo for text generation, our method enables better-aligned outputs without modifying the underlying model or even requiring logit access. We demonstrate the effectiveness of QAlign on mathematical reasoning benchmarks (GSM8K and GSM-Symbolic) using a task-specific RM, showing consistent improvements over existing test-time compute methods like best-of-n and majority voting. Furthermore, when applied with more realistic RMs trained on the Tulu 3 preference dataset, QAlign outperforms direct preference optimization (DPO), best-of-n, majority voting, and weighted majority voting on a diverse range of datasets (GSM8K, MATH500, IFEval, MMLU-Redux, and TruthfulQA). A practical solution to aligning language models at test time using additional computation without degradation, our approach expands the limits of the capability that can be obtained from off-the-shelf language models without further training.