Каузальное внимание с опережающими ключами
Causal Attention with Lookahead Keys
September 9, 2025
Авторы: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu
cs.AI
Аннотация
В стандартном причинно-следственном внимании запрос, ключ и значение (QKV) каждого токена остаются статичными и кодируют только предшествующий контекст. Мы представляем CAuSal aTtention with Lookahead kEys (CASTLE) — механизм внимания, который постоянно обновляет ключи каждого токена по мере разворачивания контекста. Мы называем эти обновленные ключи "опережающими ключами", поскольку они принадлежат более ранним позициям, но интегрируют информацию из токенов, появляющихся позже относительно этих позиций, при этом строго сохраняя авторегрессивное свойство. Хотя механизм кажется последовательным, мы выводим математическую эквивалентность, которая позволяет избежать явного материализации опережающих ключей на каждой позиции и обеспечивает эффективное параллельное обучение. На бенчмарках языкового моделирования CASTLE стабильно превосходит стандартное причинно-следственное внимание на различных масштабах моделей, снижая перплексию на валидации и улучшая производительность на ряде последующих задач.
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.