Elamur Research⌕
☰
home / papers / 2026 / TARo: Token-level Adaptive Routing for LLM Test-time Alignment

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

Paper · 2026 · Arushi Rai, Qiang Zhang, Hanqing Zeng et al.

Summary

Test-time alignment steers a frozen LLM during decoding with a smaller reward model, avoiding expensive post-training such as RLVR. It has mostly targeted preference alignment, and methods like GenARM mix base and reward distributions with a fixed weight α. That weight must be tuned by hand, and its best value shifts across tasks, domains and model scales. TARo applies this approach to reasoning in two steps. First, it trains a reasoning-aware reward model on step-wise math preference pairs (Math-StepDPO-10K), using a token-level reward formulation. Second, it adds a lightweight learnable router. At each decoding step, the router reads the output logits of both models, not their hidden states, and outputs a mixing weight α_t. Because the router depends only on logits, it can be moved to other backbones. On MATH500, TARo beats the base model by up to 22.4% and GenARM by up to 8.4%. It also helps on out-of-domain MedXpertQA and AlpacaEval. A router trained with small models transfers to larger ones without retraining. This makes adaptive, inexpensive reasoning control possible for frozen LLMs.

Key points

  • The reward model is trained with a Bradley–Terry-style preference loss on Math-StepDPO-10K. Each pair shares a correct reasoning prefix and then splits into one valid and one wrong next step.
  • The router outputs α_t ∈ (0,1) and combines logits as (1−α_t)·z_base + α_t·z_reward. It has two input variants: full concatenated logits, or Top-K logits plus learnable token-index embeddings.
  • Evaluation covers MATH500 (math), MedXpertQA (out-of-domain clinical reasoning) and AlpacaEval (instruction following), with backbones including Qwen2.5-3B and Llama-3.1-8B-Instruct.
  • Inference cost was measured on 8×H100:
    • The Top-K router keeps throughput close to GenARM's.
    • The full-logit router slows decoding for large-vocabulary models such as Llama3.1-8B, though per-query overhead stays small.
  • The authors list several limitations:
    • Results depend on reward-model quality.
    • Testing is limited to math, clinical reasoning and AlpacaEval.
    • The router adds inference compute.
    • The best router configuration differs between model families.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of California, BerkeleyBerkeley · United States175 works · 227,440 citations🎓University of PittsburghPittsburgh · United States59 works · 22,324 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Arushi Rai, Qiang Zhang, Hanqing Zeng, Yunkai Zhang, Dipesh Tamboli, Xiangjun Fan, Zhuokai Zhao and Lizhu Zhang. 2026. TARo: Token-level Adaptive Routing for LLM Test-time Alignment. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.50

@inproceedings{rai2026taro,
  title = {{TARo: Token-level Adaptive Routing for LLM Test-time Alignment}},
  author = {Arushi Rai, Qiang Zhang, Hanqing Zeng, Yunkai Zhang, Dipesh Tamboli, Xiangjun Fan, Zhuokai Zhao and Lizhu Zhang},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.50},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.50}
}

References

