Self-Attention и Multi-Head Attention

Самовнимание: токены смотрят друг на друга

В прошлом уроке запросы шли от одной последовательности (слова перевода), а ключи и значения — от другой (слова оригинала). В самовнимании (self-attention) всё берётся из одной последовательности: каждый токен одновременно

• задаёт вопрос — свой запрос q_i: «какая информация мне нужна?»;
• предлагает себя — свой ключ k_i: «вот что у меня есть»;
• отдаёт содержимое — своё значение v_i: «вот что я передам тому, кто на меня посмотрит».

Все три вектора получаются из эмбеддинга токена умножением на обучаемые матрицы:

Здесь X — матрица (T, d): эмбеддинги T токенов по строкам. Выход — тоже T векторов, но теперь каждый из них знает о контексте: вектор «ключа» в предложении про родник подмешал информацию от «бил» и «холодный».

import numpy as np
import vis

rng = np.random.default_rng(0)
tokens = ["кот", "сел", "на", "тёплый", "коврик"]
T, d, dk = len(tokens), 8, 4
X = rng.normal(size=(T, d))                         # эмбеддинги (пока случайные)
Wq, Wk, Wv = (rng.normal(size=(d, dk)) for _ in range(3))

Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = Q @ K.T / np.sqrt(dk)
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)
out = weights @ V
print("Q, K, V:", Q.shape, K.shape, V.shape, " веса:", weights.shape, " выход:", out.shape)
vis.heatmap(weights, xlabels=tokens, ylabels=tokens, title="Веса самовнимания (случайные матрицы)")

Матрицы здесь случайные, поэтому и веса бессмысленные. В обученной модели W_Q, W_K, W_V подбираются градиентным спуском так, чтобы каждый токен смотрел туда, откуда ему полезно взять информацию.

Зачем три разные матрицы

Почему бы не сравнивать сами эмбеддинги: \text{softmax}(XX^\top)X? Тогда каждый токен больше всего смотрел бы на себя (скалярное произведение вектора с собой обычно самое большое), а сходство было бы симметричным: если «кот» смотрит на «сел», то и «сел» на «кот» с той же оценкой. А нужно другое: глаголу полезно найти подлежащее, прилагательному — существительное, к которому оно относится. «Что я ищу» и «что я предлагаю» — разные вещи. Отдельные W_Q и W_K позволяют выучить такие несимметричные связи, а W_V — решить, какую именно информацию передавать.

Самовнимание не видит порядок

Переставим токены — что станет с выходом?

import numpy as np
import nanotorch as nt

rng = np.random.default_rng(0)
X = rng.normal(size=(1, 5, 8))                      # (батч, T, d)
nt.manual_seed(0)
mha = nt.MultiHeadAttention(d_model=8, num_heads=1)
out = mha(nt.tensor(X)).data[0]

perm = [3, 0, 4, 1, 2]
out_perm = mha(nt.tensor(X[:, perm])).data[0]
print("выход для переставленного входа == переставленный выход:", np.allclose(out_perm, out[perm]))

Каждый токен получил ровно тот же вектор, что и раньше, — просто на другом месте. Самовнимание работает с множеством токенов: веса зависят только от содержимого векторов, а не от их позиций. «Собака кусает человека» и «человек кусает собаку» для него неразличимы. Поэтому к эмбеддингам добавляют позиционные кодировки — это первое, что мы сделаем в следующем уроке.

Цена: квадратичная сложность

Матрица весов имеет размер T \times T: каждый токен сравнивается с каждым. Для 1000 токенов — миллион весов на каждую голову каждого слоя, для 100 000 токенов — 10 миллиардов. Время и память растут квадратично с длиной текста. Это главное ограничение трансформеров, и на борьбу с ним придумано множество приёмов: эффективные реализации вроде FlashAttention, разреженное и локальное внимание, кэширование ключей и значений при генерации.

Зато у самовнимания есть огромный плюс: путь между любыми двумя токенами — один шаг, как бы далеко они ни стояли. И все позиции обрабатываются параллельно — это идеально подходит для видеокарт.

Многоголовое внимание

Одна голова внимания даёт для каждого токена одно распределение весов — один «взгляд». Но связей в тексте много сразу: синтаксис (кто что делает), кореференция («он» → «кот»), соседние слова, тема. Multi-head attention запускает h голов параллельно, у каждой — свои W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} меньшей размерности d_{\text{head}} = d_{\text{model}} / h:

Выходы голов склеиваются обратно в вектор размера d_{\text{model}}, а матрица W_O смешивает информацию от разных голов.

На практике все головы считаются одним матричным умножением: делают одну большую W_Q размера d_{\text{model}} \times d_{\text{model}}, а результат разрезают на h частей по столбцам. Поэтому многоголовое внимание стоит столько же, сколько одноголовое той же ширины.

import inspect
import numpy as np
import nanotorch as nt

nt.manual_seed(0)
mha = nt.MultiHeadAttention(d_model=16, num_heads=4)
x = nt.tensor(np.random.default_rng(0).normal(size=(2, 6, 16)))    # батч из 2 последовательностей по 6 токенов
out = mha(x)
print(mha)
print("вход:", x.shape, " выход:", out.shape)
print("веса внимания (батч, головы, T, T):", mha.attn.shape)
print("параметров:", mha.num_params(), "= 4 · (16·16 + 16)")
print()
print(inspect.getsource(nt.MultiHeadAttention.forward))
print(inspect.getsource(nt.MultiHeadAttention._split))

В _split тензор формы (N, T, d_{\text{model}}) превращается в (N, h, T, d_{\text{head}}): последняя ось режется на головы, и ось голов переезжает вперёд. Дальше scaled_dot_product_attention работает со всеми головами сразу — матричное умножение @ в NumPy применяется к последним двум осям, а остальные считаются «батчем».

Итог

• Self-attention: Q = XW_Q, K = XW_K, V = XW_V из одной последовательности; каждый токен собирает информацию из всех, и выход — контекстные векторы.
• Три разные проекции позволяют выучить несимметричные связи «что ищу» и «что предлагаю».
• Самовнимание эквивариантно к перестановкам — о порядке токенов нужно сообщить отдельно.
• Сложность O(T^2) по длине; зато любые два токена связаны за один шаг, и всё считается параллельно.
• Multi-head: h голов размерности d_{\text{model}}/h смотрят на разные связи, их выходы склеиваются и умножаются на W_O; параметров 4d_{\text{model}}^2 (+ bias).