Attention Mechanism

Механизм внимания — это операция нейронной сети, которая формирует каждый выход как взвешенное среднее входных признаков. Веса вычисляются на основе того, насколько каждый вход важен для обрабатываемого элемента. Релевантность оценивается сравнением вектора запроса с векторами ключей, а полученные оценки задают веса соответствующих векторов значений. Внимание представили в 2014 году для машинного перевода: благодаря ему декодер рекуррентной нейронной сети (RNN) мог обращаться к каждому слову исходного текста, а не сжимать всё предложение в один вектор фиксированной длины. В архитектуре Transformer, представленной в 2017 году, внимание стало основой всей архитектуры. Сегодня оно лежит в основе больших языковых моделей, зрительных трансформеров и многих моделей обнаружения объектов.

Как работает механизм внимания#

Внимание использует три обучаемые линейные проекции входных данных. Запрос (Q) описывает, что ищет текущий токен, каждый ключ (K) — что предлагает токен, а каждое значение (V) несёт передаваемое содержимое. Модель сравнивает запрос со всеми ключами, преобразует результаты сравнения в веса и возвращает взвешенную сумму значений. Когда модель машинного перевода переводит «The black cat» как «noir», запрос для этого шага получает наивысшую оценку соответствия с ключом для слова «black», хотя во французском языке прилагательное стоит после существительного.

Многоголовое внимание выполняет несколько таких вычислений параллельно — каждое на собственных проекциях меньшей размерности, — а затем объединяет результаты. Благодаря этому разные головы могут специализироваться на разных отношениях внутри одного слоя.

Внимание с масштабированным скалярным произведением#

Внимание в Transformer записывается как Attention(Q, K, V) = softmax(QKᵀ / √d_k) V. Каждый запрос сравнивается со всеми ключами с помощью скалярного произведения, а полученные оценки делятся на квадратный корень из размерности ключа d_k, чтобы они не росли вместе с размером векторов и не приводили softmax к насыщению. Softmax преобразует оценки каждого запроса в веса, сумма которых равна 1; результат — взвешенная сумма векторов значений.

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# Запросы, ключи и значения для последовательности из 4 токенов размерности 8 каждый
q, k, v = torch.randn(3, 4, 8).unbind(0)

# Сравни каждый запрос со всеми ключами, раздели оценки на sqrt(d_k) и нормализуй каждую строку с помощью softmax
weights = (q @ k.T / k.shape[-1] ** 0.5).softmax(dim=-1)

# Каждый выходной токен — это взвешенное среднее векторов значений
output = weights @ v

print("Attention weights of token 0:", [round(w, 3) for w in weights[0].tolist()])
print("Matches PyTorch SDPA:", torch.allclose(output, F.scaled_dot_product_attention(q, k, v), atol=1e-6))

Код выводит [0.121, 0.741, 0.106, 0.032] и True: токен 0 получает 74% своего результата из вектора значений токена 1, а сумма четырёх весов равна 1. Результат ручного вычисления совпадает с результатом объединённой операции scaled_dot_product_attention в PyTorch, которая на поддерживаемых GPU использует эффективные по памяти ядра, например FlashAttention.

Типы внимания#

• Самовнимание: Запросы, ключи и значения берутся из одной последовательности, поэтому каждый токен получает контекст от всех остальных. Самовнимание используется в кодировщиках Transformer и ViT.
• Перекрёстное внимание: Запросы берутся из одной последовательности, а ключи и значения — из другой. Декодер перевода обращается к исходному предложению, диффузионная модель для преобразования текста в изображение — к запросу, а запросы объектов декодера DETR — к признакам изображения.
• Каузальное внимание: маска не позволяет каждому токену учитывать последующие токены, поэтому языковая модель генерирует текст слева направо, а ключи и значения предыдущих токенов можно хранить в KV-кэше.
• Эффективное внимание: разреженное внимание, внимание со скользящим окном и линейное внимание снижают квадратичные затраты на сравнение каждого токена со всеми остальными, а деформируемое внимание выбирает лишь несколько обучаемых точек вокруг каждой опорной позиции.
• Внимание к каналам в CNN: сети Squeeze-and-Excitation перевзвешивают каналы карты признаков на основе её глобального среднего. Это недорогой вариант внимания, который лёг в основу заявки-победителя в классификации на ILSVRC 2017.

Механизм внимания и свертка#

Свёртка применяет один и тот же небольшой фильтр в каждой позиции, поэтому каждый выходной элемент видит фиксированное локальное окно, а веса не зависят от изображения. Внимание вычисляет веса с учётом содержимого и может связывать любые две позиции в пределах одного слоя. За это приходится платить затратами, растущими пропорционально квадрату числа токенов. Изображение 640×640 при шаге 32 даёт сетку 20×20, то есть 400 токенов и 160 000 пар «запрос — ключ»; сетка того же изображения при шаге 8 содержит 6 400 токенов и почти 41 миллион пар. В YOLO26 блоки внимания размещены с учётом этих затрат: блок C2PSA в конце бэкбона и один блок C3k2 с поддержкой внимания в шее выполняют многоголовое самовнимание на карте признаков 20×20 с шагом 32, а во всех слоях с более высоким разрешением используются свёртки.