Campiona, Non Cercare: Ripensare l'Allineamento al Momento del Test per i Modelli Linguistici

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

April 4, 2025

Autori: Gonçalo Faria, Noah A. Smith

cs.AI

Abstract

L'aumento del calcolo al momento del test è emerso come una direzione promettente per migliorare le prestazioni dei modelli linguistici, in particolare negli scenari in cui il fine-tuning del modello è impraticabile o impossibile a causa di vincoli computazionali o di pesi del modello privati. Tuttavia, i metodi esistenti di ricerca al momento del test che utilizzano un modello di ricompensa (RM) spesso peggiorano in qualità man mano che il calcolo aumenta, a causa dell'ottimizzazione eccessiva di quelli che sono intrinsecamente proxy di ricompensa imperfetti. Introduciamo QAlign, un nuovo approccio di allineamento al momento del test. Man mano che aumentiamo il calcolo al momento del test, QAlign converge a campionare dalla distribuzione allineata ottimale per ogni singolo prompt. Adottando i recenti progressi nel Markov chain Monte Carlo per la generazione di testo, il nostro metodo consente output meglio allineati senza modificare il modello sottostante o addirittura senza richiedere l'accesso ai logit. Dimostriamo l'efficacia di QAlign su benchmark di ragionamento matematico (GSM8K e GSM-Symbolic) utilizzando un RM specifico per il compito, mostrando miglioramenti costanti rispetto ai metodi esistenti di calcolo al momento del test come best-of-n e voto a maggioranza. Inoltre, quando applicato con RM più realistici addestrati sul dataset di preferenze Tulu 3, QAlign supera l'ottimizzazione diretta delle preferenze (DPO), best-of-n, voto a maggioranza e voto a maggioranza ponderato su una gamma diversificata di dataset (GSM8K, MATH500, IFEval, MMLU-Redux e TruthfulQA). Una soluzione pratica per allineare i modelli linguistici al momento del test utilizzando calcoli aggiuntivi senza degradazione, il nostro approccio espande i limiti delle capacità che possono essere ottenute da modelli linguistici pronti all'uso senza ulteriore addestramento.

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.