Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2509.07301 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912614673547264 |
|---|---|
| author | Song, Zhuoqing Sun, Peng Yuan, Huizhuo Gu, Quanquan |
| author_facet | Song, Zhuoqing Sun, Peng Yuan, Huizhuo Gu, Quanquan |
| contents | In standard causal attention, each token's query, key, and value (QKV) are static and encode only preceding context. We introduce CAuSal aTtention with Lookahead kEys (CASTLE), an attention mechanism that continually updates each token's keys as the context unfolds. We term these updated keys lookahead keys because they belong to earlier positions yet integrate information from tokens that appear later relative to those positions, while strictly preserving the autoregressive property. Although the mechanism appears sequential, we derive a mathematical equivalence that avoids explicitly materializing lookahead keys at each position and enables efficient parallel training. On language modeling benchmarks, CASTLE consistently outperforms standard causal attention across model scales, reducing validation perplexity and improving performance on a range of downstream tasks. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2509_07301 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Causal Attention with Lookahead Keys Song, Zhuoqing Sun, Peng Yuan, Huizhuo Gu, Quanquan Computation and Language Machine Learning In standard causal attention, each token's query, key, and value (QKV) are static and encode only preceding context. We introduce CAuSal aTtention with Lookahead kEys (CASTLE), an attention mechanism that continually updates each token's keys as the context unfolds. We term these updated keys lookahead keys because they belong to earlier positions yet integrate information from tokens that appear later relative to those positions, while strictly preserving the autoregressive property. Although the mechanism appears sequential, we derive a mathematical equivalence that avoids explicitly materializing lookahead keys at each position and enables efficient parallel training. On language modeling benchmarks, CASTLE consistently outperforms standard causal attention across model scales, reducing validation perplexity and improving performance on a range of downstream tasks. |
| title | Causal Attention with Lookahead Keys |
| topic | Computation and Language Machine Learning |
| url | https://arxiv.org/abs/2509.07301 |