Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking
Quiet-STaR：语言模型可以自我学习在发言前思考

1 Introduction 1 引言

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.
   我们将 STaR 推广到从多样化的非结构化文本数据中学习推理。据我们所知，这是第一项明确训练语言模型从文本中进行一般推理的工作，而不是在精心策划的推理任务或推理任务集合上进行。
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. On multiple tasks, we demonstrate that thinking allows the LM to predict difficult tokens better than one trained on the same web text, improving with longer thoughts.
   在多个任务中，我们证明了思考使得语言模型能够比在相同网络文本上训练的模型更好地预测困难的标记，并且随着思考时间的延长，预测能力有所提高。

2 Related Work 2 相关工作

2.1 Reasoning in Language Models
2.1 语言模型中的推理

2.2 Training Language Models to Reason
2.2 训练语言模型以进行推理

Algorithm 1: Quiet Self-Taught Reasoner (Quiet-STaR)
Input: Language model \(\theta_{0}\), training steps num_steps, sequence length \(l\), thought length
    \(t\), learning rate \(\alpha\), batch size \(b\), number of thoughts \(n_{\text {thoughts }}\), number of ground
        truth tokens used for supervising each thought \(n_{\text {true }}\)
Output: Language model \(\theta\) that generates rationales to predict future text
for \(i=0\) to num_steps do
    Sample batch of sequences \(X\) of length \(l\)
    \(h^{\text {init }} \leftarrow\) hidden_states \(\theta_{i}(X)\)
    for \(j=1\) to \(l\) in parallel using attention mask do
        \(\log p_{j: j+n_{\text {true }}}^{\text {init }} \leftarrow \operatorname{lm}\) head \(_{\theta_{i}}\left(h_{j: j+n_{\text {true }}}^{\text {init }}\right) \quad / /\) Predict next tokens
        \(T_{j} \leftarrow\) generate_tokens \(\theta_{i}\left(\left[X_{: j} ;<\right.\right.\) start_thought \(\left.\left.>\right], t, n_{\text {thoughts }}\right) / /\) Generate thought
        \(T_{j} \leftarrow\left[T_{j} ;<\right.\) end_thought \(\left.>\right]\)
        \(h_{j: j+n_{\text {true }}}^{\text {thought }} \leftarrow\) hidden_states \({ }_{\theta_{i}}\left(\left[X_{: j} ; T_{j} ; X_{j: j+n_{\text {true }}-1}\right]\right)\)
        \(\log p_{j: j+n_{\text {true }}}^{\text {thought }} \leftarrow \operatorname{lm}\) head \(_{\theta_{i}}\left(h_{j: j+n_{\text {true }}}^{\text {thought }}\right) \quad / /\) Predict next tokens \(\mathrm{w} /\) thought
        \(w_{j: j+n_{\text {true }}} \leftarrow\) mixing_head \({ }_{\theta_{i}}\left(h_{j: j+n_{\text {true }}}^{\text {thought }}, h_{j: j+n_{\text {true }}}^{\text {init }}\right)\)
        \(\log p_{j}^{\text {talk }} \leftarrow w_{j: j+n_{\text {true }}} \cdot \log p_{j: j+n_{\text {true }}}^{\text {init }}+\left(1-w_{j: j+n_{\text {true }}}\right) \cdot \log p_{j: j+n_{\text {true }}}^{\text {thought }} \quad / /\) Mix logits
        \(\mathcal{L}_{j}^{\mathrm{NLL}} \leftarrow-\log p_{j: j+n_{\text {true }}}^{\text {talk }}\left(X_{j+1: j+n_{\text {true }}+1}\right)\)
        \(r_{j}=\log p_{j: j+n_{\text {true }}}^{\text {talk }}\left(X_{j+1: j+n_{\text {true }}+1}\right)-\log \bar{p}_{j: j+n_{\text {true }}}^{\text {talk }}\left(X_{j+1: j+n_{\text {true }}+1}\right)\)
        \(\nabla_{\theta} \mathcal{L}_{j}^{\text {REINFORCE }} \leftarrow-r_{j} \mathbb{1}\left[r_{j}>0\right] \cdot \nabla_{\theta} \log p_{\theta_{i}}\left(T_{j} \mid\left[X_{: j} ;<\right.\right.\) start_thought> \(\left.]\right)\)
        \(\nabla_{\theta} \mathcal{L}_{j} \leftarrow \nabla_{\theta} \mathcal{L}_{j}^{\mathrm{NLL}}+\nabla_{\theta} \mathcal{L}_{j}^{\text {REINFORCE }}\)
    \(\theta_{i+1} \leftarrow \theta_{i}-\alpha \sum_{j=1}^{l} \nabla_{\theta} \mathcal{L}_{j} \quad\) // Update model parameters
