Paper: Think before you speak: Training Language Models With Pause Tokens (ICLR-2024)

Reviewer comment (Evaluation):
It seems mysterious and problematic that sometimes more pause tokens can lead to worse performance.

Author response:
This is an interesting question, **but we believe that this is not mysterious if we view it from the fundamental lens of _overfitting_.** Recall that if a model is inundated with a lot of features, it can fail to generalize well as it is unable to pick up on the general patterns in those features. In our case, if we look at each intermediate vector input to some hidden layer as a feature, when we add $p$ pause tokens, we add $p$ new features. When we add an excessive number of pause tokens, the subsequent layer may suffer from overfitting. Of course, the value of $p$ for which overfitting starts to happen, must vary from task to task, depending on the task’s complexity.

Reviewer comment (Reproducibility):
How does one randomly inject pause token during pretraining?

Author response:
We add pause tokens at uniformly randomly selected positions in input sequences. We then  trim the appended sequence back to the original length. We do this so as to have a fair comparison with standard pretraining without pause tokens: this way, we preserve the same computational cost during pretraining.

Reviewer comment (Experiments):
Could there be more than one pause token in a row during pretraining?

Author response:
Yes, there can be some cases  where many pause tokens come in a row. In fact, it is an interesting direction for future work to come up with better ways of adding pause tokens during pretraining, as the consecutive pause tokens can be more helpful. This is because at inference, we do add a consecutive chunk of pause tokens at the end of prompt.

Reviewer comment (Reproducibility):
Is there an upperbound on how many pause tokens can be in a row?

Author response:
Yes, there can be some cases  where many pause tokens come in a row. In fact, it is an interesting direction for future work to come up with better ways of adding pause tokens during pretraining, as the consecutive pause tokens can be more helpful. This is because at inference, we do add a consecutive chunk of pause tokens at the end of prompt.

Reviewer comment (Experiments):
Have you tried to prompt LLM with filler words that typically represent pause or thinking? Like "let me think a bit ...." etc?

Author response:
What we tried was appending periods (“.”) as the pause tokens, as discussed in Section 4.3. We doubt that even prompting with words like “let me think a bit” will allow the model to leverage additional computational pathways within the transformer, as our work clearly shows that one needs to pretrain the model accordingly for the same.

Reviewer comment (Novelty):
Probably the closest idea to pause tokens is adaptive computation time (ACT) [1], originally proposed for recurrent neural networks with the exact same motivation: to let the model "think" before generating a difficult token. The idea of ACT translates easily to Transformers, and many subsequent works in the last 7 years (e.g. [2,3]) use ACT for various transformer types, in both sequence and depth dimensions. If authors choose to also overview ACT in related work, there are many more studies on adaptive computation time beyond the few referenced. Furthermore, unlike Pause Tokens, ACT offers a way of teaching the model to automatically decide how many extra turns it needs to solve a problem.

Author response:
We thank the reviewer for their time and providing detailed feedback. The reviewer appreciated the “simplicity of our algorithm” and “comprehensive results” and raised some valid concerns about missing related work, which we address below. 

> My main problem with the paper is that it ignores several very similar works from the early years of Transformer.

Thank you for pointing us to these lines of works, which we have now cited (Appendix K), along with a few other papers as well, which we had indeed missed in our original version. **However, we would like to strongly emphasize why our work is orthogonal to (and complements) adaptive compute approaches, and why it is unreasonable to compare against them as baselines:** 

## Comparison with Adaptive Computation Time (ACT): 

1.  **Ours is an extremely simple change that appends the same number of dummy tokens to all inputs, without adapting to the input or cleverly selecting any positions.** (during pre-training we randomly insert these tokens). In contrast, ACT [3][4][5] and related methods enjoy heavily-perfected adaptive machinery such as trained gating, halting, probing and routing mechanisms that allow them to adaptively allocate compute based on the input (and the model’s confidence etc.,). **Our current exploration is not concerned with such complex input-adaptiveness in any way.** The core question we ask is a much simpler one as you note — does appending dummy tokens even do anything? — for which we surprisingly find positive evidence. This, we believe is a significant novel finding in its own right. While we analyzed some basic varying of inference-time compute (Section 5.2) for the sake of ablations, we note that these again are done independent of the input i.e., we report accuracy, appending the same $p$ tokens to all inputs, and vary $p$ universally). 

**If** one compelled us to view this as comparable to adaptive compute techniques, we still argue that this is orthogonal and if at all, complementary: 

