Think before you speak: Training Language Models With Pause Tokens
Authors: Sachin Goyal
Ziwei Ji
Ankit Singh Rawat
Aditya Krishna Menon
Sanjiv Kumar
Vaishnavh Nagarajan
Venue: ICLR 2024 poster (2024)
Keywords: LLM training and inference
Downstream finetuning
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We explore delaying model's prediction of next token, by appending (learnable) pause tokens to allow increased inference-time computations .

Abstract: Language models generate responses by producing a series of tokens in immediate succession: the $(K+1)^{\rm th}$ token is an outcome of manipulating $K$ hidden vectors per layer, one vector per preceding token. What if instead we were to let the model manipulate say, $K+10$ hidden vectors, before it outputs the $(K+1)^{\rm th}$ token? We operationalize this idea by performing
 training and inference on language models with a (learnable) $\textit{pause}$ token, a sequence of which is appended to the input prefix. We then delay extracting the model's outputs until the last pause token is seen, thereby allowing the model to process extra computation before committing to an answer. We empirically evaluate $\textit{pause-training}$ on decoder-only models of 1B and 130M parameters with causal pretraining on C4, and on downstream tasks covering reasoning, question-answering, general understanding and fact recall. Our main finding is that inference-time delays show gains when the model is both pre-trained and finetuned with delays. For the 1B model, we witness gains on 8 of 9 tasks, most prominently, a gain of $18\\%$ EM score on the QA task of SQuAD, $8\\%$ on CommonSenseQA and $1\\%$ accuracy on the reasoning task of GSM8k. Our work raises a range of conceptual and practical future research questions on making delayed next-token prediction a widely applicable new paradigm.

Published as a conference paper at ICLR 2024
THINK BEFORE YOU SPEAK: TRAINING LANGUAGE MODELS WITH PAUSE TOKENS
Sachin Goyal
Ziwei Ji
Ankit Singh Rawat Google Research, NY mikktsrawat@gpogle.com
Machine Learning Department Carnegie Mellon University sachingo@and rew. cmu.edu
Google Research, NY ziweij riweijigooggle.oo i@google. com
Aditya Krishna Menon Google Research, NY adi adiityaamenod@goolee.com com
Sanjiy Kumar Google Research, NY sanjivk@google.co com
Vaishnavh Nagarajan Google Research, NY vaishnavh@gpogle.com com
ABSTRACT
Transformer- language models generate responses by producing series of tokens in immediate succession: the (K + 1 token is an outcome of manipulat- ing K hidden vectors per layer, one vector per preceding token. What if instead we were to let the model manipulate say, K + 10 hidden vectors, before it outputs the (K + 1)th token? We operationalize this idea by performing training and in- ference on language models with (learnable) pause token, a sequence of which is appended the input prefix. We then delay extracting the model's outputs until the last pause token is seen, thereby allowing the model to process extra computa- tion before committing to an answer. We empirically evaluate pause-training on decoder-only models of 1B and 130M parameters with causal pretraining C4, general under-
and on downstream tasks covering reasoning, question-answeree standing and fact recall. Our main finding is that inference-time delays show gains on our tasks when the model is both pretrained and finetuned with delays. For the 1B model, witness gains eight tasks, most prominently, gain of 18% EM score on the QA task of SQuAD, 8% on CommonSenseQA and 1% accuracy on the reasoning task of GSM8k. Our work raises a range of conceptual and prac- tical future research questions on making delayed next-token prediction widely applicable new paradigm.
INTRODUCTION
Transformer-base causal language models generate tokens one after the other in immediate succes- sion. To generate the (K 1)sh token, the model consumes the K previous tokens, and proceeds layer by layer, computing K intermediate vectors in each hidden layer. Each vector in itself is the output of a module (consisting of self-attention and -er--llyee--preptoonssss operating on the pre- vious layer's output vectors. However sophisticated this end-to-end process may be, it abides by peculiar constraint: the number of operations determining the next token limited by the number of tokens seen so far. Arguably, this the most natural design choice when the Transformer was first conceived by Vaswani et al. (2017). But in hindsight, one may wonder whether for some inputs, the (K + 1)th token demands K + . Transformer operations in each layer (for M > 0), which cannot be met by the arbitrarily constrained K operations per layer. This paper explores one way to free the Transformer of this arbitrary per-layer computational constraint.
The approach we study is to append dummy tokens into a decoder-only model's input, thereby de- laying the model's output. Specifically, append one or more copies of <pause> as a sequence to the input. We simply ignore the model's cor- responding outputs until the last <pause> token seen, after which we begin extracting its response.
select a (learnable) pause token (denoted <pause>) and
Work done in part as a Student Researcher at Google.
1 "Corrsponning Authors

---

Published as a conference paper at ICLR 2024
Ignore Output
Ignore Output
Ignore Output
Targets
25+
25+
4 äs
29
Layer 2
Layer 1
Inputs
5
1
25
is
52+
4=
pausse
ppause>
spause
25+
is
(a) Standard inference and finetuning
(b) Pause-inference and finetuning
Standard vs. pause-inference (and finetuning). We consider a downstream task where, given prefix, the decoder-only model (bidirectionally) attends to all of the prefix to generate its target answer. The rounded squares denote one Transformer operation (a self-attention and MLP) 2-layer Transformer. Any Ignore Output denotes that during inference, the corresponding out- put token is not extracted and thus, not fed back autoregressively: during finetuning, this output not backpropagated through. The connecting lines denote some (not all) of the "computational pathways" within the model. Specifically, we visualize only those pathways that begin at a specific token the prefix (here arbitrarily chosen be "4 is") and end an output token (here arbitrarily chosen be "25+"), All differences between the two settings highlighted in color. (a) In stan- dard inference (finctuning,, the model's output is extracted immediately upon sceing the last prefix token. (b) In pause-inference (and mause-finetunigg ally specified number of <pause> tokens. This introduces new computational pathways (the colored lines) between the prefix token and the output token of interest.
Figure I:
this is initiated only after appending manu-
Crucially, we consider injecting such delays not just at inference, but also during downstream fine- tuning (see Fig 1) and pretraining (see Fig 2, which provides additional technical details).
A-priori, it is unclear what this simple change would bring about in practice. Optimistically, Transformer may take advantage of a "wider" computational pathway induced by the delay. A more mundane outcome though would be that the model simply skips any delays introduced by the <pause> tokens. After all, neither do the <pause> tokens provide any additional information during inference, nor are there sufficiently many new parameters (barring the few embedding parameters of the single <pause> token) that encode any additional information from training data. Worse still, these uninformative tokens may drown out informative signals, and hurt the model.
Partial answers this question be found in the literature, motivated somewhat differently. To un- derstand where the benefits of chain-of-thoughe (Wei et al., 2022) come from, Lanham et al. (2023) append dummy thoughts in the form of periods ('...), but only during inference. This, they report, does not help. Presumably, an off-the-shelf model may not have learned to utilize the new compu- tational pathways offered by the inference-time delay. Burtsev et (2020) learn with prepended dummy tokens, with the orthogonal motivation of adding memory (rather than extending computa- tion). They train with these tokens only on the target task, and observe minimal performance gains.
What then can we hope for when injecting (appended) delays on all stages of training and inference? Our work empirically evaluates this, and other key questions that come up when training the Trans- former with delays. For this, we study pause-training model, trained on C4 (Raffel et al., 2020) and finctuned on nine downstream tasks spanning extrac- tive question answering, reasoning, general understanding and fact recall. In summary, we make the following contributions:
1B and 130M parameter decoder-only
(1) We pose the question of what happens we delay a model's 's answer generation, and how can we execute these delays? We design way: training with dummy <pause> tokens. Accordingly, we design a pause-injected pretraining, downstream finetuning, and inference procedure.
(2) We find that variety of downstream tasks, training models with <pause> tokens during both pretraining and downstream finetuning, exhibits clear gains compared to standard end-to-em training and inference. Most notably, for the IB model, in the SQuAD extractive question-
2

---

Published as a conference paper at ICLR 2024
answering task, this approach improves the exact match score by 18%. Similarly observe 8% gains on the general understanding task of CommonSense QA and 1% accuracy gain on the reasoning task of GSM8K over the standard model'ss accuracy of 7.5%.
(3) On the flip side, when we introduce <pause> tokens only in the downstream finetuning stage (on the standard pretrained model), we find that the gains show up in far fewer instances, and are relatively mild. In some instances, we even find clear drop in performance.
(4) We also conduct series of key ablations: (a) We find that appending <pause> tokens largely better than prepending them, (b) perhaps unsurprisingly, for any downstream task, there is corresponding optimal number of <pause> tokens, and (c) when decreasing the number of inference-time <pause> tokens, we find graceful degradation of performance even though pause-training does not explicitly train for such robustness.
Overall, our work explores the new paradigm of delayed next-token generation in Transformer mod- els, and finds that there are benefits to this simple change, provided the change is implemented both during pretraining and finetuning. Our preliminary step here inspires variety of conceptual and practical future research questions, ranging from understanding how Transformer delays work mechanistically, making pause-training more generally applicable for practice.
PRELIMINARIES
We briefly outline the next-token prediction process in a standard causal decoder-only language model (details in $A). Consider vocabulary V and an input PI:K E y of K tokens. Let f denote a Transformer- baased language model, from which we sample the next token as PK+1 ~ S f(P1:)) To achieve this, internally, each layer / [1, L] of the Transformer produces intermediate vector v corresponding to each input token. The next token, i.e., PK+1 is then sampled from a distribution inferred from the last vector in the last layer, Vk"". (L)
On a high level, each layer in the above process can be represented as function T D D Its input a matrix K vectors, VI.K- = [Vi...., Vx], and likewise, the output, Vi,x- This mapping itself involves two key (parameterized) modules. The first is the attention module A which takes as inputs two matrices Vkey. Vvalue E RDx (for any N) and a "query" vector v E R to produce an output vector in DR, This is followed by a feedforward module fFF : RR - RD Then, given the inputs Vk, and given the layer-norm module PN R + RR,, the outputs v/, for k < K can be expressed as:
ag iIN o Attn (Vis,, Visk Va. + Vx) v iN (DrF (ak) + aa))
(2)
Observe here that the th output v1 is obtained by manipulating exactly the k previous hidden embeddings in the same layer, Vi
3 PAUSE-TRAINING
In the current paradigm of language models, we compute exactly К embeddings Vi.... Vx in each layer. before generating the (K + +1)th token, PK+1- Our premise that this limit of K operations an arbitrary one. Instead, we wish to expend more than K operations towards producing the next to- ken, PK+1- While something to this effect could achieved by increasing the number of attention heads in each layer, we are interested in an orthogonal approach that introduces hardly any parame- ters into the network. The idea to synthetically increase the input sequence length by appending M dummy tokens to the input, thus delaying the model's next response by M tokens of input. In effect, this M-token-delay lets the model manipulate an additional set of M intermediate vectors VK+1..... VK+M before committing to its next (output) token, PK-1- Intuitively, these vectors could provide a richer representation of the input (e.g., by attending differently). thus resulting in a better next token from the model. We visualize this in Figure 1.
3.1 LEARNING AND INFERENCE WITH THE <pause> TOKEN
A simple choice for the dummy tokens are special characters such as !.' or "#', as Lanham et al. (2023) chose for inference. But to prevent the model from confounding the role of delays with the
3

---

Published as a conference paper at ICLR 2024
Targets
kid
playzng
Soccer
kid
is
Ignore Output
playing
Ignore Output
Layer
Layer
Inputs
The
is
playing
The
kid
spause>
playind
spause>
(a) Standard pretraining
(b) Pause-pretraining
Figure 2: Standard vs. pause-pretraining We consider pretraining based on causal language mod- eling, where each token is predicted given all preceding tokens in the sequence, using unidirectional self-attention. Here, we visualize the computational pathways beginning from the token "is" on the input side of the decoder-only model, Figure for a guide on how to follow this visualization. (a) In standard pretraining, we compute the model's 's loss each output token, and backpropagate through it. (b) In cuuse-pretraining
subsequent token "soccer" on the output side. Please
insert multiple copies of <pause> tokens at uniformly random locations in the input. However, we do not apply a loss on the model predict these tokens, as indicated by each corresponding Ignore Out- put flags. This introduces new computational pathways connecting the input token and the output token of interest.
role the above characters play in natural language, we choose a single <pause> token residing outside of the standard vocabulary. To impose multi-token delays, simply repeat this token. Building this core idea, below we discuss our specific techniques for pause-pretraining and pause-finttuning.
Pretraining with the <pause> token The sequences in our pretraining data do not come with an annotation of which suffix constitutes the answer, since every input token also functions as target output. Thus, it is impossible to execute the simple delaying strategy of appending dummy tokens before extracting the answer. Therefore, for a given pretraining sequence P1.N, we insert multiple <pause> tokens (say Mpt many) uniformly random locations obtain a pause-injected sequence, PI:N+ Mge- We visualize this in Figure 2b. We then train the model with the standard next-token prediction loss this pause-injected sequence, while ignoring any loss term that corresponds to predicting the pause tokens themselves. Formally, let S Signore = {k: pb+1 <pause>) denote the positions where the next token is <pause> token. Then, for the decoder-only language model f. the pause-training loss is given by:
N+
C C (f.PLN4M,.)
Signnre
where LCE denotes the cross-entropy loss. Observe that the loss is skipped over indices in Signore- The rationale is that, we only want to use the <pause> tokens as a way of enforcing a delay in the model's computation; demanding that the model itself produce these tokens, would only be pointless distraction. Finally, as standard, we update the parameters both the model and of all the tokens, including those of the <pause> token. We term this mause-pretraiine (Algorithm l).
Finetuning with the <pause> token In downstream finetuning, we are given prefix P1.V an- notated with a target t1.r. Here, we append Mr copies of the <pause> token to P1:N. to cre- ate our new prefix, P1.N+Ma- We visualize how this introduces new computational pathways in Figure 1. As before, we ignore the model's outputs until the last <pause> token is seen. We apply the standard next-token prediction loss on the target with the new prefix, thus minimizing Eh Ccr(t++1. ((PINN Ma tial)), where [:] denotes the concatenation operation. Note that for any given downstream task, we fix Mj to be the same across all inputs for that task. We again

---

Published as a conference paper at ICLR 2024
update both the parameters of the model, and that of the whole vocabulary, including the <pause> token, as is standard. We term this finsunnin (Algorithm 2).
Pausing during inference During inference the downstream task, append Minf <pause> tokens to the prefix and as always, we ignore the output of the model until the last <pause> token (Figure 1). We term this pause-inference (Algorithm 3).
3.2 VARIANTS OF PAUSE- TRAINING
While pause tokens can incorporated in either pretraining or finetuning, in our study, we will consider all combinations of this. Our hope is to identify if there any differences in how each stage of pause-training affects inference-ting performance. In total, we study four techniques:
1. Standard Pretraining and Standard Finetuning (StdPT_StdFT).
2. Standard Pretraining and Pause-Finetuning SStdPT.PauseFF. We train with <pause> tokens only during downstream finetuning. If this technique helps, it would promise a practically viable approach for pause-training off-the-shell models.
3. Pause- Pause-Pretraining and Standard Finetuning (PausePT StdFT): Here we introduce <pause> tokens during pretraining, but abandon downstream. This is purely for analytical purposes (See 84.3). PauseeFT: We inject delays into both stages.
4. Pause-Pretraining and
Unless stated otherwise, we use the same number of pauses at inference as finetuming (Mit Mu).
4 EXPERIMENTS
Our main experiments broadly aim address two questions:
(1) Does delaying the model computation via pausing help (hopefully, due to the wider computa- tional flow), have effect (since the tokens provide new information, and substantially new parameters are added) or hurt (perhaps, by distracting the model with stray tokens)?
(2) If at all these delays have any effect, is there difference in performance when inject it into the pretraining stage versus finetuning stage versus both?
4.1 EXPERIMENT SETUP
We consider decoder-only models of size IB and 130M for main experiments. For our ablations, we stick to the IB model. Both the standard and pause models are pretrained on the C4 English mixture (Raffel et al., 2020), using the causal next token prediction objective for total of 200B tokens (slightly more than 1 epoch on C4). For pause-pretraining, we insert the <pause> token randomly at 10% of the sequence length (2048) positions, and trim the now-longer sequence to its original sequence length. We then conduct pause- pretraining and standard- saundard-pettainingg for the same number of total tokens (200B). We use single <pause> token embedding, effectively increasing the parameter count by 1024 (the token embedding size), a quantity that is dwarfed by the I billion total parameter count (the token constitutes a 10- fraction of the model size).
Since we expect different downstream tasks to benefit from a different number of finetuning <pause> tokens Mr. we run finetuning with Mj. (and likewise Mnr)) set to 10 and 50 and report the best of these two for our consolidated results. However, we provide the values for both Mn. E 510,50) in Appendix E, in addition to a more finegrained ablation this hyperparameter Section For all the downstream finetuning experiments, we report mean and standard deviation over 5 runs (with the randomness purely from the finetuning stage). We tune the learning rate and batch size for standard end-to-end training, and use the best hyperparameter for all other training variants well. We share all the hyperparameters in Appendix H.
4.2 DOWNSTREAM DATASETS
We consider nine varied downstream tasks: (a) reasoning (GSM8k (Cobbe al., 2021)), (b) ex- tractive question answering (SQuAD (Rajpurkar et al., 2016), CoQA (Reddy et al., 2019)), (c) gen- eral understanding (CommonSenseQA (Talmor et al., 2019), PhysicalIQA (Bisk et al., 2020)), (d)
5

