Think before you speak: Training Language Models With Pause Tokens (2310.02226v3)

Abstract: LLMs 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 LLMs 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.

• The paper demonstrates that incorporating pause tokens during both pretraining and finetuning yields significant gains, including an 18% increase in SQuAD EM.
• The methodology leverages a delay mechanism to expand the model's computational width by allowing extra attention operations beyond the input length.
• The study emphasizes that an optimal number of pause tokens is crucial, as excessive or missing tokens can degrade performance or reduce robustness.

Training LLMs with Pause Tokens: A Paradigm for Delayed Next-Token Prediction

Introduction and Motivation

This work introduces a novel approach to LLM (LM) training and inference by incorporating "pause tokens"—special, learnable tokens that are appended to the input sequence to intentionally delay the model's output generation. The central hypothesis is that the standard causal Transformer architecture imposes an arbitrary constraint: the number of per-layer operations available to compute the next token is limited by the number of tokens seen so far. By appending $M$ pause tokens, the model is afforded $K+M$ hidden vectors per layer (for $K$ input tokens), potentially enabling richer intermediate representations and more expressive computation before outputting the next token.

The method is operationalized by introducing a unique <pause> token, appended multiple times to the input during both pretraining and finetuning. The model is trained to ignore outputs corresponding to these tokens, only producing meaningful outputs after the final pause. This approach is evaluated on decoder-only models of 1B and 130M parameters, pretrained on C4 and finetuned on a diverse set of downstream tasks.

Methodology: Pause-Training and Inference

Pause-Token Mechanism

The pause-training paradigm consists of three stages:

1. Pause-Pretraining: Insert $M_{\rm pt}$ <pause> tokens at random positions in the pretraining sequence. The model is trained with the standard next-token prediction loss, but loss terms for predicting <pause> tokens are omitted. This ensures the model learns to utilize the additional computation without being distracted by the need to predict the pause tokens themselves.
2. Pause-Finetuning: For downstream tasks, append $M_{\rm ft}$ <pause> tokens to the input prefix. The model is trained to predict the target sequence only after the last pause token, again ignoring outputs for the pause tokens.
3. Pause-Inference: At inference, append $M_{\rm inf}$ <pause> tokens to the input, and ignore model outputs until the last pause token is seen.

This approach is compared against several baselines: standard pretraining and finetuning (StdPT), pause-finetuning only (StdPT), pause-pretraining only (PausePT), and pause-pretraining plus pause-finetuning (PausePT).

(Figure 1)

Figure 1: Downstream performance for a 1B model. Injecting delays in both stages of training (PausePT) outperforms standard end-to-end training (StdPT) on a wide variety of tasks (except HellaSwag). In contrast, introducing delays only in the finetuning stage provides only lukewarm gains, and even hurts in GSM8k.

Empirical Results

Main Findings

• PausePT (Pause-Pretraining + Pause-Finetuning) yields consistent improvements across a range of tasks for the 1B model, with the most notable gains being an 18% increase in SQuAD EM, 8% in CommonSenseQA, and 1% in GSM8k accuracy over the standard baseline.
• Pause-finetuning alone (StdPT) provides only mild or inconsistent gains, and in some cases, degrades performance.
• Pause-pretraining alone (PausePT) without downstream delay offers limited benefits, indicating that both pretraining and finetuning with pauses are necessary to realize the full advantage.
• Filler tokens (e.g., periods) as delays do not confer benefits, corroborating prior findings that models must be explicitly trained to utilize such delays.

(Figure 2)

Figure 2: Downstream performance of pause-training on a 130M decoder-only model. On six out of nine tasks, PausePT outperforms StdPT, but gains are less pronounced than for the 1B model.

Ablation Studies

• Optimal Number of Pauses: Each downstream task has an optimal $M_{\rm ft}$; excessive pauses can degrade performance, likely due to overwhelming the self-attention mechanism.
• Robustness to Inference-Time Delay Mismatch: Pause-trained models degrade gracefully when the number of inference-time pauses differs from training, but performance collapses if no pauses are provided at inference.
• Appending vs. Prepending: Appending pauses is generally superior, but prepending still outperforms the baseline, suggesting positional biases induced by pause-pretraining.

  Figure 3: Varying finetuning delay: There exists an optimal number of <pause> tokens for each downstream dataset; gains diminish or reverse beyond this point.

  

Figure 4: Zero-shot evaluation of pause-pretrained models. Zero-shot inference with pause tokens gives gains on some tasks, but absolute accuracies remain low for small models.

Theoretical Analysis

The authors formalize the intuition that pause tokens increase the "computational width" of the Transformer. In standard inference, the number of parallel operations per layer is limited by the input length $K$. By appending $M$ pause tokens, the model can perform $K+M$ parallel computations per layer, potentially enabling it to solve tasks that require more independent operations than the input length allows.

A key theoretical result (Theorem~\ref{thm:pause}) demonstrates that, under reasonable assumptions about the representational capacity of the attention-feedforward block, there exist tasks that a 2-layer Transformer can solve if and only if it uses pause tokens. This is because the number of distinct operations the model can implement is bottlenecked by the input length in standard inference, but can be expanded to match the parameter count of the attention module with pause tokens.

Computational and Practical Considerations

• Parameter Efficiency: The addition of a single learnable pause token increases the parameter count negligibly (e.g., 1024 parameters for a 1B model).
• FLOPS and Latency: Pause tokens increase the number of attention operations per layer, but do not add sequential depth. Thus, with sufficient parallelism, wall-clock overhead is minimal compared to adding layers or attention heads.
• Comparison to Chain-of-Thought (CoT): While both pause-inference and CoT increase computational width, CoT also increases computational depth via autoregressive generation of intermediate steps. CoT does not require pretraining modifications, whereas pause-inference does.

Limitations and Open Questions

• Generalization: Gains are not universal; some tasks (e.g., HellaSwag) do not benefit, and the optimal number of pauses is task-dependent.
• Accessibility: The requirement for pause-pretraining limits immediate applicability to existing pretrained models.
• Robustness: Pause-trained models are not robust to zero-delay inference; performance collapses if no pauses are provided.
• Scaling: Preliminary results suggest that larger models benefit more from pause tokens, contrary to the hypothesis that smaller models would benefit due to increased effective capacity.
• Future Directions: Open questions include developing methods to make pause-training effective for standard pretrained models, determining the optimal number of pauses adaptively, and extending the approach to other architectures and pretraining objectives.

Implications and Future Developments

The pause-training paradigm challenges the conventional design of causal LLMs by decoupling the number of per-layer operations from the input length. This has both theoretical and practical implications:

• Theoretical: It motivates a re-examination of the relationship between model capacity, computational width, and input sequence length in Transformer architectures.
• Practical: It suggests a new axis for model improvement—computational width via input manipulation—distinct from parameter scaling or architectural changes.

Potential future developments include adaptive pause mechanisms (varying the number of pauses per input), integration with adaptive compute methods, and exploration of pause tokens in encoder-decoder or multi-modal architectures.

Conclusion

Pause-training introduces a simple yet effective modification to the Transformer training and inference pipeline by leveraging learnable pause tokens to expand computational width. Empirical results demonstrate significant gains on a variety of tasks when pauses are introduced during both pretraining and finetuning. Theoretical analysis supports the intuition that pause tokens unlock otherwise inaccessible representational capacity. While the approach has limitations—most notably, the need for pause-pretraining and lack of robustness to zero-delay inference—it opens a promising direction for both the mechanistic understanding and practical enhancement of LLMs.