# Segment-Based Attention Masking for GPTs

arXiv: https://arxiv.org/abs/2412.18487 · HTML: https://mlsanity.com/paper/2412.18487/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 26 references that are in the map.
- First posted on arXiv 24 Dec 2024.

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [DoRA: Weight-Decomposed Low-Rank Adaptation](https://mlsanity.com/paper/2402.09353/) (2024)
- [Unified Language Model Pre-training for Natural Language Understanding and Generation](https://mlsanity.com/paper/1905.03197/) (2019)

## Built on by (most cited first, 3 of 3)

- [AutoSOTA: An End-to-End Automated Research System for State-of-the-Art AI Model Discovery](https://mlsanity.com/paper/2604.05550/) (2026)
- [Back to Bytes: Revisiting Tokenization Through UTF-8](https://mlsanity.com/paper/2510.16987/) (2025)
- [Dual-objective Language Models: Training Efficiency Without Overfitting](https://mlsanity.com/paper/2512.14549/) (2025)

## Abstract (arXiv)

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.