分析下面论文的内容， 结合qwen2大模型实现论文中的方案；要求代码能够训练和推理；Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking
Eric Zelikman
Stanford University
Nick Haber
Stanford University
Georges Harik
Notbad AI Inc
Yijia Shao
Stanford University
Noah D. Goodman
Stanford University
Varuna Jayasiri
Notbad AI Inc
Abstract
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%→10.9%) and CommonsenseQA (36.3%→47.2%) and observe a perplexity improve- ment 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.
“Life can only be understood backwards; but it must be lived forwards.”
— Sren Kierkegaard
1 Introduction
Much of the meaning of text is hidden between the lines: without understanding why statements appear in a document, a reader has only a shallow understanding. Moreover, this has been repeatedly shown to be true for LMs as well, in the contexts of tasks ranging from commonsense reasoning to theorem proving to programming (Wei et al., 2022b; Nye et al., 2021; Zelikman et al., 2022; 2023a; Kojima et al., 2022). Reasoning about implications of text to predict later text has consistently been shown to improve LM performance on a variety of tasks, but methods for allowing LMs to learn from their reasoning (e.g., Zelikman et al. 2022) have focused on solving individual tasks or predefined sets of tasks (e.g., Wei et al. 2021b). These works rely on carefully curated datasets to provide either specific reasoning tasks or in some cases, the reasoning itself. We instead ask, if reasoning is implicit in all text, why shouldn’t we leverage the task of language modeling to teach reasoning?
1
arXiv:2403.09629v2 [cs.CL] 18 Mar 2024
   OriginalText
4 LM 2 LM
+ LM 1 LM 2 LM = LM 5 LM 4 LM \n
        Thought0
Thought1
Thought2 Thought3
1) Think
Thought4
Thought5
Thought7
 2) Talk Thought6
Helps Prediction
Reward
Harms Prediction
Discard
3) Learn
LM
                  Sampled LM 2 LM + LM 2 Thought START
END
         Sampled LM 4 LM + LM Thought START
1
END
   Figure 1: Quiet-STaR. We visualize the algorithm as applied during training to a single thought. We 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). We apply REINFORCE, as in STaR, to increase the likelihood of thoughts that help the model predict future text while discarding thoughts that make the future text less likely (learn).
In particular, the Self-Taught Reasoner (STaR, Zelikman et al. 2022) showed that LMs can bootstrap their reasoning ability on question-answering (QA) datasets by sampling rationales to attempt to answer questions, training on rationales if they led to a correct final answer, and then repeating this to iteratively solve more difficult problems. Yet, training from curated QA datasets limits the scale and generalizability of the rationales. QA datasets, especially high-quality ones, require thoughtful curation and will inherently only ever cover a subset of reasoning tasks. Thus, we extend STaR – instead of the LM learning to reason on particular tasks like mathematical QA, we train an LM to generate reasoning that helps it infer future text from a large internet text corpus. As a result, we allow the LM to learn from the diverse tasks present in language (Weber et al., 2021). This builds on an intuition essential to the current language modeling paradigm, namely, that ”language models are unsupervised multitask learners” (Radford et al., 2019). Thus, as in STaR, we leverage the LM’s pre-existing reasoning ability to generate rationales and train the LM on them with a REINFORCE-based reward (Williams, 1992). We refer to this technique as Quiet-STaR, as it can be understood as applying STaR “quietly”, training the model to think before it speaks.
Broadly, Quiet-STaR proceeds by generating rationales after every token to explain future text (think), mixing the future-text predictions with and without rationales (talk), and then learning to generate better rationales using REINFORCE (learn). We apply Quiet-STaR to Mistral 7B (Jiang et al., 2023) using the web text datasets OpenWebMath (Paster et al., 2023) and Colossal Clean Crawled Corpus (C4, Raffel et al. 2020). We find that, even without dataset-specific fine-tuning, Quiet-STaR results in improvements to zero-shot direct- reasoning abilities on CommonsenseQA (36.3%→47.2%) and GSM8K (5.9%→10.9%), and that these improvements consistently increase with the number of tokens used in the LM’s internal thoughts. Lastly, we qualitatively investigate patterns in the generated rationales.
In solving this task, we make the following contributions:
1. We generalize STaR to learn reasoning from diverse unstructured text data. To our knowledge, this is the first work explicitly training LMs to reason generally from text, rather than on curated reasoning tasks or collections of reasoning tasks.
2. We propose and implement a parallel sampling algorithm that makes our training procedure scalable, generating rationales from all token positions in a given string.
3. We introduce custom meta-tokens at the start and end of each thought to allow
the LM to learn that it should be generating a rationale and when it should make a
prediction based on that rationale.
4. We apply a mixing head to retrospectively determine how much to incorporate the
next-token prediction from a given thought into the current next-token prediction.
5. We show that a non-myopic loss, including multiple tokens ahead for language
modeling, improves the effect of thinking.
6. Onmultipletasks,wedemonstratethatthinkingallowstheLMtopredictdifficultto-
kens better than one trained on the same web text, improving with longer thoughts. 2

   11% 10% 9% 8% 7% 6%
