PaperlayerView exact revision on arXiv

Source-linked AI summary

Implicit Chain of Thought Reasoning via Knowledge Distillation

Yuntian Deng, Kiran Prasad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, Stuart Shieber

arXiv:2311.01460v1cs.CLcs.AIcs.LG

TL;DR

Language models commonly generate explicit chain-of-thought, but the paper investigates whether hidden states can support reasoning without producing intermediate words. It distills explicit teacher reasoning into vertical hidden-state computation and finds direct solutions on difficult math tasks, with speed closer to no-CoT inference.

  • Problem

    Explicit chain-of-thought delays final answers, motivating methods that exploit its reasoning benefits while directly producing the answer.

  • Method

    A teacher’s explicit reasoning states are distilled into an emulator and student that perform implicit vertical reasoning through hidden states.

  • Results

    The method directly solves up to five-digit-by-five-digit multiplication and reaches 22% accuracy on GSM8k without explicitly generating intermediate steps.

  • Takeaways & Limitations

    Implicit chain-of-thought enables direct-answer reasoning on tasks previously difficult for models without explicit chain-of-thought, while offering faster inference for long reasoning tasks.

  • Takeaways & Limitations

    Implicit chain-of-thought sacrifices the transparency and interpretability provided by explicit intermediate reasoning steps.

Abstract

from arXiv · show

To augment language models with the ability to reason, researchers usually prompt or finetune them to produce chain of thought reasoning steps before producing the final answer. However, although people use natural language to reason effectively, it may be that LMs could reason more effectively with some intermediate computation that is not in natural language. In this work, we explore an alternative reasoning approach: instead of explicitly producing the chain of thought reasoning steps, we use the language model's internal hidden states to perform implicit reasoning. The implicit reasoning steps are distilled from a teacher model trained on explicit chain-of-thought reasoning, and instead of doing reasoning "horizontally" by producing intermediate words one-by-one, we distill it such that the reasoning happens "vertically" among the hidden states in different layers. We conduct experiments on a multi-digit multiplication task and a grade school math problem dataset and find that this approach enables solving tasks previously not solvable without explicit chain-of-thought, at a speed comparable to no chain-of-thought.

1 INTRODUCTION

The paper proposes implicit chain-of-thought reasoning, which distills a teacher’s explicit reasoning into vertical hidden-state computation so students can answer directly. Experiments show direct solutions on difficult arithmetic tasks, with faster generation than explicit reasoning.

  • Motivation: The approach targets the time cost of generating chain-of-thought by letting models perform intermediate computation internally rather than outputting every reasoning step.Transformer computation can proceed horizontally through generated words or vertically through hidden states across layers.
  • Approach: Implicit chain-of-thought compiles a teacher’s autoregressive reasoning states into predicted layer-wise hidden states used by a student to produce the final answer.This replaces horizontal reasoning through generated tokens with vertical reasoning through transformer layers.
  • Contributions: The work shifts from teacher-forcing explicit reasoning to teacher-teaching, distilling a teacher’s internal reasoning into a student’s implicit computation.The authors identify faster generation, reasoning distillation, and improved direct math responses as the main contributions.
  • Approach: The method combines mind-reading the teacher, thought emulation, and end-to-end coupling and optimization.The student first uses teacher states, an emulator predicts those states from the input, and the combined system is optimized jointly.
  • Results: 22% accuracy on GSM8k is achieved without explicitly generating intermediate steps, while five-digit-by-five-digit multiplication becomes directly solvable with GPT-2 Medium.The method addresses both synthetic multiplication and grade-school math problems.

2 EXPLICIT, IMPLICIT, AND NO CHAIN-OF-THOUGHT REASONING

