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

Quite-STaR is an extension of Self-Taught Reasoner (STaR). Quiet-STaR allows LLMs to learn by generating rationales or thoughts at every token position to interpret the subsequent text, thereby improving prediction performance.
This paper proposes a token-wise parallel sampling algorithm, using meta-tokens to indicate the start and end of thoughts. The generated thoughts consistently increase with the number of tokens used in the LM’s internal thoughts. After continued pretraining with Quiet-STaR LLM, the zero-shot performance improved on GSM8K (from 5.9% to 10.9%) and CommonsenseQA (from 36.3% to 47.2%), while also reducing the difficulty of handling hard tokens in natural text.

Introduction

  • Many implicit meanings are hidden in the text, and understanding this hidden information is very important for LLMs. Consider this scenario: when reading a complex mathematical proof, although the final answer may be clear, the key lies in understanding the middle steps between the initial conditions and the conclusion. Similarly, in conversations, we rely on intermediate assumptions and background knowledge to understand what is being said, and this implicit reasoning ability is something that earlier works like STaR lack.
  • Previous research has shown that reasoning helps improve language model performance across various tasks. Existing methods (such as STaR) learn reasoning techniques by inferring rationales (the reasoning step) from question–answer pairs. However, STaR is limited to specific tasks and requires the existing existing answer–rationale pairs, which limits the generality and scalability of reasoning. Based on STaR, Quiet-STaR enables LLMs to generate rationales/thoughts for any text they produce, making the reasoning process more general and broadly applicable.

Basic Method

📌

Problem Statement

This study introduces an auxiliary rationale variable between every pair of tokens, optimizing a language model with parameter so that it can generate intermediate thoughts (or reasoning rationales).
The model uses tokens from to together with their rationales/thoughts to predict the next to tokens, rather than only the immediate next token. Implementing Quiet-STaR faces several challenges:
  • Generating rationales for every token is computationally expensive, and LLMs themselves cannot simultaneously produce or leverage such internal reasoning.
  • Quiet-STaR not only needs to predict the next word, but also must consider longer-term dependencies within the text.
📌

Quiet-STaR Overview

Quiet-STaR is divided into three parts: think, talk, and learn.
  1. Think: Generate a thought based on the current token and previous tokens. Then, use the thought token and the original text tokens to predict subsequent tokens.
  1. Talk: Smoothly balance the logits from next-token prediction with or without thoughts, preventing overly severe out-of-distribution (OOD) issues at the beginning.
  1. Learn: Finally, use the REINFORCE RL algorithm (See Section 3.2.9 Policy Gradient Algorithm) based on policy gradients to optimize model parameters, thereby increasing the probability of outputting useful thoughts.
Quiet-STaR generate thoughts, in parallel, following all tokens in the text (think). The model produces a mixture of its next-token predictions with and without a thought (talk). Then apply REINFORCE to increase the likelihood of thoughts that help the model predict future text while discarding thoughts that make the future text less likely (learn).
Quiet-STaR generate thoughts, in parallel, following all tokens in the text (think). The model produces a mixture of its next-token predictions with and without a thought (talk). Then apply REINFORCE to increase the likelihood of thoughts that help the model predict future text while discarding thoughts that make the future text less likely (learn).
📌

1. Think and Parallel Rationale Generation:

A key challenge of Quiet-STaR is efficiently generating rationales for every token position in the input sequence. This requires each token to perform an independent forward pass, which becomes computationally difficult for long sequences. To address this, the paper uses a masking mechanism to achieve efficient parallel computation:
  • Suppose our goal is to generate a 2-step thought at every token position
    • notion image
  • To generate the first thought token, an attention mask is used for causality
    • notion image
  • When the first step thought to is generated, append the thought to the end of the sentence, and create a new attention mask matrix
    • notion image
  • For the next step, the second step thought and is appened again to the sequence, and used to generate the third step thought
    • notion image
This process repeats step by step. Each thought token requires conditioning on the mask of its corresponding text token.
  • At each token position, multiple thoughts are generated.
  • Each thought is marked with <|startofthought|> and <|endofthought|> tokens to indicate its beginning and end, and is learned in a trainable way.
📌

2. Talk: Mixing Reasoning with Basic Prediction

When starting training from a pretrained model, the generated thoughts are initially out of distribution, which harms performance.
To enable a smoother transition toward reasoning, this paper introduces a learned mixing weight between thought-based and non-thought LM predictions. Specifically, this is implemented as a three-layer MLP mixing head.
  • Given the hidden state of the thought-ending token and the corresponding hidden state of the original text token, the mixing head outputs a scalar weight.
  • This weight is then applied to interpolate between the LM’s prediction logits generated with thoughts and without thoughts.
This method helps mitigate distribution shift problems during fine-tuning.

Optimization and Loss Function Design

Optimize Start-of-Thought and End-of-Thought Tokens:
<|startofthought|> and <|endofthought|> tokens serve as meta-tokens that control the model to generate rationales/thoughts. Optimizing the representation of these tokens is critical, especially the <|startofthought|> token. In this study, the embeddings of these two tokens are initialized as em dashes (“−−−”), since em dashes in text often indicate pauses or thoughts. This leverages prior knowledge, and during training, the embedding parameters of these tokens are further updated.
The loss function has two parts:
  1. NLL loss: using Non-myopic Scoring and Teacher-forcing
  1. REINFORCE loss

1. NLL loss: Non-myopic Scoring and Teacher-forcing