---

Published as a conference paper at ICLR 2024
GSM8K
SQUAD V1
CommonSenseQA
LAMBADA
WebQA
19.5
16.5
Accuracy
54
18.0
8.0
48
16.5
7.5
38.0
15.0
13.5
7.0
27.5
13.5
Naturalu
CoQA
PhysicalIOA
He llaSwag
31.5
38.0
Training Algorithns
26
31.0
StdPT_ StdFT
25
30.5
73.8
37.6
StdPT_ PauseFT
PausePT StdFF
24
PausePT PauseFT
23
73.2
36.8
Figure 3: Downstream performance for 1B model. Injecting delays in both stages of training (PausePT.PauseFT) outperforms the standard end-end training StdPT_StdFT wide variety of tasks (except HellaSwag). In contrast, introducing delays only in the finetuning stage provides only lukewarm gains, and even hurts in GSMSk.
long term context recall (LAMBADA (Paperno et al., 2016)), (c) natural language inference (Hel- laSwag (Zellers et al., 2019)), and (f) fact recall (WehQuestions (Berant et al., 2013), Natural Ques- tions (Kwiatkowski et al., 2019)). HellaSwag and PhysicalIQA are scoring tasks. We note that our implementation of CommonSenseQA is as decoding task, and hence we report Exact Match (EM) scores. Detailed dataset description is in Appendix G.
4.3 RESULTS: EFFECT OF PAUSE-TRAINING
We report the performance of the four approaches in $3.2 on all our downstream tasks for our 1B model Figure 3, and for our 130M model in Appendix B. We discuss zero-shot results SD.
The benefit of mouse-pretraining followed by pause- mouse-finetuning (PausePT PauseFT), Our first core finding that there are clear gains when <pause> tokens are introduced during both pretraining and finetuning (PausePT_Pausee PPasseFa), across majority of the tasks consider. Overall, this outperforms the standard baseline (StdPT StdFT) on eight tasks for the IB model, and on six tasks for the 130M model (Appendix Fig 5) albeit to varying extents. Most prominently, for the 1B model the SQuAD guestion-answering task, PausePT PaasseFF improves over StdPT_StdFT by an 18% EM score. Similarly, we observe upto 8% gains on the general understanding task of CommonSenseQA. On the reasoning task of GSM8k, PausePT_PauseFT gives an accuracy of 8.5% compared to 7.5% of the standard baseline. Similar gains are observed in other tasks like long-term context understanding (LAMBADA) and also on fact recall tasks like WebQA and NaturalQuestion.
The lukewarm effect of pause-finetuning a canndard-peetainned model In contrast to the above observations, introducing delays only during downstream finetuning gives mixed results. While there are gains on about benchmarks, they are com- paritively less. On the remaining, the performance mirrors (or is worse than) standard training.
Isolating the benefits of pause-pretraining independent of downstream delay (PausePT StdFT), The gains in PausePT PauseFT may come not only from inference-time delays, but also from better representations learned by pautraprera.ining. To isolate the latter effect, we consider PausePT_StdFf where we do not inject delays in the downstream task. Here the gains clear only in two tasks (CoQA and PhysicallAA). Thus, we conclude that mause-pretraining improves the representation for few downstream tasks; conversely, in most tasks, the gains of PTuseTT must come from well-learned delayed computations executed at inference-tinn
Filler characters <pause>; For completeness,
also report results for inference StdPT StdFT models, delayed with 10 or 50 periods (t.)). Corroborating the observations of Lanham al. (2023), find gains in doing this (Figure 4a).

---