No-CoT models map inputs directly to answers, explicit-CoT models generate intermediate steps first, and implicit-CoT models process those steps internally. Implicit reasoning therefore preserves intermediate computation during training while avoiding explicit reasoning tokens at inference.

  • 2.1 NO CHAIN-OF-THOUGHT REASONING: No CoT directly parameterizes Pθ(y | x), mapping an input x to output y without exposing intermediate reasoning steps.For 12 × 3, the model directly predicts 36.
  • 2.2 EXPLICIT CHAIN-OF-THOUGHT REASONING: Explicit CoT models factor reasoning as Pθ(z | x)Pθ(y | x, z), generating intermediate steps z before the final output y.For 12 × 3, the model generates 6 + 30 before producing 36.
  • 2.3 IMPLICIT CHAIN-OF-THOUGHT REASONING: Implicit CoT trains on intermediate steps but uses internal hidden states to produce the final answer without explicitly generating those steps at test time.Its inference diagram matches No CoT, although the model internally processes predicted reasoning states.
  • 2.3 IMPLICIT CHAIN-OF-THOUGHT REASONING: Implicit CoT is presented as a middle ground: it retains intermediate reasoning during training but omits explicit intermediate outputs during testing.The approach aims to combine internal reasoning with direct answer generation.
  • 2.3 IMPLICIT CHAIN-OF-THOUGHT REASONING: The proposed three-step strategy trains a student to read teacher states, an emulator to predict them vertically, and a coupled system to optimize end-to-end.The teacher remains fixed during the initial student and emulator stages, while the combined system is subsequently optimized.

3 APPROACH TO IMPLICIT CHAIN OF THOUGHT REASONING

The approach distills a teacher’s explicit chain-of-thought computation into hidden states that a student can use to answer directly. It extracts layerwise teacher states, emulates them from the input, and jointly optimizes the resulting system, with mixture modeling for multiple reasoning paths.

  • The three-stage method trains a student to use teacher states, an emulator to predict those states vertically, and a combined model to optimize final-answer prediction end-to-end.The stages are called Mind-Reading the Teacher, Thought Emulation, and Couple and Optimize.
  • Teacher hidden states are arranged across layers and intermediate tokens, then diagonal or evenly spaced layerwise vectors are selected to represent the reasoning trajectory.A dynamic spacing value can adapt selection to the number of intermediate tokens.
  • The student receives the selected teacher vectors after the input and is trained to produce the final answer while the teacher remains fixed.At test time, the teacher vectors are unavailable, motivating the emulator.
  • The emulator predicts one teacher-state vector per layer from the input using mean squared loss, with a transformer-compatible layerwise design.A special token separates reasoning from output so lower-layer states can contribute to later answer tokens.
  • When multiple reasoning paths exist, mean squared prediction can average incompatible states, so the method models teacher states as a mixture of components.Component supervision uses the intermediate token at the selected position, and a low-temperature softmax approximates differentiable pathway selection.
  • The combined system can learn an internal reasoning process that diverges from the teacher and can be optimized without intermediate-reasoning annotations.The mixture pathway selection is made differentiable with a low-temperature softmax.

4 EXPERIMENTAL SETUP

The experiments evaluate implicit chain of thought on challenging multiplication and grade-school mathematics tasks. They use augmented training data, compare against no-CoT and explicit-CoT baselines, and retain the original test sets.

  • 4.1 DATA: Experiments cover four-digit and five-digit multiplication from BIG-bench and grade-school math problems from GSM8K.The multiplication settings are selected because they are difficult under no chain of thought.
  • 4.1 DATA: Synthetic data augmentation is used for both tasks because preliminary experiments indicated that implicit chain of thought requires a large training set.The original test sets are unchanged.
  • 4.1 DATA: For GSM8K, GPT-4 generates 400k additional problems in the dataset’s format, while multiplication equations are randomly sampled without overlap with BIG-bench.The augmented GSM8K collection is called GSM8K-Aug.
  • 4.2 MODELS AND BASELINES: The study compares implicit reasoning with no-CoT and explicit-CoT training using GPT-2 Small, Medium, and Large, plus few-shot ChatGPT and GPT-4 comparisons in an appendix.Explicit reasoning increases generated-token counts substantially in the augmented datasets.
  • 4.2 MODELS AND BASELINES: The implicit-CoT models are GPT-2 Small and GPT-2 Medium, with normalized teacher states, an MLP for mind-reading, and an LSTM with self-attention for thought emulation.The mixture approach is used for GSM8K-Aug but not multiplication because multiplication’s intermediate steps are unique given each input.

