Образец, не ищи: Переосмысление адаптации языковых моделей на этапе тестирования

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

April 4, 2025

Авторы: Gonçalo Faria, Noah A. Smith

cs.AI

Аннотация

Увеличение вычислительных ресурсов на этапе тестирования стало перспективным направлением для повышения производительности языковых моделей, особенно в сценариях, где тонкая настройка модели невозможна или нецелесообразна из-за вычислительных ограничений или приватности весов модели. Однако существующие методы поиска на этапе тестирования с использованием модели вознаграждения (RM) часто теряют в качестве по мере масштабирования вычислений из-за чрезмерной оптимизации несовершенных прокси-метрик вознаграждения. Мы представляем QAlign — новый подход к согласованию на этапе тестирования. По мере увеличения вычислительных ресурсов QAlign сходится к выборке из оптимального согласованного распределения для каждого отдельного запроса. Используя последние достижения в методах Монте-Карло с марковскими цепями для генерации текста, наш метод позволяет получать более согласованные выходные данные без изменения базовой модели или даже доступа к логитам. Мы демонстрируем эффективность QAlign на задачах математического рассуждения (GSM8K и GSM-Symbolic) с использованием RM, специфичной для задачи, показывая устойчивые улучшения по сравнению с существующими методами, такими как best-of-n и мажоритарное голосование. Более того, при применении с более реалистичными RM, обученными на наборе данных предпочтений Tulu 3, QAlign превосходит прямое оптимизацию предпочтений (DPO), best-of-n, мажоритарное голосование и взвешенное мажоритарное голосование на разнообразных наборах данных (GSM8K, MATH500, IFEval, MMLU-Redux и TruthfulQA). Наш подход предлагает практическое решение для согласования языковых моделей на этапе тестирования с использованием дополнительных вычислений без ухудшения качества, расширяя пределы возможностей, которые можно получить из готовых языковых моделей без дополнительного обучения.

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.