Что такое attention sink и почему он возникает в длинных контекстах?

Краткий тезис

Attention sink — это феномен, при котором модель Transformer уделяет непропорционально много внимания первым нескольким токенам последовательности (особенно токену начала строки BOS). Причина кроется в свойствах softmax: он не может обнулить малые значения, поэтому модель «сбрасывает» избыточное внимание на всегда присутствующие tokens|начальные токены. Это приводит к эффекту «lost in the middle» — информация из середины контекста забывается. Для борьбы используют sliding window attention, attention|gated attention и модификации нормализации.

1. Термин: Attention sink (сток внимания)

Attention sink — это термин, введённый в статье Efficient Streaming Language Models with Attention Sinks (2023). Он описывает ситуацию, когда модель концентрирует значительную долю внимания на первых токенах последовательности, даже если они не несут полезной информации. Эти токены становятся «раковиной» (sink), в которую стекается избыточное внимание.

Почему это проблема

• Модель тратит ёмкость внимания на фиксированные токены, а не на содержательные части контекста.
• При удлинении контекста внимание к середине падает, что ухудшает качество ответов на вопросы, требующие информации из середины.

2. Причина: свойства softmax

Softmax — функция активации, превращающая произвольные числа (логиты) в распределение вероятностей:

[ [text](/wiki/text){softmax}(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}} ]

Ключевое свойство: сумма всех вероятностей равна 1, и ни одно значение не может быть строго нулевым (если только логит не равен (-\infty)). Даже для нерелевантных токенов softmax выдаёт маленькое, но ненулевое значение.

Почему это порождает sink

• В длинных последовательностях модель должна распределить внимание между сотнями и тысячами токенов.
• Если какой-то токен не важен, модель всё равно вынуждена дать ему немного внимания, чтобы сумма сошлась.
• Первые токены (особенно BOS) присутствуют всегда, поэтому модель «учится» сбрасывать на них этот остаточный вес. Это дешёвый способ удовлетворить ограничение softmax.

3. Механизм возникновения

Рассмотрим процесс на примере decoder-only трансформера (например, LLaMA, GPT).

1. Начало последовательности добавляется специальный токен <s> (BOS).
2. Обучение модель вычисляет attention scores для всех пар токенов. Для токенов из середины контекста scores могут быть низкими, но softmax всё равно распределяет ненулевые веса.
3. Формирование sink модель быстро обнаруживает, что проще всего «сбросить» лишнее внимание на токен <s>, так как он всегда на первой позиции и не меняется.
4. Усиление в процессе обучения градиенты закрепляют это поведение — attention к началу растёт, а к середине падает.

Визуализация attention карты (heatmap) для длинного контекста:

• Первый столбец (или строка) ярко-жёлтый — высокие веса на BOS.
• Остальные токены — бледные, с небольшими всплесками на локально важных позициях.

4. Последствия: lost in the middle

Lost in the middle — эффект, при котором модель хуже использует информацию из середины длинного контекста, чем из начала или конца. Attention sink — одна из причин этого явления.

Экспериментальные данные (из статьи Liu et al., 2023):

• При длине контекста > 1K токенов точность ответов на вопросы из середины падает на 10–20% по сравнению с вопросами из начала.
• Attention sink объясняет до 40% этого падения (остальное — позиционное кодирование и кэш KV).

Таблица влияния позиции на recall

Данные условны, иллюстрируют тренд.

5. Экспериментальные наблюдения

В статье Efficient Streaming Language Models with Attention Sinks авторы показали:

• Attention sink возникает во всех обученных трансформерах, независимо от архитектуры (GPT-2, LLaMA, Falcon).
• Величина sink растёт с увеличением длины контекста — модель вынуждена всё больше «сбрасывать» внимание на начальные токены.
• Удаление BOS токена не решает проблему: sink переходит на следующий токен (например, первое слово предложения).
• Sink можно измерить как долю внимания, приходящуюся на первые 4 токена. В длинных контекстах она может достигать 20–30% от всего распределения.

6. Методы борьбы с attention sink

6.1 Sliding window attention

Идея ограничить поле внимания окном фиксированного размера (например, 2048 токенов). Токены за пределами окна не участвуют в вычислении attention.

Плюсы

• Устраняет sink, так как начальные токены могут не попасть в окно.
• Линейная сложность по длине контекста (O(n) вместо O(n²)).

Минусы

• Теряется способность «видеть» далёкие зависимости.
• Для задач, где важны глобальные связи (например, суммаризация), может ухудшить качество.

