W2S-AlignTree: Weak-to-Strong Inference-Time Alignment for Large Language Models via Monte Carlo Tree Search

Authors

  • Zhenyu Ding Xi'an Jiaotong University
  • Yuhao Wang Xi'an Jiaotong University
  • Tengyue Xiao Xi'an Jiaotong University
  • Haoying Wang Xi'an Jiaotong University
  • Guojun Ma Tsinghua University
  • Mingyang Wan Xi'an Jiaotong University
  • Caigui Jiang Xi'an Jiaotong University
  • Ning Ding Xi'an Jiaotong University

DOI:

https://doi.org/10.1609/aaai.v40i36.40306

Abstract

Large Language Models (LLMs) demonstrate impressive capabilities, yet their outputs often suffer from misalignment with human preferences due to the inadequacy of weak supervision and a lack of fine-grained control. Training-time alignment methods like Reinforcement Learning from Human Feedback (RLHF) face prohibitive costs in expert supervision and inherent scalability limitations, offering limited dynamic control during inference. Consequently, there is an urgent need for scalable and adaptable alignment mechanisms. To address this, we propose W2S-AlignTree, a pioneering plug-and-play inference-time alignment framework that synergistically combines Monte Carlo Tree Search (MCTS) with the Weak-to-Strong Generalization paradigm for the first time. W2S-AlignTree formulates LLM alignment as an optimal heuristic search problem within a generative search tree. By leveraging weak model's real-time, step-level signals as alignment proxies and introducing an Entropy-Aware exploration mechanism, W2S-AlignTree enables fine-grained guidance during strong model's generation without modifying its parameters. The approach dynamically balances exploration and exploitation in high-dimensional generation search trees. Experiments across controlled sentiment generation, summarization, and instruction-following show that W2S-AlignTree consistently outperforms strong baselines. Notably, W2S-AlignTree raises the performance of Llama3-8B from 1.89 to 2.19, a relative improvement of 15.9% on the summarization task.

Published

2026-03-14

How to Cite

Ding, Z., Wang, Y., Xiao, T., Wang, H., Ma, G., Wan, M., … Ding, N. (2026). W2S-AlignTree: Weak-to-Strong Inference-Time Alignment for Large Language Models via Monte Carlo Tree Search. Proceedings of the AAAI Conference on Artificial Intelligence, 40(36), 30521–30529. https://doi.org/10.1609/aaai.v40i36.40306

Issue

Section

AAAI Technical Track on Natural Language Processing I