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

Arushi Rai; Qiang Zhang; Hanqing Zeng; Yunkai Zhang; Dipesh Tamboli; Xiangjun Fan; Zhuokai Zhao; Lizhu Zhang

cs.CLMar 19, 2026

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

Authors: Arushi Rai, Qiang Zhang, Hanqing Zeng, Yunkai Zhang, Dipesh Tamboli, Xiangjun Fan, Zhuokai Zhao, Lizhu Zhang

Organizations: Meta · University of Pittsburgh · University of California, Berkeley

Abstract

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.

Figures & tables

Figure 1: Performance on MATH500 (accuracy) and AlpacaEval (length-controlled win rate) for the state-of-the-art test-time alignment approach (GenARM) under different mixing coefficients $\alpha\in[0,1]$ . An $\alpha=0$ corresponds to decoding solely from the base model, while $\alpha=1$ uses only the reward model.

Figure 2: Learnable token-level router design. At each LLM decoding step $t$ , the base and reward models produce logits $z^{\text{base}}_{t}$ and $z^{\text{reward}}_{t}$ . The logits are passed as input to Feature Concat , which either (i) concatenates logits, or (ii) concatenates logits plus learnable token-index embeddings (as discussed in § 3.3 ). The router consumes the concatenated feature and outputs a routing weight $\alpha_{t}\in(0,1)$ . The guided distribution $(1-\alpha_{t})\,z^{\text{base}}_{t}+\alpha_{t}\,z^{\text{reward}}_{t}$ is then used for sampling next token. This design makes the router portable across base model scales and families.

Llama-3.1-8B Instruct Base + DeepSeek Distill Llama-8B Reward
Method	MATH500	MedXpertQA	AlpacaEval	Average
UniR Kim et al. (2025)	51.2	–	–	–
Llama-3.1-8B Instruct (Base)	32.0	13.0	17.3	20.8
DeepSeek Llama-8B (Reward)	47.2	11.6	2.8	20.5
GenARM Xu et al. (2024) *	49.2	11.2	10.8	23.7
TARo (Ours)	54.4	13.2	20.8	29.5

Table 1: Performance across reasoning (MATH500, MedXpertQA) and instruction-following (AlpacaEval) benchmarks. Reward models are trained on Math-StepDPO-10K described in § 3.2 . Results for UniR are taken from its original paper, * indicates evaluated with our Math-StepDPO-10K trained reasoning reward model.

Figure 3: Weak-to-strong generalization of learned router on reasoning. Learned router and reasoning reward model are not retrained for this scale.

High $\alpha$ Tokens	Low $\alpha$ Tokens
rightarrow	period
Exp	-coordinate
Step	profit
cases	factors
critical	shaded
Part	squares

Table 2: Tokens in top 0.1% and bottom 0.1% of

\alpha

from generated responses to MATH500 responses; filtered out tokens with less than 50 occurrences and shorter than 2 characters (difficult to interpret). Tokens on the left show strong reward model influence, reflecting mathematical operators, formatting, and reasoning scaffolds (e.g. “cases”, “Step”). Tokens on the right are dominated by the base model, largely consisting of vocabulary from the problem context. Base Model = Qwen2.5-3B; Reward Model = Qwen2.5-0.5B.

Router Input Feature	MATH500 Accuracy (%)
Reward hidden state only	51.2
Base + reward hidden states	49.6
Reward logits only	52.6
Base + reward logits	54.4

Table 3: Ablation on router input features for Llama-3.1-8B Instruct with DeepSeek Distill Llama-8B reward. Results are reported on MATH500. See Appendix C for results on all three benchmarks.

Model	Variant	MATH	MedX
Qwen	$K$ = 100	60.2	7.7
Qwen	$K$ = 1,000	60.2	7.3
Qwen	$K$ = 10,000	61.8	7.4
Qwen	$K$ = 100 + $\lambda_{\text{entropy}}=0.5$	59.8	7.3
Qwen	$K$ = 1,000 + $\lambda_{\text{entropy}}=0.5$	64.8	9.8
Llama	$K$ = 1,000 + $\lambda_{\text{entropy}}=0.5$	50.6	14.5

