PaperlayerView exact revision on arXiv

Source-linked AI summary

Think before you speak: Training Language Models With Pause Tokens

Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, Vaishnavh Nagarajan

arXiv:2310.02226v3cs.CLcs.AIcs.LG

TL;DR

The paper asks whether language models can use extra computation before committing to the next token. It implements this with learnable pause tokens during pretraining, finetuning and inference, finding broad gains when delays are learned at both training stages, while noting important scope and cost limitations.

  • Problem

    The paper examines whether the arbitrary limit of K per-layer operations can be relaxed when producing the (K+1)th token.

  • Method

    The approach appends learnable <pause> tokens and delays output extraction, incorporating pauses during pretraining, downstream finetuning and inference.

  • Results

    Training with pauses during both pretraining and finetuning outperforms standard training on eight tasks for the 1B model and six tasks for the 130M model.

  • Takeaways & Limitations

    Delayed next-token generation can benefit multiple downstream tasks when models learn to use pauses during both pretraining and finetuning.

  • Takeaways & Limitations

    The study does not establish generalization beyond 1B and 130M models, encoder-decoder architectures, or other pretraining mixtures and objectives, partly because pause-pretraining is expensive.

Abstract

from arXiv · show

Language models generate responses by producing a series of tokens in immediate succession: the $(K+1)^{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)^{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.

1 INTRODUCTION

The paper asks whether language models can benefit from extra computation before producing the next token, implementing delays with learnable pause tokens during training and inference. Gains are strongest when delays are used in both pretraining and finetuning, improving performance across most evaluated tasks.

  • Motivation: The paper challenges the constraint that producing the next token receives only as many per-layer operations as tokens seen so far.It asks whether some inputs require more than K operations to determine the (K+1)th token.
  • Approach: Pause-training appends one or more learnable <pause> tokens and delays output extraction until the final pause token.The delay lets the model process additional hidden vectors before generating its response.
  • Findings: 18% EM improvement appears on 1B-model SQuAD, alongside up to 8% gains on CommonSenseQA and a 1% GSM8k accuracy gain.The GSM8k result is reported relative to the standard model’s 7.5% accuracy.
  • Findings: Introducing pauses only during downstream finetuning produces fewer and milder gains, with clear performance drops in some instances.This contrasts with the broader benefits observed when pauses are introduced during both pretraining and finetuning.
  • Ablations: The ablations find that appending pauses is generally better than prepending them, each task has an optimal pause count, and reducing inference pauses degrades performance gracefully.The robustness to fewer pauses is not explicitly trained.

2 PRELIMINARIES

A causal decoder-only Transformer predicts the next token from hidden vectors computed for preceding tokens. Each layer maps a sequence of token representations through attention and feedforward modules, with the next-token distribution read from the final representation.

  • Next-token prediction: Given K input tokens, the model produces one intermediate vector per token in every Transformer layer before sampling the next token.The next-token distribution is inferred from the last vector in the final layer.
  • Layer computation: Each Transformer layer maps a matrix of K hidden vectors to another matrix of K hidden vectors.The layer operates on the representations corresponding to the input sequence.
  • Layer computation: The attention module combines key, value and query representations, followed by a feedforward module and layer normalization.These parameterized modules form the layer’s token-representation transformation.
  • Causal structure: The kth output representation is computed using exactly the k preceding hidden embeddings available at that layer.This causal structure limits each position to information from earlier positions.

3 PAUSE-TRAINING

Pause-training synthetically lengthens the input with a learnable pause token, delaying response extraction while preserving ordinary next-token prediction on informative outputs. The procedure applies distinct pause mechanisms to pretraining, finetuning and inference.

  • Core mechanism: Appending M dummy pause tokens gives the model M additional intermediate vectors before producing the next response.This increases the number of parallel computations available in each layer without substantially adding parameters.
  • Core mechanism: A single out-of-vocabulary <pause> token is repeated to create multi-token delays, avoiding confusion with ordinary punctuation.The pause token’s parameters are learned with the rest of the model vocabulary.
  • Pause-pretraining: During pause-pretraining, pauses are inserted at uniformly random locations and loss terms for predicting pause tokens are omitted.The model retains the standard next-token cross-entropy objective for non-pause positions.
  • Pause-finetuning: During pause-finetuning, Mft pauses are appended to each task prefix, outputs are ignored until the final pause, and loss is applied to the target sequence.The same pause count is fixed across inputs within each downstream task.
  • Pause-inference: During downstream inference, Minf pauses are appended to the prefix and response extraction begins only after the final pause token.Unless otherwise stated, the inference and finetuning pause counts are equal.
  • Experimental setup: The study compares standard and pause training across pretraining, finetuning and inference combinations.These combinations test whether delays must be learned at particular stages to affect inference-time performance.

4 EXPERIMENTS

The experiments evaluate pause-training across nine downstream tasks and show that gains depend strongly on introducing delays during pretraining and finetuning, especially for the 1B model.

  • Main results: 8 tasks: PausePT PauseFT outperforms StdPT StdFT for the 1B model, while the 130M model improves on 6 tasks.The improvements vary across tasks, and the 1B comparison excludes HellaSwag from the set of improved tasks.
  • Main results: 18% EM score: PausePT PauseFT improves over StdPT StdFT on SQuAD for the 1B model.CommonSenseQA gains reach 8%, while GSM8k accuracy rises from 7.5% to 8.5%.
  • Main results: About 5 benchmarks: StdPT PauseFT shows smaller gains, while other results match or fall below standard training.The figure caption specifically notes that finetuning-only delays can hurt GSM8k.
  • Main results: PausePT StdFT shows clear gains on only CoQA and PhysicalIQA, suggesting most gains in PausePT PauseFT involve delayed computations at inference.The paper separately attributes some benefits to improved representations from pause-pretraining.
  • Ablations: Adding periods as inference-time filler tokens produces no gains for standard models.This reproduces the reported observation that untrained inference-time delays do not help.
  • Interpretation: The authors conclude that pause-pretraining appears crucial for delays to help during downstream inference.They conjecture that standard pretraining biases models toward quick computation.
  • Experimental caveat: PausePT PauseFT has a computational advantage during finetuning and inference, while pause-pretraining sees only 90% as many meaningful tokens as standard pretraining.The comparison equalizes total tokens seen during pretraining by inserting dummy pause tokens.

5 ABLATIONS: WHERE AND HOW MANY <pause> TOKENS TO USE

The ablations show that pause placement and count matter: appended pauses are preferred, each dataset has an optimal finetuning count, and performance is robust to moderate inference-time shifts but not zero delay.

  • Pause count: Each downstream dataset has an optimal finetuning pause count, beyond which gains diminish.On GSM8k, 10 pauses are optimal and accuracy returns to baseline when the count reaches 50; SQuAD has a different optimum.
  • Inference robustness: PausePT PauseFT remains above baseline when inference uses half as many pauses as training.StdPT PauseFT is reported as even more robust to varying inference-time pause counts.
  • Inference robustness: Zero inference-time pauses can catastrophically break PausePT PauseFT performance, despite reasonable behavior with as few as 2 pauses.Designing models robust to zero-delay inference is left as future work.
  • Pause placement: Appending pauses remains more optimal than prepending them, although prepended pauses still improve over standard end-to-end training for PausePT PauseFT.The result suggests positional embeddings affect how readily the learned delays are used.

6 DISCUSSION AND KEY OPEN QUESTIONS

The discussion frames pause tokens as expanding the computation available to Transformers without simply increasing parameters, while identifying theoretical and practical limits and contrasts with chain-of-thought.

  • Computational interpretation: M pause tokens increase computational width from K to K+M parallel computations when producing the next token.The authors hypothesize that this wider pathway may support more precise and diverse lower-layer representations.
  • Computational interpretation: Pause tokens can relieve the gap between raw parameter capacity and implementation capacity, which is bottlenecked by the number of input tokens.The paper presents this as a key theoretical insight about why delays may unlock existing representational capacity.
  • Theory: If K is much larger than L, some tasks with L < N < K independent computations are implementable by a 2-layer Transformer if and only if it uses pause tokens.This is stated as an informal version of Theorem J.1.
  • Scaling question: Smaller models may benefit less because they may lack enough raw capacity to exploit the additional computational pathways.This conjecture aligns with the theorem’s assumption that parameter count K is sufficiently large.
  • Scaling question: At a fixed inference-time FLOPS budget, comparable gains from more heads or layers would require substantial parameter and FLOPS expansion.The authors present pause tokens as notable because they expand computation without the same scale of expansion in those resources.
  • Comparison with chain-of-thought: Chain-of-thought delays the final answer with meaningful tokens and additionally increases computational depth through autoregressive generation.With M reasoning tokens and L layers, the final token arises from roughly M · L sequentially composed operations.

7 RELATED WORK

Related work connects pause-training to dummy-token memory methods and adaptive-compute Transformers, while distinguishing its focus on delayed computation rather than memory or input-dependent computation.

  • Memory tokens: Prior memory-token work prepends dummy tokens during training and target-task inference, but reportedly yields minimal gains on smaller scratch-trained and pretrained BERT models.The paper describes this as echoing its mixed results for the comparable standard-pretrained, finetuning-delayed setup.
  • Adaptive computation: Adaptive-compute methods vary inference computation by input, whereas pause-inference uses the same number of pause tokens across inputs.The authors view pause-inference as a potential basis for future adaptive-compute methods rather than as adaptive compute itself.
  • Additional literature: The paper points readers to Appendix K for additional related work.

8 CONCLUSION, LIMITATIONS AND FUTURE WORK

Pause-training extends next-token prediction by training models to use appended dummy pause tokens for additional inference-time computation. The paper finds benefits across tasks when pauses are used during pretraining and finetuning, while identifying broad scope and theoretical limitations.

  • Conclusion: Pause-training trains models with dummy <pause> tokens so they can harness additional inference-time computation.The approach extends immediate next-token prediction by delaying output extraction until the pause sequence is processed.
  • Conclusion: Training with <pause> tokens during both pretraining and downstream finetuning improves performance across a variety of tasks.The supported conclusion is conditional on using delays in both training stages.
  • Limitations: Pause-training is not claimed to benefit every downstream task because each task has an optimal number of <pause> tokens, which may be zero.The paper therefore treats the delay count as task-dependent rather than universally beneficial.
  • Limitations and future work: The expense of pause-pretraining limits evaluation across additional model sizes, encoder-decoder architectures, and alternative pretraining mixtures or objectives.The paper also leaves rigorous theoretical understanding and several algorithmic extensions for future work.

B ADDITIONAL DOWNSTREAM FINETUNING RESULTS

On the 130M model, PausePT PauseFT outperforms standard training on six of nine downstream tasks, with gains across reasoning, general understanding, and fact recall but not SQuAD.

  • Downstream performance: Six of nine tasks improve with PausePT PauseFT over StdPT StdFT on the 130M decoder-only model.Reported gains occur on GSM8k, CommonSenseQA, LAMBADA, WebQA, NaturalQA, and minimally PhysicalIQA; SQuAD does not improve.
  • Downstream performance: PausePT PauseFT clearly outperforms StdPT StdFT on GSM8k, CommonSenseQA, LAMBADA, WebQA, and NaturalQA.These include reasoning, general understanding, and fact-recall tasks.
  • Downstream performance: SQuAD gains observed for the 1B model do not appear for the 130M model.The result contrasts with the larger-model outcome reported elsewhere in the paper.

C PREPENDING VS APPENDING PAUSE TOKENS

The additional results compare pause-training configurations and model sizes, showing broad but task- and scale-dependent benefits for PausePT PauseFT over standard training.

  • Prepending versus appending: Prepending pause tokens still beats StdPT StdFT for PausePT PauseFT, but is slightly worse than appending them on some benchmarks including GSM8k and SQuAD.For StdPT PauseFT, prepending and appending produce equal numbers of wins and losses.

D ZERO-SHOT EVALUATION

Pause-pretrained models show some zero-shot gains with pause tokens, while downstream results indicate task-specific optimal delays and graceful tolerance to several inference-time delay shifts.

  • Zero-shot evaluation: Zero-shot pause inference gives gains on tasks such as GSM8k and HellaSwag for pause-pretrained models.The comparison evaluates 0, 10, and 50 appended pause tokens against a standard pretrained 1B model.
  • Zero-shot evaluation: Zero-shot accuracies remain quite low because the experiments use a small 1B parameter model.The paper qualifies the reported zero-shot gains with this absolute-performance caveat.
  • Inference-time delay shifts: Pause-trained models degrade gracefully when inference-time pauses differ from finetuning delays, across Minf ∈[5, 25] in the tested settings.An exception is GSM8k with Mft = 50, and PausePT PauseFT drops drastically when inference uses no delay.
  • Finetuning delay: Downstream finetuning typically has an optimal number of pause tokens that depends on the task.The number of finetuning pauses is denoted Mft.

F ROBUSTNESS TO VARYING NUMBER OF INFERENCE TIME PAUSES

Pause-trained models tolerate moderate changes in the number of inference-time pause tokens, with performance degrading gracefully rather than abruptly. Removing the delay entirely, however, causes a drastic performance drop.

  • Performance degrades gracefully when pause-trained models use different numbers of inference-time pauses than during finetuning.This robustness persists even when the number of tokens seen is halved.
  • Halving the number of inference-time pause tokens still produces graceful performance degradation in pause-trained models.
  • No inference-time delay causes a drastic performance drop in pause-trained models.

G DOWNSTREAM DATASET DESCRIPTION

The paper evaluates pause-training across six downstream datasets spanning reasoning, question answering, understanding, and fact recall, while also analyzing its computational and theoretical implications. These analyses compare pause tokens with conventional capacity expansion and characterize when delayed computation can overcome input-length bottlenecks.

  • Downstream datasets: The evaluation includes GSM8k, a reasoning dataset containing 8.5k grade school math word problems.
  • Downstream datasets: SQuAD V1 evaluates reading comprehension using Wikipedia-based questions.
  • Downstream datasets: CommonSenseQA tests commonsense knowledge through decoding, with results reported using Exact Match scores.
  • Downstream datasets: LAMBADA tests last-word prediction from long contexts, while Web Questions evaluates recall of commonly asked web questions.
  • Downstream datasets: PhysicalIQA evaluates physical commonsense reasoning and the ability to understand interactions with the world.
  • Efficiency analysis: Pause-inference is presented as more FLOPS- and wall-clock-efficient than adding layers or attention heads, though CoT remains stronger in performance.For SQuAD, 10 pause tokens produced an 18% EM increase, while an equivalent FLOPS budget allowed at most two additional layers.
Loading 2310.02226v3…