Language Models as Implicit Tree Search

Thanks for your comments and suggestions.

Q1. The concept of Implict Tree search (ITS) is not clearly stated in the main part, and authors should explain what the difference between explicit and implicit tree search and what the advantages of implicit tree search in the introduction section

R1: The concept ITS is derived from the theoretical result in (Grill et al. 2020), which verfiies that given each intermediate node in MCTS, if the intermediate node is treated as a state $s$ in RL, then the search process can be approximated by a state-specific stochastic policy $\overline{\pi}(\cdot|s)$ (Theorem.3.3). With such regards, for each AlphaZero-like MCTS, it always holds a stochastic variant defined by $\overline{\pi}(\cdot|s)$ across all possible intermediate nodes in the MCTS. However, $\overline{\pi}$ can not be universally optimized by gradient descent. Instead, given each state $s$, $\overline{\pi}(\cdot|s)$ can only be obtaind by solving a state-specific problem via dichotomic search. To this end, $\overline{\pi}(\cdot|s)$ can not be used as explicit search like MCTS and instead, only applicable for policy distillation to update the Q value function in AlphaZero-like MCTS (see Eq.7).

The difference between MCTS (i.e., explicit tree search) and ITS mainly refers to three parts.

First, MCTS is defined and executed by UCB in Eq.7, while ITS is defined through the set of policy optimization objectives (Eq.8) across all possible intermediate nodes in MCTS.

Second, the execution of MCTS relies on the empirical visit distribution $\frac{1 + n(s,a)}{|\mathcal{A}{\sf MCTS}| + \sum_{a'\in \mathcal{A}_{\sf MCTS}}n(\boldsymbol{s},\boldsymbol{a}')}$. The empirical visit distribution suffers from the cold-start search problems and its expressivness is always lower than the continuous distrbituion. It implies that ITS is more powerful than MCTS.

Third, ITS can not be explicitly used to search like MCTS due to its complexity to obtain the universal form of $\overline{\pi}$. Instead, ITS was mostly used in policy distillation to enhance Q-learning.

The introduction would be updated to present more details of ITS.

Q2. The figure 2 in the appendix should be put in the main part, so that readers can have a clear high-level intuitive understanding at the begining; Line 137, a typo Convnetional

R2: Will rearrange the figures, and fix all typos in the next version.

Q3. How to implement an implicit tree search (ITS) process?

R3: In the original ITS defined in (Grill et al. 2020), it can not be implemented as the search strategy. However, in this paper, we establish a connection between ITS and DPO. So the universal ITS policy $\overline{\pi}$ can be approximated using a second LLM $\varphi$ by solving either Step-Synchronous IT-PO (Eq.16) or Step-Asynchronous IT-PO (Eq.18), alongside the LLM $\theta$ obtained through DPO-like algorithms. Once LLM $\varphi$ is properly trained, the ITS process can be implemented either by directly using $\varphi$ for next-token prediction or by employing our newly proposed decoding strategies: ITS-$\alpha$ and ITS-rollout.

Thanks for your comments and concerns. We attempt to address the concerns by discussing the writing quality and the causes of marginal performance gain.

Q1: Writing quality.

R1: We response from three aspects.

First, we appreciate the critique to phrases and grammar, and commit that our writing can be largely improved to encourage the readability. In contrast to complicated terms, particularlly, our abstract and introduction, opt for more accessible and illustrative expressions to enhance comprehension. Here are the argues to some cases found by Reviewer U4U8:

Case 1.1 "... remains enabling LLM-based reasoning as AlphaZero": AlphaZero (Silver et, al. 2017b) is the famous Chess game Go system combining MCTS and value-based RL. Our IT-PO exactly uses the stochastic tree-search strategy derived from the consistent AlphaZero-like MCTS pipeline in (Feng, et al. 2023).

Case 1.2 "It's also concerning... as a universal approximator": Policies separatly learned by state-specific objectives (Eq.8) can be uniformly approximated by the same LLM.

