Causale aandacht met vooruitziende sleutels
Causal Attention with Lookahead Keys
September 9, 2025
Auteurs: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu
cs.AI
Samenvatting
In standaard causale aandacht zijn de query, key en value (QKV) van elk token statisch en coderen ze alleen de voorafgaande context. We introduceren CAuSal aTtention with Lookahead kEys (CASTLE), een aandachtmechanisme dat de keys van elk token voortdurend bijwerkt naarmate de context zich ontvouwt. We noemen deze bijgewerkte keys lookahead keys omdat ze tot eerdere posities behoren, maar informatie integreren van tokens die later verschijnen ten opzichte van die posities, terwijl het autoregressieve eigenschap strikt behouden blijft. Hoewel het mechanisme sequentieel lijkt, leiden we een wiskundige equivalentie af die expliciet materialiseren van lookahead keys op elke positie vermijdt en efficiënte parallelle training mogelijk maakt. Op taalmodelleringsbenchmarks presteert CASTLE consistent beter dan standaard causale aandacht over verschillende modelschalen, waarbij de validatieperplexiteit wordt verlaagd en de prestaties op een reeks downstream taken worden verbeterd.
English
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.