[Paper Review] ReST meets ReAct: Self-Improvement for Multi-Step Reasoning LLM Agent

Renat Aksitov, Sobhan Miryoosefi+11 more|arXiv (Cornell University)|Dec 15, 2023

Topic Modeling4 citations

TL;DR

This paper proposes a self-improving ReAct-style LLM agent for multi-step reasoning by combining AI feedback with a ReST-like iterative fine-tuning process. Using synthetic trajectories from a large model, the method enables distillation into small models with performance comparable to the teacher model after just two training iterations, achieving high accuracy with 100x fewer parameters.

ABSTRACT

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.

Motivation & Objective

• To address the challenge of improving multi-step reasoning LLM agents without relying on costly human-annotated trajectories.
• To enable continuous self-improvement of agents through iterative synthetic data generation and AI feedback, avoiding end-to-end differentiable training.
• To demonstrate effective distillation of a large, high-performing agent into smaller models while preserving performance on complex, compositional question-answering tasks.
• To establish a proxy evaluation framework using the Bamboogle and BamTwoogle datasets for auto-eval of long-form, grounded answers.
• To explore whether self-critique and iterative refinement can enhance agent robustness and reasoning quality in process-based, tool-augmented LLM agents.

Proposed method

• The agent follows a ReAct-style workflow, interleaving chain-of-thought reasoning with search tool calls and observations to answer complex questions.
• Synthetic trajectories are generated via self-improvement loops: in the 'grow' phase, the current policy samples full reasoning trajectories; in the 'improve' phase, trajectories are filtered and ranked using LLM-based reward feedback.
• The method employs growing-batch reinforcement learning with AI feedback, where the policy is iteratively fine-tuned on progressively larger synthetic datasets derived from prior agent behavior.
• Self-critique steps are integrated into the agent’s reasoning flow to verify answer relevance and grounding in retrieved snippets, improving final output quality.
• The final distilled model is trained on synthetic data generated during the self-improvement process, enabling performance close to the original large model despite having two orders of magnitude fewer parameters.
• Evaluation relies on automated metrics using the Bamboogle and BamTwoogle datasets, with multiple agent runs per question to account for stochasticity and ensure robustness.

Experimental results

Research questions

• RQ1 Can a ReAct-style LLM agent be effectively improved through iterative self-improvement using only AI feedback and synthetic data?
• RQ2 To what extent can synthetic data generated from a large model’s reasoning traces enable distillation into smaller models with comparable performance?
• RQ3 Does self-critique improve reasoning quality in multi-step agents, and is this effect consistent across model sizes?
• RQ4 How many iterations of self-improvement are needed to achieve non-trivial performance gains, and is there a saturation point?
• RQ5 Can the ReST-like self-improvement framework be extended to multi-tool or unseen-tool scenarios, and what modifications are needed?

Key findings

• After just two iterations of the self-improvement loop, a fine-tuned small model achieves comparable performance to the original large model on the Bamboogle benchmark, with a mean accuracy of 78.5% and standard deviation of 2.1% over 10 runs.
• The distilled small model achieves performance on par with the teacher model despite having two orders of magnitude fewer parameters, demonstrating effective knowledge distillation from synthetic trajectories.
• Self-critique steps provide a small but positive performance boost, with the effect being more pronounced in larger models and independent of the self-improvement process.
• The method enables reliable auto-evaluation of agent trajectories using PaLM 2-L for scoring, which aligns well with human raters and supports robust hyperparameter analysis.
• The ReST-like iterative training process significantly improves agent performance without requiring human-annotated data, relying solely on AI feedback and synthetic data generation.
• The approach shows promise for scaling to multi-tool environments, though future work is needed to determine if it generalizes to unseen tools and whether additional architectural changes are required.