PaperlayerView exact revision on arXiv

Source-linked AI summary

Distilling System 2 into System 1

Ping Yu, Jing Xu, Jason Weston, Ilia Kulikov

arXiv:2407.06023v3cs.CLcs.AI

TL;DR

System 2 techniques can improve LLM outputs, but their higher inference cost and latency limit production use. This paper studies self-supervised distillation of System 2 outputs into direct System 1 generations, finding that several approaches maintain or improve performance at lower computational cost.

  • Problem

    System 2 techniques can improve accuracy through explicit reasoning, but their higher inference cost and latency mean production systems mostly use System 1 generations.

    Source: §1 · p. 1

  • Method

    The paper applies System 2 methods to unlabeled examples, filters predictions using unsupervised consistency checks, and fine-tunes System 1 to match selected predictions.

    Source: §1 · p. 2

  • Results

    Across diverse settings, distillation sometimes improves on System 2 performance and produces predictions at a fraction of its computational cost.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    For System 2 Attention, Rephrase and Respond, and Branch-Solve-Merge, distillation offers a lower-cost route to direct outputs without intermediate generations.

    Source: §5 · p. 9, §6 · p. 9

    Effectiveness varies by task and training dataset, and this method did not effectively distill Chain-of-Thought.

    Source: §6 · p. 9

Abstract

from arXiv · show

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

The introduction proposes compiling System 2 outputs into System 1 through unsupervised distillation, aiming to retain reasoning gains without generating intermediate steps at inference.

  • System 1 directly generates responses without intermediate tokens, whereas System 2 uses deliberate, effortful reasoning that generates intermediate tokens.Source: §1 · p. 1
  • Explicit reasoning can improve accuracy, but its greater inference cost and latency help explain why production systems mostly rely on System 1 generations.Source: §1 · p. 1
  • On unlabeled examples, the method applies System 2, filters outputs using unsupervised quality checks such as self-consistency (Wang et al., 2022), then fine-tunes System 1 without intermediate steps.Source: §1 · p. 2
  • Across four System 2 approaches and five tasks, distillation often improved System 1, sometimes surpassed its teacher, and reduced compute, but failed on complex chain-of-thought math.Source: §1 · p. 2

2 Related work

The related work contrasts fast and deliberate reasoning in humans and language models, then distinguishes this work’s same-model distillation objective from standard teacher–student approaches.

2.1 System 1 and System 2 in Humans

  • Human System 1 recognizes patterns and familiar symbols quickly, whereas System 2 supports complex symbolic problems; practice can make tasks automatic.The learning process is described as progressing toward unconscious competence, or using a skill without conscious thought.Source: §2.1 · p. 2

2.2 System 1 and System 2 Models

2.3 (Standard) Distillation

  • Unlike standard teacher–student distillation, this work applies the same language model with or without intermediate reasoning to improve System 1 rather than generate internal thoughts.Prior work distills larger-model thoughts into smaller students; Zhang et al. (2024) distill Tree-of-Thought into CoT, while Deng et al. (2023b, 2024) explore vectorized thoughts or gradually removing steps, which still does not match explicit CoT.Source: §2.3 · p. 3

3 Distilling System 2 into System 1

The section defines System 1 as direct generation and System 2 as generation that uses intermediate outputs, potentially through multiple calls. System 2 distillation uses unlabeled inputs to create and optionally curate targets for fine-tuning a System 1 model, aiming to retain improved responses with lower inference cost.

3.1 Setup: System 1 and System 2 models

  • System 1 produces output directly, whereas System 2 generates intermediate tokens before its final response and may use multiple prompts, branching, iteration, or search.System 2 examples include Chain-of-Thought (Wei et al., 2022), System 2 Attention (Weston and Sukhbaatar, 2023), Rephrase and Respond (Deng et al., 2023a), and Branch-Solve-Merge (Saha et al., 2023).Source: §3.1 · p. 3

