Recursion of Thought: A Divide-and-Conquer Approach to Multi-Context Reasoning with Language Models
Year: 2023

Abstract:
Generating intermediate steps, or Chain of Thought (CoT), is an effective way to significantly improve language models’ (LM) multistep reasoning capability. However, the CoT lengths can grow rapidly with the problem complexity, easily exceeding the maximum context size. Instead of increasing the context limit, which has already been heavily investigated, we explore an orthogonal direction: making LMs divide a problem into multiple contexts. We propose a new inference framework, called Recursion of Thought (RoT), which introduces several special tokens that the models can output to trigger context-related operations. Extensive experiments with multiple architectures including GPT-3 show that RoT dramatically improves LMs’ inference capability to solve problems, whose solution consists of hundreds of thousands of tokens.

heading: Future Work
text: Nonetheless, the impressive experimental results suggest that the multi-context paradigm of RoT might play an important role in future LMs. In the supplementary file, we provide our PyTorch (Paszke et al. , 2019) implementation that can fully reproduce the experiments. 2 Related Work
Scratchpad (Nye et al. , 2021) is one of the earliest approaches demonstrating that fine-tuning lan-
623
guage models to produce CoT can largely improve reasoning accuracy. In the paper, the authors also mention the confined context size as a major hurdle to scaling their method. More recently, it has been found that sufficiently large pre-trained language models can be induced to produce CoT, by simply tuning the prompt (Wei et al. , 2022; Kojima et al. , 2022). Several concurrent works extend CoT prompting to decompose complex problems into smaller problems (Dua et al. , 2022; Zhou et al. , 2022; Khot et al. , 2022). Although these works also share the principle of divide and conquer like RoT, they mostly focus on improving the reasoning accuracy of relatively small problems whose solutions usually can fit in a single context. On the other hand, we focus on solving problems that the solutions are orders of magnitude longer than the context size. More detailed description of related work can be found in Appendix A. 3 Recursion of Thought (RoT)
3. 1 Inference
We start with how an RoT-trained LM performs at test time. RoT is a model-agnostic framework, whose only requirement is that the model can infer p(xi+1∣X1∶i), the probability of the next token xi+1 given a sequence X1∶i = [x1;. ;xi]. For recursive context control, we introduce the following special tokens: GO , STOP , and THINK. GO and STOP respectively mark the start and end of a problem sequence. They can be nested inside another GO -STOP pair to indicate a subproblem. THINK initiates a recursion procedure. RoT teaches a model how to use these tokens so that it can perform divide-and-conquer problem-solving. We formulate each inference context of a QA problem, denoted X , as the following concatenation:
X = [Q;Qsub,1;Asub,1;. ;Qsub,N ;Asub,N ;A] (1)
where Q and A are the main question and answer sequence, and Qsub,∗ and Asub,∗ are those of the top-level subproblems. Although a subproblem can have smaller, lower-level subproblems recursively, only the top-level subproblems remain in an RoT context. During inference, a model is given Q and tasked to generate the rest. Questions (Q and Q
sub,∗) start with a GO token, and answers (A and A
sub,∗) end with a STOP token. In the base cases, contexts do not have (Qsub,∗, Asub,∗) pairs. Figure 1 presents an example of solving 408 + 351 for better understanding. The pseudocode and more detailed illustrations can be found in Appendix B and E. RoT starts by initializing the context X with the original question Q (i. e. , GO 4 0 8 + 3 5 1 = in Figure 1). Then, similar to CoT, the model solves multiple subproblems (generating Q
sub,∗ and Asub,∗) before producing the final answer. However, there is a key difference: instead of producing a sub-answer directly, the model outputs the THINK token. This special token triggers a recursive process that separates the sub-question in a new context. If the new context is a base case (i. e. , X
2, X4, and X5), the answer is produced directly. Otherwise, the model recursively solves more subproblems. If enough subproblems are solved, the model generates the final answer ending with a STOP. Once an answer is returned to the previous context, it replaces the THINK token, and the generation continues. For tail recursion, where the last subquestion’s answer becomes the final answer, we additionally introduce the TAIL token. If TAIL is used in the place of a GO token in the last subquestion Qsub,N , its answer Asub,N is treated as the final answer A. 3. 2 Training
Currently, we train RoT in a supervised manner, using ground truth (GT) intermediate steps that include when to output the special tokens. The GTs are constructed following the standard procedures developed for humans. For example, the procedures for arithmetic problems are borrowed from elementary school math. More details can be found in Appendix H. We leave training RoT with less supervision as a future work. Each training example is constructed as a pair of a ground truth context sequence X and the corresponding target sequence Y. An example and the pseudocode for creating a target sequence are presented in Figure 3 and Algorithm 2 in Appendix D. Overall, Y is a copy of X except for the parts corresponding to Q and Asub,∗. Since the question Q is always given in a context, Q is replaced by special PAD tokens, which are excluded from the loss function. Each subproblem’s answer Asub,n is replaced by a THINK token followed by several PAD s that fill in the rest to make sure ∣X∣ = ∣Y ∣. This way, the model is trained to output THINK instead of the first token of Asub,n. Since the whole A sub,n will be returned from the recursive process
outputs the THINK token when it needs to generateAsub,∗, the answer of a subproblem. The THINK token triggers a recursive process and is later replaced by the answer returned from the process. and replace the THINK during inference, we do not need a training signal for the rest of Asub,n. Given a pair (X,Y ), the training objective is defined as follows:
L = −∑ i I[yi+1 ≠ PAD] log p(yi+1∣X1∶i) (2) where I is the indicator function that excludes PAD s from training. Its form is almost identical to the standard LM objective: LLM = −∑i log p(xi+1∣X1∶i), which is to predict the next token given previous tokens. Therefore, any sequence model is trained in the standard way, i. e. , end-to-end via stochastic gradient descent. 4 Experiments
4. 1 Baselines We compare RoT with two baselines. The first one is to output an answer directly from a question, which we call Without Thought (WT). The other one is to generate all the intermediate steps before the answer without recursion (Nye et al. , 2021), which we refer to as Chain of Thought (CoT; not to be confused with the CoT prompting (Wei et al. , 2022)). We construct the ground truths for CoTs by unraveling the same recursive process which we design for RoT, into a single context sequence (see Appendix F for examples). Therefore, the number of tokens to generate while solving a problem is the same for both CoT and RoT (if we do not count the THINK tokens). However, the sizes of the individual contexts of CoT are far longer than those of RoT due to the recursively nested subproblems, limiting the range of solvable problems. Refer to Appendix M for a more detailed analysis of the context sizes. For a fair comparison, we train these
baselines and do not use any prompting technique. When evaluating, we consider a problem to be correctly solved only if all the intermediate steps and the answer are correct. 4. 2 The Reasoning Problems
To evaluate the reasoning capabilities, we test four basic arithmetic tasks and four algorithmic tasks: addition, subtraction, multiplication, division, longest common subsequence, longest palindromic subsequence, 0-1 knapsack, and matrix chain multiplication. The details can be found in Appendix G. We choose these tasks because we can easily increase the problem difficulty while being able to get ground truth solutions. Therefore, we can test problems whose solution contains hundreds of thousands of tokens. All problems are formulated in pure sequence modeling, without any external programs (e. g. , calculator) involved. 4. 3 Experiments with GPT-3
Using the OpenAI API, we fine-tune GPT-3 for each reasoning task in §4. 2 for 10K steps with a batch size of 256. The results are presented in Figure 2a, and the technical details are described in Appendix I. Each point in the graphs represents one experiment at a certain problem difficulty. We report the accuracy on a test set of 1K unique problems randomly sampled as explained in Appendix G. To the best of our knowledge, the problems at this scale (e. g. , 48-digit addition/subtraction and 16-digit multiplication/division) have never been solved by any LM without the help of external programs. For reference, Minerva (Lewkowycz et al. , 2022) achieves around 80% accuracy on 10-digit addition and 20% on 18-digit addition. Results. Even WT fine-tuning cannot make GPT3 deal with such a level of complexity, while CoT is not applicable due to the context limit of 2048. The green dotted lines mark the maximum difficulty that can be handled by CoT under the context limit. On the other hand, RoT enables the GPT-3 to achieve near-perfect scores in every experiment. As presented in Appendix M, solving each problem requires up to tens of thousands of tokens. Without any architectural change, RoT makes GPT-3 handle these extremely complex problems. 4. 4 Experiments with Tiny Language Models Recent research on reasoning has been mostly focused on extremely large pre-trained LMs. In this section, we show an interesting result that RoT can make even tiny models, without any pre-training, perform convoluted reasoning procedures. We test the two basic sequence model architectures: a Transformer (Vaswani et al. , 2017) with 536K parameters and an LSTM (Hochreiter and Schmidhuber, 1997) with 272K parameters. These models are more than a million times smaller than the recent 540B-parameter PaLM (Chowdhery et al. , 2022). The context limit is set to 2048 for the Transformer and 512 for the LSTM. By virtue of their small sizes, we conduct far more extensive experiments than GPT-3, which are presented in Figure 2b and Figure 2c. For each experiment, we train a randomly initialized model and evaluate it on a test set of 30K unique problems. We repeat each experiment eight times and report
the average and standard deviation of the accuracies. With the tiny Transformer, we experiment to the extent that even humans would find daunting. For example, we test addition/subtraction up to 64 digits and multiplication/division up to 32 digits. Note that a 32-digit number cannot even fit into the 64-bit integer datatype. Throughout the experiments, we observe consistent patterns: (i) WT’s accuracy drops most quickly as the problem difficulty increases, (ii) CoT achieves near-perfect accuracy, but it can only be applied to simple problems due to the context limit, (iii) RoT achieves near-perfect accuracy and can be scaled up to extremely complex problems. Despite the small sizes, RoT makes the Transformers master all types of extremely complex problems. We do not test more difficult problems mainly because the evaluation becomes too costly, not because RoT is incapable of learning them