Such writing style aims to provide a straightfoward sight of our work before diving into our formulations. It may also cause the unexpected confusion, which is promised to fix.

Second, though the clarity can be improved, we do not agree with the judge "unclear motivation and lack of sufficient background" to our work. Our motivation is intuitivly explained in the first paragraph: LLM trained for preference alignment hardly achieves reasoning without value modeling or RL, in particular, LLM-based reasoning with MCTS. To this, through analyzing the connection between MCTS and its stochastic variant (i.e. implicit tree search in our paper) , we propose IT-PO, the first alignment framework that simultaneously achieve LLM-based reasoning without RL and value modeling. Then we provide the prelimary and warm-up sections (2.5 pages) to elaborate our technical motivation and background. The comments in Reviewers MF7d, REkf, zco2, and also your summary, captured the motivation of this paper.

Beyond this, we found that some unjustified comment misinterpret our claim, e.g.""inherits all advantages of DPO" without introducing drawbacks is unprecise" (we did not claim its escape from the DPO drawbacks.)

Third, we did provide some intuitions to our theoretical parts. In Figure.2, we have illustrated the theoretical connection across MCTS, ITS derived from (Grill et, al. 2020b), and ITS learned by IT-PO. And each formula was always followed by its source and interpretation, e.g., 131-136 (right) for Theorem 3.3; 175-181 (left) for Lemma 3.4; 207-214(right), 234-237(left) for Lemma 4.1; 263-269 (left) for Theorem 4.2, etc. We commit that more illustrative examples can be introduced to improve the understanding of Theorem 4.3 (and will be refined), while the line of proofs in the paper still can be followed by the majority of reviewers.

In general, we acknowledge the manuscript's writing deficiencies and promise for revision, yet we believe the aforementioned evidences insufficient to rate it below the ICML bar.

Q2: The causes of marginal performance gain.

R2: Our following experiments further explain, even refine the marginal gain:

Insufficient iterations of policy distillation. In Sec.4.3, we know IT-PO trains LLM $\varphi$ along with cDPO with respect to fine-tuning LLM $\theta$, which constructs the alternative optimization between $\varphi$ and $\theta$. While in our implementation, we only set the iter=5 in AnthropicHH, and iter=1 in mathematical reasoning tasks (consistent with MCTS-based baselines). With regards to the analysis to Eq.19, increasing the iter number supposes to benefit the IT-PO performance. To this, we reset the IT-PO's evaluation with iter=7 in AnthropicHH, and iter=2 in GSM8K (we also reset iter=2 for MCTS for a fair comparison in Path@1). The results below support our claim.

More appropriate experimental setup. Based on the analysis in Sec.2, ITS could be promising in the cold-start search senarios, or more complex reasoning problems. This claim is supported by the results in ProofWriter (R2 in Reviewer MF7d, gain 6%$\sim$17%) and MATH (R2 in Reviewer REkf).

Q3. Missing numbers in table 2

R3: In GSM8K, the numbers in Equal-Token are consistent with those in Path@1 using tree-search baselines (see (Feng, et al. 2023)).

Thanks for your comments and suggestions.

Q1. The limitation due to the demand of two language models.

R1: It is a good question since the requirement of LLM $\theta$ and $\varphi$ is inevitable in IT-PO. Despite so, the implementation can partially mitigate the limitation. First, despite the alternative training between $\theta$ and $\varphi$, we may only use one of them for inference, in particular, we may decode with either $\theta$ or $\varphi$ while the LLMs are derived from the IT-PO for token-level alignment. Second, since LLM $\varphi$ is the key to achieve reasoning, we may resort to the small-size LLM to reduce the training burden. The experiment of MATH in R2 verified this claim.

Q2. Larger language models as the base of IT-PO;

Q3. Evaluation on more diverse and complex reasoning scenarios.