5 RESULTS

Implicit CoT solves multiplication and grade-school math tasks that no-CoT models struggle with, while generating answers much faster than explicit CoT. Performance depends on model scale and remains below explicit CoT overall.

  • 96% accuracy on 5 × 5 multiplication with GPT-2 Medium under implicit CoT, versus 2% without CoT.Table 3 defines accuracy as exact-match final-answer accuracy.
  • 22% accuracy on GSM8K-Aug with implicit CoT exceeds the best GPT-2 no-CoT result of 17%.
  • GPT-2 Small reaches 97% on 4 × 4 multiplication but only 10% on 5 × 5, while GPT-2 Medium reaches 96% on 5 × 5.The authors suggest sufficient layers may be needed for the intermediate calculations.
  • 73% normalized throughput with implicit CoT on 5 × 5 multiplication exceeds explicit CoT’s 14% and approaches no-CoT speed.Throughput is measured with batch size 1 and normalized by the corresponding no-CoT model.
  • Implicit CoT still lags explicit CoT by a large margin, possibly because the models were pretrained for horizontal reasoning and had only 24 layers.

6 ANALYSIS

Analysis shows that how teacher hidden states are compressed and how the emulator and student are optimized materially affects implicit CoT accuracy and interpretability.

  • Taking Different Subsets as Teacher’s Thought Process: 100.0% validation accuracy from diagonal hidden-state elements on 4 × 4 multiplication exceeds first-column, top-row, and bottom-row alternatives.The alternatives achieve 29.9%, 84.4%, and 57.6%, respectively.
  • Mixture: 20.2% GSM8K-Aug validation accuracy with mixture components exceeds 11.2% without the mixture approach.The mixture approach is described as crucial because multiple reasoning pathways exist.
  • Coupling & Optimization: 21.9% validation accuracy after further optimization exceeds 9.4% from coupling the emulator and student without optimization.Fixing the emulator while optimizing only the student reduces accuracy to 13.0%.
  • Mixture: After optimization, predicted mixture components are no longer interpretable, despite resembling the teacher’s intermediate reasoning steps beforehand.Directly using the pre-optimization mapped words yields 9.4% accuracy.

7 RELATED WORK

Related work situates implicit CoT alongside scaling-based arithmetic improvements and knowledge distillation, but distinguishes it by distilling horizontal reasoning into vertical hidden-state computation.

  • Emergent Capabilities: Prior work reports that increasing model and data size improves direct performance on multi-step reasoning tasks, including GSM8K without CoT.Wei et al. (2022a) observed accuracy rising from below 5% to around 7% as training FLOPs increased from 10^21 to 10^24.
  • Knowledge Distillation: The paper uses knowledge distillation to transfer a teacher’s horizontal reasoning process into vertical reasoning across hidden states.This differs from traditional uses such as model compression and non-autoregressive translation.

8 LIMITATIONS

The paper identifies three limitations: implicit CoT is opaque, depends heavily on the teacher’s reasoning process, and currently underperforms explicit CoT.

  • Lack of Transparency and Interpretability: Implicit CoT sacrifices human interpretability because its reasoning occurs internally in hidden states rather than transparent intermediate steps.
  • Reliance on the Teacher’s Thought Process: The initial strategy relies heavily on the teacher’s thought process instead of allowing the model to develop its own reasoning trajectory.
  • Performance Discrepancies: Current implicit CoT results still lag behind explicit CoT, although the authors describe the method as an initial step with room for optimization.

9 CONCLUSION AND FUTURE WORK