Published as a conference paper at ICLR 2024
GSMBK
SQuAD V1
GSMBK
BrMKk:PauseFPPausee M,
SQuAD Souu0:StdPPPP Stt.rPasssPPT...oco Pa MM
Start
8.4
Accuracy 1
Accuracy
8.0
7.6
38
7.2
5WA GAPT
25
25 50 Num Inference Pauses
2
Filler Tokens as Pause
Num Finetune Pause
Num Inference Pauses
(a)
(b)
(c)
(d)
Figure 4: Key Ablations (55): (a) We compare pause-trained model vs. standard model delayed
using filler periods (...'). As in Lanham et al. (2023), the filler periods do not give any gains.
(b) There exists an optimal number of finetuning <pause> tokens (Mn) for given downstream dataset beyond which gains diminish. (c) and (d) We test the robustness of pause-trained models to varying number inference time <pause> tokens (setting Min not equal to Mn). which exposes the model to a serious test-time distribution shift. Pause-training degrades gracefully to shifts as wide as Maaf E [5, 25] for Ma = 10 both for (c) PausePT PauseFT and (d) StdPT. PauseFT.
Thus, to the core question of our exploration - whether delays help. hurt or do nothing at all we find that the answer depends on when these delays are introduced. Specifically, cause-pretraining appears crucial for delays to help in downstream inference-time. We conjecture that standard- pretrained model has strong biases that prevent it from fully realizing benefits inference-time delays e.g., standard pretraining biases the model to be "quick" in its computations.
As a concluding note, we remind the reader that the PausePT.Pauserf model has (de- liberately injected) computational advantage compared to StdPT. during finetuning and in- ference. However, there no computational advantage during cuuue-pretraining since we equalize the number of tokens seen. In fact, this only results in slight statistical disadvantage: the pause- pretrained model sees only 90% the (meaningful) pretraining tokens that the standard model sees, the remaining 10% are dummy <pause> tokens.
Remark:
ABLATIONS: WHERE AND HOW MANY <pause> TOKENS TO USE
NUMBER OF <pause> TOKENS DURING FINETUNING Recall that we append Mt. copies of (the same) <pause> token to the prefix during finetuning. We find that for each downstream dataset, there is corresponding optimal value of Mr. For example, on GSM8k, 10 <pause> tokens optimal with accuracy reducing to that of baseline as <pause> tokens are increased to 50 (See Fig- ure 4b), while for SQuAD. 10 is sub-optimal (see Appendix E). Possibly, for each dataset, there exists certain threshold of <pause> tokens beyond which the self-attention mechanism becomes overwhelmed.
ROBUSTNESS TO A VARYING NUMBER OF INFERENCE-TIME PAUSES So far, we have set the inference-time delay be the same what was seen during finetuning (Mnt MR). Next, we examine what happens if we vary Miat during inference. Note that this presents severe test-time distribution shift as we provide no supervision for the model until the last <pause> token (the Mi one) is seen. Thus the model may very well output garbage if we begin eliciting response that is either premature (Mir < M+) or belated (Mit MR). Yet, practice, we find (somewhat) graceful behavior.
First, PauseeT_Pausee PauseFT model is robust to a wide test-time shift the number of <pause> tokens (see Figure 4c and Appendix F): the performance remains above the baseline even if the inference- time pause tokens are half that of training-ting in computational constraints. Relatively, the StdPT PauseFF model (wherein inject delays only during finetuning) is even more more robust (see Fig and also Appendix F).
This is desirable in case real-time fluctuations
Now, ideal robustness criterion would be that, in the absence of any <pause> tokens, the pause- finetuned model performs just as well as a standard finetuundd model. Unfortunately. this isn't 't the case for any of models. In fact, for PausePT PauseFT, providing zero delay during inference

---

Published as a conference paper at ICLR 2024
breaks its performance spectacularly (Figure 4c and also Appendix F), even if the model behaves reasonably with as few as 2 inference-time <pause> tokens. The design of zero-delay delay-robust pause- trained models is thus an important question for future work.
APPENDING VS. PREPENDING PAUSES In our main experiments, chose to append <pause> tokens since it is the most natural format for a general setting e.g. in long-text-generatinn as in a conversational agent, one would append <pause> tokens to the current text rather than deploying the tokens all at once at the beginning of the conversation. Furthermore, when there is unidirec- tional attention, prepending these tokens should make no difference. Nevertheless, in our down- stream tasks which use bidirectional attention on the prefix, it makes sense to consider prepending <pause> tokens. We investigate this Table in Appendix C. Most importantly, we find that, for PausePT.PauseFT,, even prepending the <pause> token performs improves over standard end-to-end training. However, appending is still the more optimal choice. This indicates that pause- quuus-pretaining induces considerable biases in how readily the delays are used based on their positional embeddings.
DISCUSSION AND KEY OPEN QUESTIONS
Enhanced computational width. One hypothesis as to why Transformer delays can help is that increases expressivity by increasing the computational width. To produce the (K +1)h token, stan- dard inference involves a computational depth of L (corresponding to the sequential computation of L layers), and computational width of K (corresponding layer). With M <pause> tokens however, we perform K + MM parallel computations. We hypothesize that this additional width helps certain downstream tasks. Take for example, comprehension-base question-ansse answering tasks. Here, having greater number of attention units per layer, would per- mit finer distribution of attention across various parts of the supporting context (where the answer resides). We speculate that this would allow the lower layers to extract more precise and diverse representations, which a higher layer can more reliably aggregate to produce final answer.
the parallel K computations per
We formalize this in Theorem J.1 (stated informally below). Our key theoretical insight is that the attention module can have high "raw" parameter-count-basee capacity, but low "implementation capacity": the number of operations implemented for given input, which is bottlenecked by the number of tokens. Pause tokens can help relieve this bottleneck and tap into the raw representational capacity. We hope this preliminary result can inspire further discussions how to formalize the Transformer's 's implementation capacity and differentiate it from the raw parameter count.
Theorem 6.1. (informal; see Theorem J.1) Assume that the attention module has sufficiently many parameters (K) that it is much larger than the number of inputs tokens (L). Then there are tasks that involve many independent computations N, where N > L (but N < K), such that a 2-layer Transformer can implement the task if and only if it uses pause tokens.
Computational expansion along with parameter expansion. How do the gains offered by <pause> tokens vary with parameter count? An immediate hypothesis would be that for smaller models, delays become more beneficial they provide much-needed capacity increase in an oth- erwise sapacity-deprived model. But preliminary comparison between our two model sizes sur- prisingly suggests the opposite. Intuitively, we conjecture that smaller model may simply not have enough raw capacity to implement a variety of distinct computations to utilize the new pathways introduced by <pause> tokens. This intuition is also echoed in Theorem 6.1 where smaller models would break our assumption that parameter count K is large enough.
Computational expansion vs parameter expansion. There are trivial ways to extend the next- token computation process, say via more attention heads or more layers. For a fixed inference-time FLOPS budget, can these give similar gains as pause tokens? In Appendix I, we argue why this is not the case attaining similar gains requires significant parameter expansion and significant expansion of the FLOPS budget. Thus, it is both practically and theoretically remarkable that pause tokens can yield gains despite negligible addition to the parameter count.
Pause-inference vs. Chain-of- Thouugh,,, is worth contrasting the above computational advantage with that enjoyed by tf-hhugh (CoT) prompting (Wei et al., 2022). Here, one prompts the model to generate a long series of intermediate reasoning steps before producing the final answer. Thus, COT also corresponds to greater computational width, by way of delaying its final answer
8

---

Published as a conference paper at ICLR 2024
(albeit with meaningful tokens). However, CoT has vital added advantage: it also increases the computational depth to a significant extent (Feng et al., 2023, Theorem 3.3 and 3.4). In particular, each (meaningful) delay token generated by CoT is autoregressively generated by the model. Thus, there are M such tokens and I. layers, the final token arises out of roughly M I. sequentially composed operations. Thus, CoT has computational depth that is larger by a mulriplicative factor of M, compared to cause-inference. Finally, we note that one major advantange of CoT is that it does not seem to require any special modifications in the pretraining objective.
RELATED WORK
Multiple works have proposed using dummy tokens as a ways of introducing memory into the model. Closest to our work is Burtsev et al. (2020) who prepend these tokens (rather than append them) and crucially, introduce them only during training and inference on the target tasks. On smaller scratch-trained models (with parameter counts of 10M, 65M and 277M) and a pretrained BERT model (109M), this reportedly gives minimal gains. This echoes our own mixed results for the comparable StdPT_PauseFT variant, and the fact that our smaller model shows gains on fewer datasets. Orthogonally, Sukhbaatar et al. (2019) introduce a large set of "global" memory tokens (of the order of ten million new parameters) across all layers. However, these are independent of the input, and only act as keys and values (not as queries). For vision transformers (ViTs), AdaTape (Xue et al., 2023) appends tokens from a learnt memory bank. Concurrent work in ViTs (Darcet et al., 2023), use multiple different "register tokens" appended to the image patches. In the context of recurrent networks, closest to is Herel & Mikolov (2023) who train with "thinking" tokens for extra compute for the target task and find small perplexity gains on reasoning tasks. Grave al. (2016) use a memory cache as a way to attend to past hidden vectors.
MEMORY TOKENS
Adaptive Compute Transformers Adaptive compule methods provide users with flexibility in inference time depending on the input to the model. Although adaptive compute is not the objective of this paper (after all, we use the same number of pause tokens across all inputs), pause-inference can be viewed as a potential basis for a future adaptive compute method. The Universal Transformer (Dehghani et al., 2019) adaptively increases serial computations by repeating certain layers. Graves (2017) explored a similar recurrence over hidden units in the context of RNNs. A future adaptive- compute version of pause-inferene would fall in much less explored paradigm that varies the number of input tokens, akin to AdaTape in ViTs (Xue et al., 2023).
We discuss other relevant works in detail in Appendix K.
8 CONCLUSION, LIMITATIONS AND FUTURE WORK
tonseetrainig takes a step beyond the paradigm of "immediate" next-token prediction in language models. The key idea is to train models with (dummy) <pause> tokens so that the model can learn to harness the additional inference-tine time computation. This can improve performance on variety of tasks, if we train with <pause> tokens both during pretraining and downstream finetuning.
However, by extension of the fact that every downstream task has an optimal number of <pause> tokens, we do not claim that pause-training should benefit every downstream task. Some tasks may simply be better off with zero <pause> tokens. The most important limitation though is that the expense of mause-pettraiiing comes in the way of making this idea more widely accessible. Consequently, we do not study how the gains generalize across more model sizes (beyond 1B and 130M), to encoder-decoder architectures, or other pretraining mixtures and objectives. Next, while we have laid out some preliminary theory for why <pause> tokens may be beneficial, we leave rigorous understanding for future study. We also leave open variety of follow-up algorithmic questions: pause-training with multiple different <pause> tokens, better determining the number of <pause> tokens (perhaps using model confidence), inducing robustness to shifts in delays, and so on. But the most pressing next step would be to find ways to make delays helpful directly on standard pretrained model. Overall, we hope that our work opens up many avenues for theoretical and practical work in the paradigm of delayed next-token prediction.
Acknowledgements: We would like to thank Kaifeng Lyu, Srinadh Bhojanapalli, Yongchao Zhou, Nikunj Saunshi, and Seungyeon Kim for helping set up the initial codebase and for their guidance.

---

Published as a conference paper at ICLR 2024
REFERENCES
Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder, 2021.
Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing Freebase from question-answer pairs. Proceedings of the 2013 Conference on Empirical Methods Natural Language Processing, pp. 1533-1544. Association for Computational Linguistics, 2013.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer. In NeurIPS, 2022.
Mikhail Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov. Memory transformer. arXiv preprint (rXXi:20006.15222) 2020.
Haw-Shiuan Chang, Ruei Yao Sun, Kathryn Ricci, and Andrew McCallum. Multi-cls BERT: an efficient alternative to traditional ensembling. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, pp. 821 2544. Association for Computational Linguistics, 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers, 2023.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers, 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLIT 2019, Volume I (Long and Short Papers). Association for Computa- tional Linguistics, 2019.
Devvrit, Sncha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi, Sham Kakade, Ali Farhadi, and Prateek Jain. Matformer: Nested transformer for elastic inference, 2023.
Cristóbal Eyzaguirre, Felipe del Río, Vladimir Araujo, and Álvaro Soto. Dact-bert: Differentiable adaptive computation time for an efficient bert inference, 2021.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. abs/2305. 15408, 2023. doi: 10.448550/axxv..23055 15408. URL https:t////ooo org/10../00. 48550/arXiv 2305 15408.
Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with continuous cache, 2016.
Alex Graves. Adaptive computation time for recurrent neural networks, 2017.
Karen Hambardzumyan, Hrant Khachatrian, and Jonathan May. WARP: Word-level Adversarial ReProgramming, In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers). Association for Computational Linguistics, August 2021.
David Herel and Tomas Mikolov. Thinking tokens for language modeling. 2023.
Wittawat Jitkrittum, Neha Gupta, Aditya Krishna Menon, Harikrishna Narasimhan, Ankit Singh Rawat, and Sanjiv Kumar. When does confidence-basee cascade deferral suffice?, 2023.
10

