Sampleer, Zoek Niet: Een Heroverweging van Testtijd-Uitlijning voor Taalmodellen

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

April 4, 2025

Auteurs: Gonçalo Faria, Noah A. Smith

cs.AI

Samenvatting

Het vergroten van de rekentijd tijdens het testen is naar voren gekomen als een veelbelovende richting voor het verbeteren van de prestaties van taalmodelen, met name in scenario's waar het finetunen van het model onpraktisch of onmogelijk is vanwege rekenbeperkingen of private modelgewichten. Bestaande zoekmethoden tijdens het testen die gebruikmaken van een beloningsmodel (RM) gaan echter vaak in kwaliteit achteruit naarmate de rekencapaciteit toeneemt, vanwege de overoptimalisatie van inherent onvolmaakte beloningsproxies. Wij introduceren QAlign, een nieuwe aanpak voor uitlijning tijdens het testen. Naarmate we de rekentijd tijdens het testen opschalen, convergeert QAlign naar het bemonsteren van de optimaal uitgelijnde verdeling voor elke individuele prompt. Door recente vooruitgang in Markov-keten Monte Carlo voor tekstgeneratie te benutten, maakt onze methode beter uitgelijnde uitvoer mogelijk zonder het onderliggende model aan te passen of zelfs toegang tot logits te vereisen. We demonstreren de effectiviteit van QAlign op benchmarks voor wiskundig redeneren (GSM8K en GSM-Symbolic) met behulp van een taakspecifiek RM, en laten consistente verbeteringen zien ten opzichte van bestaande methoden voor rekentijd tijdens het testen, zoals best-of-n en meerderheidsstemming. Bovendien, wanneer toegepast met meer realistische RM's getraind op de Tulu 3 voorkeursdataset, presteert QAlign beter dan directe voorkeursoptimalisatie (DPO), best-of-n, meerderheidsstemming en gewogen meerderheidsstemming op een diverse reeks datasets (GSM8K, MATH500, IFEval, MMLU-Redux en TruthfulQA). Als een praktische oplossing voor het uitlijnen van taalmodelen tijdens het testen met behulp van extra rekenkracht zonder kwaliteitsverlies, breidt onze aanpak de grenzen uit van de mogelijkheden die kunnen worden behaald met kant-en-klare taalmodelen zonder verdere training.

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.