Saved in:
Bibliographic Details
Main Authors: Katz, Shahar, Ringel, Liran, Romano, Yaniv, Wolf, Lior
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2412.18487
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910762036887552
author Katz, Shahar
Ringel, Liran
Romano, Yaniv
Wolf, Lior
author_facet Katz, Shahar
Ringel, Liran
Romano, Yaniv
Wolf, Lior
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2412_18487
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Segment-Based Attention Masking for GPTs
Katz, Shahar
Ringel, Liran
Romano, Yaniv
Wolf, Lior
Computation and Language
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.
title Segment-Based Attention Masking for GPTs
topic Computation and Language
url https://arxiv.org/abs/2412.18487