---

Published as a conference paper at ICLR 2024
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ra- manujan, William Howard- -nyder,, Kaifeng Chen, Sham Kakade, Prateek Jain, and Ali Farhadi. Matryoshka representation learning. 2022.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Al- berti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wey Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: a benchmark for question answering research. Transactions of the Association of Computational Linguistics, 2019.
Jack Lanchantin, Shubham Toshniwal, Jason Weston, Arthur Szlam, and Sainbayar Sukhbaatar. Learning to reason and memorize with self-notes, 2023.
Tamera Lanham, Anna Chen, Ansh Radhakrishnan, Benoit Steiner, Carson Denison, Danny Her- nandez, Dustin Li, Esin Durmus, Evan Hubinger, Jackson Kernion, Kamilè Lukošiûte, Karina Nguyen, Newton Cheng, Nicholas Joseph. Nicholas Schiefer, Oliver Rausch, Robin Larson, Sam McCandlish, Sandipan Kundu, Saurav Kadavath, Shannon Yang, Thomas Henighan, Tim- othy Maxwell, Timothy Telleen-Lawton. Tristan Hume, Zac Hatfield-Dodds, Jared Kaplan, Jan Brauner, Samuel R. Bowman, and Ethan Perez. Measuring faithfulness in chain-of-thouge rea- soning. 2023.
Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference Empirical Methods in Natural Language Pro- cessing, EMNLP 2021, 7-11 November, 2021. Association for Computational Linguistics, 2021.
Xiang Lisa Li and Percy Liang. Prefix-tuming: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the Ithh International Joint Conference on Natural Language Processing (Volume 1: Long Papers). Association for Computational Linguistics, 2021.
Xiao Liu, Yanan Zheng, Zhengxiao Du Ming Ding, Yujie Qian, Zhilin Yang, and Jie Tang. Gpt understands, too, 2021.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P- tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks, 2022.
Yang Liu. Fine-tune bert for extractive summarization, 2019.
Lajanugen Logeswaran, Ann Lee, Myle Ott, Honglak Lee, Marc' Aurelio Ranzato, and Arthur Szlam. Few-shot sequence learning with transformers, 2020.
Aman Madaan and Amir Yazdanbakhsh. Text and patterns: For effective chain of thought, it takes two to tango, 2022.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought, 2023.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7- 2016, Berlin, Germany, Volume Long Papers. The Association for Computer Linguisties, 2016
Guanghui Qin and Jason Eisner. Learning how to ask: Querying LMs with mixtures of soft prompts. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Lin- guisties, 2021.
11

---

Published as a conference paper at ICLR 2024
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pp- 2383-2392. The Association for Computational Linguistics, 2016.
Siva Reddy, Danqi Chen, and Christopher D. Manning. CoQA: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7, 2019.
Timo Schiek and Hinrich Schütze. It's not just size that matters: Small language models are also few-shot learners, 2021.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling, 2022.
Roy Schwartz, Gabriel Stanovsky, Swabha Swayamdipta, Jesse Dodge, and Noah A. Smith. The right tool for the job: Matching model and instance complexities, 2020.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory, 2019.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le, Ed Chi, Denny Zhou, and Jason Wei. Challenging big-bench tasks and whether chain-of-thought can solve them. In Findings of the Association for Computa- tional Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023. Association for Computational Linguistics, 2023.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. CommonsensseQA: A ques- tion answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Con- ference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers). Association for Computational Lin- guistics, 2019.
John Thickstun. The transformer model in equations. University of Washington: Seattle, WA, USA, 2021.
Miles Turpin, Julian Michael, Ethan Perez, and Samuel R. Bowman. Language models don't 't always say what they think: Unfaithful explanations in tfain-of-thoughee prompting, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, pp- -008- 2017.
Boshi Wang, Sewon Min, Xiang Deng, Jiaming Shen, You Wu, Luke Zettlemoyer, and Huan Sun. Towards understanding chain-of thain-of-thought prompting: An empirical study of what matters. In Proceedings of the 6t Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023. Association for Computational Linguistics, 2023a.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoe V. Lc, and Denny Zhou. Chain-of-thought prompting clicits reasoning in large language models. In NeurlPS, 2022.
Fuzhao Xue, Aixin Sun, Hao Zhang, Jinjie Ni, and Eng Siong Chng. An embarrassingyy simple model for dialogue relation extraction. In IEEE International Conference on Acoustics. Speech and Signal Processing, ICASSP 2022. IEEE, 2022.
12

---

Published as a conference paper at ICLR 2024
Fuzhao Xue, Valeri Likhosherstov, Anurag Arnab, Neil Houlsby, Mostafa Dehghani. and Yang You. Adaptive computation with elastic input sequence. In International Conference on Machine Learning, ICML 2023. Proceedings of Machine Learning Research. PMLR, 2023.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023,
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah D. Goodman. Star: Bootstrapping reasoning with reasoning. NeurlPS, 2022.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Zexuan Zhong, Dan Friedman, and Danqi Chen. Factual probing is [MASK]: Learning vs. learning to recall. In Proceedings of the 2021 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguisties, 2021.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoe V. Le, and Ed H. Chi. Least-to-most prompting en- ables complex reasoning large language models. In The Eleventh International Conference on Learning Representations, ICLR 2023. OpenReview.net, 2023.
13

---

