Kausale Aufmerksamkeit mit vorausschauenden Schlüsseln

Causal Attention with Lookahead Keys

September 9, 2025

papers.authors: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu

cs.AI

papers.abstract

Bei der standardmäßigen kausalen Aufmerksamkeit sind die Query-, Key- und Value-Vektoren (QKV) jedes Tokens statisch und kodieren nur den vorhergehenden Kontext. Wir stellen CAuSal aTtention with Lookahead kEys (CASTLE) vor, einen Aufmerksamkeitsmechanismus, der die Keys jedes Tokens kontinuierlich aktualisiert, während sich der Kontext entfaltet. Wir bezeichnen diese aktualisierten Keys als Lookahead-Keys, da sie zu früheren Positionen gehören, aber Informationen von Tokens integrieren, die relativ zu diesen Positionen später erscheinen, wobei die autoregressive Eigenschaft strikt erhalten bleibt. Obwohl der Mechanismus sequenziell erscheint, leiten wir eine mathematische Äquivalenz ab, die das explizite Materialisieren von Lookahead-Keys an jeder Position vermeidet und ein effizientes paralleles Training ermöglicht. Bei Sprachmodellierungs-Benchmarks übertrifft CASTLE durchweg die standardmäßige kausale Aufmerksamkeit über verschiedene Modellgrößen hinweg, reduziert die Validierungsperplexität und verbessert die Leistung bei einer Reihe von Downstream-Aufgaben.

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.