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

Thank you for the supportive and thought-provoking comments and questions!

As Quiet-STaR incurs significant extra computation overhead, would it be possible to achieve similar performance by spending the extra computation on scaling up the model size & training data (or even just use a larger model during the inference, noting that the extra overhead during inference is also significant)?

At least in the current framing, Quiet-STaR is presented largely as a continual training / post-training technique, so most likely the best comparison is compute-equivalent training without new parameters. The most expensive runs in our work have compute comparable to training on tens of millions of tokens, roughly a millionth of the pretraining compute (supposing Mistral 7b was trained on trillions of tokens). In contrast, continued training on OWM / C4 results in little improvement (or in some of our baselines, actually results in deteriorations – we’ll also update the figure to make it clearer that our baseline corresponds to naive continual pretraining).

As for inference cost, the question is definitely trickier. LM scaling laws suggest that a sufficiently large model trained on sufficiently much data should likely perform as well as Quiet-STaR, even on difficult tokens. However, it’s not clear how we might estimate the potential size of such a model or its training corpus. Perhaps the simplest answer is that, given a (sufficiently large) pretrained model, Quiet-STaR allows us to trade inference compute for additional performance in a general way (with other recent works like Rapaso 2024 also being exciting steps in this direction). Indeed, we hope this lays the groundwork for more adaptive computation strategies.

I doubt the argument that "We expect that this is a more difficult task, as predicting the usefulness of a thought is simpler when one has already generated the thought," which was used to justify Quiet-STaR generating many tokens before generating every additional token.

We absolutely agree! We mention this in the conclusion and limitations but dynamically predicting the amount of thinking tokens to generate is an essential future direction. This line is mostly saying that predicting the utility of a generated thought is easier when one has the generated thought (indeed, this must be true if there is variance in the utility of thoughts and a thought provides some information about its usefulness). We’ll clarify this in the revision.

Thank you for your encouraging, thorough, and positive feedback!

Thanks for the great questions and many encouraging words!

First, to the question about baselines, specifically Self-improvement (Huang et al., 2022) and V-STaR (Hosseini et al., 2024): while these are both excellent papers, they aim to solve a very different problem. Fundamentally, these techniques aim to train a language model to perform well on a given question-answer dataset, specifically by optimizing on questions from that dataset. Instead, our goal is more general: can a language model learn to generate rationales that help predict arbitrary future text?

Both are excellent and inspiring papers. Notably, Huang et al 2022 differentiates itself from STaR primarily by using majority vote instead of ground truth answers, but does not claim to outperform STaR given ground truth – in Quiet-STaR, the continuation of the text is our ground truth. As the length of predicted future text increases, the chance of multiple predictions matching also decreases, so using majority vote becomes impractical. V-STaR is also excellent (note, preprinted the month before the COLM deadline) but for us the future text is the ground truth so the verifier would just be the LM.

The idea of generating a rationale for each token may not be reasonable, as the rationale for a single token may not be contextually meaningful… Instead of generating a rationale for each single token, how about sampling thought based on a chunk of tokens?

Absolutely, using only a single token ahead would result in a much noisier reward signal – this is why we use the multi-token ahead teacher-forcing trick (see Section 4.4.2).

Figures 1, 3, and 4 are not well explained… For example, it is hard to understand the parallel generation process based on Figure 3 and the sketchy description in Section 4.2.

Thank you for this point! We would be happy to add additional explanations for these figures using the additional page for the final version. Do you have a particular change in mind for Figure 3? We tried to directly visualize the attentions of the thought tokens in both parts of the figure.

In Figure 1, the second sampled thought seems like Thought_5. Why does it appear as Thought_6?

The second sampled thought shows an example of a thought that would not help with predicting “4”, so the model should not generate thoughts like that.

Thank you for your excellent questions!

First, to the question about baselines, we need to emphasize that the key goal of this paper is to show that language models can improve their general inter-token reasoning ability. We attempted to analyze / control for this in the experiments visualized in Figure 2. In general, this method is complementary to CoT, since the model can generate an inner thought before each token of a CoT (as shown in Figure 5 and Appendix E). We will update our paper to note that we chose a base language model because it allows us to study this problem directly, without introducing side-effects from (typically unreleased) RLHF and chat-style fine-tuning strategies used to fine-tune these models.

As for why the start and end thought initializations were selected, it was indeed somewhat arbitrary. We wanted a single token to indicate the start and end of an intermediate thought, similar to what might be used in an actual text, ideally without imposing too many additional biases. We’ll add a note about this.

One simple explanation could be that the model now has more tokens to generate

Seeing as this is the setup of the pause token paper (Goyal et al 2023), which finds no improvements in a stronger “pause-finetuning” setup (they finetune directly on the downstream task, while we don’t), this explanation is less likely -- though, of course, it is an important one to consider. Notably, this concern is a key challenge for reasoning-related work, and is related to the problem of faithfulness, i.e. how do we know that the internal processing meaningfully corresponds to the generated thought tokens? We'd be happy to elaborate on this more in our limitations section in a revised final version.

How did the authors determine which tokens are difficult?

In the context of discussing Figure 7, “difficult” was used to refer to high-prediction-loss tokens, but this could be more explicit. We’ll add “(high loss)” to make this clearer.