给GPT的预填充阶段引入分段注意力掩码,提升生成效率
Segment-Based Attention Masking for GPTs
- 按输入段落分块处理,预填充时允许段内非因果访问
- 在Llama和Qwen上实现顶尖性能,无额外计算开销
- 适合追求推理速度与质量的大型语言模型应用
现代语言模型的成功很大程度上依赖于掩码因果注意力,这是生成式预训练变换器(GPT)模型的核心。尽管GPT能一次性处理完整用户提示,但因果掩码仍逐令牌施加,模拟生成过程。这在初始“预填充”阶段造成不必要的约束,此时模型正在处理输入提示并生成内部表示,尚未产生输出令牌。本文提出在预填充阶段基于已知的块结构进行注意力掩码,之后再采用传统的逐令牌自回归过程。例如,在典型聊天提示中,系统提示被视为一个块,用户提示为下一个块。每个块作为掩码单位,使得每个块的第一个令牌可非因果地访问后续令牌。随后,模型答案以传统因果方式生成。这种分段式方案不增加额外计算开销。将其集成到Llama和Qwen等模型中,均实现了最先进的性能。
原文摘要 · Abstract (English)
Modern Language Models (LMs) owe much of their success to masked causal attention, the backbone of Generative Pre-Trained Transformer (GPT) models. Although GPTs can process the entire user prompt at once, the causal masking is applied to all input tokens step-by-step, mimicking the generation process. This imposes an unnecessary constraint during the initial "prefill" phase when the model processes the input prompt and generates the internal representations before producing any output tokens. In this work, attention is masked based on the known block structure at the prefill phase, followed by the conventional token-by-token autoregressive process after that. For example, in a typical chat prompt, the system prompt is treated as one block, and the user prompt as the next one. Each of these is treated as a unit for the purpose of masking, such that the first tokens in each block can access the subsequent tokens in a non-causal manner. Then, the model answer is generated in the conventional causal manner. This Segment-by-Segment scheme entails no additional computational overhead. When integrating it into models such as Llama and Qwen, state-of-the-art performance is consistently achieved.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。