Re-Tuning: Overcoming the Compositionality Limits of Large Language Models with Recursive Tuning (2024)

Abstract: We present a new method for large language models to solve compositional tasks. Although they have shown strong performance on traditional language understanding tasks, large language models struggle to solve compositional tasks, where the solution depends on solving smaller instances of the same problem. We propose a natural approach to solve compositional tasks recursively. Our method, Re-Tuning, tunes models to break down a problem into subproblems, solve those subproblems, and combine the results. We show that our method significantly improves model performance on three representative compositional tasks: integer addition, dynamic programming, and parity. Compared to state-of-the-art methods that keep intermediate steps towards solving the problems, Re-Tuning achieves significantly higher accuracy and is more GPU memory efficient.

## introduction large language models (llm) have obtained the state-of-the-art performance on a wide set of tasks @xcite @xcite @xcite . however, recent studies @xcite @xcite show these models struggle to generalize to compositional tasks, where the solution depends on solutions to smaller instances of the same problem. an example task, integer addition, is shown in figure 1a . when calculating '1234 + 4567', we first break the problem into a smaller subproblem '234 + 567'. after obtaining the solution to this subproblem, the original problem is partially solved. similarly, to solve '234 + 567', we first sum '34 + 67'. this recursion is the fundamental operation to solve compositional tasks. however, no existing approach has explicitly captured the recursive nature of compositional tasks. in this paper, we propose a recursion-based method for llms to better solve compositional tasks. more specifically, we adopt a top-down approach to solve problems recursively. we train llms to recursively call themselves on subproblems of reduced size, recognize and solve the base case directly, and combine the solutions up the associated call stack to obtain the solution to the original problem (figure 1a ). the above procedure is referred to as recursive tuning (or re-tuning in short). the basic idea behind re-tuning is motivated by two lines of work. first, recent work @xcite @xcite show that training llms on high-quality scratchpad data, which includes intermediate steps towards solving a problem, can improve performance on certain compositional tasks such as integer addition and parity. instead of using the intermediate steps to train models, which is computationally costly, re-tuning breaks down the problems into smaller and smaller subproblems. each subproblem runs independently within its own context in the associated call stack. the solution to each subproblem is then propagated up the call stack to produce the final solution. since each level of the call stack only includes the information necessary to solve the current subproblem, models can more easily attend to the relevant context, improving the accuracy of solving each subproblem. second, our tuning process is reminiscent of recent works that incorporate tool use in llms @xcite . similar to how these models call a tool and resume generating output based on the output of the tool, with re-tuning, the models call themselves on a subproblem and resume generating after receiving the subproblem's solution. we empirically evaluate the performance of re-tuning on three representative compositional tasks: integer addition @xcite , a dynamic programming problem @xcite , and the parity problem @xcite . our results show re-tuning improves the average performance of llama 7b and llama 13b on all tasks by 37.4% and 31.9% over baseline training. compared to scratchpad training, our improvement is striking, with average improvements of 34.5% and 36.7% on llama 7b and llama 13b respectively. importantly, we show re-tuning saves significant gpu memory compared to the scratchpad method when training. we hope our results foster future research on recursive learning of large foundation models. ## approach we present re-tuning in this section. re-tuning recursively tunes llms to solve compositional tasks. specifically, the method involves (1) recursively decreasing the size of the problem, (2) solving the base case, and (3) passing the solutions up the recursion stack, solving subproblems of increasing complexity along the way. first, with re-tuning, an llm recursively calls itself on subproblems of decreasing length or complexity. for example, when adding 1234 + 5678, the llm calls itself to add 234 + 678. this call is then sent to a new context in which the llm calls itself to add 34 + 78, which is again sent to a new context where the llm calls itself to add 4 + 8. next, the base case is solved. the base cases are easy enough to be solved directly in the same context. for the integer addition problem, the base case is to add the two least-significant digits together (e.g., adding 4 + 8). finally, the subproblem solutions are passed up the recursive call stack. specifically, subproblem solutions are appended directly after the associated call in the context one level up the call stack. again, sticking with integer addition, it helps to know the sum of 4 + 8 when tasked with adding 34 + 78. as such, the llm-generated solution to 4 + 8 is appended to the context tasked with solving 34 + 78. this process of propagating subproblem solutions continues up the recursive call stack until the solution to the first recursive call is passed to the initial context, which helps to solve the initial problem. to accomplish this, we train llms to (1) generate recursive subproblems, (2) solve base cases, and (3) use the answers propagated up from these recursive calls in their computation for the problem in the current context. we do so by randomly generating a set of seed data from which we programmatically construct training instances for all three types (see figure 1a ). during generation, the model can designate some of its generated text to be a recursive call by enclosing the text between 'call: ' and '\n'. once a recursive call is made, we stop generating in the current context and prompt the model with the call in a new context. in each new context, we follow the exact same generation procedure, except for the base case where the model learns to directly output the answer rather than making another recursive call. when generation in the new context is complete, we take the subproblem solution, which is separated by the text '\nanswer: ', and append that to the context one level higher in the associated call stack. then we continue generation in the initial context. the pseudo-code for the recursive generation procedure is in figure 1 . in the integer addition example, the model only needs to generate one recursive call in each context. however, our method works more generally than this. many recursive calls may be generated in a single context. for example, the dynamic programming problem we describe below requires multiple recursive calls in the initial context to solve the problem. ## experiments we consider three tasks: integer addition, a dynamic programming problem, and the parity problem. for each task, we train 3 types of models: baseline, scratchpad, and re-tuning. the baseline models were trained to simply output the solution to the problem. the scratchpad models were trained to generate a scratchpad @xcite containing intermediate reasoning steps before generating the final solution to the problem. the re-tuning models are as described above. during evaluation, we consider both indistribution and out-of-distribution (ood) data. the in-distribution data are those with problem lengths that were seen in training and the ood data are those with problem lengths longer than seen in training. for example, on the integer addition task, the training data consists of numbers with lengths up to 15 digits. evaluation examples with 1-15 digits are considered in-distribution and examples with 16 or more digits are considered ood. we train llama 7b and 13b @xcite using low-rank adapters @xcite . see appendix a.2 for additional details on the training setup. additionally, we provide results on the smaller galactica @xcite 125m and 1.3b parameter models, in appendix b. ## analysis and further discussion in this section, we conduct additional experiments in order to better understand the effects of various mechanisms behind re-tuning. ## related work several works have explored the length generalization ability of llms on compositional problems. @xcite suggests that llms solve compositional tasks via "linearized subgraph matching" and thus fail to learn the underlying algorithm necessary to solve more complex problem instances. @xcite showed that training on a combination of in-context learning and scratchpad prompting could enable better performance. similarly, re-tuning involves training pretrained llms to make recursive calls in order to improve performance on compositional tasks. other works have studied length generalization on small, purpose-built transformer models. @xcite and @xcite showed that training small transformer models from scratch on scratchpad data could enable better length generalization. @xcite showed that transformers can achieve strong ood performance on addition by using special positional embeddings. various papers have explored the idea of llms prompting themselves or other llms, although, to our knowledge, no papers explicitly train a language model to do so. @xcite prompts a language model to break a problem down into simpler steps and then prompts itself to solve each step individually in a sequential, non-recursive manner. similar methods have been proposed as a way to improve the logical consistency of the generated responses @xcite @xcite . @xcite use a language model to generate prompts by extracting relevant information from the context. similarly, re-tuning generates a recursive call that includes the relevant information for solving a simpler subproblem. a recent topic of interest has been teaching language models to use tools @xcite @xcite @xcite , which often involve stopping the generation and waiting for the tool output before continuing with generation. re-tuning can be interpreted as teaching llms to use themselves as a tool. several papers have investigated the ability of language models on arithmetic tasks @xcite @xcite @xcite @xcite . in many cases, it was noticed that performance was significantly worse on problems longer than those seen during training. ## conclusion we study the problem of solving compositional tasks with large language models. we propose a new tuning paradigm that decomposes the original compositional problem into smaller and smaller instances of the same type, solves each, and combines the results to produce the final answer. to the best of our knowledge, our method is the first to utilize the recursive property of compositional tasks. experimental results on three representative compositional tasks demonstrate the effectiveness of our method. our method not only significantly outperforms standard training and state-of-the-art methods, especially on out-of-distribution problem instances, but is also more memory efficient during training. we hope our method can be applied to more tasks where recursive computation is inherent and computational resources are limited.