Source-linked AI summary
Distilling System 2 into System 1
Ping Yu, Jing Xu, Jason Weston, Ilia Kulikov
TL;DR
LLMs can improve responses by using costly System 2 intermediate reasoning, motivating methods that transfer this behavior into direct System 1 generations. The paper uses self-supervised filtering and fine-tuning on unlabeled data, finding successful distillation for several methods with lower inference cost, while complex chain-of-thought math remains difficult to distill.
Problem
System 2 methods can improve accuracy but usually require costly intermediate generations, creating a need to transfer their benefits into direct System 1 outputs.
Method
The method runs System 2 on unlabeled inputs, filters responses using consistency criteria, and fine-tunes System 1 to reproduce the retained targets without intermediate outputs.
Results
Across 4 System 2 approaches and 5 tasks, several methods were distilled with improved performance over System 1 and lower inference cost than System 2.
Takeaways & Limitations
System 2 distillation can transfer capabilities for bias handling, instruction clarification, and LLM evaluation into cheaper direct generations, while leaving some reasoning tasks for System 2.
Takeaways & Limitations
Effectiveness varies by task and training dataset; the study could not effectively distill chain-of-thought, and performance depends on the self-supervised filters used.
Abstract
from arXiv · showhide
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 (Wei et al., 2022), many such System 2 techniques have been proposed such as Rephrase and Respond (Deng et al., 2023a), System 2 Attention (Weston and Sukhbaatar, 2023) and Branch-Solve-Merge (Saha et al., 2023). 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 such 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.
1 Introduction
System 2 uses deliberate intermediate reasoning to improve LLM responses, but its extra inference cost limits deployment. This work distills selected System 2 behavior into direct System 1 generations using unlabeled data.
- System 2 generates intermediate thoughts to reason and plan before producing a response, whereas System 1 responds directly without intermediate tokens.
- Many System 2 techniques improve accuracy through explicit reasoning but typically incur higher inference cost and latency than System 1.
- System 2 distillation applies a System 2 method to unlabeled examples, filters reliable predictions using unsupervised quality checks, and fine-tunes System 1 without intermediate steps.
- Experiments across 4 approaches and 5 tasks show that several methods can be distilled, sometimes outperforming their System 2 teacher while using a fraction of the computational cost.
- Complex math reasoning requiring chain-of-thought was not successfully distilled into System 1.
2 Related work
Related work frames System 1 as direct, largely automatic processing and System 2 as deliberate reasoning with intermediate outputs. Prior research covers human automaticity, LLM reasoning methods, and conventional model-to-model distillation.
- System 1 is associated with recognizing patterns and making quick judgments, while complex symbolic problem-solving is described as requiring System 2 reasoning.
- In LLMs, System 1 directly produces responses but can struggle with complex symbolic reasoning because intermediate latent representations are difficult to manipulate.
- Intermediate-step prompting or scratchpads can enable language models to perform multi-step computations that they cannot perform directly.
- System 2 research includes decomposition, self-refinement, self-verification, repeated prompting, branching, and search using LLM-generated intermediate outputs.
- Conventional distillation trains a student model to mimic a separate teacher through output distributions, layer activations, or derivatives of teacher outputs.
- Earlier approaches to distilling chain-of-thought may improve performance by removing intermediate steps gradually, but still do not match explicit chain-of-thought.
3 Distilling System 2 into System 1
The paper distinguishes direct System 1 responses from System 2 procedures that generate intermediate tokens through repeated calls, branching, or search. Its distillation pipeline uses System 2 outputs from unlabeled inputs, filters them for consistency, and fine-tunes System 1 to reproduce the resulting targets.
- 3.1 Setup: System 1 and System 2 models: System 1 produces output y directly from input x, while System 2 generates intermediate tokens z before returning the final response.
- 3.1 Setup: System 1 and System 2 models: System 2 methods can involve multiple prompts, branching, iteration, and search, whereas System 1 directly calls the underlying language model on the original input.
- 3.1 Setup: System 1 and System 2 models: Chain-of-thought, Rephrase and Respond, System 2 Attention, and Branch-Solve-Merge instantiate System 2 through intermediate generations or multiple LLM calls.
- 3.2 Method: System 2 Distillation: The method applies a System 2 model to unlabeled inputs because System 2 inference is slower and unlabeled instruction-following prompts are commonly available.
- 3.2 Method: System 2 Distillation: Intermediate outputs are discarded, leaving System 2 responses as candidate targets for System 1 distillation.
- 3.2 Method: System 2 Distillation: Because System 2 responses may be noisy, the method filters examples using output self-consistency or consistency under input perturbation.
- 3.2 Method: System 2 Distillation: The filtered synthetic dataset is used for supervised fine-tuning, producing a System 1 model expected to provide performance similar to the evaluated System 2 model.
4 Experiments
Experiments test whether unsupervised System 2 distillation transfers stronger reasoning into direct System 1 generation across multiple methods and tasks. The results show substantial gains and lower inference cost in several settings, while performance remains task-dependent.
- Rephrase and Respond: 98.0% exact-match accuracy exceeds 44.5% for 2-Step RaR and 30.0% for the System 1 baseline on last-letter concatenation.The distilled model preserves System 2’s accuracy benefits while using inference costs comparable to System 1.
- Rephrase and Respond: 75.69% success rate approaches 77.2% for 2-Step RaR and exceeds 56.1% for zero-shot System 1 on coin-flip reasoning.The distilled model achieves comparable performance without executing the two-prompt System 2 program.
- Analysis and limitations: Prompt engineering raises System 1 coin-flip performance from 56.11% to 66.84%, but reliance on prompt selection remains a limitation requiring additional human effort.The distilled System 2 model instead shows consistent performance across prompt variations and lower sensitivity to prompting.
- System 2 Attention: System 2 Attention distillation retains strong performance on biased inputs while substantially reducing generated tokens relative to both the baseline and S2A.Using unsupervised self-consistency filtering is important for overall results; removing it produces inferior performance.
- Branch-Solve-Merge: Distilled BSM outperforms CoT and BSM on OASST2 while generating only four tokens, and it exceeds GPT-4-0125-preview in human agreement and consistency.BSM improves over CoT but requires increased inference time, whereas the distilled model combines stronger evaluation results with very short outputs.
- Branch-Solve-Merge: On MT-bench, distilled BSM improves agreement and inconsistency rates with fewer output tokens, though it slightly trails GPT-4-0125-preview in agreement.Per-category results show distilled BSM surpassing GPT-4-0125-preview in writing, math, and STEM but lagging in reasoning, coding, and extraction.
- Analysis and limitations: System 2 distillation performs poorly on GSM8k across decoding hyper-parameters, indicating that its effectiveness depends on the reasoning task.The paper attributes this task difference to GSM8k requiring a substantially different kind of reasoning from the other evaluated tasks.
5 Conclusion
The paper shows that System 2 reasoning can often be distilled into System 1 outputs without intermediate generations, maintaining or sometimes improving performance while reducing inference cost.
- System 2 distillation transfers useful reasoning into direct LLM outputs without intermediate generations.
- The method succeeds for System 2 Attention, Rephrase and Respond, and Branch-Solve-Merge across bias handling, instruction clarification, and LLM-as-a-Judge evaluation.
- Distilling these approaches makes them more efficient at inference time and more likely to be used by practitioners.
- The authors propose continuous training loops as a future direction for systems that reserve reasoning effort for tasks they cannot yet solve well.
6 Limitations
The approach is effective for several System 2 methods but does not generalize uniformly across tasks, and its self-supervised performance depends on data-quality filters.
- Distillation effectiveness varies with the specific task or dataset used for model training.
- Chain-of-Thought for complex reasoning could not be effectively distilled using the proposed method.
- Model performance depends on the consistency filters used to construct self-supervised training data.
- The study used self-consistency of outputs and input perturbations but did not explore alternative data-quality strategies.
A.1 Prompts
The appendix presents prompts for Rephrase and Respond, including a one-step version and a two-step version that separates rephrasing from answering.
- The 1-step RaR prompt asks the model to rephrase the question and then provide an answer in one step.
- The prompt examples illustrate how RaR restructures questions before generating responses.
- The 2-step RaR prompt for last letter concatenation separates the process into a first step and a second step.
A.2 Experiment Details
The appendix documents training settings, System 2 prompts, evaluation materials, and example Branch-Solve-Merge data and judgments.
- Model training: Llama2 70B Chat initializes supervised fine-tuning with cross-entropy loss applied only to answer tokens.
- System 2 Attention: S2A uses nucleus sampling with top-p 0.9 in both generation stages, with occasional reduction of generated answers when context limits are exceeded.
- Prompts: The appendix includes prompts for System 2 Attention and the branch, solve, and merge stages of Branch-Solve-Merge.
- Evaluation: The appendix provides experimental-details and benchmark tables for BSM, Rephrase and Respond, and Chain-of-Thought-related evaluations.
- BSM evaluation: The BSM examples show assistant responses being scored on criteria such as relevance, accuracy, clarity, helpfulness, and personalization.
- BSM evaluation: In the coffee example, Assistant A receives higher scores than Assistant B and is selected as the preferred response.