Test-time Alignment
[Paper Note] TARo: Token-level Adaptive Routing for LLM Test-time Alignment, Arushi Rai+, ACL'26 Findings, 2026.03
Paper/Blog Link My Issue
#LanguageModel #Alignment #ACL #Decoding #Findings #Routing #KeyPoint Notes #Author Thread/Post Issue Date: 2026-04-07 GPT Summary- 推論時に固定されたLLMsを用いて、トークンレベル適応ルーティング（TARo）を提案。報酬モデルにより数学的推論の一貫性信号を捉え、ルーターが基盤モデルを自動制御。TARoは推論性能を最大+22.4%向上させ、分布外の臨床推論や指示遵守を改善。再訓練なしでの一般化も可能で、堅牢な推論を実現。 Comment
元ポスト:
巨大なベースモデル全体を特定ドメインに適用するためにpost-trainingするのは大変なので、代わりに小規模なdomain-expertなRewardモデルを学習し（今回は数学のstep-wiseにlogicが正しいことをpreferenceとして与えるような学習方法を採用したようである; 3.2節）、各decoding step tにおいて、ベースモデルとRewardモデルのトークンのlogitを線形補完することで、出力トークンをガイドする。logitの線形補完において、固定されたスカラー値（e.g., 0.5など。GenARMという手法らしい）を用いる研究などが先行研究ではあるが、これはベースモデルの特定タスクにおいてベースモデルの性能を劣化させるので、本研究ではdecoding step t時点で出力されたベースモデル、Rewardモデルのlogitを入力として、FFNによって線形補完の重みα_tをdecoding step tごとに決定する（α_tを決定するネットワークをRouterと呼ぶ）。FFNは2種類のvariantがあり、双方のlogitをconcatしたものを入力するものと、top-kをサンプリングし、kごとにindexに基づいたembeddingをconcatして入力する方法の二種類がある（3.3節）。
結果としては、GenARMと比較して提案手法は有効ではあるが、ベースモデルとrewardモデルの組み合わせによっては、baseモデルよりも性能が悪化するということもありそうに見える。
またRouterはベースモデルのサイズを大きくしても、性能が転移するので再学習が不要である。