Source-linked AI summary

Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking

Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, Noah D. Goodman

arXiv:2403.09629v2cs.CLcs.AIcs.LG

TL;DR

Language models often need to infer reasoning implicit in arbitrary text, whereas prior self-taught approaches focused on curated reasoning tasks. Quiet-STaR trains models to generate and use token-level rationales during language modeling, yielding zero-shot reasoning improvements and better prediction of difficult tokens without task-specific fine-tuning. The method remains limited by substantial computational overhead and evaluation on only a 7-billion-parameter model.

  • Problem

    Prior methods learned reasoning mainly from curated tasks, limiting the scale and generalizability of rationales despite reasoning being implicit in diverse text.

  • Method

    Quiet-STaR generates rationales after each token, mixes predictions with and without thoughts, and learns useful rationales with REINFORCE, parallel sampling, learned thought tokens, and extended teacher forcing.

  • Results

    Quiet-STaR improves zero-shot direct reasoning on both GSM8K and CommonsenseQA and disproportionately helps difficult-to-predict tokens, without dataset-specific fine-tuning.

  • Takeaways & Limitations

    Quiet-STaR points toward language models that learn reasoning more generally from diverse web text rather than specializing narrowly for particular datasets.

  • Takeaways & Limitations

    The method incurs substantial token-generation overhead, lacks dynamic rationale stopping, and has been applied only to a 7-billion-parameter model.

Abstract

from arXiv · show

When writing and talking, people sometimes pause to think. Although reasoning-focused works have often framed reasoning as a method of answering questions or completing agentic tasks, reasoning is implicit in almost all written text. For example, this applies to the steps not stated between the lines of a proof or to the theory of mind underlying a conversation. In the Self-Taught Reasoner (STaR, Zelikman et al. 2022), useful thinking is learned by inferring rationales from few-shot examples in question-answering and learning from those that lead to a correct answer. This is a highly constrained setting -- ideally, a language model could instead learn to infer unstated rationales in arbitrary text. We present Quiet-STaR, a generalization of STaR in which LMs learn to generate rationales at each token to explain future text, improving their predictions. We address key challenges, including 1) the computational cost of generating continuations, 2) the fact that the LM does not initially know how to generate or use internal thoughts, and 3) the need to predict beyond individual next tokens. To resolve these, we propose a tokenwise parallel sampling algorithm, using learnable tokens indicating a thought's start and end, and an extended teacher-forcing technique. Encouragingly, generated rationales disproportionately help model difficult-to-predict tokens and improve the LM's ability to directly answer difficult questions. In particular, after continued pretraining of an LM on a corpus of internet text with Quiet-STaR, we find zero-shot improvements on GSM8K (5.9%$\rightarrow$10.9%) and CommonsenseQA (36.3%$\rightarrow$47.2%) and observe a perplexity improvement of difficult tokens in natural text. Crucially, these improvements require no fine-tuning on these tasks. Quiet-STaR marks a step towards LMs that can learn to reason in a more general and scalable way.

1 Introduction

Quiet-STaR generalizes self-taught reasoning from curated question-answering tasks to diverse unstructured text. It generates token-level rationales, uses them to improve future-token predictions, and reports zero-shot gains on reasoning benchmarks without task-specific fine-tuning.

  • Reasoning is implicit in much written text, but prior LM methods largely learned it from curated tasks or predefined task sets.
  • Quiet-STaR generates rationales after every token, mixes predictions with and without them, and improves rationale generation through REINFORCE.
  • Quiet-STaR extends STaR by training language models to generate reasoning that helps predict future text from large internet-text corpora.
  • 5.9%→10.9% on GSM8K and 36.3%→47.2% on CommonsenseQA are reported zero-shot direct-reasoning improvements without dataset-specific fine-tuning.
  • The method uses parallel sampling, learned thought-boundary meta-tokens, a mixing head, and non-myopic loss to make general reasoning training scalable.

2 Related Work

Related work includes supervised reasoning traces, self-taught reasoning, self-play-style methods, custom control tokens, and the Quiet-STaR training procedure. Quiet-STaR addresses scalability by generating rationales in parallel across token positions and learning from future-token prediction outcomes.

  • Training Language Models to Reason: Training on mined reasoning traces can work, but manual annotation is expensive, off-policy, difficult to scale, and bounded by annotator capability.
  • Quiet-STaR: Quiet-STaR samples rationales after every token in parallel, predicts future ground-truth tokens with thoughts, and updates the model using rationale-dependent losses.
  • Self-Taught Reasoning: STaR iteratively trains a language model on reasoning that leads to correct answers, enabling it to solve increasingly difficult problems.
  • Custom Tokens: Custom thought tokens differ from a learned pause token because they initialize an explicit thought rather than representing the entire pause.

3 Problem Statement

