Small Language Models Fine-tuned to Coordinate Larger Language Models improve Complex Reasoning (2023)

Abstract: Large Language Models (LLMs) prompted to generate chain-of-thought (CoT) exhibit impressive reasoning capabilities. Recent attempts at prompt decomposition toward solving complex, multi-step reasoning problems depend on the ability of the LLM to simultaneously decompose and solve the problem. A significant disadvantage is that foundational LLMs are typically not available for fine-tuning, making adaptation computationally prohibitive. We believe (and demonstrate) that problem decomposition and solution generation are distinct capabilites, better addressed in separate modules, than by one monolithic LLM. We introduce DaSLaM, which uses a decomposition generator to decompose complex problems into subproblems that require fewer reasoning steps. These subproblems are answered by a solver. We use a relatively small (13B parameters) LM as the decomposition generator, which we train using policy gradient optimization to interact with a solver LM (regarded as black-box) and guide it through subproblems, thereby rendering our method solver-agnostic. Evaluation on multiple different reasoning datasets reveal that with our method, a 175 billion parameter LM (text-davinci-003) can produce competitive or even better performance, compared to its orders-of-magnitude larger successor, GPT-4. Additionally, we show that DaSLaM is not limited by the solver’s capabilities as a function of scale; e.g., solver LMs with diverse sizes give significant performance improvement with our solver-agnostic decomposition technique. Exhaustive ablation studies evince the superiority of our modular finetuning technique over exorbitantly large decomposer LLMs, based on prompting alone.

