July 08, 2024
Large language models (LLMs) can spend extra compute during inference to generate intermediate thoughts, which helps to produce better final responses. Since Chain-of-Thought [1], many such System 2 techniques have been proposed such as Rephrase and Respond [2], System 2 Attention [3] and Branch-Solve-Merge [4]. In this work we investigate self-supervised methods to “compile” (distill) higher quality outputs from System 2 techniques back into LLM generations without intermediate reasoning token sequences, as this reasoning has been distilled into System 1. We show that several such techniques can be successfully distilled, resulting in improved results compared to the original System 1 performance, and with less inference cost than System 2. We posit that System 2 distillation will be an important feature of future continually learning AI systems, enabling them to focus System 2 capabilities on the reasoning tasks that they cannot yet do well.
Generating intermediate thoughts allows a model (or human!) to reason and plan in order to successfully complete a task or respond to an instruction. We refer to such deliberate thinking as System 2 reasoning, following its description for humans in [5], [6] and later for AI models [3], [7], [8]. In System 2 reasoning effortful mental activity is exerted, especially in situations where System 1 – more automatic thinking – is likely to make errors. In standard Large Language Models (LLMs) we thus define System 1 as application of the Transformer [9] to directly produce a response given an input, without generation of intermediate tokens. We define System 2 as any approach which generates intermediate tokens, including methods that perform search, or prompt multiple times, before finally generating a response. A battery of such System 2 techniques have been proposed, among them Chain-of-Thought [1], Tree-of-Thoughts [10], Graph-of-Thoughts [11], Branch-Solve-Merge [4], System 2 Attention [3], Rephrase and Respond [2] and more. Many of these methods are shown to produce more accurate results due to this explicit reasoning, but typically do so at much higher inference cost and latency for a response. Due to the latter, many of these approaches are not used in production systems, which mostly use System 1 generations.
For a human, the process of learning to transfer a skill from deliberate (System 2) to automatic (System 1) in psychology is referred to as automaticity, and the use of procedural memory [12]. For example, when driving to work for the first time one might typically expend conscious effort planning and making decisions to get there. After a driver repeats this route, the driving process becomes “compiled” into the subconscious [13]. Similarly, playing a sport such as tennis can become “second nature”. In this work, we explore an analogous technique for AI models. Our approach performs this compilation, which we refer to as System 2 distillation, in an unsupervised manner given a set of unlabeled examples. For each example we apply the given System 2 method, and then measure the quality of the prediction in an unsupervised manner. For example, for tasks with unique answers we apply self-consistency [14], sampling multiple times. For examples where System 2 is consistent enough, we assume this result should be distilled, and add it to the distillation pool. We then fine-tune System 1 to match the predictions of the System 2 method on the collected pool of examples, but without generating the intermediate steps. Figure 1 illustrates the overall process of distilling System 2 into System 1.
We conduct experiments across 4 different System 2 LLM approaches and 5 different tasks. We find our approach can distill System 2 reasoning into System 1 in a diverse array of settings, sometimes even improving the results over the System 2 teacher. Moreover, these predictions are now produced at a fraction of the computational cost. For example, we see successful distillation for tasks involving dealing with biased opinions or irrelevant information (System 2 Attention), clarifying and improving responses in some reasoning tasks (Rephrase and Respond), and for fine-grained evaluation of LLMs (Branch-Solve-Merge). However, we also show that not all tasks can be distilled into System 1, particularly complex math reasoning tasks requiring chain-of-thought. This is also mirrored in humans, who cannot execute some tasks without deliberate System 2 reasoning [6].
In humans, System 1 reasoning is described as being capable of recognizing patterns, making quick judgments, and understanding simple or familiar symbols. For instance, it is used to identify common traffic signs, recognize faces, or associate basic symbols with specific emotions or ideas. However, for complex problem-solving or for example manipulation of abstract symbols (like algebraic equations or logical statements), System 2 reasoning is deemed necessary [6]. In psychology the concept of automaticity describes behavior that becomes so well-practiced that it can be performed with little to no conscious thought, with an example being driving a familiar route [13]. In general, humans are said to use procedural memory to consolidate a specific task into memory, learning through practice, so that it can be later performed without conscious awareness [12]. The concept of unconscious competence is classified as a later stage of learning. Initially a person recognizes their incompetence, and consciously seeks to learn a skill until they acquire conscious competence. Finally, the aim is to utilize it without conscious thought when it is said to become, in common language, “second nature” [15].
We refer to a neural network that outputs a response directly without intermediate outputs as a System 1 model. Such a network can nevertheless compute intermediate latent representations in its layers before it outputs a response. As these states are represented as vectors they typically encode distributed knowledge, rather than discrete decisions, and have difficulty manipulating complex symbolic reasoning tasks directly [16]–[19], which is analogous to issues with System 1 reasoning in humans. Nevertheless, a vast array of tasks can be solved with success directly in this manner without intermediate generations [20].
[16] showed that the same language model that is unable to perform complex multi-step computations can perform those tasks when asked to generate intermediate steps into a “scratchpad” using either few-shot prompting or supervised training. Chain-of-thought reasoning was shown to be elicited from LLMs even using zero-shot prompting [21] as well as by supervised [17] or few-shot [1] methods. LLM pretraining allows such reasoning to be built into the model because reasoning steps in discrete symbols (text) are present in the training corpora written by humans. Such System 2 model approaches output discrete tokens which is good for making sequential correct logical reasoning steps – but obviously has a downside if the reasoning is generated incorrectly. An incorrect discrete decision is difficult to recover from, unlike latent vector-based reasoning that might more easily model a distribution.
Recently, many approaches have been proposed to execute deeper reasoning using the LLM as part of an inner loop where it generates intermediate outputs, sometimes referred to as LLM Programs [22]. These include subquestion decomposition [23], self-refinement [2], [3], [24], self-verification and asking [25]–[27], and various search techniques such as Tree-of-Thoughts and others [10], [11].
The concept of distillation is usually applied to taking separate models, a powerful teacher model (or multiple teacher models) and a less powerful student model with separate parameters. The student model is then trained to mimic the behavior of the teacher(s). Methods of distillation include training the student to have similar output distributions [28], layer activations [29] or derivatives of the target teacher outputs [30]. Earlier works considered distillation from an ensemble of multiple teacher models [28], [31]. As neural networks have become larger, distilling from a larger to a smaller network has become a common paradigm [32]. In contrast, in our work the teacher and student model are the same language model, but applied differently (either with intermediate reasoning, or not).
For chain-of-thought reasoning in particular, several distillation approaches have been considered [33]–[35]. These again follow the paradigm of distilling a separate larger model’s output into a smaller model, i.e. the student model is asked to mimic the internal thoughts of the teacher model. The work of [36], however, considers distilling a slower System 2 method (Tree-of-Thought) into a faster System 2 method (CoT), which can use the same model as student and teacher. In contrast our work’s goal is to not generate internal thoughts (to improve System 1). Some exceptions are [37], [38]. The former still uses a separate student and teacher model, but attempts to distill the intermediate thought tokens into the layers of the network by representing reasoning steps as vectors and then setting them as targets. The latter recent work attempts to distill CoT by gradually removing the intermediate steps, which can improve performance greatly compared to not doing so, but still does not match explicit CoT.
Given an input \(x\), in this work we consider the setting of a single model, in our case a large language model (LLM), that is capable of two modes of response:
System 1: Produces the output \(y\) directly. This is done by forwarding through the layers of the underlying autoregressive neural network (Transformer) to produce the output tokens.
System 2: We define System 2 models as methods that use the underlying Transformer to generate intermediate output tokens \(z\) of any kind before generating the final response tokens. This may include multiple calls (prompts).
More formally, we consider a System 2 model \(S_{\text{\uppercase{\romannumeral 2\relax}}}\) as a function that takes an LLM \(p_\theta\) and input \(x\), and can call the LLM possibly repeatedly to generate intermediate tokens \(z\) using a specific algorithm, before returning an output \(y\): \[\begin{align} \label{eq:s2} S_{\text{\uppercase{\romannumeral 2\relax}}}(x; p_\theta) \rightarrow z, y. \end{align}\tag{1}\] System 2 approaches can potentially involve multiple prompts, branching, iteration and search, all the while using the LLM to generate intermediate results for further processing. In contrast, a System 1 model only considers the original input \(x\) and calls the LLM \(p_\theta\) directly to produce an output \(y\): \[\begin{align} S_{\text{\uppercase{\romannumeral 1\relax}}}(x) = p_\theta(x) \rightarrow y. \end{align}\]
There are many existing instantiations of System 2 models. Chain-of-thought prompting only requires a single LLM prompt, but still outputs intermediate generations before a final response, typically used in math and other reasoning tasks [1].
Methods like System 2 Attention [3] and Rephrase and Respond [2] require two calls to the LLM, where in the former the first call is used to attend to the context and remove bias, and in the latter to expand on the question. The second call is then used to finally respond to the answer given the intermediate generations. Some methods are much more sophisticated for example Branch-Solve-Merge [4] which generates a plan via an LLM which branches into several more LLM calls until a final stage merges the results.
We will perform experiments with the four methods just described, but there are many other system 2 approaches, for example Tree-of-Thoughts [10], Graph-of-Thoughts [11] and more, see related work in 2.
Many System 2 methods, by their nature, are significantly slower at inference time due to multiple prompt calls and generation of intermediate tokens. The aim of System 2 Distillation is to distill all the reasoning from \(S_{\text{\uppercase{\romannumeral 2\relax}}}\) back into \(S_{\text{\uppercase{\romannumeral 1\relax}}}\) so that the direct outputs from the language model \(p_\theta(x)\) are improved. We assume a setting where the model has access to unlabeled inputs \(\mathcal{X}\) from which it can learn, in analogy to how humans learn their procedural memory without supervision. For language-based tasks, it is common to have access to instruction following prompts (inputs) as they can be collected from humans, e.g. the 1M released WildChat interactions [39] where inputs are given but correct labels are unknown. Hence this is a realistic setup.
The first step of the proposed method is to generate responses using the System 2 model over the unlabeled inputs \(\mathcal{X}\): \[\begin{align} y_{S_{\text{\uppercase{\romannumeral 2\relax}}}}^i = S_{\text{\uppercase{\romannumeral 2\relax}}}(x^i; p_\theta), ~~~\forall x_i \in \mathcal{X}. \end{align}\]
Note we discard (do not store) the intermediate outputs \(z\) from Eq. 1 . These responses \(y_{S_{\text{\uppercase{\romannumeral 2\relax}}}}^i\) can then be used directly as System 2 distillation targets for fine-tuning a System 1 model. However, they are subject to noise: some of these responses could be high quality, while others could be low quality or incorrect. For shortform QA and reasoning tasks involving a short response with a typically unique correct (but unknown) answer, we thus consider an unsupervised curation step to attempt to improve training data quality. We consider two variations which both rely on a consistency criterion:
self-consistency of outputs: we sample \(S_{\text{\uppercase{\romannumeral 2\relax}}}(x^i; p_\theta)\) a total of \(N\) times, and accept the response that is the majority vote; if there is no majority winner, we discard the example.
self-consistency under input perturbation: we perturb the input \(x^i\) in such a way that the output should not change, e.g. changing the order of multiple-choice items in the prompt, and compute \(S_{\text{\uppercase{\romannumeral 2\relax}}}\) for each perturbation; if the outputs do not agree, we discard the example.
After that, we end up with the synthetic dataset \((\mathcal{X}_{S_{\text{\uppercase{\romannumeral 2\relax}}}}, \mathcal{Y}_{S_{\text{\uppercase{\romannumeral 2\relax}}}})\), where \(\mathcal{X}_{S_{\text{\uppercase{\romannumeral 2\relax}}}}\) is a filtered subset of \(\mathcal{X}\) with targets \(\mathcal{Y}_{S_{\text{\uppercase{\romannumeral 2\relax}}}}\). The final step is then supervised fine-tuning of the LLM with parameters \(p_\theta\) using this distilled training set. We typically initialize this model from the current state \(p_\theta\) and continue training with the new dataset.
After fine-tuning we obtain an LLM \(\hat{p_\theta}\) which is a System 1 model that is expected to provide outputs and performance gains similar to the evaluated System 2 model.
We use Llama-2-70B-chat [40] as the base model for all our experiments. We require a base model of sufficient power that it can be performant as a System 2 model, but also have open weights that can be fine-tuned, hence this choice. We consider several System 2 methods, including Rephrase and Respond (RaR), System 2 Attention (S2A), Branch-Solve-Merge (BSM), and Chain-of-Thought (CoT), focusing on tasks where each method has demonstrated strong performance. For System 1, we conduct zero-shot inference using the instruction-tuned base model as a standard baseline. We report task-specific metrics for each task, and the “#Tokens” metric which measures the average number of tokens generated per input across the evaluation set. For System 2 methods this includes both intermediate token generations as well as the final output token generations. Detailed descriptions of the experimental setups are available in the Appendix 7.2.
Rephrase and Response (RaR) [2] is a System 2 method that first prompts the language model to rephrase the original question with further elaboration, and then secondly to generate a response based on the rephrased question with the aim that this provides superior output. The authors introduce two approaches, 1-step RaR and 2-step RaR, where the latter involves two separate prompts rather than a combined one as in the former, see Appendix 7.1 for specific prompts. They find that 2-step RaR significantly improves performance on several reasoning tasks that are challenging for the baseline LLM. We consider two tasks from the original paper where it performed well: the last letter concatenation task and coin flip reasoning. We then assess whether it is possible to distill this System 2 approach.
We build the System 2 distillation dataset for RaR using self-consistency of outputs. For each input, we conduct eight sampling iterations for the last letter task and eight for each stage of the coin flip task.1 We then apply a majority vote to determine the final output.
This task focuses on symbolic reasoning, requiring the model to concatenate the last letters of given words. For instance, the instruction: “Take the last letters of the words in ‘Edgar Bob’ and concatenate them.” As demonstrated in [2], this task benefits significantly from the application of the RaR method. We compiled a dataset by randomly selecting 1200 unique English words. Using this, we constructed 200 samples each for training, validation, and test.
Overall results are given in 1. The baseline System 1 model (Llama-2-70B-chat) achieves an accuracy of 30.0%, and is outperformed by the System 2 methods of 1-Step and 2-Step RaR (39.5% and 44.5%, respectively). Distilling the 2-Step RaR method back into a System 1 Llama-2-70B-chat model via our unsupervised technique, we achieve a remarkable accuracy of 98.0%. The model can effectively learn from this training data how to solve the task, in comparison to the zero-shot chat model. Distillation of Rephrase and Respond effectively inherits the advantages of both System 2 and System 1. It maintains the accuracy benefits of System 2, while its inference cost is comparable to that of System 1 (see # of generated Tokens).
To evaluate the effectiveness and necessity of our unsupervised curation step using self-consistency of outputs we conducted an ablation study by creating a distillation dataset without applying the self-consistency filter. When we distilled the System 2 model using this unfiltered dataset under the same setting, it achieved an exact match accuracy of 87.5% (with 98% for the filtered version). This comparison underscores the critical role of consistency filtering. Nevethess, in both cases constructing training data does improve results over zero-shot performance. We also attempted to distill the System 1 predictions using the same filtering technique, which results in a lower accuracy of 69.5%.
| Last Letter | Coin Flip | |||
|---|---|---|---|---|
| Acc\(\uparrow\) | #Tokens | Acc\(\uparrow\) | #Tokens | |
| System 1 | ||||
| Llama-2-70B-chat | 30.0% | 27.1 | 56.1% | 61.9 |
| Distill System 1 | 69.5% | 24.4 | 54.5% | 30.4 |
| System 2 | ||||
| 1-Step RaR | 39.5% | 106.6 | 58.5% | 158.9 |
| 2-Step RaR | 44.5% | 41.5 | 77.2% | 112.4 |
| Distill System 2 | ||||
| Distill 2-Step RaR | 98.0% | 25.5 | 75.69% | 50.3 |
This symbolic reasoning task has frequently been tested in research, including in [1] and [2]. It involves determining the final face (heads or tails) of a coin, starting from a known initial position after a series of flips described in natural language , such as “A coin is heads up. Roxas does not flip the coin. Schneiderman does not flip the coin. Is the coin still heads up?” [2] showed that even strong language models do not succeed at this task, whereas applying the RaR method improves their performance. There are 20k training examples, which we use for unsupervised learning (without labels), 3.33k validation and 1.33k test examples.
Overall results are given in 1. Llama-2-70B-chat (zero-shot) has a success rate of 56.1% on this task, while 1-Step and 2-Step RaR have success rates of 58.5% and 77.2% respectively. We thus only see a large improvement with the 2-Step method. Distilling 2-Step RaR back into a system 1 Llama-2-70B-chat via our unsupervised technique yields 75.69%. Hence, we find that our distilled System 2 model delivers performance comparable to that of System 2 (2 Step RaR), but without the need to execute the LLM program with 2 prompts (see # of generated Tokens).
The RaR method in [2] incorporates prompt engineering tricks, such as appending phrases like “Flip means reverse. Answer the Yes or No question” to the original query, which has been shown to enhance model performance. Following their approach, we evaluated model performance using different prompts, see 7. When testing the Llama-2-70B-chat model (System 1) with prompts like “Flip means reverse” and “Flip means reverse. Answer the Yes or No question,” we observed a significant improvement in performance, from 56.11% to 66.84%. This highlights the critical role of prompt selection in optimizing the performance of System 1 models. However, this reliance on prompt engineering also represents a limitation, necessitating additional human effort.
| Acc\(\uparrow\) | Acc\(\uparrow\) | ||
| Model | (biased) | (unbiased) | #Tokens |
| System 1 (Zero-shot) | 51.6% | 73.8% | 165 |
| System 2 (S2A) | 76.0% | 69.3% | 147 |
| Distill S2A | 81.3% | 78.6% | 56 |
| Distill S2A (no USC) | 78.6% | 75.3% | 58 |
We also attempted to distill the System 1 model, which gave poor performance. In this case, we also observed fluctuations in performance with different prompts. In contrast, the distilled System 2 model demonstrated consistent performance across various prompts, with a lower sensitivity to prompt variations. This consistency indicates that extensive prompt engineering might not be essential for the distilled System 2 model.
[3] proposed System 2 Attention (S2A), a method that helps to reduce models’ reasoning pitfalls such as relying on biased information in the input or attending to irrelevant context. S2A is a two-stage inference method where the first stage rewrites the input so that it does not contain undesired information such as bias or irrelevant context, and the second stage attends to the shorter rewritten context (in contrast to RaR which expands the context), see 6. In this work we verify the feasibility of distilling S2A into System 1. In particular, we focus on the SycophancyEval question answering task [41] that contains biased information in the input that is known to hurt LLM performance. We use 6668 examples from SycophancyEval as unlabeled training data, and 400 examples for evaluation, where the latter are split into biased inputs (350) and without bias (50).
We use universal self-consistency (USC) [42] to select high quality targets. Specifically, we sample 20 generations and then use the Llama-70B-chat model with a USC prompt (provided in 12) to compose a self-consistent (majority) final answer that is used as the distillation target.
The results are provided in 2, reporting average accuracy over 3 random seeds. The baseline (System 1) LLM has low accuracy on the biased portion as expected, being susceptible to biased inputs. S2A improves performance dramatically for biased inputs. System 2 distillation shows similarly strong performance as the System 2 approach. There is, however, a signification reduction in the average number of tokens used compared to both the baseline and the S2A model. This is because biased inputs tend to make the baseline LLM generate more output tokens, while S2A has to generate intermediate tokens as well. 11 shows a representative example. Finally, we show that using USC for distillation is important for overall results, by also reporting results without USC (last row), where the latter provides inferior results. This highlights the importance of the distillation data quality that is used during fine-tuning.
| OASST2 Eval | MT-bench Eval | ||||||
|---|---|---|---|---|---|---|---|
| Agreement \(\uparrow\) | Agreement \(\uparrow\) | ||||||
| System 1 | |||||||
| GPT-4-0125-preview | 44.7% | 35.5% | 4 | 68.1% | 25.6% | 4 | |
| Llama-2-70B-chat | 32.0% | 56.7% | 4 | 28.1% | 80.9% | 4 | |
| System 2 | |||||||
| CoT (GPT-4-0125-preview) | 48.7% | 28.2% | 603.7 | 73.8% | 16.2% | 548.8 | |
| CoT (Llama-2-70B-chat) | 45.2% | 37.7% | 432.6 | 58.9% | 30.8% | 411.8 | |
| BSM (Llama-2-70B-chat) | 49.1% | 30.4% | 2117.8 | 64.5% | 21.1% | 2063.1 | |
| Distill System 2 | |||||||
| Distill BSM (Llama-2-70B-chat) | 58.4% | 12.2% | 4 | 72.4% | 9.1% | 4 | |
Branch-Solve-Merge (BSM) [4], consists of three modules: branch, solve, and merge. These modules work together to break down a task into several parallel sub-tasks, each guided by specific prompts. BSM has proven effective when used in the context of an LLM acting as a judge, see 14. The method begins by prompting the language model to list evaluation metrics (branch) tailored to a given user query. Subsequently, the LLM is queried to evaluate a response based on each metric independently in parallel (solve). Finally, the scores from each branch are averaged to arrive at a comprehensive evaluation decision (merge). Notably, this method incurs an inference cost 5-6 times greater than that of a conventional (System 1) LLM evaluation approach, making it much less practical. We assess the feasibility of distilling BSM, aiming to retain its benefits while reducing computational cost.
Following [43], [44], we used the Open Assistant Dataset v2 (OASST2) [45] with turn 1 and English only data. We use queries along with two candidate responses from the OASST2 training set as inputs (19,672 examples in total). We use self-consistency under input perturbations to ensure the quality of our distillation data. Specifically, as two responses are being judged, we evaluate each sample twice with BSM - once in the original order and once in the swapped order. The winning response should remain consistent regardless of the order. We filter out samples that do not yield a consistent winner when the response order is swapped.
We evaluate our models on two popular benchmarks, the OASST2 valid set and MT-bench [46]. The OASST2 validation set comprises 273 samples, restricted to turn 1 and English language only. Evaluations of response pairs are performed in both original and swapped orders. As we trained our distilled model on the OASST2 training set, the OASST2 validation set functions as an in-distribution evaluation set, while MT-bench is more out-of-distribution. MT-bench is a popular benchmark that evaluates LLM-as-judges of other LLM’s responses when acting as helpful AI assistants conversations. It consists of instructions from 8 diverse domains e.g., writing, reasoning, math, coding, etc.
Following [46], we assessed the Agreement between model votes and human expert votes. A well-documented limitation of LLM-as-a-judge is
position bias, where a Language Model (LLM) tends to favor certain positions over others. This bias is evident as altering the position of responses in the evaluation prompt often leads to different decisions by the model. To quantify this, we not only
measure agreement but also calculate the Percentage of Inconsistent examples to assess position bias.
| Model | k=1 | k=5 | k=10 | |||
|---|---|---|---|---|---|---|
| Acc % | #Tokens | Acc % | #Tokens | Acc % | #Tokens | |
| System 1 | ||||||
| Few (8)-shot (no CoT) | 7.58% | 57 | 9.40% | 295 | 10.31% | 620 |
| System 2 | ||||||
| CoT zero-shot | 52.77% | 270 | 57.54% | 1385 | 59.44% | 2760 |
| CoT few (8)-shot | 36.39% | 297 | 54.97% | 1560 | 63.84% | 3120 |
| Distill System 2 | ||||||
| Distill CoT zero-shot | 7.13% | 18 | 7.13% | 90 | 7.35% | 180 |
3 provides results on the OASST2 dataset. Compared to baseline (System 1) LLMs, the Chain-of-Thought (CoT) method improves performance by improving agreement and reducing inconsistency rates (see prompts in Appendix). While BSM outperforms CoT, this comes at the cost of increased inference time (#Tokens). Remarkably, our distilled System 2 BSM model requires the generation of only four tokens and still outperforms both CoT and BSM. Furthermore, our distilled model based on Llama-2-70B-chat outperforms GPT-4-0125-preview, achieving higher human agreement and greater consistency.
3 also provides results on MT-bench, which serves as an out-of-distribution test. The results mirror those from the OASST2 evaluation. Both Chain-of-Thought (CoT) and BSM improve model performance but at the expense of significantly increased inference costs. Our distilled BSM model not only achieves higher human agreement and lower inconsistency rates but also requires less computational resources. Although our model slightly underperforms in agreement compared to the state-of-the-art GPT-4-0125-preview model, it was trained solely on unlabeled data from OASST2 based on Llama-2-70B-chat. Despite this, it is more consistent and inference is cheap in terms of output tokens.
Here, we further analyze the MT-Bench results in terms of Agreement by category. Figure 2 shows the per category agreement. We observe that CoT improved agreement compared to the base model
(Llama-2-70B-Chat) on all categories. BSM is better than CoT and our distilled BSM is even better than BSM. Although Distilled BSM achieves superior performance compared to the baselines across all categories, it still lags behind GPT-4-0125-preview in
reasoning, coding, and extraction. However, it surpasses GPT-4-0125-preview in writing, math, and STEM.
Chain-of-Thought (CoT) [1] has been shown to be an effective method to improve LLM’s reasoning abilities, such as for solving graduate school math problems. The LLM generates intermediate tokens that are steps (chain) of reasoning (thoughts) before it produces the final answer. We consider two variants of the approach: (i) few-shot CoT, whereby multiple [question, CoT, answer] examples from the training set are provided as part of the context followed by the question; and (ii) zero-shot, whereby an explicit instruction to think “step by step” is added to the prompt in addition to the question, see Appendix 10.
We use CoT to produce answers for questions from the training split of GSM8k [17] (which we consider unlabeled), using majority voting with \(K=10\). The resulting distillation training set consists of \(7461\) [question, answer] pairs i.e., without any intermediate reasoning steps. The accuracy of the self-supervised targets, computed for analysis purposes, is \(56.81\%\).
We report evaluation accuracy computed over the GSM8k test set with majority voting with different values of \(K\). Similarly to our previous experiments, we report the average number of predicted tokens for each method. Note that we compute this average over all generated tokens when we run majority voting to see how the increase in \(K\) affects the inference cost. We consider several baselines: System 1 and System 2 (CoT) methods evaluated with zero-shot or 8-shot input contexts. Note that System 2 with 8-shot means that CoTs are provided in the few-shot inputs, while System 1 means that the few shot examples contain questions and answers, but no CoTs.
Evaluation results are presented in 4. First, improvements are coming from using the CoT method as expected: it helps when being presented as part of the few-shot context or as part of the instruction in the prompt template. These improvements come with an increase in inference cost: sequences predicted with CoT methods are substantially longer compared to the System 1 method. Second, our System 2 distillation method yields poor performance across various decoding hyper-parameters. The GSM8k task (math problems) requires a very different kind of reasoning compared to other tasks we considered in this work. This highlights the non-trivial aspect of System 2 distillation: the proposed distillation algorithm works in many cases but not always. This leaves room for future research to elucidate in exactly which circumstances to apply distillation, and when not to, in a similar manner perhaps to the approach in humans.
Recent work has shown that complex reasoning procedures using LLMs in the inner loop, called System 2 approaches, can improve performance. In this work we have shown that in many cases it is possible to distill this System 2 reasoning into the outputs of the LLM without intermediate generations while maintaining, or sometimes even improving, performance. While not all methods can be distilled easily using our method, with Chain-of-Thought for complex reasoning being a challenging counterexample, this is possible for diverse approaches. Our method works for System 2 Attention for dealing with bias and irrelevant context, Rephrase and Respond for clarifying task instructions, and Branch-Solve-Merge for improved LLM-as-a-Judge evaluation. Pragmatically, distilling these approaches makes them more likely to be used by LLM practitioners, and they are more efficient at inference time. Looking forward, systems that can distill useful tasks in this way free up more time to spend on reasoning about the tasks that they cannot yet do well, just as humans do. Hence, we expect exploring this approach in a continuous training loop will be a fruitful research direction.
In this paper, we explored three System 2 methods—RaR, S2A, and BSM—which have been successfully distilled, yielding enhanced results compared to the original System 1 performance while incurring lower inference costs than System 2. However, the effectiveness of these methods can vary depending on the specific task or the dataset used for model training. For instance, we observed that the CoT method could not be effectively distilled back to System 1 using our method. We note that recent methods have tried alternative ways to distill CoT [37], [38].
Moreover, due to the self-supervised nature of these methods, model performance relies on the specific filters applied. In our study, we depended on a consistency criterion that includes self-consistency of outputs and self-consistency under input perturbation. Although there are multiple alternative strategies to enhance data quality in self-supervised learning, these were not explored in our research.
None
Figure 3: 1-step RaR prompt. The 1-step RaR process involves the model rephrasing the question and subsequently providing an answer, all in a single step..
None
Figure 4: 2-step RaR prompt for last letter concatenation task, step 1 (top), step 2 (down) The 1-step RaR process involves the model rephrasing the question and subsequently providing an answer, all in a single step..
None
Figure 5: 2-step RaR prompt for coin flip task, step 1 (top), step 2 (down) The 1-step RaR process involves the model rephrasing the question and subsequently providing an answer, all in a single step..
None
Figure 6: System 2 Attention prompts. We use the prompts from [3] to extract the training signal for distillation. The output after the second stage is used as the distillation target..
None
Figure 7: BSM: Branch prompt..
None
Figure 8: BSM: Solve prompt..
None
Figure 9: BSM: LLM merge prompt..
None
Figure 10: Chain-of-Thought prompt. We consider this prompt as the one that provides a formatting requirement for the model so that answer extraction without the CoT is feasible..
None
Figure 11: System 2 Attention examples. Representative examples showing the difference in conciseness and style of predicted responses from baselines and distilled System 2 model. Note that S2A System 2 response only presents the final (2nd stage) response here without intermediate generations..
None
Figure 12: Universal Self-Consistency prompt. We use USC [42] in order to find the most consistent final answer from a set sampled by the S2A model..
We use Llama2 70B Chat as the initialization for SFT training with CE loss. The loss is only applied on the answer part of the sequence. Model is trained with dropout \(0.1\), learning rate \(5.5e-6\), with warmup \(1\). Table 5 shows details about total training steps and total training tokens per step.
For S2A, in both generation stages we use nucleus sampling with top-p value \(0.9\). During distillation, for USC, in some cases the generated answers are too long and 20 do not fit in the Llama2 context. In these rare cases we reduce the answer set to 10 or select an answer randomly if 10 generated answers are still too long.
| Methods | Dataset | Total Training Steps | Total Training Tokens per Step |
|---|---|---|---|
| RaR | Last Letter Concatenation | 3 | 66k |
| RaR | Coin Flip | 100 | 66k |
| S2A | TriviaQA | 350 | 23k |
| BSM | OASST2 | 600 | 131k |
| CoT | GSM8K | 5000 | 33k |
None
Figure 13: BSM training data sample. (Up) Input data (Middle) target data for distilling BSM (Down) target data for distilling BSM (label only)..
| writing | reasoning | math | humanities | roleplay | coding | extraction | stem | |
|---|---|---|---|---|---|---|---|---|
| gpt-4-0125-preview | 65.38% | 78.79% | 73.33% | 75.17% | 69.94% | 78.57% | 76.32% | 75.51% |
| llama2-70b-chat | 48.98% | 54.89% | 60.00% | 66.67% | 58.89% | 62.17% | 48.54% | 66.67% |
| BSM | 63.08% | 64.65% | 61.67% | 70.74% | 65.64% | 70.63% | 57.02% | 76.19 |
| Distill System 1 | 53.59% | 66.00% | 54.72% | 67.11% | 62.17% | 67.73% | 43.86% | 70.07% |
| Distill System 2 | 68.46% | 67.34% | 67.78% | 74.94% | 68.30% | 70.64% | 61.69% | 75.51% |
| Distill System 2 (label only) | 70.77% | 70.71% | 76.95% | 74.50% | 68.92% | 74.34% | 61.70% | 79.59% |
| Data Input Prompt | Exact Match | Miss Match Rate | |
|---|---|---|---|
| System 1 | {question} | 56.11% | 4.65% |
| System 1 | {question} Flip means reverse. | 66.84% | 0.15% |
| System 1 | {question} Flip means reverse. Answer the Yes or No question. | 52.89% | 0% |
| 1 Step RaR | Prompt in Fig. 3 | 58.51% | 0% |
| 2 Step RaR | Prompt in Fig. 5 | 77.19% | 0% |
| Distill system 1 | {question} | 54.54% | 3.75% |
| Distill system 1 | {question} Flip means reverse. | 62.64% | 1.13% |
| Distill system 1 | {question} Flip means reverse. Answer the Yes or No question. | 63.39% | 0.60% |
| Distill system 2 | {question} | 75.69% | 0% |
| Distill system 2 | {question} Flip means reverse. | 78.92% | 0% |
| Distill system 2 | {question} Flip means reverse. Answer the Yes or No question. | 74.49% | 0% |
Figure 14 shows the overview of Branch-solve-merge. We copied figure from [4].
This approach was adopted after observing that sampling just once for the rephrase stage yielded suboptimal results.↩︎