Table of Contents
Fetching ...

Weak-to-Strong Generalization with Failure Trajectories: A Tree-based Approach to Elicit Optimal Policy in Strong Models

Ruimeng Ye, Zihan Wang, Yang Xiao, Zinan Ling, Manling Li, Bo Hui

TL;DR

This work fine-tune a strong model with trajectories of intermediate actions generated by a weak model, and construct ``trajectory trees," a hierarchical representation that organizes weak model-generated action trajectories, coupled with Monte Carlo Tree Search (MCTS) to optimize the strong model.

Abstract

Weak-to-Strong generalization (W2SG) is a new trend to elicit the full capabilities of a strong model with supervision from a weak model. While existing W2SG studies focus on simple tasks like binary classification, we extend this paradigm to complex interactive decision-making environments. Specifically, we fine-tune a strong model with trajectories of intermediate actions generated by a weak model. Motivated by the human learning process, we propose to generalize not only success knowledge but also failure experience so that the strong model can learn from failed trajectories accumulated by weak models. To effectively and efficiently elicit the potential of strong agents, we further construct ``trajectory trees," a hierarchical representation that organizes weak model-generated action trajectories, coupled with Monte Carlo Tree Search (MCTS) to optimize the strong model. Through theoretical analysis, we provide formal guarantees for the effectiveness of our method in improving W2SG performance. Our empirical evaluations demonstrate substantial improvements in reasoning and decision-making capabilities across diverse task domains, validating the scalability and robustness of our proposed framework.

Weak-to-Strong Generalization with Failure Trajectories: A Tree-based Approach to Elicit Optimal Policy in Strong Models

TL;DR

This work fine-tune a strong model with trajectories of intermediate actions generated by a weak model, and construct ``trajectory trees," a hierarchical representation that organizes weak model-generated action trajectories, coupled with Monte Carlo Tree Search (MCTS) to optimize the strong model.

Abstract

Weak-to-Strong generalization (W2SG) is a new trend to elicit the full capabilities of a strong model with supervision from a weak model. While existing W2SG studies focus on simple tasks like binary classification, we extend this paradigm to complex interactive decision-making environments. Specifically, we fine-tune a strong model with trajectories of intermediate actions generated by a weak model. Motivated by the human learning process, we propose to generalize not only success knowledge but also failure experience so that the strong model can learn from failed trajectories accumulated by weak models. To effectively and efficiently elicit the potential of strong agents, we further construct ``trajectory trees," a hierarchical representation that organizes weak model-generated action trajectories, coupled with Monte Carlo Tree Search (MCTS) to optimize the strong model. Through theoretical analysis, we provide formal guarantees for the effectiveness of our method in improving W2SG performance. Our empirical evaluations demonstrate substantial improvements in reasoning and decision-making capabilities across diverse task domains, validating the scalability and robustness of our proposed framework.

Paper Structure

This paper contains 26 sections, 1 theorem, 17 equations, 5 figures, 8 tables.

Key Result

Theorem 1

$\hat{\pi}_s^{\text{TreeDPO}}$ is the policy obtained by minimizing the TreeDPO loss over the random sampling of $\mathcal{D}_w$: for some constant $C > 0$. This implies that $\hat{\pi}_s^{\text{TreeDPO}}$ can outperform the SFT-strong baseline if the potential improvement margin $\mathcal{R}(\pi^*) - \mathcal{R}(\pi_s^{\text{SFT}})$ is sufficiently large relative to the third term (estimation er

Figures (5)

  • Figure 1: Illustration of the trajectory tree construction and the weak-to-strong generalization with trajectories. The left part explores different trajectories with a weak model. A trajectory tree is constructed by merging the same action path. Nodes on the tree represent different actions and performing different actions will lead to various subsequent paths. Then the trajectory tree is used to elicit the ability of the strong model.
  • Figure 2: Illustration of the Weak-to-Strong framework. (a) Given an instruction, the weak LLM agent interacts with the environment to collect both success and failure trajectories of actions. (b) The explored trajectories are used to construct a trajectory tree by merging prefixes of actions. We propose two methods to supervise the strong model: (c1) DPO with structural contrastive failure-success pairs of trajectories instead of random pairs; (c2) Fine-tune the strong model with Monte Carlo Tree Search.
  • Figure 3: Three agents are compared: SFT, WTS, and MCTS. X shows the time steps while Y illustrates the scores. The task ID is shown above the plot.
  • Figure 4: The performance changes with the increase of trajectory numbers.
  • Figure 5: An example of weak-to-strong generalization in the WebShop scenario. The 4 interactions from left to right are the trajectories generated by the SFT Weak Model, the SFT Strong Model, the W2S model trained with DPO, and the W2S Model trained with MCTS.

Theorems & Definitions (1)

  • Theorem 1: Performance Guarantee for W2SG via Tree-Guided DPO