Recursive Introspection: Teaching Language Model Agents How to Self-Improve

Original Paper: https://arxiv.org/abs/2407.18219

By: Yuxiao Qu, Tianjun Zhang, Naman Garg, Aviral Kumar

Abstract: 

A central piece in enabling intelligent agentic behavior in foundation models is to make them capable of introspecting upon their behavior, and reasoning, and correcting their mistakes as more computation or interaction is available. Even the strongest proprietary large language models (LLMs) do not quite exhibit the ability to continually improve their responses sequentially, even in scenarios where they are explicitly told that they are making a mistake. In this paper, we develop RISE: Recursive IntroSpEction, an approach for fine-tuning LLMs to introduce this capability, despite prior work hypothesizing that this capability may not be possible to attain. Our approach prescribes an iterative fine-tuning procedure, which attempts to teach the model how to alter its response after having executed previously unsuccessful attempts to solve a hard test-time problem, with optionally additional environment feedback. RISE poses fine-tuning for a single-turn prompt as solving a multi-turn Markov decision process (MDP), where the initial state is the prompt. Inspired by principles in online imitation learning and reinforcement learning, we propose strategies for multi-turn data collection and training so as to imbue an LLM with the capability to recursively detect and correct its previous mistakes in subsequent iterations. Our experiments show that RISE enables Llama2, Llama3, and Mistral models to improve themselves with more turns on math reasoning tasks, outperforming several single-turn strategies given an equal amount of inference-time computation. We also find that RISE scales well, often attaining larger benefits with more capable models. Our analysis shows that RISE makes meaningful improvements to responses to arrive at the correct solution for challenging prompts, without disrupting one-turn abilities as a result of expressing more complex distributions.

Summary Note

Figure 1: Recursive Introspection (RISE). Using iterative multi-round training on on-policy rollouts and supervision from a reward function, RISE trains models that are capable of improving themselves over multiple turns. At inference, we run majority voting on candidate outputs from different turns to obtain the final response.

Figure 2:Left: Problem formulation. We convert single-turn problems into multi-turn MDPs as discussed in Section 4.1. The state is given by the prompt, history of prior attempts, and optional feedback from the environment. An action is a response generated from the LLM given the state of multi-turn interaction so far. Right: Data collection. We collect data by unrolling the current model k−1 times followed by an improved version of the response, which is obtained by either (1) self-distillation: sample multiple responses from the current model, and use the best response, or (2) distillation: obtain oracle responses by querying a more capable model. In either case, RISE then trains on the generated data.

Figure 3:RISE Inference. There are two ways to query the model trained via RISE upon inference: (1) with oracle (Left): each time the model improves its response, it is allowed to check its answer against an environment and terminate early as soon as a correct answer is found; or (2) without oracle (Right): we ask the model to sequentially revise its own responses j times, and perform majority voting on all candidate outputs from different turns to obtain the final response. If the turn number j is larger than the iteration number k, the agent only keeps the most recent history with k interactions to avoid test-time distribution shift.

In the fast-paced world of artificial intelligence, especially in the domain of language models, the ability to self-improve is a coveted trait. Imagine a scenario where a language model, upon realizing its mistakes, iteratively refines its responses to achieve better accuracy. This is precisely what the recent research on Recursive IntroSpEction (RISE) aims to accomplish.

Introduction

Imagine a world where machines not only perform tasks but also learn from their mistakes and improve with each attempt. This is the essence of Recursive Introspection (RISE), a novel approach designed to enhance Large Language Models (LLMs) by enabling them to self-improve through iterative fine-tuning. This blog post delves into the transformative potential of RISE, its methodologies, findings, and implications for the future of artificial intelligence.

The Quest for Self-Improving Language Models

Modern LLMs have shown remarkable proficiency in generating human-like text, solving complex problems, and even engaging in meaningful conversations. However, their ability to learn from mistakes and improve responses sequentially has remained limited. This gap poses a significant challenge, especially for tasks requiring multiple iterations to refine answers, such as complex problem-solving and logical reasoning.

