Amostre, Não Pesquise: Repensando o Alinhamento em Tempo de Teste para Modelos de Linguagem

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

April 4, 2025

Autores: Gonçalo Faria, Noah A. Smith

cs.AI

Resumo

O aumento da computação durante o teste tem surgido como uma direção promissora para melhorar o desempenho dos modelos de linguagem, especialmente em cenários onde o ajuste fino do modelo é impraticável ou impossível devido a restrições computacionais ou a pesos de modelos privados. No entanto, os métodos existentes de busca durante o teste que utilizam um modelo de recompensa (RM) frequentemente apresentam degradação na qualidade à medida que a computação escala, devido à superotimização de proxies de recompensa que são inerentemente imperfeitos. Introduzimos o QAlign, uma nova abordagem de alinhamento durante o teste. À medida que escalamos a computação durante o teste, o QAlign converge para a amostragem da distribuição alinhada ótima para cada prompt individual. Ao adotar avanços recentes em cadeias de Markov Monte Carlo para geração de texto, nosso método permite saídas melhor alinhadas sem modificar o modelo subjacente ou mesmo exigir acesso aos logits. Demonstramos a eficácia do QAlign em benchmarks de raciocínio matemático (GSM8K e GSM-Symbolic) usando um RM específico para a tarefa, mostrando melhorias consistentes em relação a métodos existentes de computação durante o teste, como best-of-n e votação majoritária. Além disso, quando aplicado com RMs mais realistas treinados no conjunto de dados de preferências Tulu 3, o QAlign supera a otimização direta de preferências (DPO), best-of-n, votação majoritária e votação majoritária ponderada em uma variedade de conjuntos de dados (GSM8K, MATH500, IFEval, MMLU-Redux e TruthfulQA). Uma solução prática para alinhar modelos de linguagem durante o teste usando computação adicional sem degradação, nossa abordagem expande os limites da capacidade que pode ser obtida de modelos de linguagem prontos para uso sem treinamento adicional.

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.