Causal Attention with Lookahead Keys
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.
68c0e08a3912ed54cf54325d
Zhuoqing Song
68c0e08a3912ed54cf54325e
Peng Sun
68c0e08a3912ed54cf54325f
Huizhuo Yuan
68c0e08a3912ed54cf543260
Quanquan Gu
Causal Attention with Lookahead Keys
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.
CASTLE, an attention mechanism that updates keys with future context while maintaining autoregressive properties, outperforms standard causal attention in language modeling.
causal attention QKV lookahead keys autoregressive property parallel training validation perplexity downstream tasks