2. A hypothetical pause-token-based adaptive compute technique would fall under a completely different paradigm of flexible compute (by varying input-length), rather than varying depth or model-size (which is used in cascades [6] or early-exit models, including DACT-BERT[4] and the Universal Transformer [3] you pointed us to). Importantly, increasing the input-length increases the computational width of the Transformer (i.e., the number of parallel operations) which is orthogonal to increasing the computational depth via parameter count or recurrent layers. **We formalize the benefits of expanding the computational width in the now-added Appendix J, where we show that there are tasks that benefit from this, but cannot benefit from depth-expansion in existing ACT.** Consider for example, computing $f(v_1, v_2, ... v_{50} ) = (v_1+v_{32}) \cdot (v_3 + v_{44}) ... (v_{27}+v_{2})$ where $f$ involves say  $1000$ many parallel addition operations — here, under some natural conditions, repeating sequential operation like in Universal Transformers would be of no use, while adding pause tokens would help implement many of these parallel operations! 

3. On that note, adding recurrent layers like in Universal Transformers, would add much more to the wall-clock time during inference (owing to the extra sequential computations) compared to pause tokens (which only add pause-token-specific computations in parallel to existing computations).  

4. To further appreciate why adaptive compute techniques are orthogonal to ours, observe that our technique can be directly combined with existing techniques like cascades, early-exits etc., i.e., one may increase/decrease the number of pause tokens, and decide to exit from a lower/higher layer. 

5. Besides, it is worth noting that even within the ACT literature, cascade approaches and early-exit approaches treat each other as different paradigms, and do not typically compare against each other as baselines. 

To reiterate, regardless of our technical points 2, 3, 4 and 5 above, we believe that solely for point 1, our contributions should not be treated as comparable to adaptive compute time (ACT), nor do we claim to provide a competitive ACT-esque approach. **What our work explores is a fundamental question about a simple change to the Transformer.**  We agree that ACT is an important line of work that needs to be cited and discussed (Appendix K). **Unfortunately, we believe it is too reductive to reject our paper on the basis that we do not compare numbers with ACT algorithms that enjoy specialized machinery and are tangentially related.**

Reviewer comment (Experiments):
Another related (but less similar) study is [4], which introduces additional learnable tokens for attention layers. The paper itself uses an attention-only architecture, while subsequent works use the same idea for regular transformers, including a notable implementation [5]. To the best of the reviewer's understanding, [4] proposes a mathematically similar idea that is likely to have similar effects to pause training, therefore there is no reason it should not be a baseline for comparison.

Author response:
## Comparison with Memory tokens

Again, we are grateful for these excellent pointers. We had unfortunately missed this, but we have fixed this now (see Appendix K, where we have also added more works outside of your review). 

**We explain why the related Transformer works, which may appear superficially similar, are orthogonal to us.** Conceptually, ideas like Persistent Memory [7], RNN-cache [8] and kNN-LM [9] rely on introducing “global information” into the model. This:

 - introduces significant number of new token parameters (e.g., the Persistent Memory tokens require 1024 global vectors in each layer, so roughly O(10M) new parameters).

 - These vectors act only as keys and values, but not as queries. 

 - These vectors do not depend on the inputs.


In contrast, pause tokens can be thought of as providing “local” information to the model by:

- Introducing virtually no parameters besides a handful of parameters for the *single* pause token, (precisely 1 new vector, so O(1k) parameters)

- Which subsequently generate intermediate vectors that act as keys, values **and queries** 

- And due to the above important difference, the new intermediate vectors depend on the input.

Reviewer comment (Novelty):
Additionally, [6-7] (and others) propose augmenting language models with attention to additional virtual tokens to improve language model quality. These other works pursue different goals from this paper, but to the reviewer's knowledge, they are no less related than PEFT methods discussed on page 9.

Author response:
## Comparison with Memory tokens

Again, we are grateful for these excellent pointers. We had unfortunately missed this, but we have fixed this now (see Appendix K, where we have also added more works outside of your review). 

**We explain why the related Transformer works, which may appear superficially similar, are orthogonal to us.** Conceptually, ideas like Persistent Memory [7], RNN-cache [8] and kNN-LM [9] rely on introducing “global information” into the model. This:

 - introduces significant number of new token parameters (e.g., the Persistent Memory tokens require 1024 global vectors in each layer, so roughly O(10M) new parameters).

 - These vectors act only as keys and values, but not as queries. 

 - These vectors do not depend on the inputs.


In contrast, pause tokens can be thought of as providing “local” information to the model by:

- Introducing virtually no parameters besides a handful of parameters for the *single* pause token, (precisely 1 new vector, so O(1k) parameters)