R2-R3: We provide the experiment to simultaneously address the concerns 2 and 3. It is because larger LLMs demand a more difficult task setup for the evaluation, and sufficiently complex reasoning task also requires a base model strong enough to support IT-PO. Specifically, we consider the reasoning-based evaluation on MATH and construct the training set integrated with the training splits of GSM8K and MATH. Due to our computation resource limited in the rebuttal period, comphrensive results based on 70B-size LLM are hardly achieved. Instead, we resort to other larger base LLMs in the Qwen family, i.e., Qwen1.5-32B. Besides, we introduced LLAMA3.1-8B as a small LLM, aiming to verify our claim in R1. Then our evaluated LLMs ($\theta$, $\varphi$) in IT-PO are defined Qwen1.5-32BX2, (Qwen1.5-32B,LLAMA3-8B), then employed the decoding strategies ITS-$\alpha$ and ITS-rollout, respectively. For a comparison, we further train Qwen1.5-32B with MCTS to provide the baselines MCTS-$\alpha$ and MCTS-rollout, where the search depth = 4 and search width = 3. In our implementation, all IT-PO models are trained with the hyper-parameters consistent with GSM8K in SFT, cDPO, and IT-PO, while the search depth and breadth for training are not limited in MATH (i.e., LLM $\varphi$ trained as a stochastic policy without restriction). For decoding, all tree-search baselines use the search depth and width setup in GSM8K. The results are presented as

In terms of the greedy CoT results in LLAMA3-8B (20.5), some observations can be found in the table. First, we note that the MCTS strategy derived from is unreliable in MATH, probably due to the conflict between the problem complexity and limited search width and depth during training, while ITS does not suffer from this problem. Second, with a weaker model (LLAMA3-8B) for LLM $\varphi$, IT-PO still enable the improvement of LLM-based reasoning.

Q4. Confusion between T-DPO and TDPO.

R4: Will fix it in the next version.

Q5. Quantify the inference-time overhead of IT-PO compared to alternatives like standard DPO and MCTS-based approaches.

