ReST meets ReAct: Self-Improvement for Multi-Step Reasoning LLM Agent | Summary

Contents

Summary
1 Introduction
2 Background: Search Agent
3 Methods
4 Evaluation
5 Experiments
6 Discussion
7 Related Work
8 Conclusion
Appendix
Brief Thoughts

This article explains the key points of ReST meets ReAct: Self-Improvement for Multi-Step Reasoning LLM Agent.

2023-12-15 (arXiv)
Aksitov, Renat, Miryoosefi, Sobhan, Li, Zonglin, Li, Daliang, Babayan, Sheila, Kopparapu, Kavya, Fisher, Zachary, Guo, Ruiqi, Prakash, Sushant, Srinivasan, Pranesh, et al.
Google Research, Google DeepMind, Google
Paper

Summary

Answering complex questions with a large language model (LLM) requires coordinating searches and reasoning steps, while external retrieval prevents direct end-to-end training through differentiation. After two training iterations, the paper’s small PaLM 2-S search agent achieves 69.7% automatic accuracy on Bamboogle, a benchmark requiring two linked pieces of evidence, compared with 70.3% for the prompted large PaLM 2-L teacher.
The method combines ReAct, which interleaves reasoning, actions, and observations, with Reinforced Self-Training (ReST), which alternates generating sequences of agent steps, called trajectories, and training on selected steps. A large model generates trajectories, an instruction-tuned model ranks candidate actions, and fine-tuning improves the large agent while transferring its behavior to smaller models through distillation.
Human evaluation on BamTwoogle, a separate 100-question test benchmark, improves from 68.0% for prompted PaLM 2-L to 74.0% for its second-generation version, while both smaller versions score 63.0%. The results cover 125 Bamboogle questions, 100 BamTwoogle questions, three model sizes, one search tool, and two training iterations.

1 Introduction

The paper addresses the cost of obtaining human annotations for complete reasoning trajectories. It proposes improving a search agent through synthetic trajectories and stepwise AI feedback without requiring task-answer labels during trajectory collection or training-data filtering.

ReST’s “grow” stage samples trajectories from the current model; its “improve” stage trains on selected data. Here, an LLM directly ranks actions instead of using a separate reward model trained on human preferences.
The same large-model trajectories train smaller agents, allowing the experiments to study both iterative improvement and transfer of the teacher’s behavior.

Bamboogle automatic accuracy across model sizes and training stages, with means and standard deviations over 10 runs.

2 Background: Search Agent

Search Agent is a ReAct-style system with Reflexion, implemented here as self-checking and answer revision. It repeatedly decides whether more information is needed, searches the web, and summarizes retrieved snippets; after terminating the search loop, it drafts an answer and performs separate relevance and grounding checks, where grounding means support from the retrieved snippets.

The relevance check asks whether the draft addresses the original question. The grounding check asks whether the retrieved evidence supports the answer.
Figure 2 separates the workflow into five types of model calls: decision, summarization, answer generation, relevance self-check, and grounding self-check.

Search Agent state machine showing search, summarization, answer generation, and two self-check stages.

3 Methods

The method collects completed search trajectories and converts individual reasoning steps into training examples. Repeated generation, action ranking, and fine-tuning form the improvement loop; training smaller models on the same examples provides distillation.

3.1 Prompting

The authors manually write few-shot prompts, meaning prompts containing example inputs and outputs, for all five reasoning-step types. These prompts use Python code so that model outputs can be parsed and integrated with the search workflow.

The model continues the supplied code with a structured action. Among the pretrained models tested, only PaLM 2-L consistently handled this format well enough to generate the initial trajectories.
The smaller agents therefore learn from the large model’s trajectories rather than independently establishing the workflow from prompts.

3.2 Implementation Details

Experiments use pretrained and fine-tuned PaLM 2 base models in extra-small (XS), small (S), and large (L) sizes. Each step usually samples at temperature T = 0.5, which controls generation randomness, and continues with the candidate having the lowest perplexity, a measure of how unpredictable a sequence is under the model: four candidates for L and S, and 16 for XS.

An internal Google Q&A API, the interface to the search tool, supplies snippets, typically the top three per query. Each trajectory permits at most 10 searches.
REMAINING_SEARCHES tracks the search budget, and PAST_ACTIONS stores the accumulated trajectory.

3.3 Input DATA

The initial question pool contains 500 randomly selected training questions from each of HotpotQA, Eli5, Eli5-askH, and Eli5-askS, totaling 2000 questions. HotpotQA requires combining information from multiple documents; the remaining datasets provide open-ended questions from r/explainlikeimfive, r/askhistorians, and r/askscience.

Only questions are used. Dataset answers do not filter trajectories, and these datasets’ validation splits are not used for hyperparameter selection or performance assessment.

3.4 Fine-Tuning