- Which subsequently generate intermediate vectors that act as keys, values **and queries** 

- And due to the above important difference, the new intermediate vectors depend on the input.

Reviewer comment (Evaluation):
Answering the above questions in the current form would require the authors to reformulate too much of the paper content and message, necessitating a full additional review, leading the reviewer to currently recommend rejection. However, the reviewer encourages resubmission once the authors can properly analyze how PT compares to earlier methods.

Author response:
**We believe our work can complement and inspire new ACT approaches in the future. However, we hope you agree that reducing this paper to a new ACT-like approach [that should compete with other ACT approaches], and rejecting it on those grounds, is not well-warranted. We are afraid this characterization is unfortunately not accurate (as our approach does not claim to be, nor is as sophisticated as, an ACT approach) and unfortunately too reductive (as it ignores the value in our fundamental finding that adding a simple change can be helpful). We sincerely hope you will reconsider your score for our work in light of this discussion, and we once again thank you for pointing these important lines of work to us which we've now discussed in detail in the paper.**

Reviewer comment (Novelty):
The reviewer could raise the score within this review cycle if either:  
1. the authors prove that the reviewer misunderstood a significant chunk of the paper and it is not in fact related to ACT family and others, or  
2. the authors prove that all of those questions can be answered trivially, without running many additional experiments or changing much of the paper.

Author response:
**We believe our work can complement and inspire new ACT approaches in the future. However, we hope you agree that reducing this paper to a new ACT-like approach [that should compete with other ACT approaches], and rejecting it on those grounds, is not well-warranted. We are afraid this characterization is unfortunately not accurate (as our approach does not claim to be, nor is as sophisticated as, an ACT approach) and unfortunately too reductive (as it ignores the value in our fundamental finding that adding a simple change can be helpful). We sincerely hope you will reconsider your score for our work in light of this discussion, and we once again thank you for pointing these important lines of work to us which we've now discussed in detail in the paper.**

Reviewer comment (Experiments):
On page 7, the authors found that the optimal number of pause tokens depends on the specific task, e.g., GSM8k or SQuAD. The reviewer asks whether the authors tried to train the model to dynamically decide the number of pause tokens, essentially allowing the model to predict pause tokens at a certain penalty (like ACT) or spread a certain fixed budget of pause tokens between different positions (like CTC training). Since the baselines are able to do the same thing, it would be interesting to check which approach (PT or ACT or CTC) is better at deciding where to put the pause tokens.

Author response:
**Determining the number of pause tokens**: We determined the number of pause tokens by simply choosing between 10 or 50. We believe there is value in making the first step in the “pause tokens” direction simple so we can be confident that this core idea itself is promising. 
Complementing this with more sophisticated ACT-type techniques is an important direction for future work!

Reviewer comment (Theory):
Regardless of the empirical gains, we need more theoretical insights into why and how "pause tokens" work during pre-training and fine-tuning. There is not enough motivation behind this trick. We need to understand why we need to "delay a model's answer generation." There are a few intuitions, but are not well-articulated and convincing enough. The reason to answer this question is necessary because the community can benefit if the pause tokens are so important to replace normal autoregressive LLMs.

Author response:
## Theoretical intuition behind Pause tokens (added in Section J)
>  There is not enough motivation behind this trick. We need to understand why we need to "delay a model's answer generation." There are a few intuitions but are not well-articulated and convincing enough.

We completely agree with the concern that the intuition was originally fuzzy. To address this, **we have now added a rigorous and formal analysis of when and why pause tokens can help – please see Appendix J, leading to Theorem J.1.** To summarize for your convenience, we make two formal insights:

1. We formally show that pause tokens are helpful in problems that require a number of parallel operations that are much larger than the number of input tokens. For example, assume the model is given as input a sequence of $L$ numbers $v_1, v_2, ... v_L$, and it must compute a polynomial $(v_1 + v_2) \cdot (v_3 + v_1) ... (v_4 + v_8)$, that involves $N$ parallel additions, where $N \gg L$. A standard transformer can only implement $L$ operations in parallel, while appending pause tokens can allow implementing all $N$ operations.
We argue that a similar case can arise in natural language settings e.g., consider an evidence-based multiple-choice question task with $C$ choices and $E$ pieces of evidence. Solving this would require $C \cdot E$ operations that individually corroborate each choice against each piece of evidence. But $C \cdot E$ scales much higher than the input sequence length which would scale as $O(C + E)$. 

