Zhuokai Zhao
2026
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
Findings of the Association for Computational Linguistics: ACL 2026
Arushi Rai | Qiang Zhang | Hanqing Zeng | Yunkai Zhang | Dipesh Tamboli | Xiangjun Fan | Zhuokai Zhao | Lizhu Zhang
Findings of the Association for Computational Linguistics: ACL 2026
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.
Mixture-of-Minds: Multi-Agent Reinforcement Learning for Table Understanding
Yuhang Zhou | Mingrui Zhang | Ke Li | Mingyi Wang | Qiao Liu | Qifei Wang | Jiayi Liu | Fei Liu | Serena Li | Weiwei LI | Mingze Gao | Abhishek Kumar | Xiangjun Fan | Zhuokai Zhao | Lizhu Zhang
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Yuhang Zhou | Mingrui Zhang | Ke Li | Mingyi Wang | Qiao Liu | Qifei Wang | Jiayi Liu | Fei Liu | Serena Li | Weiwei LI | Mingze Gao | Abhishek Kumar | Xiangjun Fan | Zhuokai Zhao | Lizhu Zhang
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)
Understanding and reasoning over tables is a critical capability for many real-world applications. Large language models (LLMs) have shown promise on this task, but current approaches remain limited. Fine-tuning based methods strengthen language reasoning; yet they are prone to arithmetic errors and hallucination. In contrast, tool-based methods enable precise table manipulation but rely on rigid schemas and lack semantic understanding. These complementary drawbacks highlight the need for approaches that integrate robust reasoning with reliable table processing. In this work, we propose MIXTURE-OF-MINDS, a multi-agent framework that decomposes table reasoning into three specialized roles: planning, coding, and answering. This design enables each agent to focus on a specific aspect of the task while leveraging code execution for precise table manipulation. Building on this workflow, we introduce a self-improvement training framework that employs Monte Carlo Tree Search (MCTS) rollouts to generate pseudo-gold trajectories and optimize agents with reinforcement learning (RL). Extensive experiments show that MIXTURE-OF-MINDS delivers substantial gains, reaching 62.13% on TableBench and surpassing GPT-o3-mini. These results demonstrate the promise of combining structured multi-agent workflows with RL to advance table understanding.
2025
DISCO Balances the Scales: Adaptive Domain- and Difficulty-Aware Reinforcement Learning on Imbalanced Data
Yuhang Zhou | Jing Zhu | Shengyi Qian | Zhuokai Zhao | Xiyao Wang | Xiaoyu Liu | Ming Li | Paiheng Xu | Wei Ai | Furong Huang
Findings of the Association for Computational Linguistics: EMNLP 2025
Yuhang Zhou | Jing Zhu | Shengyi Qian | Zhuokai Zhao | Xiyao Wang | Xiaoyu Liu | Ming Li | Paiheng Xu | Wei Ai | Furong Huang
Findings of the Association for Computational Linguistics: EMNLP 2025
Large Language Models (LLMs) are increasingly aligned with human preferences through Reinforcement Learning from Human Feedback (RLHF). Among RLHF methods, Group Relative Policy Optimization (GRPO) has gained attention for its simplicity and strong performance, notably eliminating the need for a learned value function. However, GRPO implicitly assumes a balanced domain distribution and uniform semantic alignment across groups—assumptions that rarely hold in real-world datasets. When applied to multi-domain, imbalanced data, GRPO disproportionately optimizes for dominant domains, neglecting underrepresented ones and resulting in poor generalization and fairness. We propose Domain-Informed Self-Consistency Policy Optimization (DISCO), a principled extension to GRPO that addresses inter-group imbalance with two key innovations. Domain-aware reward scaling counteracts frequency bias by reweighting optimization based on domain prevalence. Difficulty-aware reward scaling leverages prompt-level self-consistency to identify and prioritize uncertain prompts that offer greater learning value. Together, these strategies promote more equitable and effective policy learning across domains. Extensive experiments across multiple LLMs and skewed training distributions show that DISCO improves generalization, outperforms existing GRPO variants by 5% on Qwen3 models, and sets new state-of-the-art results on multi-domain alignment benchmarks.
2024
AutoPRM: Automating Procedural Supervision for Multi-Step Reasoning via Controllable Question Decomposition
Zhaorun Chen | Zhuokai Zhao | Zhihong Zhu | Ruiqi Zhang | Xiang Li | Bhiksha Raj | Huaxiu Yao
Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)
Zhaorun Chen | Zhuokai Zhao | Zhihong Zhu | Ruiqi Zhang | Xiang Li | Bhiksha Raj | Huaxiu Yao
Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)
Recent advancements in large language models (LLMs) have shown promise in multi-step reasoning tasks, yet their reliance on extensive manual labeling to provide procedural feedback remains a significant impediment. To address this challenge, in this paper, we propose a novel self-supervised framework **AutoPRM** that efficiently enhances the fine-tuning of LLMs for intricate reasoning challenges. Specifically, **AutoPRM** first decomposes complex problems into more manageable subquestions with a controllable granularity switch, then sequentially apply reinforcement learning to iteratively improve the subquestion solver. Additionally, we propose context-guided decoding to avoid reward tampering and guide the subquestion solver towards the solution of the holistic problem. Extensive experiments show that **AutoPRM** significantly improves performance on mathematical and commonsense reasoning tasks over SOTA. More encouragingly, **AutoPRM** can be easily integrated with other orthogonal reasoning pipelines.
Search
Fix author
Co-authors
- Xiangjun Fan 2
- Lizhu Zhang 2
- Yuhang Zhou (周宇航) 2
- Wei Ai 1
- Zhaorun Chen 1
- Mingze Gao 1
- Furong Huang 1
- Abhishek Kumar 1
- Weiwei LI 1
- Ke Li 1
- Ming Li 1
- Serena Li 1
- Xiang Li 1
- Fei Liu 1
- Jiayi Liu 1
- Qiao Liu 1
- Xiaoyu Liu 1
- Shengyi Qian 1
- Arushi Rai 1
- Bhiksha Raj 1
- Dipesh Tamboli 1
- Mingyi Wang 1
- Qifei Wang 1
- Xiyao Wang 1
- Paiheng Xu 1
- Huaxiu Yao 1
- Hanqing Zeng 1
- Mingrui Zhang 1
- Qiang Zhang 1
- Ruiqi Zhang 1
- Yunkai Zhang 1
- Jing Zhu 1
- Zhihong Zhu 1