In simple terms, thoughts may not be useful for every token prediction. The model’s reward should not only depend on the next token following a thought, but also depend more on the following content. Therefore, this paper supervises training by using the current token and the real tokens from multiple future steps together.
The teacher-forcing technique is represented by the dashed lines in the figure: some existing tokens are inserted (though in practice the LM generates them itself, with mixing head labels applied in between).
notion image
This method helps the model better learn to generate useful thoughts during training, thereby improving its ability to predict future tokens. Specifically, consider the following example:
  • Dataset sentence: 2 * 3 = 6 is true. If the model inserts a thought after the = sign, the sequence becomes:
    • 2 * 3 = <|startofthought|> 3 + 3 = 6 <|endofthought|> 6 is true
  • Generate a thought: After the token =, the model generates the thought span 3 + 3 = 6 between <|startofthought|> and <|endofthought|>. There is no ground truth for these tokens, since they are latent.
  • Teacher-forcing the future ground-truth tokens: After the thought ends, the next real token in the dataset is 6.
    • Compute logits with-thought (conditioned on the thought span) and no-thought (ignoring it).
    • Interpolate the two with a mixing weight.
    • Use teacher forcing: the label is the true token 6, not the model’s own sample.
    • This is similar to the training method used in the model’s pretraining step, where the source text’s next token is used to supervise the current token.
  • Non-myopic scoring (cross-entropy over multiple future tokens): Instead of stopping at 6, the model also looks ahead to the next few ground-truth tokens (is, true).
    • Predict is (label = true token is).
    • Predict true (label = true token true).
    • Cross-entropy loss is applied at each of these positions.
The key is that although the tokens inside the thought (3 + 3 = 6) have no direct labels, the losses on 6, is, and true backpropagate through the thought span. This way, the thought tokens are indirectly trained to be more helpful for predicting the real text.

2. REINFORCE loss: Optimization Objective

Use REINFORCE to optimize the probability of generating rationales. Given the previous tokens and a specific rationale, the model considers the likelihood of the next tokens (ahead tokens). To reduce variance, each token generates multiple rationales. Define the reward advantage function of each rationale/thought as the difference between:
  • The mixed likelihood of predicting the next tokens given the current thought, and
  • The baseline likelihood (averaged across multiple thoughts).
Update the language model parameters with the REINFORCE loss term, to increase the probability of generating thoughts whose performance is above average:
notion image
After we know the two parts of the loss, let’s make a brief conclusion:
  1. NLL loss (normal cross-entropy)
      • is just the negative log-likelihood of the real future tokens (like 6 is true).
      • This is the teacher-forcing supervised signal.
  1. REINFORCE loss
      • For each sampled thought , compute its reward .
      • Reward = improvement in the likelihood of future real tokens compared to a baseline.
      • Update the probability of generating this thought:
        • If , increase the probability of that thought.
        • If , decrease probability. However, in practice, Quiet-STaR removes the data when to stable training.

Experiments and Results

All experiments are based on the Mistral 7B base model
🔬
1. Downstream task performance improvement
  • On CommonsenseQA and GSM8K datasets, Quiet-STaR significantly improved the model’s zero-shot reasoning ability, with accuracy increasing by 10.9% and 5.0% respectively. As the number of reasoning tokens increased, the performance gain became even better.
  • Training on the C4 dataset also improved model performance on the above benchmarks, though the improvement was smaller.
Generalization Results
Generalization Results
🔬
2. Distribution improvement
  • On average, the model’s ability to predict arbitrary tokens improved only slightly, but for more difficult tokens, the improvement was significant, indicating that reasoning particularly helps with challenging tokens.
  • Reasoning is especially beneficial for predicting tokens that require recalling related information, such as the beginning of definitions or steps in proofs.
Most tokens are unaffected, indicated by the highest density at the (0, 0) position. For easy token (where the original token prediction loss is around 0 to 5), the improvement is hard to distinguish. But for difficult tokens (original token prediction loss > 15), the improvement is significant.
Most tokens are unaffected, indicated by the highest density at the (0, 0) position. For easy token (where the original token prediction loss is around 0 to 5), the improvement is hard to distinguish. But for difficult tokens (original token prediction loss > 15), the improvement is significant.
🔬
3. Relationship with Chain-of-Thought
  • Chain-of-Thought prompting and Quiet-STaR are orthogonal but complementary. Quiet-STaR allows the model to reason at every token, generating more structured and useful reasoning chains.
  • In GSM8K testing, majority-vote accuracy improved from 40.6% to 47.7%.
Blue curve Quiet-STaR is a 8-thought-token-trained Mistral model
Blue curve Quiet-STaR is a 8-thought-token-trained Mistral model

Conclusion

Limitations:
  1. It is unclear whether these techniques are effective when training a model completely from scratch, since the experiments were only applied to the Mistral 7B base model.
  1. Quiet-STaR introduces significant computation before generating each token, as it produces many intermediate tokens.
Conclusion:
  1. Significance and outcomes of the method
      • Quiet-STaR enables models to learn from hidden reasoning tasks, improving downstream reasoning performance.
      • The generated rationales/thoughts have meaningful interpretability.
  1. Future research directions
      • One could try aggregating multiple thoughts to further improve the prediction of future tokens.
      • Another direction is to explore learning when to reason, rather than generating a thought at every token.
 
 
 
 
Prev
CivitAI’s Payment Issue
Next
Awesome-AI-Tutorials
Loading...
Article List
LLM Learning Roadmap
✨ Awesome-Anything
🖼️ Digital Image Processing
🍃 LLM Components
🌱 LLM Pre-training
☘️ LLM Post-Training
🍀 LLM Popular Models
🪴 LLM Applications
🌿 LLM Optimization
🌾 LLM Compression
🌵 LLM Hands-on Practice
🌴 LLM Must-read Papers
🌳 LLM Q&A
🐝 VLM Image Encoders
📝 MISC.