[Paper Review] Retroformer: Retrospective Large Language Agents with Policy Gradient Optimization

Weiran Yao, Shelby Heinecke+13 more|arXiv (Cornell University)|Aug 4, 2023

Topic Modeling4 citations

TL;DR

Retroformer introduces a policy gradient-based framework that fine-tunes a retrospective language model to improve large language agents by analyzing past failures and generating actionable feedback. By learning from environment rewards without accessing the actor model’s parameters, it boosts task success rates by 18% in 4 attempts on HotPotQA, outperforming baselines that lack gradient-based refinement.

ABSTRACT

Recent months have seen the emergence of a powerful new trend in which large language models (LLMs) are augmented to become autonomous language agents capable of performing objective oriented multi-step tasks on their own, rather than merely responding to queries from human users. Most existing language agents, however, are not optimized using environment-specific rewards. Although some agents enable iterative refinement through verbal feedback, they do not reason and plan in ways that are compatible with gradient-based learning from rewards. This paper introduces a principled framework for reinforcing large language agents by learning a retrospective model, which automatically tunes the language agent prompts from environment feedback through policy gradient. Specifically, our proposed agent architecture learns from rewards across multiple environments and tasks, for fine-tuning a pre-trained language model which refines the language agent prompt by summarizing the root cause of prior failed attempts and proposing action plans. Experimental results on various tasks demonstrate that the language agents improve over time and that our approach considerably outperforms baselines that do not properly leverage gradients from the environment. This demonstrates that using policy gradient optimization to improve language agents, for which we believe our work is one of the first, seems promising and can be applied to optimize other models in the agent architecture to enhance agent performances over time.

Motivation & Objective

• To address the lack of environment reward optimization in existing language agents, which often rely on non-differentiable verbal feedback.
• To enable gradient-based learning in language agents by training a retrospective model to analyze failures and generate actionable insights.
• To develop a plug-in module that improves agent performance without requiring access to or gradient computation through the main actor LLM.
• To investigate whether policy gradient optimization can enhance reasoning and planning in language agents across diverse environments and tasks.
• To overcome the limitations of frozen LLMs in self-reflection by training a specialized retrospective model for precise credit assignment and root cause analysis.

Proposed method

• The framework uses a retrospective model $M_r$ that is fine-tuned via policy gradient optimization using rewards from the environment, treating the actor LLM as part of the environment.
• The retrospective model generates textual feedback summarizing the root cause of prior failures and proposing a revised action plan for the next attempt.
• The method avoids backpropagation through the actor LLM by treating it as a black-box environment, enabling gradient updates only on the retrospective model.
• The retrospective model is trained on offline experience trajectories, using rewards from environments like HotPotQA to optimize its ability to generate effective, structured reflections.
• The model learns to produce two-part responses: a reflection on failure and a new, actionable plan, emerging this structure without explicit supervision.
• The approach is modular and agnostic, allowing integration with any cloud-based LLM such as GPT or Bard, and can be extended to other components like memory or summarization modules.

Experimental results

Research questions

• RQ1 Can policy gradient optimization improve the reasoning and planning capabilities of large language agents through retrospective feedback?
• RQ2 Does a fine-tuned retrospective model outperform a frozen LLM in generating actionable, credit-assignment-aware feedback for iterative task improvement?
• RQ3 To what extent can a plug-in retrospective module enhance agent performance without modifying the actor LLM or propagating gradients through it?
• RQ4 How does the structure of the retrospective feedback (e.g., separation of reflection and plan) affect agent performance and learning efficiency?
• RQ5 Can the proposed method generalize across diverse environments and tasks, particularly in high-dimensional state-action spaces like tool-use in question-answering?

Key findings

• Retroformer improves the success rate of search-based question-answering tasks in HotPotQA by 18% within 4 attempts, significantly outperforming baseline methods.
• The fine-tuned retrospective model learns to perform precise credit assignment, identifying that the agent failed due to forgetting its goal during chain-of-thought reasoning.
• Unlike frozen LLMs that produce repetitive, unhelpful reflections, the reinforced model generates structured feedback with distinct reflection and action plan sections.
• The reinforced model avoids hallucination by removing redundant or irrelevant content such as 'Next trial:' in its outputs, leading to cleaner and more actionable prompts.
• The method achieves a 53% success rate in 5 trials on a benchmark, surpassing the 50% success rate of a state-of-the-art method using a much larger LLM (GPT-3)
• The policy gradient approach enables the retrospective model to learn emergent behaviors such as root cause analysis and structured response formatting through trial-and-error with environment feedback.