Teaching Language Models to Self-Improve through Interactive Demonstrations (2024)

Abstract: The self-improving ability of large language models (LLMs), enabled by prompting them to analyze and revise their own outputs, has garnered significant interest in recent research. However, this ability has been shown to be absent and difficult to learn for smaller models, thus widening the performance gap between state-of-the-art LLMs and more cost-effective and faster ones. To reduce this gap, we introduce TriPosT, a training algorithm that endows smaller models with such self-improvement ability, and show that our approach can improve LLaMA-7B’s performance on math and reasoning tasks by up to 7.13%. In contrast to prior work, we achieve this by using the smaller model to interact with LLMs to collect feedback and improvements on *its own generations*. We then replay this experience to train the small model. Our experiments on four math and reasoning datasets show that the interactive experience of learning from and correcting its *own* mistakes is crucial for small models to improve their performance.

## introduction large language models @xcite together with techniques such as fewshot prompting @xcite and chain-of-thought (cot) prompting @xcite have been shown to be effective in achieving strong performance on various downstream language tasks. more recently, a new way to adapt llms to downstream tasks has captured the attention of many researchers, namely to further enhance the llm's downstream task performance by asking the llm to provide feedback on its own generations and then use the feedback to revise its outputs @xcite @xcite . this process is often called "self-improvement", and has proven to be an effective technique to make the llm's generations more diverse, more precise, or more faithful to a given piece of knowledge @xcite @xcite . however, @xcite ; huang et al. @xcite found that the ability to generate critical feedback or to self-improve is hardly evident in smaller models foot_0 . similarly, @xcite found that fine-tuning smaller models (e.g. 7-13b) with self-improvement demonstrations from llms can still fail on tasks such as math, reasoning, and factuality. following these previous works, we performed a similar study on two math and reasoning tasks in figure 1 . we compared the accuracy of the final answer generated by prompting method codex (175b) llama+ft (7b) input q: ((4 -0 * -1 * 8 + 6) + (0 + 4 + 9 * 5 * -2)) = attempt: a 175b codex @xcite to self-improve, with prompting or training a llama-7b model to self-improve using demonstrations from codex @xcite . in figure 1 , we surprisingly find that smaller models performed worse using prior selfimprovement-related methods than simply training on ground-truth step-by-step rationales (+ft). by comparing the generated solutions from codex-175b and llama-7b, we find that smaller models, such as llama-7b, not only make more mistakes, but also different types of mistakes compared to an llm ( motivated by this finding, we propose tripost, a training algorithm that can more effectively train a small model to learn from its mistakes, generate feedback, and improve its performance on math and reasoning tasks. tripost is an iterative algorithm consisting of three stages: interactive trajectory editing, data post-processing, and model training. similar to the exploration stage in reinforcement learning, tripost first creates improvement demonstrations using the small model to interact with the expert llms or relevant python scripts. then, tripost postprocesses the collected data by filtering out failed improvement attempts, and then re-balances the dataset to disincentivize the model from trying to self-"improve" when it is not needed. finally, tripost replays the post-process dataset @xcite , and trains the smaller model using weighted supervised learning. tripost repeats entire the process several times. we evaluate our approach on four maths and reasoning datasets from the big-bench hard @xcite collection, and find that tripost-trained models can use its learned selfimprovement ability to improve their task performance. we also find that tripost-trained models achieve better in-domain and out-of-domain performance than models trained using just the ground truth step-by-step rationales and trained using direct llm demonstrations @xcite . this paper makes the following contributions: ## approach tripost is an algorithm that trains a small language model to self-improve by learning from its own mistakes. each iteration of tripost consists of three stages. on a high level, we first collect feedback: ...... update: ...... feedback: ...... filter question attempt: feedback: update: question attempt: feedback: split re-balance edited trajectory feedback: ... update: ... attempt ... feedback: ... update: ... ... question attempt: feedback: update: a set of improving trajectories by using a smaller model m θ to interact with llms. we use m θ to generate initial attempts and then use a feedback module fbk and an improvement module imp to edit parts of the m θ generated attempts. this creates a trajectory that includes attempts generated by the small model, with feedbacks and improvements tailored to the small model's capability (figure 2 ). next, we post-process the collected trajectories by 1) using scripts and other heuristics to filter out failed "improvement" attempts; and 2) re-balancing the dataset using both directly correct attempts and the improving trajectories. finally, we use weighted supervised learning to train a smaller model m θ using the post-processed data. we provide an overview of our algorithm in figure 2 , and detail each of the three stages in section 2.2, section 2.3, and section 2.4, respectively. ## analysis to investigate the factors that can influence how tripost-trained models learned to attempt selfimprovement, we focus our analysis on the multistep arithmetic and logical deduction datatset. we also mainly study tripost with @xmath0.43, which has both a measurable self-improvement contribution and improvement in its task performance (see ## conclusion we introduce tripost, a training algorithm that distills the ability to self-improve to a small model and help it achieve better task performance. tri-post first creates improving trajectories using interactions between a smaller model and an llm, then post-process the collected trajectories, and finally train the smaller model to self-improve using weighted sl. we evaluated tripost on four math and reasoning tasks from the big-bench hard collection and found that it can help small models achieve better task performance. in our analysis, we find that 1) the interactive process of learning from and correcting its own mistakes is crucial for small models to learn to self-improve and 2) learning to always generate a useful feedback and a corresponding improvement can be much harder than learning to directly generate a correct answer. these findings suggest that other data formats, beyond the traditional (input, answer) pair, could be better suited for training a language model to solve a downstream task. we believe this also opens new possibilities for future work to leverage llms to improve the performance of smaller, faster models.