Published as a conference paper at ICLR 2024
A PRELIMINARIES: TRANSFORMER
Consider a vocabulary V and an input PL.K VK of K tokens, and an T--ayer decoder-only language model. The /'th layer of the Transformer produces one intermediate vector for each token here, denoted by Vg E TRP for k = l,..., K. We first describe this operation before outlining the end-to-end next-token generation process.
(t)
Consider Transformer (Vaswani al., 2017) block T(-): : RRRDD Kx that operates over sequence of K intermediate vectors. The block is defined by H many sets of four matrices each, W Wquery, w-(h), wr(h) key" value and W:(h) E P (for h 1.... each denoting an attention head), and single parameterized feedforward module fFP : RD_ + RR. Let 4L _R_ + RD denote the layer-norm operation. Given the input vectors V1.K RP*K the output Ví.x of the Transformer block (-) can be expressed in the following steps. For all k < K,
H
(W Vink) ( W VDttn
ak = P
V&
+
Σ (wilh)T
w((h) value softmax
query V
(4)
h=1
(2) V/A = & (#FF(a)) +a).
(5)
Here, the first step computes K different self-attention outputs by attending to all K input vectors, while the second step individually processes each attention output via a feedforward network and other normalization components produce the final output of the block. Note that here have assumed a unidirectional attention mechanism; for a bidirectional mechanism over the whole K- length prefix, one simply needs to replace V1 with Vi.K in the above computation.
Given this block, the Transformer generates the next token as follows. Let Proken V + RD and position N + IR" denote the toben-emn and position-embeddding layers. With an abuse of notation, let the token unembedding layer be denoted as Φ-1 tooenn which maps from TR. to a probability vector 1 Let (-) denote the 1 Transformer layer. Then, the Transformer commits the following operations in sequence to arrive at the (K + 1)th token.
v roken (Pk) + pposition (k)
(6)
v-( ( (vilk=1) (1 VI E [1, L]
(7)
PK+1 Tb Loken (V
(L)
For a more detailed mathematical exposition of the Transformer model, we refer the reader to Thick- stun (2021).
B ADDITIONAL DOWNSTREAM FINETUNING RESULTS
We first report the downstream finetuning performance for the 1B model in Table (numbers cor- responding to Figure 3 in $4.3). Further, Figure report downstream performance various tasks for a 130M decoder-only model. Again we observe that BausePT_Pauserf clearly outperforms standard training baseline (StdPT.StdFT) on GSM8K, CommonSenseQA,, LAMBADA and on our fact recall tasks like WebQA and NaturalQA. However, surprisingly, we do not observe gains on SQuAD, contrast to the gains observed 1B model. Overall, we see an improvement six tasks for the smaller model (one of which is PhysicalIQA where the gain is minimal..
C PREPENDING VS APPENDING PAUSE TOKENS
In Section 5, we discussed the effect prepending the pause token in comparison to the default approach of appending them to the end of prefix. Table 2 compares the approaches. As stated before in Section 5, for the PausePT_PauseFT training algorithm, we observe that prepending the pause tokens still outperforms the baseline but is (slightly) worse than appending the pause tokens on some benchmarks like GSM8k and SQuAD. For StdPT PauseFT however, we see mixed results with equal number of wins and losses between the prepending and appending.
14

---

Published as a conference paper at ICLR 2024
GSMBK
SQUAD VI
CommonSense0A
LAMBADA
PhysicalIOA
Rouge1
39.0
28
37.5
27
28.0
67.4
26
12
67.2
17.5
36.0
WebQA
NaturalQA
CoQA
l l LaSwag
27.08
38.4
Training Algorithms
25.5
15
26.75
StdPT_StdFF StaFT
24.0
12
26.50
29.6
PausePT_Sted StdFT
StaPT PauseFT
22.5
26.25
29.2
PausePT_ PauseFT
21.8
Figure 5: Downstream performance of pause-training on 130M decoder-only model. We find six of nine downstream tasks, the pause-pretrained and pause-finetuned model [PausePT.PauseFT outperforms standard training (StdPT.StdFT) For example, on the reasoning task of GSM8k, we observe 3% gains in Rougel scores (we compare Rougel as the final accuracy was too low to be meaningful for our 130M model). Similarly on the general understanding task of CommonSenseQAA, we observe upto 8% gains. We note here that we solve CommonsenseQA
a 130M decoder-only model.
decoding task and not rank classification task, and hence report the Ex- act Match scores. We also highlight that while pause-trainined on the IB model showed significant gains SQuAD, they disappear here.
Dataset
Metrie
StdPT,StdFT
StdPT_ PauseFT
PauSePT_StdFT
PausePT, PauseFT
10
50
50
7.5
6.5 +
0.5
8.5 + 0.9
0.3
GSM8k
Rougel
42.3 = 0.5
41.7 = 0.7
41.2 1.3
43.5 4 0.1
44.2 +.2 0.2
44.1 = 0,2
SQuAD
EM
36.4 2.5
36.6 + 2.2
3.2
38.4 2.9
51.7 + 2.3
55.9 = 1.0
CommonSense QA
FM
26.9
28.7
2.7
34.8
0.8
LAMBADA
EM
16.4 - 1.7
18.4
18.5 0.6
1.111
18.8 0.1
18.5 1 0.2
Web Questions
EM
13.7 2.1
4.4
12.4 2.6
1.50
13.8 4: 3.7
16.0 +: 1.6
Natural Questions
EM
23.6
24.3
23.9
2..5
24.9
26.9 0.4
CoQA
29.9
0.5
30.3 - 0.5
31.1
31.3
31.6 0.5
PhysicalIQA
F1
73.3 0.2
73.9
74.0 0.2
7..1
7.11
74.2 0.2
HellaSwag
37.8 0.1
37.9 0.2
37.8
0.2
0.1
37.7
37.8 0.2
Table Downstream performance on various tasks for the 1B decoder-only model. We observe that PausePT_Pauser outperforms the standard training baseline 8 of the tasks considered in this work. See $4.3 and Figure 3 for further details.
D ZERO-SHOT EVALUATION
In 84.3, we showed the efficacy of prute-pretrained models when finetuned with pause tokens. How- ever, we note that we witness gains even in the zero-shot setting, where the model not finetuned. In Figure 6, we compare the zero-shot accuracy of the standard pretrained IB model with that of pretrapre version. For the pause- pretraieed model, we perform evaluation with 0, 10 and 50 <pause> tokens appended to the prefix. Observe that muuse-pretraining gives gains tasks like GSM8k and HellSwag. However, we note here the absolute value our zero-shot accuracies quite low, experiment with small IB parameter model.
15

---

Published as a conference paper at ICLR 2024
Dataset
Metric
StdPT_StdFT
StdPT Pause_P
PausePT_PPuseeFT
Prepending
Appending
Prepending
Appending
GSM8K
Acc.
7.5 + 0.5
8.0 1.0
6.9 + 1.0
8.0 + 0.4
8.5 0.9
SQuAD
EM
36.4 2.5
35.0 1.5
40.2 3.2
44.0 3.2
55.9 1.0
CommonQA
EM
26.9 2.9
31.0 + 1.3
28.8 1.5
34.5 1.0
34.8 1.2
Lambada
EM
16.4 = 1.7
17.8 - 0.4
18.5 - 0.6
18.0 - 1.1
18.8 + 0.1
PhysicalIQA
73.3 + 0.2
74.0 = 0.3
74.0 0.3
74.2 74.2
74.2 - 0.2
NaturalQ
EM
23.6 1.2
24.1 + 0.6 24.3 1.4
25.7 0.9
26.9 + 0.4
Table Prepending vs. appending the pause tokens ($5). We observe that prepending the pause tokens still outperforms the standard training baseline of StdPT StdFT, but is suboptimal to append- ing the <pause> tokens for PausePT.PauseFTT training algorithm. However, for StdPT_Paused both have equal number wins and losses.
GSMBK
Squad V1
PhysicalIQA
WebQuestions
Hel HellaSwag
0.8
0.25
27.5
Rouge1
2.4
60
27.0
1.6
0.4
0.10
56
8.0
0.85
0.0
CoQA
Training Algorithas
0.3
StdPT
PausePT, inference
PausePT, 10 inference Pause
PausePT, 5e inference
Figure 6: Zero-shot evaluation of pause-pretrained models. Zero-shot inference with pause to- kens on a pause-pretrained model gives gains on tasks like GSMSK and HellaSwag. However, we note that zero accuracies are quite low, experiment with small IB parameter model.
E VARYING NUMBER OF PAUSE TOKENS Mst
GSMBK
SQUAD V1
CommonSense0A
55
8.4
34
8.0
5
Accura
32
7.6
30
48
Num Finetune Pause
Num
Finetune Pause
Num
Finetune Pause
Figure 7: Varying finetuning delay: We examine the effect of varying the number <pause> tokens used in downstream finetuning (Mn, $5) on the performance. Typically, we observe that there exists optimal number of <pause> tokens as expected for each dataset.
16

---

Published as a conference paper at ICLR 2024
SQuAD:Paissin Pwa.-aaaaPPPPseFTT - 10
SQUAD:SIdPPI StddosasPPPPPPFT.... Me 50
58
50
50
30
30
0 30
30
28
20
20
0
25 50
5
10
25
2
Num Inference Pauses
Nun Inference
Pauses
Num Inference Pauses
Num Inference
Pauses
(a) SQuAD
PSMEKPPausePPPauseFT
GSM8K: SWNKK:ttPPPPuseFF.
Accuracy
Accuracy
Accuracy
Accuracy
6
6
0
2
5
10
25
50
2
5
10
25
50
10 25
50
0
2
5
10
25 50
Num Inference Pauses
Num Inference Pauses
Num Inference Pauses
Num
Inference Pauses
(b) GSM8k
30
25
25
25
25
52.7 12 2dF7
20
20
20
20
25 59
2
Num Inference
Pauses
Num Inference Pauses
Num Inference Pauses
Num Inference Pauses
(c) CommonSenseQA
Figure 8: Varying inference-time delays: We test the robustness of pause-trained models to varying number of inference time <pause> tokens (setting Mine equal to MR). which exposes the model to a serious test-time distribution shift (55). Pause-training degrades gracefully to shifts as wide as apart
Mint [5 25] for Mr. 10 and Mr. 50 both for PausePT FauseFT and from GSM8k wherein there is a drop for Ma. 50. In each row, the first and the third column considers the PausePT.PauseFT model for Mj. set to 10 and 50, respectively. Likewise, the second and the fourth column show the same for
model.
In Figure 7, we study the effect of varying the number of pause tokens used during downstream finetuning (z) on the downstream performance. We refer the reader for further details. Again we observe that there exists an optimal number of pause tokens to be used during downstream finetuning, depending on the task.
F ROBUSTNESS TO VARYING NUMBER OF INFERENCE TIME PAUSES
Recall in $5 and Figure 4c we observed that pause-training is robust to using a different number of inference time pauses compared to that used during finetuning (i.e. Minf A MR). We present additional results regarding the same in Figure Sa, Figure 8b and Figure 8c. Again, we observe that the performance degrades gracefully for the pause- pause-trained models, number tokens seen. However, we still find a drastic drop in performance when no delay given during inference for the PausePT PauseFT model.
with shifts that halve the
17

---

Published as a conference paper at ICLR 2024
Dataset
Learning Rate
Warmup Steps
Finetuning Steps
Batch Size
SQuAD
1.0E-04
100
10000
256
GSM8k
1.0E-04
200
20000
16
HellaSwag
5,0E-06
100
1000
16
PhysicalIQA
1.0E-06
50
600
32
CoQA
5.0E-05
75
3500
16
CommonSenseQA
5.0E-05
100
4000
16
LAMBADA
5.0E-05
40
2800
16
WebQuestions
5.0E-04
200
2000
16
NaturalQuestions
1.0E-04
100
5000
256
Table 3: Downstream finetuning hyperparameters for the 1B model.
G DOWNSTREAM DATASET DESCRIPTION
We finetune and evaluate the pretrained models (both standard and pause-pretine presse--reraannd on the follow-
ing datasets:
1. GSM8k: A reasoning task with 8.5k grade school math word problems (Cobbe et al., 2021).
2. SQuAD V1: comprenocomppeeensiioo task based on Wikipedia (Rajpurkar et al., 2016).
3. CommonSensseQA: Requires different types of commonsense knowledge to choose the cor- rect answer (Talmor et al., 2019). Our implementation of CommonSenseeQA is a decod- ing task, and hence we report Exact Match (EM) scores.
4. LAMBADA: Text-understanding task requiring last-word prediction based on long con- text (Paperno et al., 2016).
5. Web Questions: A fact-recall dataset of commonly-asked questions on the web (Berant et al., 2013).
6. PhysicallQA: A physical commonsense reasoning dataset, which test the ability to under- stand interactions with the world (Bisk et al., 2020).
7. Natural Questions:
QA task which requires answering fact-based questions from Wikipedia article pages (Kwiatkowski et al., 2019). Since we use the closed-bood version of this dataset (no access to helpful context), this is fact-recall task.
8. HellaSwag: Next-sentence prediction task based on common-sena inference (Zellers et al., 2019).
9. CoQA: Question-answeriigg task based on a context (Reddy et al., 2019).
H HYPERPARAMETERS: DOWNSTREAM FINETUNING
We share all the hyperparameters for downstream finetuning in Table (IB model) and Table 4 (130M model). We also provide the decoder-only architecture details for the two models considered in this work in Table 5.
18

---

Published as a conference paper at ICLR 2024
Dataset
Learning Rate
Warmup Steps
Finetuning Steps
Batch Size
SQuAD
1.00E-04
400
40000
16
GSM8k
1.00E:04
75
7500
16
CommonSenseQA
5.00E-05
100
6000
16
LAMBADA
5.00E-05
40
1400
16
WehQuestions
5.00E-04
200
2000
16
NaturalQuestions
5.00E-04
100
5000
256
CoQA
1.00E-04
75
3500
16
PhysicalIQA
1.00E-06
50
600
32
HellaSwag
1.00E-06
100
1000
16
Table 4: Downstream finetuning hyperparameters for the 130M model.
Model
130M
IB
Parameters
136,237,056
1,345,085,440
Transformer Layers
12
24
Attention Heads
12
32
Embedding Dimension
768
2048
Hidden Dimension
3072
8092
Table 5: Architecture details for the models considered in this work
Algorithm 1: Pause-pretraining
Pretraining with Pause
Inputs: Pretraining dataset Ppt, decoder-only model fe. number of <pause> tokens Mpt to insert, pause token <pause>
PI:N Dpt
/ Input Sequence corpus /* Insert Mpt pause tokens randonly in the
PI:N+ +M, raandm.inseer(pp... <pause>, Mpt) original input sequence PI:N. extending length by Mar x/
Signore = [k E [0, + Mpt 1|: pk-1 <pause>}
/* Identify the set of pasitions where the
is <pauses +/
C fe, PI,N+ M,.) = E Lk=1,k@S CCE(PK+1 (o(p))))
/ Next token prediction
error excludes targets are pause (nodel made learn predict pause itself; + 0 = 1 - VoLrauseT (fo, PM:N+M,,))
/* Update the rodel
19

---

Published as a conference paper at ICLR 2024
Algorithm 2: Pause-finetuning
Stage 2: Finetuning with Pause
Downstream labeled dataset Dr. pretrained model fo, number of <pause> tokens Ma to insert, pause token <pause>
Inputs:
PLAN frrr ~ D
/ Sangle prefix and target */ / Append prefix Mr pauses
PI:N+ Concat I:N; [<pause>] x Mu]
LpausePT (Je, P1:N+M, (I:T) Σk (CE(te+1, fo (C (1x))) / Next token
prediction error on targets
0 = A - VoLPauseFT (fe. PI:N Ma 11:1)
Algorithm 3: Pause-inference
Stage 3: Inference with Pause
Inputs: Prefix P1:N: finetuned model fe, number of <pause> tokens Min to insert, Pause token <pause>
Ti:N-Marl = [PL:N: [<pause>] x Minl]
/ Append Mist pauses to prefix +/
PN+Ma++1 [o auto-regressive fashion *
/* Predict the next token in the sequence, and continue in
20

---

Published as a conference paper at ICLR 2024
I INFERENCE TIME COST OF PAUSE TOKENS
One way to assess the inference-time compute efficiency of a method is by estimating the number of Floating Point Operations per Second (FLOPS) requires. related, but independent metric, is the Wall Clock Time as it affects the latency of deployed systems. Below, we analyze how efficient pause- couse-inference along these two metrics in this section. Broadly, we make two arguments:
1. Pause-inference offers a more FLOPS-efficient way of increasing performance, as com- pared to other natural ways of expanding the number of attention operations in Trans- former, such as by adding layers or by adding attention heads.
2. Pause-inference is also wall-clock-efficiee compared the above techniques as it virtually introduces no overhead. When compared to CoT,, pause-inference provides a computation- ally more granular and cheaper way to improve performance (although still upper-bounded by CoT in terms of performance).
I.1 PAUSE TOKENS ALLOW FOR MORE EFFICIENT USE OF FLOPS
We frame our FLOPS-efficiency analysis as follows. Consider introducing pause tokens during inference in given Transformer. How many additional FLOPS does this require? If we spent the same budget of additional FLOPS to introduce more attention operations via other techniques namely, via an appropriate number of additional layers or additional attention heads - would we find a similar improvement in quality, in terms of metrics like accuracy?
Concretely, we use a running example of the 1B model used in this paper, whose architecture details are provided in Table 5. Specifically, we have number of transformer layers as 1 = 24, input em- bedding dimension as = 2048, per attention head embedding dimension = h/a, where a = 32 100 to-
is the number of attention heads. We also anchor analysis for input prompts with kens, which represents the average prompt length for many tasks considered in this work. For our analysis, rely supporting lemmas deferred to Section 1.2.1.
FLOPS-efficiency of pause tokens vs & additional layers: In the context of the downstream task of SQuAD, appending p = 10 pause tokens yields an 18% increase in EM score. Applying these values (n = 100,p p = 10) and / = 24 in the FLOPS-efficiency analysis of Theorem I.3, we can deduce that if we were to allocate the same FLOOS budget to adding more layers to the Transformer stack,
at most add 2 layers. This enhancement corresponds count, expanding the model from a 1B parameter model to a 1.1B parameter model. However, in practice, when scaling the parameter count, significant performance improvements (such an 18% increase in EM score) are typically observed only when the parameter count is scaled by much larger factors. Thus we argue that in this case, pause tokens provide a more inference-timeFLL. efficient way of increasing performance. Conceptually, this underscores the fact that pause tokens introduce an alternative dimension to representation capacity, distinct from the traditional approach of scaling the parameter count.
modest 10% rise in parameter
Comparing FLOPS with increase in attention heads: In the standard Transformer implementa- tion used in practice, when one increases the attention heads (a), although the number of attention operations increases, the per-attention-heal embedding dimension (d) proportionally gets reduced keep the overall embedding dimension constant (h). Thus, there is effectively no change in the number of FLOPS. In contrast, adding pause tokens increases the number of attention operations, while keeping the per-attention-hen head embedding dimension d constant. Therefore, for fair com- parison, we consider the case where we increase a, while keeping the per-attention-head dimension fixed (and we fix to be equal to h).
Then, from our analysis in Theorem 1.4 we have that for an input of length n 100 and 32 attention heads, appending p- = 10 pause tokens is equivalent to increasing the number of attention heads by A = 3. However, increase the attention heads by 3 in the IB model configuration adds only 48M parameters (we add Wg, WR. W., W proj E hh per additional attention head), bringing the model to parameter count of 1. 1.048B from IB. This, argue cannot account for any significant performance improvement equivalent to the improvements seen under pause- toassetraining proposed in this paper.
21

---

Published as a conference paper at ICLR 2024
1.2 PAUSE TOKENS DO NOT ADD SEQUENTIAL COMPUTE
Comparing pause tokens with adding layers or attention heads Recall that <pause> tokens are added as a part of the input prompt, where each token processed in parallel. Thus, <pause> tokens do not add extra serial computations. If there are sufficiently many parallel threads available, the wall clock overhead from pause tokens would be a minimal percentage of the time required for standard inference. However, in contrast, increasing the number of transformer layers increases the length of sequential operations, causing the wall-clock time to increase proportionally to the fraction of layers added. Note that adding attention heads should have a similar effect as adding pause tokens, as they too introduce parallel, not sequential operations.
Comparison with -hainno (CoT) prompting Recall that CoT involves auto- regressively decoding long sequence of tokens involved in the model's reasoning. This requires significant wall clock time cost, proportional to Opt)), if p is the number of intermediate reason- ing tokens and is the number of layers. In stark contrast, pause tokens do not add extra wall-clock time. Furthermore, CoT prompting offers little flexibility in how large p can be. Pause-inference on the other hand, offers a more direct way for manipulating the number of pause-tokent (even if, in its current version, this adaptivity is not robust beyond a point).
1.2.1 SUPPORTING LEMMAS FOR ESTIMATING FLOPS EFFICIENCY
In Lemma 1.1, we present the facts about FLOPS required for basic vector and matrix calculations. Subsequently, in Lemma 1.2, we compute the overall FLOPS required for an end-to-end Transformer computation. Finally, in Theorem 1.3 and L4, we derive how different kinds of parameter expan- sions in the model compare to adding pause tokens, in terms of FLOPS efficiency. Specifically, The- orem 1.3 establishes the number of layers one needs to add to model to realize the same number of FLOPS as adding p pause tokens. Theorem I.4 presents a similar result for adding attention heads. Lemma L.1. (FLOPS for vector and matrix calculations) The number flops required to compute:
1. the dot product V1 V2 where V1, V2 E R is Odd)
2. the matrix multiplication M1 M2, where M1 Roxb and M2 E Rbx the FLOPS is (aa xaxxxx, c).
Proof. For a dot product between vi and V2. both of dimension d. the total number of FLOPS is given by the sum of multiplications and additions required. Specifically, it involves d multiplications and d additions, totaling to 2d FLOPP. For the ease of calculation, we approximate this as 2d.
For matrix multiplication of M1 E Raabb and M2 E Raxe, each element of the resulting matrix is computed by taking the dot product of row from Mj and a column from M2, which requires 2b FLOPS Since there are a X such dot products to compute for the entire matrix multiplication, the total FLOPS amount to (26) x a X c. However, for simplicity, we only consider the multiplicative operations, the FL.OP reduce to X b x
Lemma I.2. (FLOPS for one end-to-end Transformer computation) Consider an layered decoder only language model, where we denote imput embedding dimension as h, number of attention heads as a and per-attention-heur embedding dimension as d. We assume feed-forward hidden dimension be 4h and finally let denote the length of inpul sequence. Then the total FLOPS are given as:
F(n, h 1) = (Anadh + 2an2 d + 8nh?))
We note that standard Transformer implementations assume d = h/a, aa i.c. the per-attention-head embedding dimension decreases the number of attention heads are increased. However, we treat these are three independent hyperparameters for greater flexibility in our analysis.
Proof. Let us consider the various per-layer operations in decoder-only model step-by-step and count their FLOPS:
1. q, k, v vector computation: Given input token z E R²_ for the query vector computation, we have Q' = Waz'j E [1,a], where Wi E Rdxh, The same extends for key and value vector computations. Thus total FLOPS required is 3nadh.
22

---

Published as a conference paper at ICLR 2024
2. Self-attention: Given Q Raxd and K R QK incurs n2d flops. The obtained a softmar( QK") va E g is multiplied by RRad costing another nid flops. Note that
for simplicity, we ignore that FLOPS from softmax or the division by v operation as they are negligible. Thus, the total FLOPS ann'd d + nd 2an2d.
3. Combining anttithead-atteetton:: The MHA projection matrix concatenates all the outputs from individual attention heads above and projects them to output of dimension h. For simplicity, we ignore the FLOPS from the skip connection as adds only a relatively minimal number. Thus total FLOPS x h ad nodh.
4. Feed-forward network: This adds another Snh² FLOPS. Again, for simplicity ignore the FLOPS from the skip connection.
Combining the FLOPS from each of the sub-parts above, we have:
FLOPS (4nodh + 2an + 8
(10)
Theorem I.3. (FLOPS for adding layers vs. P pause tokens) Consider a layer decoder only model, with denoting the input embedding dimension and d denoting the per attention head em- bedding dimension. Let n be the length of initial prompt. Then, under the assumption that hidden embedding dimensions are much larger then the prompt sequence length and the appended pause tokens ie. d, h >> p: the additional FLOPS from pause tokens is less than that from k additional layers of transformer if n > pl/k.
Proof. From Lemma 1.2 we have that increase in FLOPS due to k additional transformer layer is give by:
FAI=K = a( F(n, a,
Fai=k = (4nadh + 2an + 8nh)
(1I)
Similarly, increase in FLOPS due to pause tokens is given by:
Fan=p = F(n + p, h, a, - F(n, h, a,
Pan=p (4pudh 2nd((n + p)² n²)) 8pli2)1
FAn=p (4adh + 2nd p) 8h2)
(12)
Now,
FAn=p < FAI=K
(2adh + od(2n + p) + 4h )pl < (2adh + and 4h2 1h²)
(2adh ad(2n + p) 4h2) pl
n (2adh + and + 4h22)::
pl
[assuming h, d >>>>> n.
(13)
Next, we derive how many attention heads can be added, to be FLOP-equivalent to adding pause to- kens. Note that in the standard transformer implementation, increasing the attention heads decreases the per attention head embedding dimension (i.e. d h/a). In contrast, adding pause tokens, in- creases the number of attention computations while keeping the per-attention- consider setting where we increase the number of attention heads,
dimension fixed.
Thus for a fair comparison, while keeping the per-attention- -ead dimension fixed. Specifically, we consider the case where per attention head embedding dimension is fixed to be the same of input embedding (i.e. d = ht.
23

---

Published as a conference paper at ICLR 2024
Theorem 1.4. (FLOPS for adding k attention heads p pause tokens) Consider decoder only language model, with the per-attention
embedding dimension d, fixed to be same as the input embedding dimension h. Let n be the length of initial prompt. Then under the assumption that hidden embedding dimension is much larger then the prompt sequence length and the appended pause tokens Le. d, >> n. p; the additional FLOPS from p pause tokens is less than that from k. additional attention head, if a (a+2)p
Proof. From Lemma 1.2, we have:
F(n, h, d, a, = (4anh² + 2an?h + 8nh?)1 where d = h.
(14)
Now, increase in FLOPS due to l additional attention head is given by:
F Fa==k F(n, a k,1) F(n, h 1)
Faa=k- (4knh² + 2k)
(15)
Similarly, increase in FLOPS due to pause tokens given by:
Fan=p = F'n + P. h a, - F(n, h, a,
Fan=p = (4aph² + 2ah((n + p)² - n²²) + 8ph2
Pan=p (4ah 2a(2n p) + 8h)
(16)
Therefore, have, for:
Fan=p < FAa=1
(4ah 4an + 2ap + 8h)hpl (Anh² 222
((2n + 4)h + 2an + ap)p < (2h + n kn
(2a + 4)h + 2an + ap
ni
(2h + n))k
(a 2)p
n2
[assuming h, d > pl
(17)
k
J THEORETICAL INTUITION
This section formalizes broad class of problems where appending pause tokens during inference can enhance expressivity and thus be helpful. Our formalization identifies two core insights:
1. Pause tokens can be critical to solve tasks that require a large number of independent par- allel computations that exceed the number of input tokens. For example, consider task where the input is a sequence of L numbers U1, , UL,
and the target is a polynomial of the form (vi + 02) (v) 03) -. (us + U2). the number of addition operations required (N) scales much larger than the total number of input tokens (L), (and so N = (()) we argue that (under some natural capacity constraints), standard inference fails as is bottle- necked in terms of its "implementational capacity": it can conduct only O(L) operations in parallel. Pause-inference however is relieved of this bottleneck.
2. The citention-eeeffowwart block any layer has "untapped" representational capacity that is independent of the inpuf length- which pause-inferent taps into. Specifically, note that the iitettion-ceedforwaadd block can implement many different operations, one for each intermediate vector it generates at each positional index. But crucially, the number of possible such operations (say, K) scales with the parameter count of the block. This quantity is independent of and in practice, is much larger than the input sequence length. Unfortunately, standard inference can only help realize at most L such operations (where T. & K), while nause-inference can tap into K different such operations.
24

---

Published as a conference paper at ICLR 2024
Combining the above two insights, main result stated informally is that, given fixed (2-layer) Transformer architecture, (a) if the underlying task requires parallel operations, where N exceeds the number of imput tokens L, and (b) as long N is not much larger than the parameter count K of the erention-feedforwarr block, pause-inference can solve tasks that standard inference cannot.
We formalize the above insights in the form of assumptions stated in abstract setting (in order to be as general as possible). We emphasize that the crux of our argument lies within these assumptions themselves, rather than the proof of our theorem. Thus, our main result here should be viewed as identifying precisely what assumptions are required for pause-inference to help.
J.1 UNDERLYING TASK
We consider an abstract set of tasks that require a first step that involve multiple parallel operations, following by simple aggregation step to arrive at the solution:
Assumption J.1. (structure of the underlying task) Given the vocabulary space V, let be generic 2-ary operator on V For input sequence length L, consider corresponding function class FL that corresponds to all functions f Vl. + V that require applying N operations independently following by generie aggregation operation Jager VN + V:
s:ve FL 5 v Bin..... iN.. LI,
(18)
U1,,, vL.) Guger (vi, Unj (11,g vig)..... ... (ViN Vix
(19)
Nindependent o operations.
Examples. This structure covers a broad range of examples.
As simple mathematical example, this covers learning polynomials of the form (x1 +22)-
(x3 + x44 (z) + x3).
As natural language example, consider multi-choice question-answer task with C choices given along with E pieces of evidence in the context. One can then imagine that each v; corresponds
piece of evidence, and each v a choice. We then require N C-E operations that compare each of the given choices against each of the given pieces of ev- idence. A final aggregation step would select the choice for which there exists piece of evidence the most confidently corroborates the choice.
J.2 (TIGHT) UPPER BOUNDS ON THE TRANSFORMER CAPACITY
the attention-feedforraadd module had, say, an infinite or exceedingly large capacity, the Trans- former would be able to trivially express the solution to any task. It is only when these modules have finite capacity - as they do in practice that we expect additional operations introduced by the pause tokens be helpful. Correspondingly, we state this intuitive assumption multiple "tight upper bound" assumptions. Fach assumption states that the modules in a Transformer can represent objects of certain complexity, but none any more complex than that.
Our first such assumption is in how much information can be represented by each intermediate vector. Specifically, we assume that each vector can precisely capture one token in V along with the positional index of the token (akin to positional embeddings injected into each token in practice). The precise form by which this information represented as a vector is abstracted away for our discussion (c.g., it may be in one- hot form). Also note that our argument can be extended to settings where each intermediate vector could potentially represent more tokens, we discuss this at the end of the section.
Assumption J.2. (information bandwidth limit of intermediate Transformer operations) We as- sume that the i I'th intermediate vector in any given layer can be represented as (u, V N.
Next, we assume finite limit on the class of functions that each intermediate Transformer operation represent. To state this, let ((u1, 1), (u2, 2),.... (wz, T.)) be the outputs of an intermediate
25

---

Published as a conference paper at ICLR 2024
layer (corresponding to L-length input sequence). For convenience, ignoring the residual and layer-norm blocks, let the i' i'th output of the next layer be expressed as:
ΦFF (O (u, u, (u,, 0))) = (ul, i)
(20)
where the first two arguments represent the keys and values, and the third argument the query for the i' i'th intermediate operation in the considered layer. Note that ΦrF and @Atta are parameterized modules that can implement finite set of functions. We assume what this set of functions consists of:
Assumption J.3. (representational limits of intermediate Transformer operations) We assume that for each index i E N, opr (erttsum(::' (t,)) can represent exactly one of two types of functions:
A single operation. Specifically, we assume the self-attention operation can select two indices as @Ann u. (u,, i)) = (We(i), Uy)())). where v, : N + N come from some finite set of "index-selecting" functions P. We then assume ΦFF can implement as ΦrF (ur(s) ())(s) = (b(()) o ())()))
The aggregating function gagg as OFF (fr(dnnn (u, u, (1,, i))) (gaggr (u), (u), ur.), i).
We explain why the above aub-assumptions are both reasonable and can hold simultaneously. First, we argue why it reasonable that each intermediate Transformer operation can implement a limited number of operations, but not any more. Assume that the model needs to represent u/y5 =1 ug- This requires the model to pay attention to the query's positional index 35, and then select the values at two different positional indices 1 and 3. Selecting these two values can be implemented by 2 self-attention heads operating independently. A subsequent feedforward network can then operate on concatenated input (ui, ua). Crucially, implementing any further operations, would require more attention heads. Thus, it is reasonable to assume a (tight) limit on the number of attention operations.
Note that the above assumption can simultancously hold with the assumption that the Transformer operation can implement the aggregation function Jaggr. This because Jaggr does not require preferentially selecting any positional indices: all inputs are aggregated equally. Thus, we only need a single self-attention head that applies equal weight to all the values.
The last and arguably most insightful assumption stipulates a tight upper limit on what the self- attention module can represent. Specifically, observe that above we assume that the self-attention module can implement a finite set of functions P. which help select "indices" v(i) and v'ii). We assume that there is a limit to this set of index selextsel functions P. Our key insight is that this limit is purely determined by the parameter count of the module, and is therefore independent of the input- sepueence-ennthh L. In practice, the size of the set P is much larger than the number of tokens L, and thus corresponds the untapped capacity in the model. This capacity is however bottlenecked by the input length L, which determines how many operations in P are executed in standard inference.
Assumption J.4. (the "untapped" capacity of self-attention operation is independent of input length) We assume that for some K >> L, the self-attention module in each layer has at least 2K log K parameters, and can hence implement the index-selecting functions v, v' :K] - [K] to be any of the K* many mappings possible. In other words, P = 21. . (1......... (1.2....K)
From our assumptions, we derive the following result on what pause-inference can implement, which standard inference cannot, given 2-layer Transformer (we discuss extensions to larger ar- chitectures in a subsequent remark):
J.3 and J.4, standard inference on a 2-layer Transformer with no pause tokens can only represent the function class FN for N < L, where N denotes the number of parallel operations required by the function class, and L denotes the length of the imput sequence. In contrast, a 2-layer Transformer with N L appended dummy tokens can represent the function class FN for any large N 2 as long as N < K, where K scales with the parameter count of the self- atteeattentoon module as in Assumption J.4.
Theorem J.1. Under Assumptions J.1, J.2,
The key insight is that the self-attention module has the representational capacity to implement K different operations, capacity that is independent of, and much larger than the number of
26

---

Published as a conference paper at ICLR 2024
input tokens. However the standard Transformer that sees only L tokens, only allows the model to realize L of these. We hope this serves as a preliminary way of formalization the notion of the "implementation capacity" of a Transformer vs. the notion of raw parameter capacity.
Proof. Under Assumption J.2 that each intermediate vector can only represent one token from V. the first layer in the Transformer will have to implement all the N o operations to represent any given f E FL. Therefore, at the ''th index, the model would have to instantiate v1() ig and (i) = ju (as defined in Assumption J.3), so that the corresponding Transformer operation can compute 2/2 V/ However, by Assumption J.3 (which states that each intermediate Transformer operation can implement only one a operation) and by how standard inference in a Transformer is defined (which allows only many operations as there are input tokens), the first layer can only compute at most L many o operations. Hence, a Transformer with standard inference can only represent FN for N S L. On the other hand, from Assumption J.4, we know that the self-attention operator can implement as many as K such operations, where K > L. Thus as long as N S K, with N L dummy tokens, the Transformer can implement FN for N >> L.
Remark J.2. (breaking the information bandwidth assumption) In Assumption J.2, recall that we assume that in each intermediate vector, we are able to communicate precisely one token from V. If this assumption were to break, then standard inference would be able to implement a larger class of functions where N >> L. However, it would still fall short of what pause inference can do. Specifically, imagine that each layer passes on its computed output, and also all computed oulputs from the previous layer. Thus, if the intermediate vector can represent upto N/L. tokens, then the Transformer could divide the N required a operations over N/T. many layers, each layer performing L operations in parallel. This would however be slower as it requires series of N N/L L computations. In contrast, pause-inference requires only set of parallel computations, and also a meagre information bandwidth of token per intermediate vector.
Remark J.3. (recurrent depth) An alternative way to exploit the untapped capacity of the attention- feedforward modules, is to repeat these operations sequentially along the depth of the model as done in Dehghani et al. (2019). This strategy would be helpful in tasks that require recursion. However, our class of problems do not involve recursion of any form. Thus, if these recurrent layers simply repeat the same operations over and over, we may not enjoy any advantages. However, one may argue that, perhups repeating the same layers over and over somehow implements different oper- ations in each repetition. In this case, we can make an argument similar to the previous remark. Specifically, to fit a task that requires N parallel operations, we would need a model that has an in- formation bandwidth of N / L, and applies a recurrence of N L layers corresponding to N/ L serial operations to compute the desired function. This should again be contrasted with muus-infereence which requires only set of parallel computations, and also meagre information bandwidth of token per intermediate vector.
K ADDITIONAL RELATED WORK
Input-only tokens. The idea of using tokens that occur only as an input has found its use in various forms, most commonly as <cls> (Chang et al., 2023; Liu, 2019; Devlin et al., 2019), <sep> or <mask> in BERT (Devlin et al., 2019) and in a line of work adding memory to transformers (Burtsev et al., 2020; Bulatov et al., 2022; Darcet et al., 2023).
Chain-of- thoutht (CoT) prompting and role of intermediate tokens. CoT prompting has been shown to significantly improve the reasoning capabilities of large language models (Wei et al., 2022: Nye al., 2021: Lanchantin al., 2023; Suzgun et al., 2023; Zelikman et al., 2022; Zhou et al., 2023; Wang et al., 2023b; Yao et al., 2023). Consequently. there has been a surge of interest in understanding the source of these CoT prompting gains. Feng al. (2023); Merrill & Sabharwal (2023) theoretically argue that CoT aids by increasing the computational expressivity of the Trans- former. Other empirical works (Turpin et al., 2023; Wang et al., 2023a; Madaan & Yazdanbakhsh, 2022) have shown that the generated intermediate reasoning steps can be unfaithful, not represent- ing the model's 's true reasoning process. Wang et al. (2023a) empirically show that even incorrect reasoning steps can preserve 80%- 90% of the performance gains. However, Lanham et al. (2023) find that simply replacing the reasoning steps with filler tokens is unhelpful. As we argue, the model needs to be primed to process such tokens help its computation.
27

---

Published as a conference paper at ICLR 2024
Lightweight finetuning techniques. Pause- fouus--fnnetuing bears resemblance to orthogonal line of work on lightweight finetuning and ensembling techniques (Liu et al., 2022: Li & Liang, 2021; Lester et al., 2021; Hambardzumyan et al., 2021; Qin & Eisner, 2021; Logeswaran et al., 2020; Liu et al., 2021; Zhong et al., 2021; Schick & Schüizze, 2021; Xue et al., 2022; Chang et al., 2023). Lightweight finetuning is concerned with parameter-efficient techniques that do not update the model's weights, and instead update a series of multiple distinct learnable tokens (prepended the input). While pause-training uses (single) learnable token too (appended the goal and effects are significantly different. First, pause-training is not intended for parameter- efficient finetuning. Infact, pause-training tunes slightly more parameters than standard finetuning. Next, in terms of the effect, while pause-training hopes to outperform standard finetuning as it is a less constrained technique, lightweight finetuning typically cannot, as is a more constrained technique. Finally, note that pause-training crucially benefits from introducing the <pause> tokens during pretraining, while lightweight methods do not affect pretraining in any way.
the input),
Adaptive compute. In the literature, there are two major paradigms of adaptive compute. In the cascading paradigm, one selects between models of varying sizes (Jitkrittum et al., 2023; Kusupati et al., 2022; Devvrit et al., 2023). Another standard approach towards adaptive compute is layerwise adaptive compute within a single model, called early-exiting (Schuster et al., 2022; Schwartz et al., 2020; Eyzaguirre al., 2021; Banino et al., 2021).
is worth noting that while adaptiveness of the input token length (Xue et al., 2023) helps expand the (parallel) computational width, aarly-exi/layer-recureee type of methods help expand the (serial) computational depth of the model. As formalized in Theorem J. J.1, the expanded computational width critical in range of problems.
28


Reviews:
[{"review_id": "ObpCr3lUcj", "reviewer": "Reviewer_SV1s", "summary": "This paper shows that pretraining and tuning are necessary for consistent gains with pause tokens, which give transformers the ability to slightly decouple the amount of FLOPs from the number of tokens. The method inserts pause tokens at random positions during pretraining, and appends $M_{ft}$ and $M_{inf}$ pause tokens at fine-tuning and inference.\n\nExperiments show that pause tokens improve performance for a variety of tasks.\nAblations show that, surprisingly, more pause tokens is not always better.", "questions": "Typos\n1. Section 3.1: tokenresiding, the this\n2. 4.3: standard-pretraing\n\nComments\n1. A recent paper [1], that shows that extra scratch space improves the computational power of transformers, could help motivate this paper.\n\n[1] Merrill, W., & Sabharwal, A. (2023). The Expressive Power of Transformers with Chain of Thought. ArXiv, abs/2310.07923.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper successfully gets pause tokens working, which were previously thought not to work. The writing and presentation are clear, as are the experiments and results.", "weaknesses": "The paper is convincing and demonstrates some gains can be found via pause-token training. One weakness is that the method does not compare directly to Chain of Thought (CoT).\n\nIn contrast to pause tokens, which require both pretraining and fine-tuning, CoT is an inference-time only method that potentially requires extra human annotations. Another difference is that CoT has the ability to perform variable-length computation, as opposed to the fixed number of pause tokens added at inference time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wml6FzLMiq", "reviewer": "Reviewer_orpq", "summary": "This paper introduces an approach to the training and inference of language models by incorporating a concept of learnable \"pause tokens.\" Language models generate responses token-by-token, where each subsequent token is influenced by the preceding ones. However, this research explores the impact of allowing the model to manipulate more hidden vectors by adding a sequence of learnable \"pause\" tokens to the input prefix before outputting the answer. The key idea is to allow the model extra computation time before committing to an answer. Experiments were conducted on models with different parameters and various tasks, which suggest that incorporating delays during inference time shows gains when the model is pre-trained and fine-tuned with these delays. Notably, significant gains were observed in tasks like question-answering on the SQuAD dataset, CommonSenseQA, and reasoning tasks.", "questions": "1. How do you determine the number of “pause” tokens?\n2. How is the generalization ability of “pause” tokens, for example, pretraining with “pause” tokens and inference on tasks without fine-tuning? Are these tokens tasks specific?\n3. 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.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper offers an intriguing exploration into the behavior of language models by increasing computation at the level of hidden states, rather than at the text level, like Chain-of-Thought. It presents the concept of the language model \"thinking\" through the generation of intermediate hidden states. \n2. The introduction of a \"pause token\" is a novel approach that enables this deeper computational process to enhance the model’s performance on various tasks. There are a few useful observations. One is that pause tokens are necessary for both pre-training and fine-tuning stages. \n3. The comprehensive experimentation involving various models and tasks contributes to the robustness of the paper. A strong aspect of the paper is its detailed ablation studies, which effectively dissect the influence of the \"pause\" token, thereby providing valuable insights into delayed next-token prediction.", "weaknesses": "1. 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. \n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3KelV0MmaE", "reviewer": "Reviewer_zLe3", "summary": "The paper presents and analyzes Pause Training - a technique for training transformer language models while inserting virtual \"pause\" tokens between real tokens. The model is then trained normally, except that the training loss does not apply to predictions at pause tokens, essentially letting the model do whatever it wants as long as it predicts the next true token correctly. The intent behind pause training is to let the model do additional computation (\"thinking\") before predicting the next real token. Authors consider several scenarios where pause tokens are introduced during pretraining or finetuning, for a range of 130M-1B model sizes, and evaluate several alternatives such as introducing extra real tokens (e.g. dots). Overall, authors observe the greatest gains from pause training when it is introduced at both pretraining  and finetuning stage, and find significantly smaller effect for finetuning only. The paper reports several additional experiments, including an ablation analysis of how many pause tokens are needed. Finally, authors draw several parallels to modern prompting and PEFT methods such as chain of thought and prompt tuning.", "questions": "On page 7, authors found that the optimal number of pause tokens depends on the specific task, e.g. GSM8k or SQuAD. Did you try to train the model to dynamically decide the number of pause tokens? Essentially, allowing 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 traning)?\nThe reason I'm asking is that, since the baselines are able to do the same thing, it would be curious to check which approach (PT or ACT or CTC) is better at deciding where to put the pause tokens.\n\n\n\nMinor typo: authors use SquAD as a dataset name. To the best of my knowledge, the original paper[8] refers to this dataset as SQuAD, with uppercase Q.\n\n[8] https://arxiv.org/abs/1606.05250", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 1, "strengths": "1. The paper presents a comprehensive experimental analysis of a (relatively simple) phenomena. Authors consider multiple downstream tasks, multiple model sizes, and tries several baselines. Overall, I believe that the quality of experiments is one of the stronger sides of this paper.\n\n2. The paper focuses on a very simple algorithm, meaning that it can be easily implemented in most frameworks and libraries.\n\n3. The paper is generally well written , both conceptual explanations and experiments were easy to follow. Discussions and parallels drawn were interesting to read.", "weaknesses": "My main problem with the paper is that it ignores several very similar works from the early years of Transformer.\n\nProbably the closest idea to pause tokens is adaptive computation time (ACT) [1]. ACT was 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 (and I encourage them to), please note that there are *many* more studies on adaptive computation time and I only referenced a few of them from the top of my head. Furthermore, unlike Pause Tokens, ACT offers a way of teaching model to automatically decide how many extra turns it needs to solve a problem.\n\nWith that in mind, I believe that a proper analysis of Pause Training should at least compare it to a temporal (regular) ACT for the same model with equivalent tuning budget.\n\n* [1] https://arxiv.org/abs/1603.08983\n* [2] https://arxiv.org/abs/1807.03819\n* [3] https://arxiv.org/abs/2109.11745\n\nAnother related (but less similar) study is [4], that introduces additional learnable tokens for attention layers. The paper itself uses attention-only architecture, while subsequent works use the same idea for regular transformers, including a notable implementation[5]. To the best of my understanding, [4] proposes a mathematically similar idea that is likely to have similar effects to pause training, therefore there is no reason it shouldn't be a baseline for comparison. Additionally, [6-7] (and others) also propose to augment 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 best of my knowledge, they are no less related than PEFT methods discussed on page 9.\n\n* [4] https://arxiv.org/abs/1907.01470\n* [5] https://github.com/lucidrains/x-transformers#augmenting-self-attention-with-persistent-memory\n* [6] https://arxiv.org/abs/1911.00172\n* [7] https://arxiv.org/abs/1612.04426\n\n\nOverall, I believe that the paper presents an interesting idea which could offer a simpler and more efficient alternative to ACT-related methods, but it could also be inferior to them, and there is no way of knowing that from the experiments in the submitted draft. I would recommend authors to consider adressing the following questions for revised version of the paper:\n- how does PT compare to ACT variants in terms of quality?\n- how does PT compare to learned extra tokens (e.g. [4])?\n- which approach is better at deciding where to put extra tokens?\n- for modern language models, PT appears more efficient than a reasonably optimized ACT. How much faster is it in terms of training speed?\n- which method is better for very few vs very many pause tokens?\n(and similar)\n\nUnfortunately, in the current form, I believe that answering those questions would require authors to reformulate too much of the paper content and message, and therefore neccesitate a full additional review. Therefore, I currently recommend rejection, but I encourage authors to resubmit a revised version once they can properly analyze how PT compares to earlier methods.\nThere are two circumstances in which I could raise the score within this review cycle:\n\n1. authors prove that I misunderstood a significant chunk of the paper content and it is NOT in fact related to ACT family & others\n2. authors prove that all of those questions can be answered trivially, without running many additional experiments or changing much of the paper", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xne17xImrv", "reviewer": "Reviewer_cnas", "summary": "This paper presents a method to inject pause tokens into language models during pretraining and fine-tuning, such that the language model can spend more tokens/compute before outputting the final answer. Specifically, pause tokens are randomly injected into sequence during pretraining and appended to questions during fine-tuning. Extensive experiments on several datasets demonstrate the robust gain induced by this method.", "questions": "- How does one randomly inject pause token during pretraining? Could there be more than one pause token in a row during pretraining? Is there an upperbound on how many pause tokens can be in a row?\n- Have you tried to prompt LLM with filler words that typically represent pause or thinking? Like \"let me think a bit ....\" etc?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written and easy to follow.\n- The method is novel and interesting.\n- The experiments are extensive and demonstrate clear gains.", "weaknesses": "- It seems mysterious and problematic that sometimes more pause tokens can lead to worse performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)