2. As another key insight, we argue that the self-attention/feedforward module has an “untapped” representational capacity that grows with its parameter count, say K. Importantly, this quantity is independent of the input length L. While standard inference only allows this module to realize L different operations, appending O(K) pause tokens can help realize as many K different operations.

Reviewer comment (Experiments):
Adding the "pause token" brings new challenges and headaches. A limitation of the "pause token" approach lies in its necessity to be integrated into both the pre-training and fine-tuning phases, potentially hindering its broader applicability. The introduction of the "pause" token consumes a significant portion of the model’s input window space—approximately 10% of the sequence length—resulting in a loss of some training data information. Additionally, the number of "pause" tokens acts as a hyperparameter, exerting a considerable influence on task performance. This number is not consistent but varies across different tasks, adding a layer of complexity and variability to the model's application and performance.

Author response:
## Challenges with Pause Tokens 

> Need to incorporate pause tokens during pretraining: 

While we agree that this is a hindrance, **this is in fact the key discovery of our exploratory paper which was hitherto unknown**. Further, some recent papers like [1] also argue that post hoc learning of chain-of-thought style reasoning may not work very well if only introduced in a post-hoc manner, as the default next token prediction pretraining objective (without learning to rely on the reasoning during training) does not model the same. 

> Pause tokens consume model’s input window space: 

There appears to be a misunderstanding here: we deliberately fixed the window length same as standard training for the sake of fairness of computational costs for both models. In fact, this indeed further highlights another interesting aspect about pause training – it is able to outperform standard training even while being at a disadvantage due to less number of “real” tokens seen!


> Number of pause tokens as the hyper-parameter: 

We agree that this hyperparameter can be a nuisance.  However, we note that this was varied (and that too, over a grid of two values!) only in the downstream finetuning, and not during pre-training.

[1] Merrill, W., & Sabharwal, A. (2023). The Expressive Power of Transformers with Chain of Thought. ArXiv, abs/2310.07923.

Reviewer comment (Experiments):
The introduction of pause tokens undoubtedly adds a computational overhead to the training and inference processes. The paper could benefit from a more detailed analysis of this trade-off, especially in terms of computational resources and time. Assessing the scalability of this approach in larger models or in resource-constrained environments would be beneficial. Providing benchmarks on computational costs and suggesting ways to optimize this aspect could enhance the practicality of the approach.

Author response:
## Assessing computational costs  (added in Section I)
> adds a computational overhead to the training and inference processes… The paper could benefit from a more detailed analysis of this trade-off,

You have raised some important practically-relevant questions regarding the compute-efficiency of pause-inference. **We have conducted an extensive analytical study of this in Appendix I (and Theorem I.3 and Theorem 1.4) in terms of FLOPS-efficiency and wall-clock-compute-efficiency.**  For your convenience, here are our main takeaways: 

- **Pause tokens are FLOPS-efficient compared to increasing parameter count:** If we were to add 10 pause tokens to our 1B model, how many extra FLOPS does it consume? Subsequently, if we allocate the same extra FLOPS budget to other ways of increasing the number of attention operations, how much accuracy gains would we see? Specifically, we consider (a) adding layers and (b) adding attention heads as ways to increase the number of attention operations. We find that for our architecture, adding 10 pause tokens at inference, over an initial prompt of length 100 tokens, becomes equivalent to (a) adding 2 layers (on top of 24 pre-existing layers) or (b) adding 3 attention heads per layer (in addition to 32 pre-existing heads). Both these approaches increase the parameter counts from 1B to (a) 1.1B and (b) 1.048B respectively, which is a modest increase in parameter count that cannot promise the gains we see in our 1B model with pause tokens. 

- **Pause tokens do not add extra serial compute:** Pause tokens are consumed in parallel with the prefix, with minimal wall-clock overhead. However, adding layers increases wall-clock time linearly with the number of layers. Similarly, other techniques like CoT introduce significant wall-clock overhead since they involve auto-regressively decoding each intermediate token (a wall-clock time of $pl$, where $p$ is the number of tokens and $l$ is the number of layers).

Reviewer comment (Experiments):
How do you determine the number of “pause” tokens?

Author response:
> How do you determine the number of “pause” tokens?

**This is a practically important question! In our current** exploratory work, we wanted to analyze the benefits of a simple no-frills approach of appending dummy tokens. Here, we simply choose between appending 10 or 50 tokens in the downstream task.   However, in the future, we could leverage the progress in the fields of adaptive computation and utilize some kind of confidence measure or gating units to introduce a variable number of pause tokens per input. Please note that we do not claim to propose any adaptive compute methods in the current submission.

