Divide-or-Conquer? Which Part Should You Distill Your LLM ? (2024)

Abstract: Recent methods have demonstrated that Large Language Models (LLMs) can solve reasoning tasks better when they are encouraged to solve subtasks of the main task first. In this paper we devise a similar strategy that breaks down reasoning tasks into a problem decomposition phase and a problem solving phase and show that the strategy is able to outperform a single stage solution. Further, we hypothesize that the decomposition should be easier to distill into a smaller model compared to the problem solving because the latter requires large amounts of domain knowledge while the former only requires learning general problem solving strategies. We propose methods to distill these two capabilities and evaluate their impact on reasoning outcomes and inference cost. We find that we can distill the problem decomposition phase and at the same time achieve good generalization across tasks, datasets, and models. However, it is harder to distill the problem solving capability without losing performance and the resulting distilled model struggles with generalization. These results indicate that by using smaller, distilled problem decomposition models in combination with problem solving LLMs we can achieve reasoning with cost-efficient inference and local adaptation.

## introduction large language models (llms), such as gpt-4 (openai, 2023), demonstrate exceptional abilities in solving knowledge-intensive tasks like open domain qa (odqa) @xcite , math @xcite , science @xcite and autonomous agents @xcite @xcite . however, the use of gigantic llms with hundreds of billions of parameters can be costly during inference, particularly when the reasoning chain generated is lengthy. additionally, due to the opaque nature of * work done during internship at apple these black box llms, they offer limited adaption options. there is a need to use cheaper and more flexible models to leverage the power of these black box llms for local adaptation and cost-efficient inference. distilling the large llms would seem like a reasonable strategy, but it often results in a significant drop in performance for downstream tasks @xcite . previous studies @xcite have indicated that effectively addressing such tasks requires the model to proficiently perform two essential capabilities simultaneously: 1) planning and decomposition, which involves breaking down complex objectives into smaller, more manageable subgoals to facilitate efficient handling of intricate tasks; and 2) execution and solving, which involves memorizing vast amounts of knowledge from extensive web training data and effectively recalling this information when needed to execute the problem-solving process. the first capability, decomposition, typically requires the model to engage in self-reflection on the input query and generate a chain-of-thoughts (cot)style reasoning chain @xcite to tackle the problem. usually, these two abilities are intertwined in a single monolithic model throughout the problem-solving process @xcite . in this paper, we first investigate whether it is possible to decouple the decomposition and solving capabilities, and how to distill these capabilities into smaller models for faster inference. we then test several hypotheses: 1) is distilling decomposition easier than distilling solving? decomposition primarily relies on semantic understanding and query parsing, while solving requires more domain expertise and knowledge. for example, decomposing the query "who is older, messi or ronaldo?" into "how old is messi?", "how old is ronaldo?", and "who is older?" only requires text comprehension, whereas solving the task necessitates memorization, retrieval, and utilization of in-figure 1 : reasoning with a long thought chain using the black box llm can be expensive and inflexible. we propose to dissect the decomposition and solving of the task, and distill only the decomposition capability to a less costly and more flexible student model, while still maintaining the original performance. a natural question arises: is it possible to distill only the long reasoning chain, which accounts for most of the inference cost, but is relatively easier to distill? to this end, we propose and evaluate the distillation of only the decomposition capability from the llm. we empirically verified our hypothesis using a teacher model, gpt-3.5-turbo, and two student models, vicuna-13b @xcite and @xcite , on qa, mathematics, and compositional datasets @xcite @xcite . ## decoupling decomposition and solving as shown in figure 1 , a common approach to solving a reasoning task using an llm involves directly generating a response to the instruction and question. this is referred to as the single-stage model. the conventional method for llm, known as the chain of thought (cot), instructs the model to "think step by step," allowing the model to take more computational steps for difficult tasks. however, cot-style reasoning has limitations as it often struggles to generalize to problems beyond the scope of the in-context examples. to address this drawback, the most notable work is the least-to-most approach @xcite , where the model breaks down the original question into subquestions and answers them sequentially. these approaches have shown improved performance compared to cot. for qa tasks, typically, the next subquestion is less dependent on the answer to the previous subquestions. conveniently, we propose a static strategy similar to hugginggpt @xcite , where in the first decomposition stage several decomposed subquestions are first generated to decompose the primary question. in the second solving stage, these subquestions are then answered one by one to obtain the final answer. we refer to this line of models as the two-stage models. interactive vs static process note that an interactive and dynamic process could be beneficial for certain reasoning tasks. in our experiments with math and qa datasets, the decomposition and solving stages are more independent, thus we did not observe gain by switching to an interactive process. our primary focus lies in understanding the impact of distilling task decomposition and solving capabilities, rather than finding the optimal framework. using a static approach would enable us to have a clearer separation of the decomposition and solving. the distilled decomposer can also potentially be integrated into more dynamic reasoning processes, enabling iterative solving and refinement based on intermediate outputs. we include a comparison in section 6.1 to illustrate the benefits of adopting a static process. ## distill the decomposition capability generating decomposed questions can be computationally expensive when the reasoning chain is long while using a black box llm. moreover, it is challenging to optimize or customize the decompo-sition process as it is performed by the black box model. our proposal aims to address these issues by utilizing a smaller trainable student model, as a drop-in replacement for the large black box llm for decomposition. to achieve this, we distill the decomposition capability from the teacher llm, referred to as t . ## ablations we provide an extensive evaluation of various instructions, and an exploration into the influence of the number of demonstrations in appendix b. ## conclusion our investigation provides a fine-grained examination of the llm's capability on reasoning tasks, by disentangling the decomposition and solving aspects. although both capacities are vital for reasoning, we demonstrate that decomposition is less dependent on specific knowledge and thus easier to distill compared to distilling solving capabilities, regardless of the availability of ground truth labels. additionally, the distilled decomposer shows strong generalization abilities across different tasks, datasets and executor/solvers. for future work, it would be interesting to train universal decomposer models using data from various tasks, and explore the use of reinforcement learning to further enhance the decomposer, leveraging the signal from the solver outcome. another possible direction for future work is to assess the effectiveness of our method in other long-horizon planning tasks, including llm-powered agent, tool use, and multiturn decision making.