Atención Causal con Claves de Anticipación
Causal Attention with Lookahead Keys
September 9, 2025
Autores: Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu
cs.AI
Resumen
En la atención causal estándar, la consulta, clave y valor (QKV) de cada token son estáticos y codifican únicamente el contexto precedente. Introducimos CAuSal aTtention with Lookahead kEys (CASTLE), un mecanismo de atención que actualiza continuamente las claves de cada token a medida que se desarrolla el contexto. Denominamos a estas claves actualizadas como claves de anticipación (lookahead keys) porque pertenecen a posiciones anteriores pero integran información de tokens que aparecen más tarde en relación con esas posiciones, preservando estrictamente la propiedad autoregresiva. Aunque el mecanismo parece secuencial, derivamos una equivalencia matemática que evita materializar explícitamente las claves de anticipación en cada posición y permite un entrenamiento paralelo eficiente. En benchmarks de modelado de lenguaje, CASTLE supera consistentemente a la atención causal estándar en diferentes escalas de modelos, reduciendo la perplejidad de validación y mejorando el rendimiento en una variedad de tareas posteriores.
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.