# Thought, Ahead Tokens 24, 12
16, 8 12, 4 10, 4
8, 4 Baseline
10 20 30 40 50 60 70 80 90 100 Training Step
 GSM8K Accuracy (%)
       # Thought, Ahead Tokens 46% 24, 12
 CommonsenseQA Accuracy (%)
 44% 42% 40% 38% 36% 34%
16, 8 12, 4 10, 4
8, 4 Baseline
       10 20 30 40 50 60 70 80 90 100 Training Step
(a) GSM8K (b) CommonsenseQA
Figure 2: Generalization Results. We evaluate the extent to which the model trained with Quiet-STaR generalizes to directly answering problems that require reasoning. The left plot (a) shows the zero-shot accuracy on GSM8K, while the right plot (b) shows the zero-shot accuracy on CommonsenseQA, without any fine-tuning. In both plots, the x-axis represents training steps, and each line corresponds to a different number of thinking tokens used during Quiet-STaR training. The y-axis measures the zero-shot direct accuracy on the respective datasets. We also include an inference normalized version of this plot in Figure 6.
2 Related Work
2.1 Reasoning in Language Models
There have been many works on training and exploiting language models to solve difficult tasks by first training them to reason through them. For example, Rajani et al. (2019) demonstrated that a pre-trained language model fine-tuned to output on human reasoning traces before answering multiple-choice commonsense reasoning questions outperformed one trained directly on answers. Shwartz et al. (2020) demonstrated that language models, when provided with some scaffolding, can generate these helpful chain-of-thought solutions without additional supervision. Later, Nye et al. (2021) demonstrated that “scratchpads” required less scaffolding when the language models were more capable, a result later reinforced by Wei et al. (2022b), emphasizing informal tasks, and further strengthened by Kojima et al. (2022), demonstrating this behavior could be accomplished zero-shot. Most recently, Wang & Zhou (2024) showed further that for commonsense-question answering, one could force a language model to leverage chain-of-thought reasoning by preventing it from emitting any valid answer tokens unless it was confident. However, once again, these approaches only work for a question-answer dataset, and Wang & Zhou (2024) relies on heuristics to identify when the model has output answer tokens. Somewhat like TRICE (Phan et al., 2023), we use the relative improvements in the log-likelihood of the target text across rationales as an estimate of quality, but we simply subtract the mean reward and do not incorporate more complex control variates.
2.2 Training Language Models to Reason
One direction that researchers have used to train language models to reason or improve their reasoning is training the language model on mined reasoning traces or reasoning-like data (Rajani et al., 2019; Wei et al., 2021a; Lewkowycz et al., 2022; Chung et al., 2022; Gunasekar et al., 2023). Although this approach has been demonstrated to be effective, it comes with drawbacks. It requires either manual annotation, which is sensitive to the capability of the annotators and is off-policy for the language model (i.e., the distribution of reasoning is not text that the language model would otherwise likely have generated). This approach is also expensive, difficult to scale, and provides no clear path to solving problems harder than those that the annotators are capable of solving.
Another direction for teaching reasoning relies on a language model’s own generated reasoning, which can be seen as building on a large body of literature on self-play (Silver et al., 2017; Anthony et al., 2017; Polu & Sutskever, 2020). These include methods such as the
3

  Algorithm 1: Quiet Self-Taught Reasoner (Quiet-STaR)