Each completed trajectory is split into its component reasoning steps to form a fine-tuning mixture. All experiments use full fine-tuning, which trains the model’s parameters; the authors conduct as many experiments as possible with XS because larger-model fine-tuning is more expensive.

3.5 RANKING ”REWARD” MODEL

An instruction-tuned PaLM 2-L, trained to follow instructions, receives the current input, candidate outputs, and ranking guidance, then selects the highest-ranked action for training. This replaces the original ReST procedure’s threshold-based filtering with relative action rankings.

Ranking mainly occurs off-policy: it replaces the action used as a training target after trajectory collection. Actual trajectory continuation still uses the lowest-perplexity candidate.
The paper relates this selection mechanism to RAFT, a method that fine-tunes on responses selected through reward rankings.

3.6 Iterative Self-Improvement

The improvement loop starts with prompted PaLM 2-L, collects trajectories for the fixed 2000-question pool, reranks actions while constructing training examples, and fine-tunes a model of the same size. Once the new model performs better, it generates the next iteration’s trajectories; smaller models can also be trained on each iteration’s data.

Each outer iteration uses one improvement stage. The question pool remains fixed while the generated trajectories change.

4 Evaluation

Evaluation measures whether the agent answers compositional questions correctly, meaning questions that require combining multiple pieces of information. Bamboogle supports repeated automatic evaluation and model selection, while BamTwoogle is reserved for final testing.

4.1 Bamboogle

Bamboogle contains 125 two-hop questions, each originally selected to resist direct Google search while having both required pieces of evidence available in Wikipedia. Because Search Agent produces long-form answers with references, correctness is assessed by humans or a separate PaLM 2-L base-model call instead of exact string matching.

Figure 3 illustrates the evaluation problem: a paragraph identifying the golden eagle and reporting 320kph must be compared with the short reference answer, 320 km/h.

A Bamboogle question, a cited long-form model answer, and its short reference answer.

4.2 Auto-Eval

Auto-eval, the model-based correctness assessment, asks whether the generated answer implies the reference answer. Across the evaluated agents, its scores correlate with human scores at Pearson 0.98, measuring linear association, with p = 6.6 × 10⁻⁸, and Spearman 0.83, measuring rank agreement, with p = 0.0015; the usual evaluation averages 10 trajectories per Bamboogle question.

Bamboogle auto-eval selects temperature T = 0.5 and checkpoints, saved training states, at step 9K for XS, 5K for S, and 3.5K for L. It also informs decisions about additional iterations, trajectory counts, and self-checks.
Bamboogle questions are excluded from trajectory-generation training data, and prompts are not tuned on them. Repeated model selection nevertheless makes Bamboogle function as a validation set.

4.3 BamTwoogle

BamTwoogle is a handcrafted test set of 100 questions designed to address overfitting concerns from repeated Bamboogle evaluation. Its questions were checked to require at least two steps, whereas search-engine changes have made some original Bamboogle questions answerable with one search.

Questions: Most require two searches or reasoning steps, and some require three or four. Each is manually checked to ensure its answer does not appear on the first page of Google results.
Expected answers: Answers should be unambiguous, stable over time, accommodate alternative proper-name forms, and preferably use Wikipedia articles without relevant quality warnings as factual references.

5 Experiments

The experiments compare a manually filtered pilot baseline, two generations of self-improvement data, and distillation into XS and S agents. Automatic evaluation reports means and standard deviations over 10 Bamboogle runs; final human evaluation covers one run per model on each benchmark.

5.1 Pilot

The pilot uses 500 trajectories generated from HotpotQA and Eli5 questions, selecting actions by minimum perplexity without model reranking. Human reviewers remove about 30% of training examples for problems such as unhelpful queries, empty reasoning, or incomplete summaries; this manually filtered mixture serves as a fine-tuning baseline.

5.2 Self-Improvement and Self-Distillation

The first generation uses 2000 trajectories from prompted PaLM 2-L; the second uses 8000 trajectories from the first-generation fine-tuned L model, repeating each of the same 2000 questions four times. Each generation’s mixture trains XS, S, and L models.

Table 1 reports Bamboogle auto-eval of 54.4 ± 3.6%, 61.9 ± 1.9%, and 74.0 ± 3.3% for first-generation XS, S, and L, increasing to 65.9 ± 2.6%, 69.7 ± 1.3%, and 76.1 ± 1.3% for the second generation. Prompted L scores 70.3 ± 3.5%.
Table 2’s human scores on Bamboogle are 68.8% for prompted L and 67.2%, 68.0%, and 74.4% for second-generation XS, S, and L. BamTwoogle scores are 68.0%, 63.0%, 63.0%, and 74.0%, respectively.
Figure 4 compares answers before and after self-checking. The improvement across training generations remains visible in both conditions.

