Sample, Don't Search: Rethinking Test-Time Alignment for Language Models
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.
67f54949f24949194dc4d771
Gonçalo Faria
claimed_verified
2025-04-08T20:06:31.707Z
63062d2641bf1fbadaec2e9d
https://cdn-avatars.huggingface.co/v1/production/uploads/63062d2641bf1fbadaec2e9d/_EjwM6W0X7snDa7cj3ikw.png
Gonçalo Faria
user
graf
67f54949f24949194dc4d772
Noah A. Smith
Sample, Don't Search: Rethinking Test-Time Alignment for Language Models
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.
QAlign improves test-time language model performance by aligning outputs without modifying the model, using advancements in Markov chain Monte Carlo for text generation and task-specific reward models.
test-time computation language model performance reward model QAlign test-time search methods Markov chain Monte Carlo text generation GSM8K GSM-Symbolic Tulu 3 preference dataset direct preference optimization (DPO) best-of-n majority voting weighted majority voting IFEval MMLU-Redux TruthfulQA