룩어헤드 키를 활용한 인과적 어텐션

초록

표준 인과적 어텐션(causal attention)에서 각 토큰의 쿼리(query), 키(key), 값(value)(QKV)는 정적이며 선행 컨텍스트만을 인코딩합니다. 우리는 컨텍스트가 전개됨에 따라 각 토큰의 키를 지속적으로 업데이트하는 어텐션 메커니즘인 CAuSal aTtention with Lookahead kEys(CASTLE)를 소개합니다. 우리는 이러한 업데이트된 키를 룩어헤드 키(lookahead keys)라고 명명하는데, 이는 이전 위치에 속하지만 해당 위치에 상대적으로 나중에 나타나는 토큰들의 정보를 통합하면서도 엄격하게 자기회귀적 속성을 보존하기 때문입니다. 이 메커니즘이 순차적으로 보이지만, 우리는 각 위치에서 룩어헤드 키를 명시적으로 구현하지 않고도 효율적인 병렬 학습을 가능하게 하는 수학적 동등성을 도출했습니다. 언어 모델링 벤치마크에서 CASTLE은 모델 규모에 걸쳐 표준 인과적 어텐션을 지속적으로 능가하며, 검증 퍼플렉서티(validation perplexity)를 감소시키고 다양한 다운스트림 작업에서 성능을 향상시켰습니다.

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.

룩어헤드 키를 활용한 인과적 어텐션

Causal Attention with Lookahead Keys

초록

Support