TARo: Token-level Adaptive Routing for LLM Test-time Alignment

📄 arXiv: 2603.18411v1 📥 PDF

作者: Arushi Rai, Qiang Zhang, Hanqing Zeng, Yunkai Zhang, Dipesh Tamboli, Xiangjun Fan, Zhuokai Zhao

分类: cs.CL, cs.AI, cs.LG

发布日期: 2026-03-19

💡 一句话要点

提出TARo：一种Token级自适应路由方法，用于LLM测试时对齐，提升推理能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 测试时对齐 自适应路由 推理能力 奖励模型

📋 核心要点

  1. 现有LLM推理能力强，但需昂贵的后训练。测试时对齐方法虽轻量，但主要针对偏好对齐，推理能力提升有限。
  2. TARo的核心在于训练奖励模型捕捉逻辑一致性，并使用token级路由器自适应地控制奖励模型对基础模型的指导。
  3. 实验表明，TARo在推理、临床推理和指令遵循方面均有显著提升，且能推广到不同大小的模型，无需重新训练。

📝 摘要（中文）

大型语言模型（LLM）展现出强大的推理能力，但通常需要昂贵的后训练才能达到高性能。最近的测试时对齐方法提供了一种轻量级的替代方案，但主要针对偏好对齐而非推理。为了弥合这一差距，我们提出了Token级自适应路由（TARo），它完全在推理时引导冻结的LLM进行结构化推理。具体来说，我们首先在逐步数学轨迹上训练奖励模型，以捕获细粒度的逻辑一致性信号，然后引入一个可学习的token级路由器，该路由器自动控制奖励模型对基础模型的指导。大量实验表明，TARo显著提高了推理性能，比基础模型提高了高达+22.4%，比现有的token级测试时对齐方法提高了+8.4%，同时还提高了分布外临床推理（MedXpertQA）和指令遵循（AlpacaEval）。此外，TARo还可以从小型主干网络推广到大型主干网络，而无需重新训练，从而将测试时对齐从偏好优化扩展到稳健的跨领域推理。

🔬 方法详解

问题定义：现有的大型语言模型虽然具备强大的推理能力，但是往往需要耗费大量的计算资源进行后训练才能达到理想的性能。而现有的测试时对齐方法虽然计算开销较小，但是主要集中在偏好对齐方面，对于提升模型的推理能力效果有限。因此，如何利用轻量级的测试时对齐方法来提升LLM的推理能力是一个亟待解决的问题。

核心思路：TARo的核心思路是利用奖励模型来指导LLM的推理过程，并引入token级别的自适应路由机制，使得奖励模型的指导能够根据不同的token进行动态调整。这样做的目的是为了让模型在推理过程中更加注重逻辑一致性，从而提升推理的准确性。通过token级别的路由，可以更精细地控制奖励模型的干预，避免过度干预导致模型性能下降。

技术框架：TARo的整体框架主要包含三个部分：基础LLM、奖励模型和token级路由器。首先，使用逐步数学轨迹训练奖励模型，使其能够评估推理过程中的逻辑一致性。然后，token级路由器根据当前token的状态，决定是否以及如何将奖励模型的输出融入到基础LLM的推理过程中。在推理时，基础LLM生成token，奖励模型评估该token的逻辑一致性，token级路由器根据评估结果调整基础LLM的输出，最终生成下一个token。

关键创新：TARo的关键创新在于引入了token级别的自适应路由机制。与传统的测试时对齐方法不同，TARo不是简单地将奖励模型的输出直接加到基础LLM的输出上，而是通过一个可学习的路由器来控制奖励模型的干预程度。这种token级别的自适应性使得模型能够更加灵活地应对不同的推理场景，从而提升推理的准确性和鲁棒性。

关键设计：TARo的关键设计包括奖励模型的训练方式和token级路由器的结构。奖励模型通过在逐步数学轨迹上进行训练，学习评估推理过程中的逻辑一致性。Token级路由器通常采用一个小型神经网络，输入包括当前token的embedding和奖励模型的输出，输出是一个权重，用于控制奖励模型的干预程度。损失函数的设计需要平衡推理的准确性和鲁棒性，避免过度依赖奖励模型导致模型泛化能力下降。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TARo在多个数据集上取得了显著的性能提升。在数学推理任务上，TARo比基础模型提高了高达+22.4%，比现有的token级测试时对齐方法提高了+8.4%。此外，TARo还在分布外临床推理（MedXpertQA）和指令遵循（AlpacaEval）方面取得了显著的提升，表明其具有很强的泛化能力。这些实验结果充分证明了TARo的有效性和优越性。

🎯 应用场景

TARo具有广泛的应用前景，可以应用于数学推理、临床诊断、代码生成等多个领域。通过在测试时对齐，可以显著提升LLM在这些领域的性能，而无需进行昂贵的后训练。此外，TARo的跨领域泛化能力使得其可以应用于新的领域，具有很高的实际价值。未来，可以进一步研究如何将TARo应用于更复杂的推理场景，并探索更有效的奖励模型训练方法。

📄 摘要（原文）

Large language models (LLMs) exhibit strong reasoning capabilities but typically require expensive post-training to reach high performance. Recent test-time alignment methods offer a lightweight alternative, but have been explored mainly for preference alignment rather than reasoning. To bridge this gap, we propose, Token-level Adaptive Routing (TARo), which steers frozen LLMs toward structured reasoning entirely at inference time. Specifically, we first train reward models on step-wise mathematical traces to capture fine-grained logical consistency signals, then introduce a learnable token-level router that automatically controls the guidance of the reward model to the base model. Extensive experiments show that TARo significantly improves reasoning performance by up to +22.4% over base model and +8.4% over existing token-level test-time alignment methods, while also boosting out-of-distribution clinical reasoning (MedXpertQA) and instruction following (AlpacaEval). Furthermore, TARo also generalizes from small to large backbones without retraining, extending test-time alignment from preference optimization to robust, cross-domain reasoning.