Attention Causale avec Clés d'Anticipation
Causal Attention with Lookahead Keys
September 9, 2025
papers.authors: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu
cs.AI
papers.abstract
Dans l'attention causale standard, les requêtes, clés et valeurs (QKV) de chaque token sont statiques et n'encodent que le contexte précédent. Nous introduisons l'attention CAuSal avec clés prospectives (CASTLE), un mécanisme d'attention qui met continuellement à jour les clés de chaque token à mesure que le contexte se déroule. Nous appelons ces clés mises à jour des clés prospectives car elles appartiennent à des positions antérieures tout en intégrant des informations provenant de tokens qui apparaissent plus tard par rapport à ces positions, tout en préservant strictement la propriété autorégressive. Bien que le mécanisme semble séquentiel, nous dérivons une équivalence mathématique qui évite de matérialiser explicitement les clés prospectives à chaque position et permet un entraînement parallèle efficace. Sur les benchmarks de modélisation du langage, CASTLE surpasse systématiquement l'attention causale standard à différentes échelles de modèles, réduisant la perplexité de validation et améliorant les performances sur une gamme de tâches en aval.
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.