Quiet-STaR inserts an auxiliary rationale between each pair of observed tokens and optimizes the model to use these intermediate thoughts when predicting later text. The rationale is an internal variable intended to support difficult computations while preserving language-modeling objectives.

  • Quiet-STaR introduces an auxiliary rationale variable between each pair of observed tokens in a sequence.
  • The model is optimized to generate intermediate thoughts that help predict future text conditioned on the observed prefix and rationale.
  • Intermediate rationales can decompose challenging computations into smaller steps, although an optimal language model would have no theoretical advantage from them.

4 Quiet-STaR

Quiet-STaR trains language models to generate rationales for future text at every token, combine rationale-based and base predictions, and optimize rationale usefulness. Its scalable procedure uses parallel generation, learned thought-boundary tokens, teacher forcing, and REINFORCE.

  • 4 Quiet-STaR: Quiet-STaR generates rationales in parallel across input tokens, marking each rationale with learned start-of-thought and end-of-thought tokens.For an input sequence, it generates multiple rationale candidates at every token position.
  • 4.2 Parallel Generation: Parallel sampling avoids separate forward passes for every token by using cached computations and attention masks that preserve each continuation’s token path.Each generated token attends to its generating tokens and itself, but not to tokens on other counterfactual paths.
  • 4.3 Mixing Post-Rationale and Base Predictions: A mixing head combines predictions made with rationales and without them, allowing the model to transition from initially out-of-distribution thoughts.The mixing weight is computed from the end-of-thought hidden state and the original text-token hidden state.
  • 4.4.2 Non-myopic Scoring and Teacher-forcing: Teacher forcing scores multiple future ground-truth tokens while inserting start and end tokens, extending supervision beyond the immediately next token.The number of future tokens included in the loss is a hyperparameter, and the loss also depends on the mixing weight.
  • 4.4.3 Objective: Excluding negative rewards from the REINFORCE loss made training more stable, although the authors note that this choice may introduce bias.This is an optimization choice rather than a claim that negative rewards are absent.
  • 4.4.3 Objective: REINFORCE increases the likelihood of rationales whose mixed predictions outperform the average across generated rationales for the same token.The resulting gradients update the language-model parameters and thought-token embeddings; excluding negative rewards improved training stability but may introduce bias.

5 Experiments and Results

Quiet-STaR improves zero-shot reasoning on GSM8K and CommonsenseQA, with gains concentrated on difficult tokens and increasing with rationale length. It also supports silent reasoning during explicit chain-of-thought generation and produces qualitatively varied internal thoughts.

  • Downstream Performance: 10.9% improvement on CommonsenseQA and 5.0% boost on GSM8K over the base model were achieved without dataset-specific fine-tuning.Performance increased with the number of rationale tokens used during training.
  • Downstream Performance: 5.9% →8.1% on GSM8K and 36.3% →42.6% on CommonsenseQA were obtained after training Quiet-STaR on C4, with smaller gains than on OpenWebMath.The C4 setup used Mistral 7B with 16 thought tokens and 4 true tokens ahead.
  • Improvement Distribution: Quiet-STaR improves difficult-to-predict tokens disproportionately, despite little average improvement in predicting arbitrary tokens.The pattern suggests that some tokens benefit more from additional thought than others.
  • Improvement Distribution: Inspection suggests that thoughts particularly help tokens requiring recalled information, such as theorem names or the start of a proof step.These examples concern tokens where relevant contextual knowledge supports continuation.
  • Downstream Performance: Quiet-STaR applies silent internal thoughts at every token while generating explicit chain-of-thought reasoning, making the two approaches complementary.The figure uses a model trained with 8 thought tokens for zero-shot GSM8K chain-of-thought.
  • Qualitative Analysis: Helpful thoughts include recalling domain knowledge, proposing near-continuations, and reasoning while reading a question before its final answer.Examples span chemistry, proof structure, and CommonsenseQA.

6 Limitations

The evaluation is limited to a pretrained 7-billion-parameter model, and Quiet-STaR introduces substantial computational overhead. The implementation also lacks dynamic control over when rationales are generated or ended.

  • Scope: Quiet-STaR was evaluated only on a 7 billion parameter model, and its behavior when training a model from scratch remains unknown.The authors expect applying the techniques to a better model could yield larger gains.
  • Compute: Quiet-STaR generates many tokens before each additional output token, creating substantial computational overhead.The authors refer to compute-adjusted performance results and identify dynamic thought allocation as future work.
  • Compute: The current implementation cannot dynamically predict whether to generate or end a rationale.The authors propose using a threshold based on the predicted usefulness of a thought as a possible extension.

7 Conclusion

Quiet-STaR trains language models to reason from diverse web text rather than narrowly specialized datasets, improving downstream reasoning while generating qualitatively meaningful rationales. The authors present it as a step toward more general and scalable reasoning, with dynamic compute allocation and thought ensembling as future directions.

  • Conclusion: Quiet-STaR trains on reasoning tasks implicit in diverse web text instead of narrowly specializing for particular datasets.The conclusion frames this as a route toward more robust and adaptable language models.
  • Conclusion: Quiet-STaR improves downstream reasoning performance while generating qualitatively meaningful rationales.The conclusion presents these results as evidence of the approach’s promise.
  • Future Directions: Future directions include ensembling thoughts and dynamically allocating compute when thought is predicted to be useful.These directions depend on improving how the model predicts the usefulness of internal thought.