54 references, 23 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ralph Allan Bradley and Milton E. Terry. 1952. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39:324. find
  2. Stephen Casper, Xander Davies, Claudia Shi, Thomas Krendl Gilbert, Jérémy Scheurer, Javier Rando, Rachel Freedman, Tomasz Korbak, David Lindner, Pedro Freire, and 1 others. 2023. Open problems and fundamental limitations of reinforcement learning from human feedback. arXiv preprint arXiv:2307.15217. find
  3. Souradip Chakraborty, Soumya Suvra Ghosal, Ming Yin, Dinesh Manocha, Mengdi Wang, A. S. Bedi, and Furong Huang. 2024. Transfer q star: Principled decoding for llm alignment. In The Thirty-eighth Annual Conference on Neural Information Processing Systems. find
  4. Angelica Chen, Sadhika Malladi, Lily H Zhang, Xinyi Chen, Qiuyi Zhang, Rajesh Ranganath, and Kyunghyun Cho. 2024. Preference learning algorithms do not learn preference rankings. Advances in Neural Information Processing Systems, 37:101928– 101968. find
  5. Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. 2023. Ultrafeedback: Boosting language models with high-quality feedback. ArXiv, abs/2310.01377. find
  6. Hao Cui, Zahra Shamsi, Gowoon Cheon, Xuejian Ma, Shutong Li, Maria Tikhanovskaya, Peter Norgaard, Nayantara Mudur, Martyna Plomecka, Paul Raccuglia, and 1 others. 2025. Curie: Evaluating llms on multitask scientific long context understanding and reasoning. arXiv preprint arXiv:2503.13517.
  7. Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language mo… find
  8. DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Jun-Mei Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiaoling Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, and 179 others. 2025. Deepseek-r1: Incentivizing reasoning capability i…
  9. William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR, abs/2101.03961. find
  10. Leo Gao, John Schulman, and Jacob Hilton. 2022. Scaling laws for reward model overoptimization. In International Conference on Machine Learning.
  11. Etash Kumar Guha, Ryan Marten, Sedrick Scott Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean-Pierre Mercat, Trung Vu, Zayne Sprague, Ashima Suvarna, Ben Feuer, Liangyu Chen, Zaid Khan, Eric Frankel, Sachin Grover, Caroline Choi, Niklas Muennighoff, Shiye Su, and 31 others… — OpenThoughts: Data Recipes for Reasoning Models
  12. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, and 1 others. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948.
  13. Zhenyu Hou, Pengfan Du, Yilin Niu, Zhengxiao Du, Aohan Zeng, Xiao Liu, Minlie Huang, Hongning Wang, Jie Tang, and Yuxiao Dong. 2024. Does rlhf scale? exploring the impacts from data, model, and method. arXiv preprint arXiv:2412.06000. — [arXiv:2412.06000](https://arxiv.org/abs/2412.06000) find
  14. James Y. Huang, Sailik Sengupta, Daniele Bonadiman, Yi an Lai, Arshit Gupta, Nikolaos Pappas, Saab Mansour, Katrin Kirchoff, and Dan Roth. 2024. Deal: Decoding-time alignment for large language models. In Annual Meeting of the Association for Computational Linguistics.
  15. Maxim Khanov, Jirayu Burapacheep, and Yixuan Li. 2024. Args: Alignment as reward-guided search. ArXiv, abs/2402.01694.
  16. Jaemin Kim, Hangeol Chang, Hyunmin Hwang, Choonghan Kim, and Jong Chul Ye. 2025. Universal reasoner: A single, composable plug-and-play reasoner for frozen llms. ArXiv, abs/2505.19075. — [arXiv:2505.19075](https://arxiv.org/abs/2505.19075) find
  17. Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang, Xiangru Peng, and Jiaya Jia. 2024. Step-dpo: Step-wise preference optimization for long-chain reasoning of llms. ArXiv, abs/2406.18629.
  18. Bolian Li, Yifan Wang, Ananth Y. Grama, and Ruqi Zhang. 2024a. Cascade reward sampling for efficient decoding-time alignment. ArXiv, abs/2406.16306. — [arXiv:2406.16306](https://arxiv.org/abs/2406.16306) find
  19. Chen Li, Weiqi Wang, Jingcheng Hu, Yixuan Wei, Nanning Zheng, Han Hu, Zheng Zhang, and Houwen Peng. 2024b. Common 7b language models already possess strong math capabilities. ArXiv, abs/2403.04706. find
  20. Dengchun Li, Yingzi Ma, Naizheng Wang, Zhengmao Ye, Zhiyuan Cheng, Yinghao Tang, Yan Zhang, Lei Duan, Jie Zuo, Cal Yang, and Mingjie Tang. 2024c. Mixlora: Enhancing large language models fine-tuning with lora-based mixture of experts. Preprint, arXiv:2404.15159. find
  21. Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. 2022. Branch-train-merge: Embarrassingly parallel training of expert language models. Preprint, arXiv:2208.03306. find
  22. Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Alpacaeval: An automatic evaluator of instruction-following models. https://github.com/tatsu-lab/alpaca_eval. ↗
  23. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023a. Let’s verify step by step. arXiv preprint arXiv:2305.20050.
  24. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023b. Let’s verify step by step. ArXiv, abs/2305.20050.
  25. Zi-Yan Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. 2025. Understanding r1-zero-like training: A critical perspective. ArXiv, abs/2503.20783.
  26. Qwen, :, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, and 25 others. 2025. Qwen2.5 technical report. Preprint, arXiv:2412.15115. find
  27. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, and 1 others. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
  28. Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen tau Yih, Jason Weston, and Xian Li. 2024. Branch-train-mix: Mixing expert llms into a mixture-of-experts llm. Preprint, arXiv:2403.07816. — [arXiv:2403.07816](https://arxiv.org/abs/2403.07816) find
  29. Llama Team. 2024. Introducing llama 3.1: Our most capable models to date. https://ai.meta.com/ blog/meta-llama-3-1/. Meta AI Blog. ↗
  30. Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, and 1 others. 2025. Kimi k1. 5: Scaling reinforcement learning with llms. arXiv preprint arXiv:2501.12599. — Kimi k1.5: Scaling Reinforcement Learning with LLMs
  31. Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar. 2024. Gsm-symbolic: Understanding the limitations of mathematical reasoning in large language models. arXiv preprint arXiv:2410.05229.
  32. Chunlin Tian, Zhan Shi, Zhijiang Guo, Li Li, and Chengzhong Xu. 2024. Hydralora: An asymmetric lora architecture for efficient fine-tuning. Preprint, arXiv:2404.19245. — [arXiv:2404.19245](https://arxiv.org/abs/2404.19245) find
  33. OpenAI. 2024. Gpt-4 technical report. https://cdn.openai.com/papers/gpt-4.pdf. Accessed: 2025-10-05. ↗
  34. OpenAI. 2024. Learning to reason with llms. https://openai.com/index/ Accessed: learning-to-reason-with-llms/. 2025-05-01. ↗
  35. Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke E. Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Francis Christiano, Jan Leike, and Ryan J. Low… — Training language models to follow instructions with human feedback
  36. Birong Pan, Yongqi Li, Weiyu Zhang, Wenpeng Lu, Mayi Xu, Shen Zhou, Yuanyuan Zhu, Ming Zhong, and Tieyun Qian. 2025. A survey on training-free alignment of large language models. arXiv preprint arXiv:2508.09016. find
  37. Zhenting Qi, Hongyin Luo, Xuliang Huang, Zhuokai Zhao, Yibo Jiang, Xiangjun Fan, Himabindu Lakkaraju, and James Glass. 2024. Quantifying generalization complexity for large language models. arXiv preprint arXiv:2410.01769. — [arXiv:2410.01769](https://arxiv.org/abs/2410.01769) find
  38. Peng-Yuan Wang, Tian-Shuo Liu, Chenyang Wang, Yi-Di Wang, Shu Yan, Cheng-Xing Jia, Xu-Hui Liu, Xin-Wei Chen, Jia-Cheng Xu, Ziniu Li, and 1 others. 2025a. A survey on large language models for mathematical reasoning. arXiv preprint arXiv:2506.08446. find
  39. Wenxuan Wang, Zizhan Ma, Meidan Ding, Shiyi Zheng, Shengyuan Liu, Jie Liu, Jiaming Ji, Wenting Chen, Xiang Li, Linlin Shen, and 1 others. 2025b. Medical reasoning in the era of llms: A systematic review of enhancement techniques and applications. arXiv preprint arXiv:2508.00669. — [arXiv:2508.00669](https://arxiv.org/abs/2508.00669) find
  40. Juncheng Wu, Sheng Liu, Haoqin Tu, Hang Yu, Xiaoke Huang, James Zou, Cihang Xie, and Yuyin Zhou. 2025. Knowledge or reasoning? a close look at how llms think across domains. arXiv preprint arXiv:2506.02126. find
  41. Jiancong Xiao, Ziniu Li, Xingyu Xie, Emily Getzen, Cong Fang, Qi Long, and Weijie J Su. 2025. On the algorithmic bias of aligning large language models with rlhf: Preference collapse and matching regularization. Journal of the American Statistical Association, (just-accepted):1–21. find
  42. Nuoya Xiong, Yuhang Zhou, Hanqing Zeng, Zhaorun Chen, Furong Huang, Shuchao Bi, Lizhu Zhang, and Zhuokai Zhao. 2026. Token-level llm collaboration via fusionroute. arXiv preprint arXiv:2601.05106. — [arXiv:2601.05106](https://arxiv.org/abs/2601.05106) find
  43. Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, and Sumitra Ganesh. 2024. Genarm: Reward guided generation with autoregressive reward model for test-time alignment. arXiv preprint arXiv:2410.08193. — [arXiv:2410.08193](https://arxiv.org/abs/2410.08193) find
  44. Yuxin Zuo, Shang Qu, Yifei Li, Zhangren Chen, Xuekai Zhu, Ermo Hua, Kaiyan Zhang, Ning Ding, and Bowen Zhou. 2025. Medxpertqa: Benchmarking expert-level medical reasoning and understanding. ArXiv, abs/2501.18362. find
  45. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025a. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  46. Chenghao Yang, Lin Gui, Chenxiao Yang, Victor Veitch, Lizhu Zhang, and Zhuokai Zhao. 2025b. Let it calm: Exploratory annealed decoding for verifiable reinforcement learning. arXiv preprint arXiv:2510.05251. — [arXiv:2510.05251](https://arxiv.org/abs/2510.05251) find
  47. Yixin Ye, Zhen Huang, Yang Xiao, Ethan Chern, Shijie Xia, and Pengfei Liu. 2025. Limo: Less is more for reasoning. ArXiv, abs/2502.03387.
  48. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, and 1 others. 2025a. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476.
  49. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma, Guangming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Wang Zhang, Hang Zhu, and 16 others. 2025b. Dapo: An open-source llm reinforcement learning sy…
  50. Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. 2023. Mammoth: Building math generalist models through hybrid instruction tuning. ArXiv, abs/2309.05653.
  51. Hanqing Zeng, Yinglong Xia, Zhuokai Zhao, Gilbert Jiang, Qiang Zhang, Jiayi Liu, Lizhu Zhang, Xiangjun Fan, and Benyu Zhang. 2025. S’more: Structural mixture of residual experts for llm fine-tuning. Preprint, arXiv:2504.06426. — [arXiv:2504.06426](https://arxiv.org/abs/2504.06426) find
  52. Qiyuan Zhang, Fuyuan Lyu, Zexu Sun, Lei Wang, Weixu Zhang, Wenyue Hua, Haolun Wu, Zhihan Guo, Yufei Wang, Niklas Muennighoff, and 1 others. 2025. A survey on test-time scaling in large language models: What, how, where, and how well? arXiv preprint arXiv:2503.24235. find
  53. Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, and 1 others. 2025a. Group sequence policy optimization. arXiv preprint arXiv:2507.18071.
  54. Chujie Zheng, Shixuan Liu, Mingze Li, Xionghui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, Jingren Zhou, and Junyang Lin. 2025b. Group sequence policy optimization. ArXiv, abs/2507.18071.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/rai-2026-taro-token-level-adaptive-routing/