## introduction in recent years, an astounding variety of text and nlp tasks have been accomplished by language * equal contribution as first author. † correspondence: tanchak@iitd.ac.in models (lms) @xcite -in essence, fitting continuous feature vectors to tokens and modeling smooth conditional distributions over thousands of token positions with multi-task objectives. the next generation of large lms (llms) such as t5, @xcite developed protean capabilities, extending to mathematical and logical ability, based on prompting and in-context learning. chain-ofthought (cot) prompting has been a key enabler @xcite . llms can solve middle-school word problems and equations reasonably well. it has also acquired the ability to invoke specialized external tools such as wolfram alpha @xcite . recent advances in llms have arisen largely from cleverly-engineered, ever-growing training data, rather than any significant change in network structure, which remains relatively regular, but with rapidly increasing network size and number of parameters. one outcome of such giant monolithic llms is unsustainable levels of hardware and energy @xcite to train them. meanwhile, neurologists and brain scientists have known, via fmri scans, inter alia, that cerebral functions are specialized and spatially localized @xcite . many recent complex reasoning challenges thrown at llms have a two-level character -the input task needs to be decomposed into subtasks, then the subtasks need to be solved, and finally, subtask solutions have to be consolidated and combined to solve the input task. existing approaches use the same llm to both decompose and solve the task, sometimes in tangled and uninterpretable ways. because the sharing of an llm across these functions cannot be closely controlled, very large models are needed for this double ability (decompose and solve) to emerge. staying entirely inside the llm regime, and avoiding the possibility of specialized tools, we figure 1 : working example of daslam on a mathematical reasoning question from the jeebench dataset @xcite . in this example, the solver lm is text-davinci-003. in step 1 , the solver is prompted to answer the question (blue textbox) and it fails to answer correctly (red textbox). a problem decomposing lm generates subproblems (violet textboxes) conditioned on the original question and the initial response of the solver in step 2 . in step 3 , the solver answers these subproblems iteratively and appends to the prompt. finally, the original problem is appended to the prompt in step 4 , and the solver answers it correctly (green textbox). ## a general overview of daslam given an lm θ and a question q as a sequence of tokens, a standard zero-shot prompting can be described as, when answering q requires multistep reasoning, one can conceptualize s as a sequence of smaller steps {s ′ 1 , s ′ 2 , • • • , s ′ n } such that the lm iteratively answers a sequence of subproblems {q ′ 1 , • • • , q ′ n } to finally reach the desired answer to the original question. eq. 1 can then be rewritten as, in daslam, we seek to alleviate this problem by offloading the task of inferring the subproblems {q ′ i } to a decomposer lm ϕ. for a more guided problem decomposition, daslam uses the answer and the steps generated by the solver lm, θ in the naive cot regime as described in eq. 1 to generate a set of subproblems { qi } @xmath0,...,n as follows: for @xmath1, n], where â0 and s 0 are the initial answer and reasoning steps, respectively generated by θ. the solver lm θ then solves the subproblem set one-by-one similar to eq. 2. however, instead of seeking to generate the final answer as a response to the last subproblem qn , we append the original question at the end and let θ answer it directly given the context generated by the subproblems, their answers, and the corresponding cots. the four stages of workflow with daslam, as described in figure 1 , can then be summarized as follows: typical lms are not suitable to reliably perform the problem decomposition stage in eq. 3. we seek to finetune the decomposer lm, ϕ for this task. specifically, we use the llama 13 billion model. instead of full lm-tuning, we use lora adapters (hu et al., 2022) for parameter-efficient training. for brevity, we will denote the adapteraugmented lm as ϕ, although only the adapter weights are being updated while training. the whole process of teaching the lm to perform the problem decomposition comprises two successive stages: (i) subproblem construction from the original question and cot, and (ii) policy optimization in conjunction with the solver lm. details of the supervised data curation required for these three steps are described in section 5. each data sample in the supervised dataset d can be conceptualized as a sequence of triplets ⟨q gold , s gold , a gold , q ′ gold ⟩, where q gold denotes the original question, s gold denotes reasoning steps in form of cot, a gold denotes the answer to the question, and q ′ gold is a sequence of subproblems generated by decomposing q gold (see section 5 and appendix a for further details on supervised data curation process). the first stage is straightforward with the decomposer lm being finetuned using language modeling objective. in the first stage, we seek to optimize the following objective: this step is somewhat similar to instruction tuning, where an lm is asked to generate some text conditioned on a context, instead of following the usual sentence completion-style behavior of the lm. intuitively, the role of this stage is to align the lm to the task of the decomposer. decomposition via policy network. the previous stage inculcates the ability to decompose a problem within ϕ. however, it is still blind to the actual errors made by the solver lm. in the next stage, we seek to make the decomposer lm work in synergy with any solver lm of choice. this solver agnosticism restrains ϕ to observe the internal representations computed by θ while solving the problem. to handle this imposed blackbox characteristics of the solver, we resort to policy gradient optimization of ϕ assuming θ to be part of the environment. we formalize the setup as follows. elements of the environment: the solver lm θ constitutes the core component of the environment. given a question q, the solver-generated cot s and the answer a as sequences of tokens define the observation of the policy. ## experimental results the tasks used to evaluate the performance of daslam contain questions that can be answered either of the three types -numerical, single correct answer mcq, and multiple correct answer mcq. daslam is better than pure prompting we start with daslam augmented with gpt-3.5 as the solver lm on math and aqua datasets (see finetuned decomposer is essential. despite being orders of magnitude smaller, a finetuned llama 13b model delivers better performance compared to gpt-3.5 as a decomposer (daslam vs. gpt-3.5 generator in feedback from the solver is important. in the preceding paragraph, we attributed the superiority of daslam over other methods to the usage of a specialized lm for problem decomposition. however, manipulating the problem decomposition upon feedback from the solver is also an important factor here. none of the existing methods does so, and therefore, remains blind towards what reasoning (and possible errors) is followed by the solver model. this is further manifested when we compare daslam with itself without the feedback module, daslam-nf. while daslam-nf is able to improve upon basic cot and other prompting methods, it falls short of a decomposer lm that has access to the initial response of the solver. ## case study to this point, we have compared the numbers produced by daslam-boosted models across different datasets. while they provide an overall assessment, deeper analyses are needed to comprehend the actual reasoning steps adopted by these different methods. figure 2 shows the reasoning steps generated by gpt-3.5 given an example problem from the math dataset with three different prompting methods -vanilla cot, l2m, and daslam. note that daslam uses the cot output to decompose the problem. both cot and l2m end up with the model an-swering incorrectly. with cot, the solver wrongly assumes that the given equation must have two real roots though it should not have any real roots. also, it mistakes the value of a 2 as a. the effect is prominent in the subproblems generated by daslam as it asks to find the value of a explicitly. furthermore, the solver lm specifically announces that @xmath2. this helps to correct the reasoning about the sign of the discriminant. with l2m, the confusion around the value of a and a 2 persists, as the solver lm substitutes a in the given equation by 49 (which is the value of a 2 ) twice in the answering process. although it substituted the correct value of a once while answering the second question, it is not explicitly declared like in daslam. we observe multiple similar failure cases with l2m. it is quite likely that prompting the model to generate the final answer after each subproblem accumulates the erroneous reasoning steps that the model falls prey to. with daslam, the reasoning steps followed by the solver remains robust throughout. it reaches the final answer much earlier (third and fourth subproblems). in the final answer, the solver simply reiterates the steps that it earlier generated to answer the subproblems. this is a common behavior that we observed across multiple problems from multiple datasets. in appendix d (see figures 3 and 4 ), we provide similar case studies on llama 13b and 33b with different prompting methods. with reduced solver capacity, the difference between least-to-most and cot generated reasoning steps further diminishes with both leading to incorrect answers; daslam, on the other hand, still guides the solver through correct steps. an interesting observation can be made by comparing how the solver behaves with cot vs. with daslam. with daslam, we do not provide any new knowledge to the solver. yet, the same model can rectify its errors made in cot response. this may point to the intuition that current llms are actually underutilized, and one can unfold even more impressive performance with cleverly composed guidance. we further provide case analyses with when daslam fails to guide the model (gpt 3.5 in this case) to successful final answers, in appendix e. while we do not find any obvious pattern of errors, one can see that the decomposer generates questions that are not readily answerable within that context. daslam does not use any method to trace back the error or generate subproblems based on the answers to the previous subproblems. this might raise such issues where the subproblems generated are not actually helping the solver in the right order. ## conclusion we challenged the design of ever-larger monolithic llms as homogeneous network structures, where diverse aspects of problem decomposition and solution are stored in a tangled and opaque manner. the formidable general-purpose problem-solving capabilities of llms are exceedingly resource-hungry, dependent on immense data engineering. inspired by brain science, we took a first step toward heterogeneity -let two different llms evolve independently and adapt to their roles of decomposing and solving complex reasoning problems. through extensive experiments on several benchmarks, we showed that such a heterogeneous network can match or exceed some of the largest contemporary llms, at a much smaller parameter count.