Learning from Mistakes via Cooperative Study Assistant for Large Language Models (2023)

Abstract: Large language models (LLMs) have demonstrated their potential to refine their generation based on their own feedback. However, the feedback from LLM itself is often inaccurate, thereby limiting its benefits. In this paper, we propose Study Assistant for Large LAnguage Model (SALAM), a novel framework with an auxiliary agent to assist the main LLM in learning from mistakes through interactive cooperation. In the gathering phase, the student assistant agent probes the main LLM, analyzes its errors, and collects the interaction in a mistake memory. During the examination phase, the study assistant provides guidelines by retrieving relevant cases to help the main LLM anticipate and avoid similar errors. We first investigate the effectiveness of a general study assistant and then customize it to provide LLM-specific guidance through imitation learning from successful guidance experiences. Our experiments on three LLMs using two challenging frameworks demonstrate that SALAM can significantly boost LLMs by an accuracy margin of up to 6.6 on BBH and 12.6 on BBQ.

## introduction large language models (llms) have demonstrated remarkable performance in a wide range of tasks @xcite @xcite . their effectiveness is further enhanced by human instructions and feedback, allowing them to better align with human intentions @xcite @xcite . furthermore, recent studies show that llms can also benefit from their own feedback to avoid mistakes, similar to human reflection @xcite . there are two main limitations to existing selfreflection methods. first, they rely on the correctness of the guidance, particularly in determining figure 1 : salam consists of two agents: a main llm and a study assistant. the main llm generates responses (in blue), while the study assistant provides guidance (in green). during the mistake-gathering phase, the main llm interacts with the study assistant, receiving feedback to refine its responses. the study assistant compares the main llm's response with the ground truth, providing guidance and collecting the mistakes made by the main llm. in the examination phase, the study assistant retrieves relevant mistakes from the mistake memory for a new query and provides guidelines without knowing the ground truth. when to terminate reflection and accept the current response. inaccurate guidance can mislead the llm by either prompting it to refine an already acceptable generation or prematurely halting the refinement of an undesired generation @xcite . prior studies have attempted to address this by utilizing additional learned discriminators and employing a value-based threshold as the termination signal @xcite @xcite , or by providing few-shot examples to encourage llms to develop their own discernment when accepting responses @xcite @xcite . however, the reliability of such criteria remains uncertain. moreover, the reflection methods are limited to the unsuccessful experiences of the query they are addressing, without acknowledging mistakes made on other queries. consequently, when confronted with a new query, the llm cannot fully utilize the experience gained from similar cases and may repeat past errors. the lack of global reflection results in an inefficient revision process. to address these challenges, we propose study assistant for large language model (salam). this novel framework introduces a cooperative agent, guiding the main llm to learn from its mistakes. it is inspired by how humans learn from their mistakes: maintaining a collection of mistakes and analyzing common misunderstandings. salam includes two cooperative agents: a main llm responsible for problem-solving, and a study assistant that collects previous error cases and provides guidance to improve the main llm's performance. the framework consists of two phases: the mistake-gathering phase and the examination phase. during the mistake-gathering phase, the llm interacts with the study assistant to receive feedback and refine its answers. simultaneously, the study assistant collects mistakes and provides guidance based on the ground truth. in the examination phase, the study assistant retrieves similar mistakes for a new query and provides guidelines to clarify misunderstandings and prevent the llm from repeating previous errors. for example, in figure 1 , the study assistant analyzes the llm's current response of '02/11/2002' compared to the ground truth of '02/12/2002' and provides the guideline 'identify the correct date from which calculations should be made' to help the llm refine its response. our proposed salam enjoys three advantages: (1) flexible: salam is a versatile framework that can be directly adapted to any llm. additionally, it has the capability to provide llmspecific guidance by fine-tuning the study assistant on the specific behaviors of the llm. (2) lightweight: in contrast to knowledge distillation, where a large teacher/teacher assistant is used to improve downstream task performance, the study assistant in salam is a small model focused on providing feedback based on mistakes. it is more cost-effective to fine-tune the small study assistant once for all downstream tasks, compared to finetuning the large llm for different complex tasks. we evaluate the effectiveness of salam on three llms: flan-t5 @xcite , gpt-neox @xcite , and llama @xcite . we use 27 tasks from bbh @xcite and bbq @xcite , which evaluate two crucial aspects of llms: reasoning ability and potential social bias. our contributions are as follows: ## related work feedback from language models large language models (llms) have exhibited a remarkable capability for providing feedback. the feedback from llms can be in the form of real numbers to evaluate the quality of the generation @xcite , or textual instruction to guide the refinement @xcite . for instance, @xcite provided feedback grounded in evidence from external knowledge. reflexion @xcite generated textual feedback utilizing trajectory history and dynamic memory with the help of fewshot examples and the signal from the environment. self-refine @xcite manually created several few-shot examples for each task to prompt llm to provide feedback. it stops reflection when it achieves the maximum iteration or exceeds the threshold value. @xcite uses responses from multiple agents and the debate between these agents as the feedback. however, @xcite finds that such intrinsic self-correction is unreliable without the correct label to determine when to stop the refinement. instead of previous instance-specific feedback, in this paper, the study assistant collects a set of mistakes to generate global feedback on model behaviors. furthermore, the study assistant used the ground truth to provide feedback, making the feedback more reliable. learning from feedback plenty of work has been done to investigate how to utilize feedback @xcite . one is to filter undesirable data based on feedback and use the filtered data to finetune the model @xcite . the other is to train a reward function and take it as the reward function in the reinforcement learning @xcite @xcite . benefiting the llms' ability to follow instructions, recent researchers add textual feedback into the prompt and directly ask models to revise their response @xcite . moreover, the feedback can be one time @xcite , or multiple times @xcite . in this work, we use feedback as the instruction of the main llm and ask it to refine its answer. teacher-student learning teacher-student learning is a knowledge distillation method to transfer knowledge from a larger teacher model to a smaller student model @xcite . the goal is to produce similar results as the powerful teacher model with fewer parameters and computational costs. the teaching assistant is an intermediate model mimicking the behavior of the teacher model and then teaches the student model. it usually has a medium size between the student and the teacher @xcite . recently, a lot of work has tried to distill knowledge in large language models to enhance the capability of small models, such as commonsense @xcite and the reasoning ability @xcite . unlike knowledge distillation, the study assistant in salam does not need a stronger capability on the downstream tasks. it is designed to analyze the output of the base model given the ground truth, providing a guideline for the base model to avoid similar mistakes. ## study assistant agent for llms in this section, we introduce salam framework. salam consists of two agents: a main llm m and a study assistant t . the m is responsible for solving the downstream tasks while the study assistant t provides text feedback for m to refine its answer. the goal of this framework is to improve m performance by the interactive cooperation between the two agents. there are two phases: in the gathering phase, the study assistant t collects mistakes while identifying common misunderstandings and providing helpful guidance for revision; the examination phase involves using m on new queries. the key difference is that in the gathering phase, the study assistant t has access to the ground truth, while in the examination phase, it does not. specifically, suppose there is a set of n queries @xmath0, q (1) , • • • , q (n ) } in the gathering phase. for each query @xmath1, the main llm m generates an initial response y 0 and the study assistant t provides text feedback a 0 based on the comparison between the current response and the ground truth y. then m generates a new response y 1 under the guidance of the feedback. there can be multiple iterations between these two agents until m gets the correct answer (or achieves the maximum iteration number l): {(y 0 , a 0 ), • • • , (y l , a l )}. these iterations are stored in the mistake memory o err . during the examination phase, given a new query q, the study assistant t retrieves the most relevant queries from o err and provides text feedback as the pre-hoc instruction for m. ## experiment we conduct experiments in two challenging benchmarks with 27 tasks: bbh and bbq, evaluating salam's ability to guide complex reasoning and reduce social biases. we further conduct comprehensive analyses from different aspects to enhance the understanding of salam. ## conclusion in this paper, we introduce a novel framework, the study assistant for large language model (salam), designed to aid llms in learning from their mistakes by interactive cooperation between the study assistant and the llm. this framework is inspired by the methods human study assistants use to support students, by identifying common errors and providing guidance. the student model sends its generations to the study assistant and refines these based on the feedback received. the study assistant identifies errors, offers feedback, and gauges its success by the student model's performance improvement. we validated the effectiveness of salam on the bbh and bbq benchmarks, showing significant improvement in the model's performance. furthermore, we use the llms' performance as the signal to further finetune the study assistant for model-specific guidance. we believe that our method offers a novel way to augment llms by the cooperation between multiple agents.