Attention Sinks

Токен-сток — это токен, на который приходится большая доля весов внимания во многих головках и слоях, хотя его содержимое не относится к запросу. Чаще всего это первый токен последовательности. Guangxuan Xiao и его коллеги из MIT, Meta AI и Университета Карнеги — Меллона дали этому эффекту название в опубликованной на ICLR 2024 работе Efficient Streaming Language Models with Attention Sinks. Они показали, что большие языковые модели (LLM), например Llama 2, MPT, Falcon и Pythia, начинают работать нестабильно, когда скользящий KV-кэш удаляет первые токены, а сохранение четырех таких токенов обеспечивает стабильную генерацию.

Почему образуются стоки внимания#

Механизм self-attention преобразует оценки для каждого запроса в веса с помощью softmax, и сумма этих весов всегда равна единице. Если в контексте нет токена, важного для конкретной головки, вес все равно должен куда-то направиться. В авторегрессионной языковой модели первые токены видны всем последующим, поэтому во время обучения они становятся стандартным местом, куда направляются веса, не нужные головке. При анализе Llama 2 с помощью StreamingLLM выяснилось, что во всех головках каждого слоя, кроме двух нижних, внимание концентрировалось на начальных токенах. Этот эффект демонстрирует логит стока, добавленный к softmax:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Без стока головка распределяет весь свой вес между шестью неважными значениями и выдает результат с нормой 0,83. При логите стока, равном 3, сток получает 77% веса, а норма результата уменьшается до 0,19 — для этого запроса головка почти ничего не делает. В моделях gpt-oss с открытыми весами от OpenAI эта функция встроена: для каждой головки внимания обучается логит стока, в который можно направить неиспользуемый вес, не влияющий на результат.

Потоковая генерация с StreamingLLM#

Модель, генерирующая текст бесконечно, со временем выходит за пределы любого контекстного окна и не может хранить каждый токен в своем KV-кэше. Простое оконное внимание, сохраняющее только последние токены, перестает работать, как только из окна выпадают первые токены. StreamingLLM постоянно хранит ключи и значения четырех начальных токенов рядом с обновляемым окном последних токенов и назначает позиции по месту в кэше, а не в исходном тексте. При этом вращательные позиционные встраивания (RoPE) применяются заново на каждом шаге:

Этот метод обеспечивал стабильное моделирование языка на последовательностях длиной от 4 миллионов токенов и работал до 22,2 раза быстрее скользящего окна, которое пересчитывает кэш для каждого нового токена. Окно по-прежнему ограничивает объем информации, которую модель может вспомнить: удаленные из середины токены исчезают навсегда, поэтому потоковый чат-бот или ИИ-агент помнит только последние сообщения. Авторы также предварительно обучили модели с одним специальным обучаемым токеном-заполнителем в начале каждого примера, который самостоятельно выполнял роль стока.

Стоки внимания и смежные понятия#

• Внимание со скользящим окном: Ограничивает контекст каждого токена соседними недавними токенами. Если применить его во время инференса к модели, обученной с полным вниманием, она перестанет работать, когда стоки покинут окно.
• Деградация контекста: Снижение точности по мере роста промпта, пока он все еще помещается в контекстное окно.

Стоки в моделях компьютерного зрения#

Визуальные трансформеры (ViT) демонстрируют похожее поведение. В работе Vision Transformers Need Registers исследователи из Meta и Inria обнаружили токены с необычно большими нормами на фоновых участках с небольшим объемом информации в контролируемых и самообучающихся ViT. Эти токены использовались для внутренних вычислений. Добавление нескольких дополнительных токенов-регистров во входные данные позволило взять на себя эту роль, удалить артефакты из карт признаков и карт внимания, а также улучшить плотное предсказание.

Модель SAM 2 от Meta, работающая в пакете Ultralytics, распределяет видеопамять по похожему принципу: хранит опорные кадры вместе со скользящим окном. Для каждого нового кадра ее модуль внимания к памяти считывает записи банка памяти, соответствующие кадрам с запросами пользователя и сохраняемые на протяжении всего видео, а также записи для шести последних кадров. Поэтому стоимость внимания на кадр остается постоянной независимо от продолжительности видео.