The paper proposes implicit chain-of-thought reasoning, distilling horizontal explicit reasoning into vertical computation among transformer hidden states. Experiments show improved direct-answer performance, while explicit chain-of-thought remains more accurate.

  • 9 CONCLUSION AND FUTURE WORK: Implicit chain-of-thought performs reasoning vertically among transformer hidden states rather than horizontally through intermediate tokens.
  • 9 CONCLUSION AND FUTURE WORK: The three-step method distills a horizontally reasoning teacher into a student and emulator trained for vertical reasoning.The steps are mind-reading the teacher, thought emulation, and coupling and optimization.
  • 9 CONCLUSION AND FUTURE WORK: On arithmetic multiplication and grade school math, the method substantially improves direct-answer performance, although explicit chain-of-thought yields higher final-answer accuracy.
  • 9 CONCLUSION AND FUTURE WORK: Future work includes end-to-end joint training with a variational autoencoder and image-modeling techniques for multimodal internal reasoning distributions.

A DATA AUGMENTATION AND PROCESSING

The experiments construct large arithmetic and grade-school-math datasets with explicit intermediate reasoning, using synthetic generation, filtering, and natural-language or equation-based steps.

  • A DATA AUGMENTATION AND PROCESSING: The multiplication data contain 808k training equations and 1k validation equations, with BIG-bench used for testing.
  • A DATA AUGMENTATION AND PROCESSING: Multiplication intermediate steps decompose products into digit-wise partial products and partial sums, including reversed partial-sum representations.
  • A DATA AUGMENTATION AND PROCESSING: The grade-school-math dataset expands GSM8K by prompting GPT-4 with five sampled examples across 80k prompts, producing 379k filtered examples.
  • A DATA AUGMENTATION AND PROCESSING: Explicit CoT uses natural-language intermediate steps, whereas the implicit-CoT teacher is trained with equations as intermediate steps.
  • A DATA AUGMENTATION AND PROCESSING: Table 4 reports exact-match answer accuracy and batch-size-one inference throughput, with daggered baselines using few-shot prompting.

B RAW RESULTS

The raw-results setup compares no-CoT, explicit-CoT, and implicit-CoT systems across arithmetic tasks and GSM8K, using prompted external baselines and several implementation choices.

  • B RAW RESULTS: GPT-4 with no CoT performs on par with GPT-2 Large finetuned with explicit CoT, although the authors suggest data contamination or emergent scale capabilities as possible explanations.
  • B RAW RESULTS: Table 4 uses ChatGPT and GPT-4 as five-shot, temperature-zero baselines, while arithmetic prompts use original numbers without reversed digits or inter-digit spaces.
  • B RAW RESULTS: The no-CoT condition directly requests final answers without decomposing multiplication problems into intermediate reasoning steps.
  • B RAW RESULTS: The explicit-CoT condition provides demonstrations containing intermediate reasoning before the final answer.
  • B RAW RESULTS: The mixture approach predicts vocabulary-sized component distributions, embeds the selected components, and processes them with an MLP and recurrent attention mechanism across layers.

D.3 COUPLE AND OPTIMIZE

Couple and Optimize makes predicted reasoning components differentiable and trains the coupled system toward final-answer prediction. Optimization improves task accuracy but makes the internal components less human-interpretable.

  • D.3 COUPLE AND OPTIMIZE: Differentiable softmax with temperature replaces nondifferentiable argmax over mixture components; the experiments fix the temperature at 0.05.
  • D.3 COUPLE AND OPTIMIZE: 22% task accuracy is reported after coupling and optimization in Table 5, rising from 11.2% before optimization.The table uses GPT-2 Medium with ∆=1.
  • D.3 COUPLE AND OPTIMIZE: For arithmetic tasks, both emulator and student are finetuned after coupling, whereas for GSM8K the student is fixed to alleviate overfitting.
  • D.3 COUPLE AND OPTIMIZE: After optimization, predicted mixture components no longer align with human-interpretable reasoning steps, suggesting a different internal reasoning process.
  • D.3 COUPLE AND OPTIMIZE: The method uses continuous hidden states for reasoning and optimizes the fully differentiable system with gradient descent.
Loading 2311.01460v1…