Attenzione Causale con Chiavi di Anticipazione

Causal Attention with Lookahead Keys

September 9, 2025

Autori: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu

cs.AI

Abstract

Nell'attenzione causale standard, la query, la chiave e il valore (QKV) di ciascun token sono statici e codificano solo il contesto precedente. Introduciamo CAuSal aTtention with Lookahead kEys (CASTLE), un meccanismo di attenzione che aggiorna continuamente le chiavi di ciascun token man mano che il contesto si sviluppa. Definiamo queste chiavi aggiornate come chiavi di lookahead perché appartengono a posizioni precedenti ma integrano informazioni provenienti da token che appaiono successivamente rispetto a quelle posizioni, preservando rigorosamente la proprietà autoregressiva. Sebbene il meccanismo appaia sequenziale, deriviamo un'equivalenza matematica che evita di materializzare esplicitamente le chiavi di lookahead in ciascuna posizione e consente un addestramento parallelo efficiente. Su benchmark di modellazione del linguaggio, CASTLE supera costantemente l'attenzione causale standard su diverse scale di modelli, riducendo la perplessità di validazione e migliorando le prestazioni su una gamma di task downstream.

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.