Think before you speak: Training Language Models With Pause Tokens

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Goyal, Sachin, Ji, Ziwei, Rawat, Ankit Singh, Menon, Aditya Krishna, Kumar, Sanjiv, Nagarajan, Vaishnavh
Format: Preprint
Published: 2023
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913322829348864
author Goyal, Sachin
Ji, Ziwei
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
Nagarajan, Vaishnavh
author_facet Goyal, Sachin
Ji, Ziwei
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
Nagarajan, Vaishnavh
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2310_02226
institution arXiv
publishDate 2023
record_format arxiv
spellingShingle Think before you speak: Training Language Models With Pause Tokens
Goyal, Sachin
Ji, Ziwei
Rawat, Ankit Singh
Menon, Aditya Krishna
Kumar, Sanjiv
Nagarajan, Vaishnavh
Computation and Language
Artificial Intelligence
Machine Learning
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.
title Think before you speak: Training Language Models With Pause Tokens
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2310.02226