Reviewer comment (Experiments):
How is the generalization ability of “pause” tokens, for example, pretraining with “pause” tokens and inference on tasks without fine-tuning? Are these tokens task specific?

Author response:
> How is the generalization ability of “pause” tokens, for example, pretraining with “pause” tokens and inference on tasks without fine-tuning? Are these tokens task specific?

**To answer your question, we have added “zero shot evaluation” results in Appendix D**, where we do not finetune the model on the downstream task, and thus the pause tokens are untouched pretraining. Here, we compare the standard and the pause pretrained model (with 0, 10 and 50 pauses at inference). While on most of our tasks, our zero-shot values of both models are too low to be meaningful, we do observe that on datasets like GSM8k (mathematical reasoning) and PhysicalIQA (general understanding), pause tokens give some gains. We believe this provides preliminary evidence that these tokens are generalizable. Future work can help in improving this generalizability.

Reviewer comment (Writing):
The title doesn't convey any meaningful information, "Think before you speak" is too casual and just for eye-catching, and "Pause tokens" are not obviously justified in the title.

Author response:
> The title doesn't convey any meaningful information, "Think before you speak" is too casual and just for eye-catching, and "Pause tokens" are not obviously justified in the title.

We will certainly rework the title in future versions.

Reviewer comment (Experiments):
One weakness is that the method does not compare directly to Chain of Thought (CoT).

Author response:
The connections to CoT are indeed conceptually interesting, which we touched upon in the Related Work and Discussion sections. We completely agree that pause-training requires pre-training while CoT doesn’t — this is a great point of comparison. We have now added this to our discussion section. However, we wish to clarify that **a direct numerical comparison with CoT would be apples-to-oranges**, and the results are predictable a-priori. We would like to explain why below.

In short, **CoT is a compute-hungry approach, while pause-training does not add (substantially) to wall-clock time:** each intermediate reasoning token in CoT is sequentially and auto-regressively decoded from the model. This naturally would give us high performance gains, but at a high compute cost. Formally, if there are $P$ intermediate CoT tokens and $L$ layers, CoT performs $P \cdot L$  **sequential** operations. Pause tokens on the other hand processes the P dummy identical tokens in **parallel**, with little excess in wall-clock time (depending on the number of parallel computational threads). Thus, by construction, pause-inference is computationally- and representationally-disadvantaged compared to CoT, rendering any numerical comparisons meaningless. 

Thus, as mentioned in the paper, we believe that in terms of accuracy, CoT is without question, an upper bound on pause-inference.  To address this discussion further, we have added relevant discussion to this in Appendix I.

Reviewer comment (Experiments):
In contrast to pause tokens, which require both pretraining and fine-tuning, CoT is an inference-time only method that potentially requires extra human annotations.

Author response:
The connections to CoT are indeed conceptually interesting, which we touched upon in the Related Work and Discussion sections. We completely agree that pause-training requires pre-training while CoT doesn’t — this is a great point of comparison. We have now added this to our discussion section. However, we wish to clarify that **a direct numerical comparison with CoT would be apples-to-oranges**, and the results are predictable a-priori. We would like to explain why below.

In short, **CoT is a compute-hungry approach, while pause-training does not add (substantially) to wall-clock time:** each intermediate reasoning token in CoT is sequentially and auto-regressively decoded from the model. This naturally would give us high performance gains, but at a high compute cost. Formally, if there are $P$ intermediate CoT tokens and $L$ layers, CoT performs $P \cdot L$  **sequential** operations. Pause tokens on the other hand processes the P dummy identical tokens in **parallel**, with little excess in wall-clock time (depending on the number of parallel computational threads). Thus, by construction, pause-inference is computationally- and representationally-disadvantaged compared to CoT, rendering any numerical comparisons meaningless. 

Thus, as mentioned in the paper, we believe that in terms of accuracy, CoT is without question, an upper bound on pause-inference.  To address this discussion further, we have added relevant discussion to this in Appendix I.

Reviewer comment (Novelty):
A recent paper [1], that shows that extra scratch space improves the computational power of transformers, could help motivate this paper.

Author response:
We thank the reviewer for pointing us to [1], which we have now cited in our work. The insights from this work are definitely relevant to us. It is worth noting that their analysis considers the CoT setting where the intermediate tokens are generated by the model itself, and fed back into the model. This enhances the “computational depth” of the model, which may be crucial to their theoretical claim. On the other hand, pause tokens enhance only the computational width of the model. In the now-added Appendix J, we formally argue why even this alone can be helpful in some tasks.