免费暂停令牌

中文摘要 AI 辅助

免费暂停令牌为语言模型提供额外计算以形成每个下一个令牌预测（如同暂停或思考令牌的作用），但该计算通过权重共享的主干在并行预测流中承载，而非作为序列中的额外令牌。在10亿参数模型上的实践中，它将下一个令牌预测的性能提升了2-3个centinats。由于暂停令牌利用现有位置而非新增位置，因此使用它是“免费”的：推理时它不增加上下文长度、不增加KV缓存，且推理浮点运算增长带来的延迟几乎可忽略，因为这并非吞吐量的活跃瓶颈。唯一的主要成本在训练阶段，与优化的预训练流水线相比，额外训练计算可低至1.14倍，同时保留大部分收益。其结果是在浮点运算量、参数数量和令牌数量均与标准下一个令牌训练的Transformer一致的情况下实现了性能提升。

英文摘要

State--prediction separation (SPS) relieves a language model's hidden state of two competing burdens---summarizing the context and predicting the next token---by splitting the forward pass into a state stream and a prediction stream. The separation works, but it is expensive: the prediction stream is a second pass over the whole backbone, costing $\sim$1.9$\times$ the pretraining FLOPs, and even more in terms of wall-clock time when using a flexible attention mask. This paper makes state--prediction separation almost free. We take the separation to its limit with a free pause token: a prediction stream that writes no keys or values at all and so rides the sequence's existing positions. It improves next-token prediction of a standard Transformer by 2-3 centinats in practice on a 1B parameter model, and because it adds no position it costs nothing at inference---no added context length, no KV cache, no decode steps, and essentially no latency, with the growth in inference flops typically irrelevant as it is not the active bottleneck on throughput. The cost is therefore entirely in training where we use four mechanisms to drive it down: a two-pass split that keeps FlashAttention kernels viable, the $w{=}0$ prediction window, a shared gated FFN that evaluates one FFN per position rather than one per stream, and phasing the separation onto the tail of the run. Together these bring the overhead versus an optimized pretraining pipeline to $1.33\times$ wall-clock while recovering ~94% of the gain compared to SPS, and to as low as $1.09\times$ along a graceful quality/compute tradeoff. Furthermore, the FFN optimization reduces the raw flops required at inference time. The result is an isoflop, isoparameter, and isotoken improvement over standard next token trained transformers.

发表机构

• Microsoft（微软）
• Cornell University（康奈尔大学）

机构由 AI 辅助整理，请以论文原文为准。