Input: Language model θ0, training steps num steps, sequence length l, thought length t, learning rate α, batch size b, number of thoughts nthoughts, number of ground truth tokens used for supervising each thought ntrue
Output: Language model θ that generates rationales to predict future text for i = 0 to num steps do
Sample batch of sequences X of length l
hinit ← hidden statesθi (X)
for j = 1 to l in parallel using attention mask do
log pinit ← lm head (hinit ) // Predict next tokens j:j+ntrue θi j:j+ntrue
Tj ← generate tokensθi ([X:j; <start thought>], t, nthoughts) // Generate thought Tj ← [Tj; <end thought>]
   hthought ← hidden statesθ ([X:j; Tj; Xj:j+n j:j+ntrue i
log pthought ← lm head (hthought ) j:j+ntrue θi j:j+ntrue
−1])
// Predict next tokens w/ thought
wj:j+n
log ptalk ← wj:j+n
) + (1 − wj:j+n
true
← mixing head (hthought , hinit
θi j:j+ntrue j:j+ntrue
j
LNLL ← − log ptalk (Xj+1:j+n
true
· log pinit
true j:j+ntrue
) · log pthought j:j+ntrue
// Mix logits
+1)
+1) − log ptalk
j j:j+ntrue
rj = log ptalk (Xj+1:j+n
j:j+ntrue
∇θ Lj ← ∇θ LNLL + ∇θ LREINFORCE
return θnum steps
Self-Taught Reasoner (Zelikman et al., 2022), which demonstrated that a language model iteratively trained on its reasoning that led to correct answers could solve increasingly difficult problems. Later work aimed to leverage additional information or assumptions such as Huang et al. (2022) which demonstrated that the algorithm proposed in STaR could still work if one assumed that the majority-vote answer was correct (although this has a lower ultimate performance). Further work has generalized the results of Zelikman et al. (2022), such as Uesato et al. (2022) which demonstrated additional usefulness to “process-based” supervision where incorrect reasoning traces were filtered, recently V- STaR (Hosseini et al., 2024) that demonstrates that training a verifier to guide generation also improves performance, as well as TRICE (Hoffman et al., 2024) which maximizes the marginal likelihood of the correct answer given several reasoning traces per problem. Finally, related work has also explored learning intermediate reasoning in the constrained setting of making mathematical statements, where statements in the model’s intermediate reasoning could be constrained to only be valid mathematical statements (Poesia et al., 2023). We include further discussion of related reasoning works in Appendix F.
2.3 Meta-tokens
Recently, a growing body of work has demonstrated the usefulness of custom tokens optimized to perform specific functions in the context of a neural network – for this reason, they have also been referred to as “function vectors.” (Todd et al., 2023). One of the original instantiations of this was prompt-tuning (Lester et al., 2021) (and relatedly prefix-tuning (Li & Liang, 2021)), where the embeddings corresponding to the tokens of a prompt could be optimized to better accomplish a task. Others have applied meta-tokens to compress long prompts (Li et al., 2023; Jung & Kim, 2023) for efficiency. Most relevant to this work, Mu et al. (2024) optimized a token such that, when the tokens after it could not attend to the tokens before it (i.e., a context compression token), it would provide sufficient information to future tokens. Although we do not focus on compression, we share the problem of learning a
true
(Xj+1:j+n ∇θLREINFORCE ←−rj1[rj >0]·∇θlogpθ(Tj|[X:j;<startthought>])
j:j+ntrue ji
+1)
true
true
jj
θi+1 ←θi−α∑lj=1∇θLj // Update model parameters
true
 4

 token that affects attention and controls complex downstream behavior. In one related work, Goyal et al. (2023) show that learning a single ”pause” token (essentially representing each token as two tokens) improves LM performance. However, unlike the thought tokens in our work, this pause token does not initialize a thought – instead, it can be seen as acting as the entirety of the thought. We find that reasoning in language is significantly more helpful.
3 Problem Statement
In this work, we introduce an auxiliary ‘rationale’ variable between each pair of observed tokens of the sequence. We then aim to optimize a language model with parameters θ with the capacity to generate intermediate thoughts (or rationales) such that
θ∗ = argmaxθEx [logpθ (xi:n|x0:i,rationaleθ (x0:i))]
Note that, in principle, this provides no advantage over an optimal language model that already correctly models the language’s distribution over strings. Yet, in practice, extensive prior work has shown that language models benefit from intermediate rationales on reason- ing tasks (Nye et al., 2021; Zelikman et al., 2022; Wei et al., 2022b). Some work has aimed to explain the effects of chain-of-thought reasoning, namely attributing it to “locality of experience” (Prystawski et al., 2024). More broadly, reasoning allows a model to decompose a challenging computation into smaller steps. In effect, we train the model to learn which decomposition and planning steps are effective in predicting future text. Also note that we formulate the objective as accurately predicting the remaining sequence, rather than only the next token. Once again, for an optimal LM these would be equivalent. However we find that the non-myopic formulation leads to a more effective loss for learning rationales.
4 Quiet-STaR
4.1 Overview
Quiet-STaR operates with three main steps (Figure 1):
1. Parallel rationale generation (think, Subsection 4.2): In parallel across n tokens xi in an input sequence x0:n, we generate r rationales of length t: ci = (ci1, . . . , cit), resulting in n × r rationale candidates. We insert learned 
 and 
 tokens to mark each rationale’s start and end.
2. Mixingpost-rationaleandbasepredictions(talk,Subsection4.3):Fromthehidden state output after each rationale, we train a ”mixing head” – a shallow MLP produc- ing a weight determining how much the post-rationale next-token predicted logits should be incorporated compared to the base language model predicted logits. This approach eases distribution shift early in finetuning, due to introducing rationales.
3. Optimizing rationale generation (learn, Subsection 4.4): We optimize the rationale generation parameters (start/end tokens and LM weights) to increase the likelihood of rationales that make future text more probable. We use REINFORCE to provide a learning signal to rationales based on their impact on future-token prediction. To reduce variance, we apply a teacher-forcing trick to include in the loss the likelihood of predicting not only the token after the thought but also later tokens.
4.2 Parallel Generation
A key challenge in Quiet-STaR is efficiently generating rationales at each token position in the input sequence. Naively, this would require a separate forward pass for each token, which becomes computationally intractable for long sequences.
We allow for highly parallel generation by first observing that an inference pass of a language model produces a probability distribution over the next tokens for all input tokens. Naturally, this allows us to sample one next token from each token in the input. If one has generated a successor from each token, it is not possible to simply continue with the original sequence. For example, imagine predicting the next token after each token of “< bos > the cat sat” one might generate “yes orange saw down” – each successor by itself is a reasonable next token
5

    Base Text
a b c
Thought Token 1
d a’ b’ c’ d’
a’ b’ c’ d’
  Thought Token 2
Thought Token 1
a’’ b’’
a
b c’’ d’’ c
a’ b’ c’ d’
Base abcd Text
Generation Paths
Parallel Inference Mask
d
    Figure 3: Parallel Generation. By constructing an attention mask that allows all thought tokens to pay attention to themselves, all preceding thought tokens within the same thought, and the preceding text, we can generate continuations of all of the thoughts in parallel. Each inference call is used to generate one additional thought token for all text tokens.
to a prefix of the sequence, but the list of tokens is a set of “counterfactual” continuations of these prefixes. We can, however, leverage these continuations to generate hidden thoughts for each observed token.
To do this efficiently, we cache each forward pass and concatenate a diagonal attention mask to the previous attention mask: each generated token now attends to all of the tokens that were used to generate it, as well as to itself (but not to token on other “counterfactual” paths). Moreover, this parallelized next-sampling token procedure can be repeated arbitrarily many times (or at least, until one runs out of memory). We visualize this procedure in Figure 3 and highlight additional ways to make this algorithm faster in Appendix B.
4.3 “Mixing” (Residual) Heads
When starting with a pre-trained model, thoughts will initially be out of distribution, and hence harm language modeling performance. To smooth the transition to thinking, we introduce a learned interpolation between the LM predictions with and without thoughts. Given the end-of-thought token’s hidden state and the hidden state of the original text token, the mixing head outputs a weight that determines the extent to which the post- thought prediction logits will be used. We use a shallow multi-layer perceptron for this head, outputting a scalar for each token. We include implementation details in Appendix A.
4.4 Optimizing Rationale Generation
4.4.1 Optimizing Start-of-Thought and End-of-Thought Tokens
The 
 and 
 tokens serve as learned meta-tokens that control the model’s rationale generation. Optimizing the representation of these tokens, especially the 
 token, is crucial but challenging due to the discrete nature of the rationale tokens. We initialize the start and end token embeddings to the embedding corresponding to the em dash, ”−−−”, which often appears in text data to denote a pause or thought. This leverages the language model’s preexisting knowledge. In addition, to allow these embeddings to be optimized more quickly, we apply a (hyperparameter) weight to the gradients of these embeddings during the update step. Intuitively, the start thought tokens can be understood as putting the model into a “thinking mode” and the end thought token can be understood as telling the model when it’s done thinking.
4.4.2 Non-myopic Scoring and Teacher-forcing
Because we do not expect thoughts to be useful in predicting every token, we would prefer the model’s reward to depend less on the exact next word in the text following the thought and more on the following semantic content. There are two primary challenges here. First, unlike in typical language modeling with transformers, only the thoughts corresponding to
6

  defgh
                         <t> </t> g h
Figure 4: Forward Pass and Teacher Forcing. We visualize a single forward pass of our algorithm. Solid lines denote language model computation, while dashed lines indicate tokens are inserted via teacher forcing, and the mixer represents the mixing head. In particular, we visualize predicting three tokens ahead. Thought generation is shown in more detail in Figure 1 and Figure 3.
a given next-token prediction receive a gradient from that prediction—a consequence of our parallel sampling strategy. We could address this by adding loss terms for future tokens by sampling the tokens before. However this would result in much higher entropy for language modeling in general and lower-quality generated text, because it would train the LM to par- tially disregard its preceding tokens. Instead, we use the parallel attention mask to compute the log probabilities of the true next tokens, applying teacher forcing by assuming the model selected the correct next ground-truth token (as implicit in normal language modeling with transformers). Note that the loss for each future token also depends on a mixing weight computed from the end thought token and the previous observed token. The number of future tokens included in the loss is a hyper-parameter. We apply the same teacher-forcing technique to insert the start and end tokens. We visualize this procedure in Figure 4.
4.4.3 Objective
We use REINFORCE to optimize the likelihoods of the rationales based on their usefullness: the log-likelihood of the ntrue true next tokens Xj+1:j+ntrue+1 under the language model given
previous observed tokens and a particular rationale (ptalk as shorthand for the mixed j:j+ntrue
prediction probabilities after thinking, see Algorithm 1). To reduce variance, we generate multiple rationale continuations for each token in the input sequence (loosely inspired by TRICE, Phan et al. (2023)). We thus define the reward rj for each rationale Tj as the difference
j:j+ntrue
r =logptalk
∇ LREINFORCE = −r · ∇ log p (T |[X ; 
]) θ j j θ θ j :j
We found it useful to exclude the negative reward from the REINFORCE loss term, as it led to more stable training, though it may introduce some bias.
This loss term encourages the model to generate rationales that improve its predictions of
future tokens compared to the average prediction across all generated rationales for that
token. The gradients from this loss are used to update both the LM parameters and the
start-of-thought and end-of-thought token embeddings, with a (hyperparameter) weight
applied to the gradients of the start-of-thought and end-of-thought token embeddings to
accelerate their optimization. By iteratively optimizing these parameters, Quiet-STaR trains
the model to generate more useful rationales throughout training. Lastly, we also include a
log-likelihood loss, LNLL, to ensure that the LM learns to optimize the talking heads and j
also receives a next-token prediction signal for the base LM head1.