サンプル、検索しない：言語モデルのためのテスト時アラインメントの再考

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は、テスト時の計算量をスケールさせることで、各プロンプトに対する最適なアライメント分布からのサンプリングに収束します。テキスト生成のためのマルコフ連鎖モンテカルロ法の最新の進展を採用することで、基礎となるモデルを変更したり、ロジットへのアクセスを必要とせずに、より適切にアライメントされた出力を実現します。タスク固有のRMを用いた数学的推論ベンチマーク（GSM8KおよびGSM-Symbolic）において、QAlignの有効性を実証し、既存のテスト時計算手法（best-of-nや多数決）を一貫して上回る改善を示しました。さらに、Tulu 3選好データセットでトレーニングされたより現実的なRMを適用した場合、QAlignは多様なデータセット（GSM8K、MATH500、IFEval、MMLU-Redux、TruthfulQA）において、直接選好最適化（DPO）、best-of-n、多数決、加重多数決を上回る性能を示しました。追加の計算量を使用してテスト時に言語モデルをアライメントする実用的なソリューションとして、本手法は、さらなるトレーニングなしで既存の言語モデルから得られる能力の限界を拡張します。

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.