Ethics Statement

The paper identifies ethical concerns about whether model-generated reasoning is faithful and whether useful reasoning may be harmful or biased. It also notes that CommonsenseQA contains biased questions and low-quality answers.

  • Model-generated rationales may not accurately represent the model’s internal processing.The paper describes this as a faithfulness concern.
  • The method provides no safeguards against harmful or biased reasoning patterns if the model finds them useful.
  • CommonsenseQA is known to contain biased questions and low-quality answers, although the paper uses it in line with prior work.
  • The paper leaves unclear how the generated rationales should be used beyond improving language modeling.

C Compute-Adjusted Plots

The compute-adjusted visualization normalizes performance curves by the number of thought and talk tokens used for training.

  • Performance curves are normalized by the number of thought and talk tokens used for training.

D Measuring the Impact of Multiple Thoughts Per Sequence and Multiple Tokens Ahead

The study evaluates how multiple sampled thoughts and predicting multiple tokens ahead affect generalization. Multiple thoughts help, while additional thoughts beyond two provide only small gains and longer-horizon supervision mainly improves rationale coherence.

  • Multiple Thoughts Per Sequence: Multiple sampled thoughts outperform computing reward solely from the loss difference, by roughly 0.5% on GSM8K and 3% on CommonsenseQA generalization.
  • Compute-Adjusted Plots: Figure 6 normalizes performance curves by the number of inference calls used.
  • Multiple Tokens Ahead: 0.3% on GSM8K and 3.1% on CommonsenseQA are the gains from predicting more than one token ahead.
  • Multiple Tokens Ahead: Additional tokens ahead beyond two did not improve performance with 12 thought tokens, but produced qualitatively more coherent rationales.The authors suggest that generation quality needs more specific future evaluation.

E.2 OpenWebMath-Trained Mistral 7B

The OpenWebMath-trained model is illustrated on arithmetic examples where generated reasoning often produces incorrect calculations, while one response gives the correct $18 result. The examples show varied rationale quality on the target problem.

  • Problem Setup: The target arithmetic problem asks for daily market earnings after subtracting three eaten eggs and four eggs used for muffins from sixteen eggs, then selling the remainder for $2 each.
  • Generated Rationales: The examples include both concise arithmetic reasoning and longer, confused attempts involving irrelevant costs, repeated steps, or inconsistent egg counts.
  • Correct Solution: $18 is the correct daily profit reported in one rationale.

E.3 Quiet-STaR-Trained Mistral 7B

This section illustrates the arithmetic reasoning traces associated with Quiet-STaR-trained Mistral 7B, using Janet’s egg-selling problem and token-level contribution visualizations. The examples show intermediate quantities leading to a final answer of $18 and introduce visualizations of where thoughts affect prediction.

  • Arithmetic reasoning example: Janet’s problem requires subtracting eggs used for breakfast and muffins from daily production before calculating market revenue.The example states 16 eggs are produced, 3 are eaten, 4 are used in muffins, and 9 remain for sale.
  • Arithmetic reasoning example: $18 is the stated total revenue from selling 9 eggs at $2 each.The example gives both the per-egg price and the resulting total.
  • Arithmetic reasoning example: The reasoning traces explicitly represent intermediate facts, arithmetic operations, and the final answer rather than only outputting the result.The traces include the egg counts, the subtraction 16 - 3 - 4, the remaining count of 9, and the final total of 18.
  • Token-level effects: Figure 7 shows that most tokens are unaffected by thoughts, while a small fraction of difficult tokens receive substantial log-probability improvements.The distribution is described as skewed toward little or no change, with disproportionate gains on hard tokens.
  • Token-level effects: Figure 8 marks helpful token-level thoughts in green and harmful ones in yellow, with opacity indicating impact.The example suggests improvements often occur where recalling relevant information helps, such as theorem names or the start of a proof step.

I Handling Instability

The paper identifies instability in learning how generated thoughts contribute to prediction and reports that this instability prevented learning with separately parameterized thinking and talking heads. The resulting design minimizes components that transform the base model’s outputs.

  • Sources of instability: The thought-to-prediction mapping is jointly learned with the thoughts, so useful thoughts may receive no learning signal when the mixing weight is zero.The paper frames this as a central source of instability and an exploration–exploitation trade-off.
  • Sources of instability: Gumbel-Softmax with a straight-through estimator produced vanishing gradients after many consecutive softmax operations.This alternative was explored but did not resolve the instability.
  • Design response: Separate thinking and talking heads were explored, but the previously identified instability prevented learning in all instances.The heads included linear layers and MLPs initialized to contribute zero residually to the base language-model outputs.
  • Design response: The final design minimizes components that can transform the language model’s outputs with and without rationales.The mixing head instead predicts a weight for combining states, functioning as an attention-like mechanism over the model’s thinking.
Loading 2403.09629v2…