Реализация (псевдокод):

def sliding_window_attention(query, key, value, window_size):
    # query, key, value: [batch, heads, seq_len, dim]
    seq_len = query.size(-2)
    attn_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)  # causal
    # sliding window mask: только диагональные полосы шириной window_size
    window_mask = (torch.arange(seq_len).unsqueeze(0) - torch.arange(seq_len).unsqueeze(1)).abs() > window_size
    mask = attn_mask | window_mask
    attn_scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(dim)
    attn_scores = attn_scores.masked_fill(mask, float('-inf'))
    attn_weights = torch.softmax(attn_scores, dim=-1)
    return torch.matmul(attn_weights, value)

6.2 Gated attention

Идея добавить обучаемый «вентиль» (gate), который регулирует, сколько внимания может быть направлено на каждый токен. Вентиль может обнулять внимание к sink-токенам.

Формула
[ [text](/wiki/text){Attention}(Q,K,V) = [text](/wiki/text){softmax}\left(\frac{QK^T}{\sqrt{d}} \odot G\right) V ]
где (G) — матрица вентилей, получаемая из дополнительного линейного слоя.

Плюсы

• Гибкость: модель сама учится, какие токены игнорировать.
• Не требует изменения архитектуры внимания.

Минусы

• Дополнительные параметры и вычислительные затраты.
• Сложнее обучать.

6.3 Нормализация внимания (attention normalization)

Идея модифицировать softmax так, чтобы он мог выдавать нулевые веса для нерелевантных токенов. Например, использовать sparse softmax (с порогом) или ReLU attention (замена softmax на ReLU + нормализация).

Пример ReLU attention
[ [text](/wiki/text){Attention}(Q,K,V) = \frac{[text](/wiki/text){ReLU}(QK^T)}{\sum [text](/wiki/text){ReLU}(QK^T)} V ]
ReLU обнуляет отрицательные логиты, что уменьшает sink.

Плюсы

• Естественное обнуление нерелевантных токенов.
• Простота реализации.

Минусы

• Может ухудшить сходимость обучения.
• Не все архитектуры поддерживают.

7. Связь с архитектурой Transformer

Attention sink — не баг, а следствие softmax-нормализации в механизме внимания. В любой архитектуре, где используется softmax (оригинальный Transformer, GPT, BERT, LLaMA), при достаточно длинном контексте возникает sink.

Влияние на RAG

• В RAG-системах контекст часто состоит из множества чанков. Attention sink может привести к тому, что модель будет игнорировать чанки из середины, даже если они релевантны.
• Это усугубляет проблему lost in the middle в RAG.

Влияние на Agentic RAG

• Агенты, которые обрабатывают длинные истории диалогов или несколько раундов retrieval, особенно уязвимы. Sink может «съедать» внимание к важным деталям из предыдущих шагов.

8. Пет-проект для закрепления

Задача Реализовать мини-трансформер с attention sink и без, сравнить перплексию на длинных текстах.

Инструменты

• Python, PyTorch
• Библиотека transformers (для baseline)
• Датасет: wikitext-2 (длинные последовательности)

Шаги:

1. Реализуйте простой decoder-only трансформер с одним слоем self-attention (softmax).
2. Обучите на коротких последовательностях (128 токенов) — sink незначителен.
3. Протестируйте на длинных последовательностях (2048 токенов) — измерьте attention weights на первых 4 токенах.
4. Добавьте sliding window attention (окно 512) и повторите тест.
5. Сравните perplexity на тестовых данных с длинным контекстом.

Ожидаемый результат

• В модели без окна attention sink составляет >15% от всего внимания, perplexity выше.
• С sliding window sink падает до <5%, perplexity улучшается на 5–10%.

Код для визуализации attention sink

import torch
import matplotlib.pyplot as plt

def plot_attention_sink(attn_weights, seq_len):
    # attn_weights: [batch, heads, seq_len, seq_len]
    sink_weights = attn_weights[:, :, :, :4].mean(dim=-1).mean(dim=(0,1))  # среднее по головам и батчу
    plt.plot(sink_weights.cpu().numpy())
    plt.xlabel("Position")
    plt.ylabel("Avg attention to first 4 tokens")
    plt.title("Attention sink visualization")
    plt.show()

9. Связь с другими вопросами

10. Навигация

• Предыдущий: 627
• Следующий: 629
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 627
• Следующий: 629
• Индекс: 00. Индекс разборов