Bamboogle automatic accuracy with and without self-critique across model sizes and training stages.

Mean Bamboogle automatic accuracy and standard deviation over 10 runs for pretrained, pilot, and self-improved agents.

Single-run human accuracy on Bamboogle and BamTwoogle for pretrained L and second-generation XS, S, and L.

5.3 Ablations

The XS ablations, comparisons that vary parts of the training setup, examine data generation, quantity, and filtering. Human-filtered pilot data scores 44.7 ± 3.1% against 47.2 ± 3.1% for unfiltered pilot data, a 2.5 percentage-point decrease; the authors hypothesize that filtering reduces examples for learning the prompt format while leaving problematic earlier steps inside PAST_ACTIONS.

With nearly equal training-example counts, first-generation data scores 54.4 ± 3.6% from 17970 examples, while second-generation data with one trajectory per question scores 63.4 ± 1.7% from 18007 examples. This supports the value of the second-generation data at approximately matched size, without isolating the individual causes of its improvement.
Increasing second-generation trajectories per question from one to two raises accuracy from 63.4 ± 1.7% to 65.6 ± 1.8%; four trajectories produce 65.9 ± 2.6%. Training-example counts are 18007, 36238, and 72424, respectively.
Comparing the initial draft with the final answer shows a self-critique benefit of roughly 0.5–1.0 percentage points for most models. Table 4 also shows an exception: first-generation XS drafts score 0.4 points higher than final answers.

Effects of filtering, data generation, trajectory counts, and training-example counts on fine-tuned PaLM 2-XS.

6 Discussion

The discussion describes the approach as process supervision, meaning training on individual workflow states and actions, combined with exploration and AI ranking. Because training does not require a correct final answer, the authors argue that useful steps can be retained from unsuccessful trajectories and that the procedure can handle questions without a single well-defined answer.

Repeated randomly sampled trajectories and large-model judging increase evaluation cost. Automatic evaluation enables detailed comparisons but still requires substantial computation.
Self-checking benefits are generally larger for larger models and do not clearly grow through self-improvement. Before self-checks, second-generation XS, S, and L score 65.6 ± 3.0%, 69.2 ± 1.8%, and 75.0 ± 1.3%.
The authors identify manually constructed prompts, small evaluations, limited models, and a single tool as limitations. Generalization to multiple or unseen tools, gains beyond two iterations, and eventual saturation remain open questions.

Bamboogle automatic accuracy before self-critique, with standard deviations and differences from final-answer scores.

The paper situates its search-based long-form answering alongside WebGPT, which learns from human demonstrations and reinforcement learning, and its agent fine-tuning alongside FireAct, another agent fine-tuning method. Its distinction is generating and selecting training steps through AI feedback without using task-answer labels for the main training process.

Human involvement remains in the few-shot prompt exemplars, benchmark construction, evaluation, and the separate pilot baseline. The paper also discusses DSP, which optimizes prompt demonstrations using labeled examples and can fine-tune agent components.
Relative to STAR and ReSTEM, self-training methods that use final-answer correctness, this method does not use correctness as its training signal. Relative to ReST and RAFT, it directly prompts a model for rankings instead of employing a dedicated reward model trained on human preferences.

8 Conclusion

The paper demonstrates an iterative synthetic-data procedure that improves a multi-step search agent and transfers much of the prompted large teacher’s performance to smaller models. The experiments establish this result for the specified PaLM 2 search workflow over two iterations.

Appendix

A.1 PROMPTS documents the structured prompt interfaces in Listings 1–8. The full reasoning prompts use nine examples for decision, six for summarization, five for answer generation, six for relevance checking, and five for grounding checking; Listings 1–6 provide first-example fragments and the shared self-check prefix.
Listing 7 supplies the full five-example auto-eval prompt, which tests whether ANSWER implies REF_ANSWER. Listing 8 supplies the ranking prompt with no examples, the only prompt that does not use the Python-code format; it ranks candidate actions for relevant, grounded answers with minimal steps. A.2 ADDITIONAL DATA contains Table 4’s draft-answer scores before self-critique, including the exception to a uniformly positive self-check effect.

Brief Thoughts

The comparison between first-generation and second-generation data at one trajectory per question is particularly informative: nearly equal training-example counts accompany a 9.0 percentage-point XS improvement. Within this experiment, changing the generated data yields a larger gain than increasing second-generation trajectory counts, although the comparison does not isolate the contribution of action ranking from changes to the trajectory generator. The smaller-model result depends on the comparison condition: S nearly matches prompted L on Bamboogle auto-eval, but XS and S both fall five percentage points below prompted L on BamTwoogle human evaluation. Reproducing the workflow is constrained by its internal search API and partially shown reasoning prompts. Strong agreement between aggregate automatic and human scores does not establish that every individual answer or citation is judged correctly.