return \(\theta_{\text {num_steps }}\)

2.3 Meta-tokens 2.3 元标记

3 Problem Statement 3 问题陈述

4 Quiet-STaR 4 安静的 STaR

4.1 Overview 4.1 概述

1. Parallel rationale generation (think, Subsection 4.2): In parallel across $n$$n$nn tokens $x_i$$x_i$x_(i)x_{i} in an input sequence $x_{0:n}$$x_{0:n}$x_(0:n)x_{0: n}, we generate $r$$r$rr rationales of length $t:c_i=(c_{i1},\dots ,c_{it})$$t:c_i=\left(c_{i1},\dots ,c_{it}\right)$t:c_(i)=(c_(i1),dots,c_(it))t: c_{i}=\left(c_{i 1}, \ldots, c_{i t}\right), resulting in $n\times r$$n\times r$n xx rn \times r rationale candidates. We insert learned $<$$<$<< startofthought $\mid >$$\mid >$∣>\mid> and $<\mid$$<\mid$<∣<\mid endofthought $\mid >$$\mid >$∣>\mid> tokens to mark each rationale’s start and end.
   并行推理生成（见第 4.2 小节）：在输入序列 $x_{0:n}$$x_{0:n}$x_(0:n)x_{0: n} 中的 $n$$n$nn 个令牌 $x_i$$x_i$x_(i)x_{i} 上并行生成长度为 $t:c_i=(c_{i1},\dots ,c_{it})$$t:c_i=\left(c_{i1},\dots ,c_{it}\right)$t:c_(i)=(c_(i1),dots,c_(it))t: c_{i}=\left(c_{i 1}, \ldots, c_{i t}\right) 的 $r$$r$rr 个推理，从而产生 $n\times r$$n\times r$n xx rn \times r 个推理候选。我们插入学习到的 $<$$<$<< 思考开始 $\mid >$$\mid >$∣>\mid> 和 $<\mid$$<\mid$<∣<\mid 思考结束 $\mid >$$\mid >$∣>\mid> 令牌，以标记每个推理的开始和结束。
2. Mixing post-rationale and base predictions (talk, Subsection 4.3): From the hidden state output after each rationale, we train a “mixing head” - a shallow MLP producing 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.
   混合后理由和基础预测（谈话，子节 4.3）：从每个理由后的隐藏状态输出，我们训练一个“混合头”——一个浅层多层感知器，生成一个权重，决定后理由预测的下一个标记的对数值应与基础语言模型预测的对数值相比应纳入多少。这种方法在微调早期由于引入理由而减轻了分布转移。
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.4 节）：我们优化推理生成参数（起始/结束标记和语言模型权重），以增加使未来文本更可能的推理的概率。我们使用 REINFORCE 根据推理对未来标记预测的影响提供学习信号。为了减少方差，我们应用教师强制技巧，将预测不仅是思考后的标记，还包括后续标记的概率纳入损失中。