RISE addresses this challenge by transforming the way LLMs are fine-tuned, enabling them to recognize and correct their mistakes over multiple turns.

This approach is inspired by techniques in online imitation learning and reinforcement learning, aiming to teach models not just what to respond with, but how to improve their responses.

Methodology: Building the Foundation for RISE

At its core, RISE employs an iterative fine-tuning process where single-turn prompts are treated as multi-turn Markov Decision Processes (MDPs). Here's a breakdown of the key methodologies used:

1. Multi-Turn Data Collection:
• RISE collects data through on-policy rollouts, where the model generates responses to a prompt over several iterations.
• Each response is evaluated, and the model is prompted to refine its answer based on feedback, creating a history of attempts.
2. Training with Reward-Weighted Regression:
• The collected data is used to fine-tune the model using a reward-weighted regression objective.
• This approach enables the model to learn from both high-quality and suboptimal responses, emphasizing improvements over sequential turns.
3. Incorporating Feedback Mechanisms:
• At each iteration, the model is evaluated using a reward function that indicates the correctness of its response.
• The process continues until the model achieves a correct response or reaches the maximum number of iterations.

Key Findings and Results

RISE has demonstrated significant improvements in the self-improvement capabilities of LLMs. Here are some of the most notable findings from the experiments:

1. Performance Boost Over Multiple Turns:
• On datasets like GSM8K and MATH, RISE significantly improved the accuracy of models like LLaMa2-7B and Mistral-7B over multiple turns.
• The LLaMa2-7B model, for instance, saw a 17.7% improvement over five turns, while Mistral-7B achieved a 23.9% improvement.
2. Scalability with Model Size:
• The benefits of RISE scale well with more capable models, indicating its potential for enhancing state-of-the-art LLMs.
• Even with large models like GPT-3.5, RISE demonstrated notable improvements, proving its scalability and robustness.
3. Effective Error Correction:
• RISE not only improves the overall accuracy but also demonstrates the ability to correct specific errors in subsequent iterations.
• This capability is crucial for tasks requiring logical reasoning and sequential problem-solving.

Implications and Real-World Applications

The implications of RISE extend far beyond academic research, offering transformative potential for various real-world applications:

1. Enhanced AI Assistants:
• AI systems equipped with RISE can provide more accurate and refined responses over multiple interactions, improving user experience and satisfaction.
• This can be particularly beneficial for customer support, where iterative problem-solving is often required.
2. Educational Tools:
• RISE can enhance educational platforms by providing more effective tutoring and feedback mechanisms, helping students learn and understand complex subjects.
3. Autonomous Systems:
• Autonomous systems, such as self-driving cars and robotics, can benefit from RISE by continuously improving their decision-making processes based on real-time feedback.
4. Healthcare:
• In healthcare, AI systems with self-improvement capabilities can assist in diagnosis and treatment planning by refining their recommendations based on patient feedback and outcomes.

Conclusion

Recursive Introspection (RISE) represents a significant leap forward in the development of self-improving language models.

By enabling LLMs to learn from their mistakes and refine their responses over multiple turns, RISE opens new avenues for more intelligent, reliable, and effective AI systems.

As we continue to explore and enhance this approach, the future of AI promises to be more adaptive and capable than ever before.

For engineers and AI enthusiasts, RISE offers a glimpse into the next frontier of machine learning, where models not only perform tasks but also evolve and improve autonomously.

The journey of recursive introspection has just begun, and its potential is boundless.

Future Directions

While RISE has shown promising results, there are still many open questions and avenues for future research.

Scaling RISE to more iterations, exploring fully online RL techniques, and integrating RISE into general instruction-tuning pipelines are exciting directions to pursue.

The journey to create self-improving AI systems is just beginning, and RISE is a significant step in that direction.