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 performance gains, like an 18% increase on the SQuAD QA task.
• The paper shows that delaying token generation allows Transformer models to perform extra computations, enhancing both memory and attention mechanisms.
• The paper highlights that tuning the number of pause tokens per task is crucial for achieving optimal performance in various downstream applications.

An Expert Analysis: Think Before You Speak - Training LLMs with Pause Tokens

The paper "Think before you speak: Training LLMs With Pause Tokens" introduces an innovative approach to training Transformer-based LLMs by incorporating pause tokens during training and inference. This technique diverges from the norm of generating tokens in immediate succession and instead suggests delaying this process to enhance the model's computational performance.

Overview

The central hypothesis of the paper is that the conventional method of producing tokens in immediate succession may constrain the computational depth of Transformer models. By adding learnable pause tokens in the input sequence, the model can utilize these tokens to process extra computations before outputting the next token. The concept is evaluated empirically using decoder-only models with sizes of 1B and 130M parameters, trained on the C4 dataset and finetuned on a range of downstream tasks.

Key Findings

1. Inference-Time Gains: The paper finds substantial improvements in performance when both pre-training and finetuning processes involve pause tokens. For instance, the 1B-parameter model shows an 18% increase in exact match (EM) score on the SQuAD QA task, an 8% increase on CommonSenseQA, and a 1% increase in accuracy on GSM8k.
2. Impact of Finetuning Alone: Introducing pause tokens only during finetuning yielded mixed results, showing benefits in fewer instances and, at times, even degrading performance.
3. Pretraining Alone Is Not Sufficient: Pause-pretrained models without pause during finetuning did not consistently offer improvements, indicating the need for delays during both stages for meaningful gains.
4. Optimal Number of Pauses: Each downstream task appears to have an optimal number of pause tokens, suggesting a nuanced dependency of performance gains on the specific configuration of pauses.

Theoretical Insights

The incorporation of pause tokens allows for a fundamental alteration in the computational pathway of the Transformer layers. The pause tokens serve as a mechanism to increase the width of computations, enabling the model to leverage additional computational steps before outputting the next token. This hypothesis is supported theoretically:

Main Theoretical Result

Given a Transformer architecture, the self-attention mechanism within each layer has a representational capacity that is bounded by the number of parameters rather than input length. The introduction of pause tokens helps in utilizing this representational capacity more effectively by offering additional computational steps. Therefore, tasks that require a larger number of parallel operations than the input tokens can be better modeled with pause tokens.

Practical Implications

1. Adaptive Compute: The technique provides a pathway for adaptive compute, wherein the inference time can be adjusted based on the system’s computational limitations, making it practical for deployment in scenarios with varying computational budgets.
2. Memory and Attention: The method offers a potential framework for enhancing memory and attention mechanisms in Transformers by introducing additional computational steps without adding significant parameters.

Future Directions

Several promising avenues for future research are identified:

• Improved Robustness: Designing models that are robust to zero-delay scenarios and varying numbers of pause tokens during inference is crucial.
• Applicability Across Models: Extending this approach to more diverse architectures, such as encoder-decoder models and larger parameter models, would test the generalizability of the findings.
• Theoretical Exploration: Further theoretical examination is required to formalize the distinction between a Transformer's raw and implementational capacities, particularly in contexts where more complex computational patterns are required.

Conclusion

The introduction of pause tokens in training Transformer-based LLMs represents a notable shift in enhancing computational capabilities. By strategically delaying token generation, the model can leverage additional parallel computational steps, thereby improving performance on various downstream tasks. This paper opens up new paradigms for future explorations in the structure and computational pathways of LLMs.