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

TL;DR

W2S-AlignTree tackles the misalignment and scalability limits of training-time alignment by moving alignment decisions to inference time. It formulates alignment as an optimal search over a generative tree, guided by a dense token-level proxy $V_{proxy}(\mathbf{x}, \mathbf{y}') = \log ( \pi_{weak}^*(\mathbf{y}'|\mathbf{x}) / \pi_{weak}^{ref}(\mathbf{y}'|\mathbf{x}) )$ derived from a weak model, and integrates this into a dual-stage Monte Carlo Tree Search with an Entropy-Aware PUCT rule. A final global re-ranking uses the full sequence-level score $r(\mathbf{x}, \mathbf{y}) = \beta \log ( \pi^*(\mathbf{y}|\mathbf{x}) / \pi_{ref}(\mathbf{y}|\mathbf{x}) )$ to select the best output, enabling fine-grained, dynamic control without changing the strong model parameters. Across sentiment generation, summarization, and instruction-following, W2S-AlignTree achieves consistent improvements over strong baselines and demonstrates robust cross-model applicability, including notable gains such as a 15.9% relative improvement on summarization for Llama3-8B. This approach offers a scalable, plug-in solution for aligning large language models with human preferences during inference, with potential extensions to online learning and multimodal settings.

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.

Figures (5)

• Figure 1: W2S-AlignTree vs. Training-Time Alignment. Unlike RLHF and DPO, W2S-AlignTree enables fine-grained inference-time alignment using weak model signals—without costly reward models or parameter updates.
• Figure 2: Dual-stage alignment process of W2S-AlignTree. (a) MCTS constructs a generative search tree where candidate chunks are proposed by the strong model and evaluated with step-level proxy values from a weak model. (b) Among all explored paths, W2S-AlignTree decides the response by globally re-ranking based on complete sequence-level alignment scores.
• Figure 3: Alignment performance across sentiment generation and summarization. W2S-AlignTree consistently outperforms strong baselines by enabling inference-time alignment with weak model guidance, achieving higher $r_{\mathrm{gold}}$ across diverse LLMs. We report mean rewards ($\pm$ standard deviations) across three random seeds. Appendix C.1 provides detailed numerical results.
• Figure 4: Hyperparameter analysis of W2S-AlignTree to chunk length $L$ and exploration coefficient $c$ across tasks. The areas with better performance are boxed.
• Figure 5: Further analysis of hyperparameter sensitivity with respect to (left) entropy constant $w \in [0.1, 1.0]$ and (right) candidate number $K \in [1, 10]$.We conduct evaluations on both GPT XL and Llama-3-8B under varying entropy settings and candidate counts.