R5: The inference-time overhead of IT-PO can be significantly different w.r.t. decoding methods. For token-level alignment, IT-PO used a single model for decoding so that the overhead is consistent of standard DPOs. For reasoning tasks, it refers to ITS-$\alpha$ and ITS-$rollout$ strategies. ITS-$\alpha$ decodes the token like beam-search and for generated sentences, it evaluates their advantages via Eq.22. The process is similar with value inference, and ITS-$\alpha$ holds the similar inference time with MCTS-$\alpha$ (In GSM8K, each token is generated in $\sim$16ms by MCTS-$\alpha$ and in $\sim$18ms by ITS-$\alpha$ using a single A100 80G). ITS-$rollout$ can be treated as test-time-training version of ITS-$\alpha$, which updates $\varphi$ by Eq.21 to execute ITS-$\alpha$ afterwards. It generates 8 (K=8 in Reviewer MF7d's R1) responses for each test prompt to achieve self-training. In our experiment, it results in the extra 30min~1h for each evaluation.

Q6. The sensitivity of the preference pairs to the hard negative response.

R6: We evaluate ITS-$\alpha$ in GSM8K with hard negative ratio ranged from 100%, 75%, 50%, 25%, 0%, along with the results ranged from 53.2, 52.9, 53.4, 51.8, 50.1, accordingly. The performance drastically drops when the ratio less than 50%.

Thanks for your comments and suggestions.

Q1. More implementation details behind theory ("they sometimes gloss...more detailed explanation").

R1: Despite Theorems 4.2, 4.3 containing numerous variables, most are derived from the data rather than serving as hyperparameters. Thus, IT-PO holds implementation consistency with its theoretical foundation except for only two details:

1.1 $\lambda_N(s^{(w)}_t),\lambda_N(s^{(l)}_t)$ indicate how many reasoning paths generated from the state $s_t$, which consists of a prompt $x$ and its responses $y^{(w)}$,$y^{(l)}$ with their contents in the previous $t-1$ steps. For each preference pair, we set $\forall t\in${$1,\cdots,T_w$}, $\lambda_N(s^{(w)}_t)=K_w$ and $\forall t\in${$1,\cdots,T_l$}, $\lambda_N(s^{(l)}_t)=K_l$, where $K_w$, $K_l$ denote how many search start from the $t$-th leaf nodes $s^{(w)}_t, s^{(l)}_t$, respectively. Note that $T_w, T_l$ change with respect to the preference pair. So IT-PO adaptively configure $K_w=\frac{K}{T_w}$, $K_l=\frac{K}{T_l}$ to balance the optimization with different response lengths in $y^{(w)}$,$y^{(l)}$ for each pair. We set $K=8$, inspired from the number of sampled responses for each prompt in many RLHF implementations.

1.2 In gradient analysis in Eq.19, the gradient of $\varphi$ consists of terms in the form $\frac{\pi{\theta}(a_t|s_t)}{\overline{\pi}{\varphi}(a_t|s_t)}$ $\nabla \log [{\overline{\pi}}(a_t|s_t)]$. Due to $\frac{\pi{\theta}(a_t|s_t)}{\pi{\varphi}(a_t|s_t)}\in(0,+\infty)$, updating the models with $\frac{\pi{\theta}(a_t|s_t)}{\overline{\pi}{\varphi}(a_t|s_t)}$ $\nabla \log [{\overline{\pi}}(a_t|s_t)]$ may suffer from exploding/vanishing gradients. To this, we used their logarithmic scaling \frac{\pi{\theta}(a_t|s_t)}{\overline{\pi}{\varphi}(a_t|s_t)}$$\rightarrow $\exp(\log(\pi{\theta}(a_t|s_t))-\log(\overline{\pi}{\varphi}(a_t|s_t)))$ to ensure the less sensitive update ratio. We also took the logarithmic scaling for the gradients in the Step-Asynchronous case (Eq.17), which will be elaborated in our next version due to our rebuttal limit.

Q2. More experiments for other reasoning tasks (Weakness and Questions For Authors).

R2: We offer evaluation based on ProofWriter [1] for deductive logical reasoning, and Chess Endgame [2] for long-term decision making. For ProofWriter, we follow [3] to generate the test set, then the rest are merged to 41,433 training instances. All training and test instances employed the prompt template in [3] that initiated the start of CoT, then we employ LLAMA2-7B as the base model and all fine-tuning methods only run for a single epoch. For Chess Endgame, we follow the experimental setup in (Feng, et al. 2023). With regards to each prompt-response pair $(x,y^{(w)})$ in ProofWriter and Chess Endgame, we find the dispreferred response $y^{(l)}$ using the same strategy in our mathematical reasoning tasks. We ensure the evaluation in the fair comparison with the CoT and LLM-based tree-search baselines : CoT-greedy, BFS-V, MCTS-$\alpha$tr, MCTS-rollout, CoT-SC-MAJ, CoT-SC-ORM, BFS-V-ORM, MCTS-ORM, whose implementations are consistent in the paper.

For simplicity, we skip the average token number metric to highlight Acc in ProofWriter and Win rate in Chess Endgame. While their results remain based on Path@1 to promise the computation efficiency, and Equal-Token to encourage the comparison in the similar scale of computation consumption cross baselines. In the table, we found that CoT variants almost fail in ProofWriter due to their performances close to random guess (33.33%). MCTS variants obtain significantly better results yet basically under-perform ITS variants with substantial gap in ACC, probably due to the cold-start effect in MCTS learned with one epoch. As for Chess Endgame, ITS variants almost solve the problems with Win rates 99.83% in Path@1 and 98.57% in Equal-Token. It proves ITS also competitive in long-horizon reasoning.

Ref:

1. Tafjord, O., et al. ProofWriter: Generating Implications, Proofs, and Abductive Statements over Natural Language (ACL-IJCNLP 2021).

2. Abdulhai, M., et al. Lmrl gym: Benchmarks for multi-turn reinforcement learning with language models. arXiv preprint

3. Pan L, et al. Logic-LM: Empowering Large Language Models with Symbolic Solvers for Faithful Logical Reasoning[C] EMNLP-2023