4.2 Parallel Generation 4.2 并行生成

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.
图 3：并行生成。通过构建一个注意力掩码，使所有思维标记能够关注自身、同一思维中的所有前置思维标记以及前面的文本，我们可以并行生成所有思维的延续。每次推理调用用于为所有文本标记生成一个额外的思维标记。
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.
为了高效地实现这一点，我们缓存每次前向传播，并将对角注意力掩码连接到之前的注意力掩码：每个生成的标记现在关注于生成它所使用的所有标记，以及它自身（但不关注其他“反事实”路径上的标记）。此外，这种并行的下一个采样标记过程可以任意多次重复（或者至少，直到内存耗尽）。我们在图 3 中可视化了这一过程，并在附录 B 中强调了加快该算法的其他方法。

4.3 "Mixing" (Residual) Heads
4.3 “混合”（残余）头

4.4 Optimizing Rationale Generation
4.4 优化推理生成

4.4.1 Optimizing Start-of-Thought and End-of-Thought Tokens
4.4.1 优化思维开始和结束标记

4.4.2 Non-myopic Scoring and Teacher-forcing
4.4.2 非近视评分与教师强制

4.4.3 Objective 4.4.3 目标

5 Experiments and Results
5 实验与结果

5.1 Downstream Performance
5.1 下游性能

5.2 Improvement Distribution
5.2 改进分配

5.3 Quiet-STaR and Chain-of-Thought
5.3 Quiet-STaR 和思维链

5.4 Examples 5.4 示例

'<s> # Magnesium reacts with nitrogen to form magnesium nitride. The chemical
    formula for this reaction is Mg+N_2-> MgN_2. What is the product, or what
    are the products, of this reaction?\n\nJan 12, 2016\n\nThe formula for
    magnesium nitride is $M {g}_{3} {N}_{2}$.\n\n#### Explanation:\n\nAs do
    many active metals, magnesium nitride can be<|startofthought|> 1 --, so the
    equation of the reaction that forms magnesium nitride is\n\n$Mg + N_2 \\to
    <|endofthought|> formed by heating the metal (fier'

An integer $n$ is odd if $n = 2k+1$ for some integer $k$.\n\nTo prove that $A =
    B$, we must show that $A \\subseteq B$ and $B \\subseteq A$. The first of
    these tends to<|startthought|> in some sense - to be the more difficult<|
    endthought|> trickiest for students

'<s> Q: Talking to the same person about the same thing over and over again is
    \(<\mid s t a r t o f\) thought \(\mid>\backslash n \backslash n(a)\) a one-to-one correlation \(\backslash n \backslash n(b)\) a one-to<|
    endofthought|> something someone can what?'

6 Limitations 6 限制

7 Conclusion 7 结论

Ethics Statement 伦理声明

Acknowledgements 致谢

References 参考文献

Appendix 附录

A Hyperparameter Choices
超参数选择

B Faster Parallel Sampling
B 更快的并行采样

C Compute-Adjusted Plots
C 计算调整图表

D Measuring the Impact of Multiple Thoughts Per Sequence and Multiple Tokens Ahead
D 测量每个序列中的多个思维和多个前置标记的影响

E Rationale-Augmented Rationales
E 理性增强的理由

E. 1 Base Mistral 7B
E. 1 基础 Mistral 7B

First Response: 首次回应：

A: Let's think step by step. How many fresh eggs remains after the three her eat for breakfast and four goes into the muffin mix for her friends. We can calculate it by 16-3-4. The amount equals 9 fresh eggs. So, the dollars it makes daily equals what
A: 让我们一步一步来思考。三个人早餐吃了多少新鲜鸡蛋，四个鸡蛋放进了给她朋友的松饼混合料中。我们可以通过 16-3-4 来计算。结果是 9 个新鲜鸡蛋。那么，它每天赚取的金额等于多少。

Second Response: 第二次回应：

Third Response: 第三个回应：

Fourth Response: 第四个回应：

Fifth Response 第五回应

E. 2 OpenWebMath-Trained Mistral 7B
E. 2 OpenWebMath 训练的 Mistral 7B

First Response: 首次回应：

Second Response: 第二次回应：