3.2 Method: System 2 Distillation

  • System 2 distillation aims to transfer reasoning into direct outputs, addressing the inference cost of multiple prompts and intermediate tokens using unlabeled inputs.Instruction-following prompts may be available without correct labels, as in the 1M WildChat interactions (Zhao et al., 2024).Source: §3.2 · p. 4
  • The method generates System 2 responses for unlabeled inputs and discards the intermediate outputs before using responses as candidate fine-tuning targets.The responses may be noisy or incorrect, motivating an unsupervised curation step for short-answer QA and reasoning tasks with typically unique correct answers.Source: §3.2 · p. 4
  • Curation keeps a majority-voted response across repeated System 2 samples or retains an example only when outputs agree across meaning-preserving input perturbations.The self-consistency procedure discards examples without a majority winner; the perturbation procedure discards examples whose outputs disagree.Source: §3.2 · p. 4
  • The curated input-target dataset is used to fine-tune the LLM, typically from its current state, yielding a System 1 model expected to approach the evaluated System 2 model's performance.Fine-tuning uses the filtered subset of unlabeled inputs and their System 2 targets.Source: §3.2 · p. 4

4 Experiments

The experiments test whether self-supervised distillation can transfer several System 2 methods into Llama-2-70B-chat generations, improving performance with lower inference cost in several tasks. CoT distillation performs poorly on GSM8k, showing that the approach does not work equally well across tasks.

4.1 Training and Evaluation Setup

  • Experiments use Llama-2-70B-chat as the fine-tuned base and zero-shot instruction-tuned Llama-2-70B-chat as the System 1 baseline, reporting task metrics and generated tokens.System 2 token counts include intermediate and final generations.Source: §4.1 · p. 4

4.2 Rephrase and Respond Distillation

  • RaR rephrases a question before responding; its 2-Step variant uses separate prompts, and distillation targets came from majority votes over eight samples per task stage.The tasks were last-letter concatenation and coin-flip reasoning, where prior work found 2-Step RaR improved performance on challenging reasoning tasks.Source: §4.2 · pp. 4–5

4.2.1 Last letter Concatenation Task

  • 98.0% accuracy: distilled 2-Step RaR exceeded 1-Step RaR (39.5%), 2-Step RaR (44.5%), and zero-shot System 1 (30.0%) on last-letter concatenation.Without self-consistency filtering, the distilled model reached 87.5% accuracy; distilling System 1 predictions with the filter reached 69.5%.Source: Table 1 · p. 5, §4.2.1 · p. 5

4.2.2 Coin Flip Reasoning Task

  • 75.69% success: distilled 2-Step RaR approached 2-Step RaR (77.2%) and exceeded zero-shot System 1 (56.1%) and 1-Step RaR (58.5%) on coin-flip reasoning.System 1 prompting raised success from 56.11% to 66.84%; unlike System 1 distillation, distilled System 2 performance was consistent across prompts.Source: Table 1 · p. 5, §4.2.2 · pp. 5–6

4.3 System 2 Attention Distillation

  • On TriviaQA, System 2 distillation achieved similarly strong performance to S2A on biased inputs, averaged over 3 random seeds, while using significantly fewer average tokens than both System 1 and S2A.Universal self-consistency selected distillation targets from 20 generations; results without USC were inferior.Source: Table 2 · p. 6, §4.3 · p. 6

4.4 Branch-Solve-Merge Distillation

  • On OASST2, distilled BSM outperformed CoT and BSM with only four generated tokens, and exceeded GPT-4-0125-preview in human agreement and consistency.Training data retained only judgments whose winning response stayed the same when candidate order was swapped.Source: Table 3 · p. 7, §4.4 · p. 7
  • On out-of-distribution MT-bench, distilled BSM improved agreement and consistency over CoT and BSM at lower inference cost, but had lower agreement than GPT-4-0125-preview.By category, it surpassed GPT-4-0125-preview in writing, math, and STEM, but lagged in reasoning, coding, and extraction.Source: Table 3 · p. 7, §4.4 · pp. 7–8, Figure 2 · p. 8

