Distilling Reasoning Capabilities into Smaller Language Models

Abstract

Step-by-step reasoning approaches like chain of thought (CoT) have proved to be very effective in inducing reasoning capabilities in large language models. However, the success of the CoT approach is fundamentally tied to the model size, and billion parameter-scale models are often needed to get CoT to work. In this paper, we propose a knowledge distillation approach that leverages the step-by-step CoT reasoning capabilities of larger models and distills these abilities into smaller models.In this work, we propose an alternative reasoning scheme, SOCRATIC COT that learns a decomposition of the original problem into a sequence of subproblems and uses it to guide the intermediate reasoning steps. We use SO-CRATIC COT to train a combination of two small distilled models: a problem decomposer and a subproblem solver. In practice, given a new problem, the two distilled models work in sync to decompose and solve complex problems. On multiple reasoning datasets (GSM8K, StrategyQA, and SVAMP), our proposed distillation strategies boost the performance of smaller models over 70% compared to the baselines. Finally, we investigate when SOCRATIC COT is an effective alternative to CoT, demonstrating cases where a much smaller model (GPT-2 large) can outperform a 10X larger model (GPT-3 6B). Our code is available here.

Democratizing Reasoning Ability: Tailored Learning from Large Language Model

“…Specifically, each round of learning consists of three key steps: (1) The student LM undergoes an "exam" on the training set for collecting mistakes which are the wrong generated rationales. Existing works (Fu et al, 2023b;Ho et al, 2023;Shridhar et al, 2023;Magister et al, 2023) merely provide the sample question for the LLM to collect annotated rationales, neglecting the importance of the student's feedback. However, the student's feedback is crucial in knowledge distillation (Fu et al, 2021;Pham et al, 2021;Ren et al, 2023).…”

“…However, these works aim to improve the instruction-following ability of smaller LMs, while the reasoning ability that we focus on is often overlooked. Some recent studies (Ho et al, 2023;Fu et al, 2023b;Shridhar et al, 2023) propose to employ LLMs to annotate rationales for training smaller student LMs towards reasoning, not considering the student's feedback to the teacher. In contrast, we exploit the potential of the black-box LLM as the teacher instead of the data annotator by proposing a multi-round learning paradigm.…”

“…4) Evaluating the correctness of generated rationale is mainly based on the final answer. Though most existing works (Zelikman et al, 2022;Ho et al, 2023;Fu et al, 2023b;Shridhar et al, 2023) in this field adopt this simple criterion, we call attention to develop more trustworthy criteria to evaluate the quality of rationales. Potential methods can be using GPT-4 (OpenAI, 2023) or a process reward model (Lightman et al, 2023) for automatic evaluation.…”

“…Recently, teaching smaller LMs towards reasoning with the help of LLMs has gained increasing attention. Most of these works (Ho et al, 2023;Magister et al, 2023;Fu et al, 2023b;Shridhar et al, 2023) can be summarized in two main steps:…”

Recurrent Neural Language Models as Probabilistic Finite-state Automata

“…Answering this question is essential whenever we require formal guarantees of the correctness of the outputs generated by an LM. For example, one might ask a language model to solve a mathematical problem based on a textual description (Shridhar et al, 2023) or ask it to find an optimal solution to an everyday optimization problem (Lin et al, 2021, Fig. 1).…”