Table 4: Ablation on Top-

K

, entropy regularization, and full-logit router designs. Accuracy (%) is reported on MATH500 (MATH) and MedXpertQA (MedX).

Question: A worker receives an annual wage of $\$ 20{,}000 $, which he always deposits into a savings account at the end of the year. By the end of the third year (when he makes the third deposit), he wants to have at least$ $66,200$ in the account to finance the purchase of a house. What is the minimal compound interest rate that the savings account must provide? Express your answer as a percentage, but do not include the percent sign.
Ground Truth: 10
	Base Model	Reward Model	Ours
Output (abridged)	Let r denote the interest rate. The worker will have a total of $20,000(1+r)^{3}$ in the account by the end of the third year. This is at least $66,200$ , so we have the inequality: $\begin{aligned} \\ 20{,}000(1+r)^{3}&\geq 66{,}200\\ (1+r)^{3}&\geq 3.31\\ 1+r&\geq\sqrt[3]{3.31}\\ r&\geq\sqrt[3]{3.31}-1\\ r&\geq 0.322\\ \\ \end{aligned}$ Thus, the smallest possible value of $r$ is 0.322, or 32.2. Final Answer: $\boxed{32.2}$ . I hope it is correct.	Let the annual interest rate be $r$ . The worker makes three deposits of $20,000 at the end of each year. After the first year:$ 20,000(1+r) $. After the second year:$ 20,000(1+r)^{2}+20,000(1+r) $. After the third year:$ 20,000(1+r)^{3}+20,000(1+r)+20,000(1+r) $. Wait, this seems incorrect$ \dots $Wait, no$ \dots $The correct amount$ \dots$ [… stuck in confusion]	Step 1: The worker deposits $20,000 each year, wants at least$ 66,200 after three years. Step 2: Total after 3 years is $20,000(1+r)^{3}+20,000(1+r)^{2}+20,000(1+r)$ . Step 3: Divide by 20,000: $(1+r)^{3}+(1+r)^{2}+(1+r)\geq 3.31$ . Step 4: Let $x=1+r$ , then $x^{3}+x^{2}+x\geq 3.31$ . Step 5: Solving gives $x\geq 1.1$ , so $r\geq 0.1$ . Final Answer: $\boxed{10}$

Table 5: Qualitative example of base (Llama-3.1 8B Instruct), reward and our model outputs (abridged to fit) for a MATH500 question. The base model oversimplifies, while the reward model understands the right plan but gets confused. Ours solves correctly by engaging in reasoning steps.

Model	Method	TPS	GPU Mem. (%)	QPS
Qwen2.5-3B	GenARM	85.6	4.29	0.0651
Qwen2.5-3B	TARo (Top-K)	85.9	5.22	0.0448
Llama3.1-8B	GenARM	80.7	6.93	0.0141
Llama3.1-8B	TARo (Full)	32.5	10.45	0.0161

Table 6: Throughput analysis on MATH500. TPS = tokens per second; QPS = queries per second. Note that these results reflect unoptimized implementations; techniques such as speculative decoding and fully sharded data-parallel inference have not been applied.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

	Larger Base	(Acc %)	(LC WR %)	(Acc %)	Avg.
Router Inputs	Scalable to	MATH500	AlpacaEval	MedXpertQA
Base + reward hidden states	✗	49.6	18.7	13.7	27.3
Reward hidden state only	✓	51.2	15.7	12.6	26.5
Reward logits only	✓	52.6	19.7	14.4	28.9
Base + reward logits	✓	54.4	20.8	13.2	29.5

Appendix

Table 7: Ablation of router input configurations on Llama3.1-8B across all three benchmarks. Scalable to Larger Base indicates whether the router can be applied to a larger base model without retraining. The greyed row is excluded from the scalable comparison. Among scalable designs, base and reward logits on average outperforms using hidden states or logits from the reward model alone. Bold indicates best result per column across all rows.

