Causal Attention with Lookahead Keys

2509.07301v1 cs.CL, cs.LG 2025-09-11
Авторы:

Zhuoqing Song, Peng Sun, Huizhuo Yuan, Quanquan Gu

Резюме на русском

## Контекст Область глубокого обучения становится все более важной в современном искусственном интеллекте. Одна из ключевых задач в этой области — понимание и моделирование зависимостей между элементами данных. Касуальное внимание (causal attention) широко используется для моделирования последовательных данных, таких как тексты или последовательности событий. Однако существующие модели имеют ограничения: они используют только предшествующий контекст для вычисления ключей (keys) и значений (values), что ограничивает модели в своей способности предсказывать будущие ситуации. Таким образом, требуется развитие методов, которые могут использовать информацию из будущих позиций для улучшения моделей. ## Метод Мы предлагаем **CASTLE** (Causal Attention with Lookahead Keys), новый метод для касуального внимания, который позволяет использовать не только предшествующую информацию, но также информацию из будущих позиций. Основной идеей является формирование "поисковых ключей" (lookahead keys) для каждой позиции, которые не только зависят от предшествующих токенов, но также учитывают информацию из будущих токенов. Для реализации этого метода мы предлагаем математическую модель, которая позволяет обновлять ключи в порядке встречи, не материализуя полностью все ключи заранее. Это позволяет эффективно использовать параллельное вычисление и сохранить авторегрессионную структуру. ## Результаты Мы проверили эффективность CASTLE на стандартных задачах моделирования языка. Модели с CASTLE показали значительные улучшения по сравнению с классическим касуальным вниманием: - **Перфекция (perplexity)**: CASTLE показал низкий уровень перфекции на текстовых данных, что указывает на более точное понимание языка. - **Задачи по применению**: Модели с CASTLE показали лучшие результаты в нескольких задачах, таких как классификация и суммирование текста. Эти результаты демонстрируют, что использование поисковых ключей улучшает способность моделей предсказывать будущие ситуации и обрабатывать более сложные задачи. ## Значимость CASTLE может применяться в различных областях, где необходимо понимание последовательных данных, например: - **Текстовые модели**: Улучшение предсказания текста и обобщения информации. - **Генерация текста**: Более точное генерирование текста в реальном времени. - **Обработка естественного языка**: Извлечение информации из текстов с учетом контекста. Преимущество CASTLE заключается в его способности обрабатывать более сложные задачи, где необходимо предсказывать будущие ситуации на основе предыдущих и будущих данных. Это может привести к переломным улучшениям

Abstract

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.

Ссылки и действия