Source-linked AI summary
ReST meets ReAct: Self-Improvement for Multi-Step Reasoning LLM Agent
Renat Aksitov, Sobhan Miryoosefi, Zonglin Li, Daliang Li, Sheila Babayan, Kavya Kopparapu, Zachary Fisher, Ruiqi Guo, Sushant Prakash, Pranesh Srinivasan, Manzil Zaheer, Felix Yu, Sanjiv Kumar
TL;DR
Human-labeled multi-step trajectories are costly and difficult to judge, complicating improvement of agents that reason and act using external knowledge. The paper applies ReST-like iterative training with AI feedback to a ReAct-style search agent and uses its trajectories for self-improvement and distillation. It reports comparable performance from fine-tuned models with up to two orders of magnitude fewer parameters than the prompted large model.
Problem
Improving process-based agents is difficult because acquiring human-labeled multi-step trajectories is costly and humans find optimal trajectories hard to judge.
Source: §1 · p. 2Method
The authors adapt ReST to a ReAct-style search agent, growing trajectories from the latest policy and improving it with LLM-based ranking or filtering.
Source: §1 · p. 2Results
Takeaways & Limitations
The authors conclude that AI-feedback ReST can iteratively build synthetic data for agent self-improvement and use that data to distill smaller multi-step agents while preserving most teacher performance.
Source: §8 · p. 9The saturation point beyond the second self-improvement iteration, including whether smaller models continue gaining or remain capped by the prompted large model, is unresolved.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Answering complex natural language questions often necessitates multi-step reasoning and integrating external information. Several systems have combined knowledge retrieval with a large language model (LLM) to answer such questions. These systems, however, suffer from various failure cases, and we cannot directly train them end-to-end to fix such failures, as interaction with external knowledge is non-differentiable. To address these deficiencies, we define a ReAct-style LLM agent with the ability to reason and act upon external knowledge. We further refine the agent through a ReST-like method that iteratively trains on previous trajectories, employing growing-batch reinforcement learning with AI feedback for continuous self-improvement and self-distillation. Starting from a prompted large model and after just two iterations of the algorithm, we can produce a fine-tuned small model that achieves comparable performance on challenging compositional question-answering benchmarks with two orders of magnitude fewer parameters.
1 INTRODUCTION
The paper addresses the difficulty of improving multi-step agents when human-labeled trajectories are costly and hard to judge. It applies ReST-style iterative training with AI feedback to a ReAct search agent, evaluating self-improvement and distillation on compositional question-answering tasks.
1 INTRODUCTION
- Human-labeled multi-step trajectories are costly to acquire because they require more data and are harder for people to judge than outcomes.The paper notes that collecting more human-labeled data is a more difficult solution for process-based systems than for outcome-based systems.Source: §1 · p. 2
- The method adapts ReST: it grows trajectories from the latest policy, then improves the policy on a fixed dataset by LLM-based ranking or filtering.The ranking step uses an LLM call rather than a reward model distilled from human preferences.Source: §1 · p. 2
- The agent is evaluated on Bamboogle and BamTwoogle, compositional questions manually verified to be unanswerable directly by a search engine.Bamboogle results also illustrate performance across model sizes and ReST iterations; smaller models are trained by distillation from data produced with the large model.Source: §1 · p. 2
- The authors report that ReST-style iterative fine-tuning improves agent reasoning traces without human-labeled training data and can distill performance into models one or two orders of magnitude smaller.Their contributions include a self-critique-enabled ReAct agent and proxy evaluation on Bamboogle and BamTwoogle.Source: §1 · p. 2
2 BACKGROUND: SEARCH AGENT
Search Agent uses web search to answer open-ended questions with long-form, attributable answers. It searches iteratively, drafts an answer, and then makes two self-revision calls.
2 BACKGROUND: SEARCH AGENT
- Search Agent is a ReAct-style agent with Reflexion that uses web search to produce long-form, explicitly attributable answers to open-ended knowledge-seeking questions.The paper describes its flow as a sequence of actions to answer diverse questions.Source: §2 · p. 2
- The agent searches until it decides it has enough information, summarizes retrieved snippets, drafts an answer, then checks relevance and grounding in two revision calls.The flow is presented as a state machine in which each blue shape is one LLM call and a distinct reasoning-step type.Source: §2 · pp. 2–3, Figure 2 · p. 3
3 METHODS
The methods prompt and run a multi-step Search Agent, collect and rank its trajectories, and iteratively fine-tune models on the resulting reasoning-step mixtures.
3.1 PROMPTING
- Only PaLM 2-L consistently follows the code-formatted few-shot prompts as valid Python, so it generates the initial reasoning trajectories.The prompts define each of the agent’s five reasoning steps; code is chosen for its structured, parseable form and natural-language comments and naming.Source: §3.1 · p. 3
3.2 IMPLEMENTATION DETAILS
- Search Agent uses PaLM 2 XS, S, and L models, samples multiple outputs per step, and continues with the lowest-perplexity sample using top-three search snippets and at most ten searches per trajectory.The number of samples is four for L and S and sixteen for XS.Source: §3.2 · p. 3
3.3 INPUT DATA
- The authors sample 500 questions from each of four datasets to create 2,000 diverse trajectory prompts, using no dataset labels or validation splits.The input datasets include HotpotQA, ELI5, ELI5-askH, and ELI5-askS.Source: §3.3 · p. 4
3.4 FINE-TUNING
- Completed trajectories are split into reasoning steps to form a fine-tuning mixture, and all experiments use full fine-tuning.The authors conduct as many experiments as possible with XS because fine-tuning costs rise sharply for larger models.Source: §3.4 · p. 4
3.5 RANKING ”REWARD” MODEL
- An instruction-tuned PaLM 2-L ranks sampled outputs for fine-tuning, differing from ReST’s human-preference reward-model filtering and resembling RAFT’s selection of high-scoring responses.LLM-based selection is mainly off-policy; on-policy rollouts use perplexity.Source: §3.5 · p. 4
3.6 ITERATIVE SELF-IMPROVEMENT
- The iterative algorithm collects trajectories from the current model, reranks steps into a fine-tuning mixture, fine-tunes a same-size model, and repeats after verifying improvement.Smaller models can also be trained on mixtures from different iterations, producing a self-distillation procedure.Source: §3.6 · pp. 4–5
4 EVALUATION
The evaluation uses Bamboogle and a held-out BamTwoogle test set, with LLM-based auto-evaluation to assess open-ended answers across repeated trajectories. BamTwoogle is designed to retain multi-step difficulty.
4.1 BAMBOOGLE
- Bamboogle is a 125-question semi-adversarial set of two-hop questions selected to resist direct Google search while keeping both evidence pieces findable in Wikipedia.The authors use it to assess the Search Agent’s tool use.Source: §4.1 · p. 5
4.2 AUTO-EVAL
- LLM-based auto-evaluation aligned with human ratings on Bamboogle trajectories across diverse agents (Pearson 0.98, p = 6.6 × 10−8; Spearman 0.83, p = 0.0015), enabling ten repetitions per question to reduce variance.Human evaluations are time-consuming and do not scale well, limiting variance reduction through repeated ratings because trajectories are stochastic.Source: §4.2 · p. 5
4.3 BAMTWOOGLE
- BamTwoogle is a complementary, slightly harder held-out test set; all its questions require at least two steps, unlike some Bamboogle questions that can be solved with one search.It is a handcrafted collection of 100 questions; most require two searches or reasoning steps, and some require three or four.Source: §4.3 · p. 6
5 EXPERIMENTS
The experiments evaluate a pilot fine-tuning mixture, a second self-improvement iteration with human evaluations, and ablations of training-data quality, quantity, filtering, and self-critique.
5.1 PILOT
- The pilot used 500 HotpotQA and Eli5 trajectories, with default best actions, and manually filtered about 30% of examples as a fine-tuning baseline.Source: §5.1 · pp. 6–7
5.2 SELF-IMPROVEMENT AND SELF-DISTILLATION
- For the second iteration, PaLM 2-L generated 8,000 trajectories from 2,000 questions repeated four times; human evaluations covered one run per model on Bamboogle and BamTwoogle.Source: §5.2 · p. 7, Table 2 · p. 7
5.3 ABLATIONS
- Fine-tuning on filtered data produced a 2.5% performance drop versus the unfiltered pilot mixture; the authors hypothesize reduced mixture size and retained bad steps may contribute.Source: §5.3 · p. 7
- Using two trajectories per question instead of one improved the fine-tuning mixture by 2.2%, while more trajectories did not significantly improve performance.Source: §5.3 · p. 7
- 9% gain from 1st-gen to 2nd-gen data at roughly equal size indicates data quality mattered more than quantity, and better data reduced evaluation-trajectory variance.Source: §5.3 · p. 7
- Self-critique produced a small positive boost of about 0.5–1.0% for most models.The effect was measured by evaluating the Answer Generation step rather than the Final Answer step on the Bamboogle trajectories.Source: §5.3 · p. 8
6 DISCUSSION
The discussion highlights process supervision without training labels, the costs of robust auto-evaluation, and a small positive effect of self-critique. It also identifies limits in the current evaluation and open questions about tool generalization and further self-improvement.
Process Supervision
- The process-based setup uses no training labels during trajectory collection and can learn from states that lead to wrong answers, including for open-ended questions without a single correct answer.It combines a state-machine agent, high-temperature exploration, AI feedback for action reranking, and state-wise fine-tuning over completed trajectories.Source: §6 · p. 8
Auto-Eval
- Robust auto-evaluation aligned with human raters is valuable for stochastic, long-form trajectories, but repeated runs and PaLM 2-L judging raise computational costs.Evaluation measures answer quality for a trajectory while accounting for variation between trajectories on the same input.Source: §6 · p. 8
Self-Critique
- Self-critique has a small positive effect on overall performance, with a larger effect for larger models and no apparent dependence on self-improvement.This contrasts with reported harm from adding self-critique in outcome-based CoT setups.Source: §6 · p. 8
Limitations and Future Directions
- Computational and time constraints limit the work to manually constructed prompts, small evaluations, a limited set of models, and one tool.These are among the limitations the authors identify.Source: §6 · p. 8
- Whether the self-improvement algorithm extends to multiple tools, improves handling of unseen tools, or requires changes to improve both remains open.The authors note that unseen-tool handling may behave like self-critique and fail to improve under ReST-like iterative training.Source: §6 · p. 8
- The saturation point beyond two self-improvement iterations, including whether smaller models continue gaining or converge to the prompted large model, remains unresolved.The authors pose these as open questions.Source: §6 · p. 8
7 RELATED WORK
The related work situates the paper among long-form web question answering, prompting and agent fine-tuning, and self-improvement methods. Its distinguishing focus is process-based agent training with synthetic data from AI feedback rather than human labels or outcome-based signals.
Long-form question answering
- Unlike WebGPT (Nakano et al., 2021), which uses many human demonstrations for imitation learning and RL, this work aims to minimize human involvement and uses only few-shot prompt exemplars as labeled demonstrations.The exemplars support the agent’s reasoning steps.Source: §7 · p. 8
Prompting and fine-tuning
- Manually designed few-shot prompts are common for language agents (Press et al. (2023); Yao et al. (2022); Shinn et al. (2023)), while DSP tunes demonstrations using labeled examples.DSP can also fine-tune specific agent components.Source: §7 · p. 9
- Agent fine-tuning is less common; unlike FireAct, this setup uses no human labels for training or filtering and instead builds synthetic data through AI-feedback self-improvement.Source: §7 · p. 9
Self-improvement
- Unlike STAR and ReSTEM, the method does not use answer correctness as a signal; unlike ReST (Gulcehre et al., 2023) and RAFT, it lacks a reward model trained on human preferences.Those four methods target outcome-based systems, whereas this work focuses on a process-based system.Source: §7 · p. 9
8 CONCLUSION
The paper concludes that ReST-like training with AI feedback can improve a multi-step reasoning agent and produce synthetic data that supports self-improvement and distillation into smaller models.
Conclusion
- ReST-like training with AI feedback is demonstrated as applicable to a multi-step reasoning LLM agent.Source: §8 · p. 9
- The iterative approach is described as a relatively simple and efficient way to build high-quality synthetic data for agent self-improvement.Source: §8 · p. 9
- Increasingly higher-quality data can also distill a multi-step agent into models several magnitudes smaller while preserving most of the large teacher’s performance.Source: §8 · p. 9
A.1 PROMPTS
The appendix lists the prompt configurations used for reasoning, self-checks, auto-evaluation, and reward-model ranking. The reasoning prompts and self-check prompts use multiple shots, while the ranking reward-model prompt is zero-shot.
Reasoning steps
- Reasoning prompts are n-shot (n > 1): decision uses 9 shots, summarization 6, and answer generation 5.Source: §A.1 · p. 11
- Relevance and grounding self-check prompts use 6 and 5 shots, respectively, and share a common prefix.Source: §A.1 · p. 11
Auto-eval
- The auto-evaluation prompt uses 5 shots.Source: §A.1 · p. 11
Ranking “Reward” Model
- The ranking reward-model prompt is 0-shot, omits the code-as-prompt format, and is intended for an instruction-tuned model.Source: §A.1 · p. 11
A.2 ADDITIONAL DATA
This section presents additional data and prompt/implementation examples for the multi-step search agent, including evaluation material and checks for grounding, answer implications, and model-output selection.
A.2 ADDITIONAL DATA
- Table 4 reports Bamboogle auto-evaluation mean accuracy and standard deviation over 10 runs, before self-critique steps.The caption gives no cell values, so it does not establish a performance comparison.Source: Table 4 · p. 11
- The decision-step fragment defines reasoning paired with an action, including Google search and termination, and illustrates searching for loveseat and vehicle dimensions.The example plans to compare Ikea Klippan loveseat dimensions with the 2019 Honda Odyssey cargo size.Source: §A.2 · p. 12
- The summarization fragment selects relevant search results and identifies missing information as a basis for a follow-up query.In the solar-light example, it rejects irrelevant results and follows a result about panel types to search for efficiency numbers.Source: §A.2 · p. 13
- The answer-generation example combines results from multiple searches to answer a question about humidifier residue and omits background that does not address it.The example attributes the residue to minerals in the water and notes calcium and magnesium as examples.Source: §A.2 · p. 14
- The self-check prompt tests whether an answer addresses the original question and is based on selected links, while the relevance example answers how to exclude a site from Google results.The example gives the -site:website_name query syntax and a Pinterest example.Source: §A.2 · pp. 15–16
Additional data
- The grounding self-check traces a multi-step answer from identifying John Glenn to finding his 1962 achievement, then checks that the answer is supported by cited links.When an initial query uses the wrong year, the example revises it to 1962 before selecting evidence and answering.Source: §A.2 · p. 17
- The auto-evaluator returns true when an answer implies the reference answer and false when it omits or fails to establish that answer.Examples include accepting a supported yes response and rejecting answers that omit the relevant person or leave the requested relationship undetermined.Source: §A.2 · p. 18
- The reward-model prompt instructs raters to filter bad actions and favor relevant, grounded answers that use minimal steps.Empty thoughts or comments are automatically considered bad.Source: §A.2 · p. 19
- Raters must choose among outputs using only the visible inputs, explain the choice specifically, and follow the required three-line format.The prompt specifies an explanation, an answer identifying the selected output, and a ranking.Source: §A.2 · p. 19