Method	MATH500	AlpacaEval	MedXpertQA	Est. TFLOPs
	(Acc %)	(LC WR %)	(Acc %)
Majority Voting ( $N$ =8)	56.4	—	10.6	113.6
GenARM	56.4	18.9	9.8	28.4
TARo	64.8	21.6	9.8	28.4

Appendix

Table 8: Comparison with majority voting on Qwen2.5-3B. TFLOPs are estimated for a single forward pass with a 2048-token context. Majority voting is not applicable to AlpacaEval (—) due to its open-ended generation format. Bold indicates best result per column.

Routing Granularity	MATH500	AlpacaEval	MedXpertQA
	(Acc %)	(LC WR %)	(Acc %)
Prompt-level	33.2	22.1	13.3
Token-level	49.6	18.7	13.7

Appendix

Table 9: Comparison of prompt-level versus token-level routing with hidden state inputs on Llama3.1-8B. Token-level routing significantly outperforms prompt-level routing on mathematical reasoning, validating the design choice of adaptive token-level

\alpha

prediction.

Inference-Time Language Model Alignment

Inference-time language model alignment steers or selects outputs during generation to better meet safety, preference, and task goals without retraining.

See all 78 papers

Training-Free Test-Time Contrastive Learning for Large Language Models
Apr 15, 2026Kaiwen Zheng, Kai Zhou, Jinwu Hu +3Test-Time Adaptation Language Model Robustness
South China University of Technology · Pazhou Laboratory
Attention-Aligned Reasoning for Large Language Models
Oct 3, 2025Hongxiang Zhang, Yuan Tian, Tianyi ZhangLLM Reasoning Structured Reasoning in Language Models
Purdue University
Route-and-Generate: KL-Regularized Router Improvement for Reasoning Routing in LLMs
Oct 21, 2026Tien-Phat Nguyen, Truong Nguyen, Thanh-Thien Le +4Adaptive Reasoning in Language Models Efficient Language Model Reasoning
Hanoi University of Science and Technology, Vietnam · Monash University, Australia · New York University, USA +1

Apr 15, 2026J/K move · Enter open · S save

Training-Free Test-Time Contrastive Learning for Large Language Models

Kaiwen Zheng, Kai Zhou, Jinwu Hu, Te Gu+2

South China University of Technology · Pazhou Laboratory

Test-Time Adaptation Language Model Robustness Domain Adaptation for LLMs Contrastive Learning

Open paper AskPDF arXiv

Large language models (LLMs) demonstrate strong reasoning capabilities, but their performance often degrades under distribution shift. Existing test-time adaptation (TTA) methods rely on gradient-based updates that require white-box access and need substantial overhead, while training-free alternatives are either static or depend on external guidance. In this paper, we propose Training-Free Test-Time Contrastive Learning TF-TTCL, a training-free adaptation framework that enables a frozen LLM to improve online by distilling supervision from its own inference experiences. Specifically, TF-TTCL implements a dynamic "Explore-Reflect-Steer" loop through three core modules: 1) Semantic Query Augmentation first diversifies problem views via multi-agent role-playing to generate different reasoning trajectories; 2) Contrastive Experience Distillation then captures the semantic gap between superior and inferior trajectories, distilling them into explicit textual rules; and 3) Contextual Rule Retrieval finally activates these stored rules during inference to dynamically steer the frozen LLM toward robust reasoning patterns while avoiding observed errors. Extensive experiments on closed-ended reasoning tasks and open-ended evaluation tasks demonstrate that TF-TTCL consistently outperforms strong zero-shot baselines and representative TTA methods under online evaluation. Code is available at https://github.com/KevinSCUTer/TF-TTCL.

Abstract

Inference-Time Language Model Alignment

Explore similar work

Training-Free Test-Time Contrastive Learning for Large Language Models

Attention-Aligned Reasoning for Large Language Models

Route-and-Generate: KL-Regularized Router Improvement for Reasoning Routing in LLMs