4.5 Chain-of-Thought Distillation

  • On GSM8k, CoT improved accuracy over System 1 but required longer generations, while System 2 distillation performed poorly across decoding settings.The distillation set contained 7,461 question-answer pairs produced using majority voting with K = 10; the paper identifies math reasoning as a distinct challenge for the method.Source: Table 4 · p. 8, §4.5 · pp. 8–9

5 Conclusion

System 2 reasoning can often be distilled into LLM outputs without intermediate generations while maintaining or improving performance, though Chain-of-Thought for complex reasoning is a challenging counterexample. The authors argue this can make effective methods more efficient and practical, and identify continuous training as a future research direction.

5 Conclusion

  • System 2 reasoning was distilled into LLM outputs without intermediate generations while maintaining or sometimes improving performance across diverse methods.The demonstrated methods are System 2 Attention for bias and irrelevant context, Rephrase and Respond for clarifying task instructions, and Branch-Solve-Merge for LLM-as-a-Judge evaluation.Source: §5 · p. 9
  • Chain-of-Thought for complex reasoning remains a challenging counterexample that is not easily distilled using the proposed method.The conclusion notes that not all System 2 methods can be distilled easily.Source: §5 · p. 9
  • Distilled approaches are more efficient at inference time and, pragmatically, more likely to be used by LLM practitioners.Source: §5 · p. 9
  • The authors identify exploring System 2 distillation in a continuous training loop as a fruitful research direction.They suggest this could free systems to spend more time reasoning about tasks they cannot yet do well.Source: §5 · p. 9

6 Limitations

RaR, S2A, and BSM were successfully distilled, improving on original System 1 performance at lower inference cost than System 2, but distillation effectiveness and self-supervised data quality have limitations.

6 Limitations

  • RaR, S2A, and BSM distillation enhanced results over original System 1 performance while incurring lower inference costs than System 2.Source: §6 · p. 9
  • Effectiveness varies by task or training dataset; CoT was not effectively distilled to System 1, while recent work explores alternative CoT distillation (Deng et al., 2023b, 2024).Source: §6 · p. 9
  • Performance depends on self-supervised data filters; this study used output and input-perturbation self-consistency criteria, leaving alternative data-quality strategies unexplored.Source: §6 · p. 9

A Appendix · A.1 Prompts

The appendix shows prompts for one-step and two-step Rephrase and Respond (RaR): the one-step process rephrases a question and answers it in a single step, while Figure 4 presents a two-step prompt for last-letter concatenation.

A.1 Prompts

  • The 1-step RaR process rephrases the question and then answers it in a single step.Source: Figure 3 · p. 12
  • The 1-step prompt asks the model to rephrase the question based on the initial inquiry.Source: §A.1 · p. 12
  • The prompt asks for two words to be separated in the revised question.Source: §A.1 · p. 12
  • The prompt requires the two words to remain unchanged from the original question.Source: §A.1 · p. 12
  • Figure 4 presents a 2-step RaR prompt for the last-letter concatenation task.Source: Figure 4 · p. 12
  • Figure 4 labels step 1 at the top and step 2 below.Source: Figure 4 · p. 12

A.2 Experiment Details

Experiments fine-tune Llama2 70B Chat with answer-only cross-entropy loss, and S2A uses two-stage nucleus sampling to generate distillation targets. The appendix also specifies training hyperparameters and handling for S2A answers that exceed context limits.

Model training

  • Llama2 70B Chat is fine-tuned with answer-only cross-entropy loss, dropout 0.1, learning rate 5.5e−6, and warmup 1.Table 5 is identified as giving total training steps and tokens per step, but those values are not included in the supplied passage.Source: §A.2 · p. 12, Table 5 · p. 13

S2A

  • S2A uses nucleus sampling with top-p 0.9 in both generation stages; when 20 answers exceed context, the set is reduced to 10 or randomly sampled from if still too long.The rare context-limit case occurs during USC distillation.Source: §A.2 · p. 12
  • For S2A distillation, the second-stage output is the target, using prompts from Weston and Sukhbaatar (2023).Figure 6 presents the System 2 Attention prompts used to extract the training signal.Source: Figure 6 · p. 12

Related papers

Loading 2407.06023v3…