Внимание

«У Мии была новая кукла. Мия всюду брала её с собой: в парк, в магазин, к бабушке. А вечером Мия уложила спать…» Кого? Любой ребёнок ответит: куклу. Но слово «кукла» прозвучало два предложения назад, и чтобы его вспомнить, нужно оглядываться на весь рассказ, а не только на последние несколько слов.

Модель из прошлой главы так не умеет: она видит ровно восемь последних токенов, и всё, что левее, для неё не существует. В этой главе мы научим Ростка оглядываться. Для этого есть механизм внимания (attention) — главное изобретение, на котором стоят все современные языковые модели, от GPT до Llama. Ничего сверхъестественного в нём нет: скалярные произведения, softmax и взвешенное среднее, и всё это вам уже знакомо по прошлым главам.

Окну не хватает памяти

Вспомним, как работала сеть из главы 6. Она брала 8 последних токенов, заменяла каждый его эмбеддингом из 48 чисел, склеивала их в один вектор длины $8 \times 48 = 384$ и пропускала через скрытый слой. У такого устройства три беды.

Первая: всё, что за окном, потеряно. Если имя героя прозвучало в начале истории, через пару предложений оно выпадает из окна — и модель уже никак не может его использовать. Дело не в том, что она плохо обучена: у неё физически нет входа, через который это слово могло бы попасть внутрь. Проверьте сами:

Вторая: окно дорого расширять. Первый слой хранит отдельный вес для каждого числа на каждой позиции окна — и так для каждого скрытого нейрона. При окне в 8 токенов и скрытом слое из 384 нейронов это $8 \times 48 \times 384 = 147\,456$ весов. Растянем окно до 512 токенов, как у Ростка, и получим $512 \times 48 \times 384 \approx 9{,}4$ миллиона весов в одном только первом слое, причём большинство из них почти ничему не научится: сигнал до них будет доходить редко.

Третья: нет общего опыта. Слово «мяч» на третьей позиции окна и на седьмой проходит через разные веса. Всё, что сеть узнала про мяч в одном месте окна, ей приходится выучивать заново для каждого другого места.

Значит, нам нужен механизм, который (1) может дотянуться до любого из прошлых токенов, как бы далеко он ни был; (2) обрабатывает все позиции одними и теми же весами; (3) сам решает, куда смотреть, исходя из содержания текста. Такой механизм и называется вниманием.

Куда смотрит Росток

Прежде чем выводить формулы, посмотрим на готовый результат. У полного Ростка 8 слоёв, и в каждом работают 6 независимых «голов» внимания. Для каждого токена голова решает, сколько взять у каждого из предыдущих токенов и у него самого, — то есть раздаёт им веса. Веса неотрицательны и в сумме дают единицу: это снова распределение вероятностей, как в первых главах, только не над словарём, а над прошлыми позициями текста.

Поиграйте с переключателями и присмотритесь к трём вещам:

• Будущее закрыто. Какой бы токен вы ни выбрали, всё правее него заштриховано: модель не может смотреть вперёд. Почему — разберём ниже, это важно.
• Разные слои смотрят по-разному. Сравните слой 1 и слой 5 на одном и том же токене. Поищите головы, которые почти всегда смотрят на соседа слева, и головы, которые дотягиваются далеко назад. В примере про «he» выберите слой 6 и голову 4: эта голова отдаёт 82% внимания «he» слову «Tom».
• Много внимания уходит на ◆. Часто голова отдаёт заметную долю метке начала — даже когда та ничего не значит. Эту странность мы объясним, когда увидим softmax в деле.

Запрос, ключ и значение

Представьте класс на уроке. Слово «she» встаёт и спрашивает: «Кто тут девочка?» Это его запрос (query). У каждого предыдущего слова на парте стоит табличка: «я имя девочки», «я цвет», «я то, что можно бросить». Это ключи (keys). Чем лучше табличка отвечает на вопрос, тем внимательнее «she» слушает это слово. А слушает оно значение (value) — то, что слово готово рассказать о себе. В итоге «she» уносит смесь рассказов, и больше всего в ней от тех слов, чьи таблички подошли лучше других.

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

Теперь то же самое на языке векторов. У каждого токена есть вектор $x$ — его эмбеддинг (или то, во что он превратился в предыдущих слоях). Из него тремя обучаемыми матрицами получаются три новых вектора:

$$q = x\,W_Q, \qquad k = x\,W_K, \qquad v = x\,W_V.$$

Это три обычных линейных слоя из главы 5. Дальше — всего три шага. Пусть токен с запросом $q$ смотрит на токены $1, \dots, t$ с ключами $k_j$ и значениями $v_j$.

1. Оценки. Похожесть запроса на каждый ключ измеряем скалярным произведением: $s_j = q \cdot k_j$. Оно велико, если векторы смотрят в одну сторону, близко к нулю, если они перпендикулярны, и отрицательно, если смотрят в разные стороны.
2. Веса. Оценки превращаем в доли знакомым softmax из главы 3: $a_j = e^{s_j} / \sum_i e^{s_i}$. Все $a_j$ положительны и в сумме дают 1.
3. Смесь. Выход — взвешенное среднее значений: $o = \sum_j a_j\, v_j$.

Вот и всё внимание. Попробуйте сами на плоскости: здесь запросы и ключи двумерные, чтобы их можно было нарисовать, а значения трёхмерные, и мы рисуем их цветом (цвет на экране — это как раз три числа: красный, зелёный и синий).

Три опыта, которые стоит поставить:

• Нулевой запрос. Все оценки равны нулю, и веса становятся одинаковыми: по 25%. Токен, которому «нечего спросить», просто усредняет всех.
• Длинный запрос. Направление то же, но вектор длиннее — оценки растут, и softmax становится резче: почти всё внимание забирает один ключ. Длина запроса работает как обратная температура из главы 3. Ползунок «Резкость» делает то же самое напрямую.
• Все ключи против запроса. Нажмите «Никто не подходит»: ключи соберутся в одной половине плоскости, а запрос отвернётся от них всех, так что все оценки станут отрицательными. Веса всё равно в сумме дают 1: побеждает наименее плохой ключ. Внимание не умеет «никуда не смотреть».

Последний опыт разгадывает загадку метки ◆ из прошлого раздела. Голове, которой на этом шаге нечего искать, всё равно нужно куда-то деть все 100% внимания. Удобнее всего сливать их в один и тот же безобидный токен, значение которого почти ничего не добавляет к смеси. Метка начала подходит для этого идеально: она есть в каждом тексте и всегда стоит на одном месте.

Вот та же арифметика на Python — с теми же числами, что и в начальном положении виджета:

import numpy as np # four earlier tokens: each has a key (how it can be found) # and a value (what it tells); values are colours, i.e. 3 numbers tokens = ["girl", "ball", "red", "ran"] K = np.array([[ 1.6, 0.9], [-1.2, 1.3], [-1.5, -0.6], [ 0.4, -1.6]]) V = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.5, 0.5, 0.5]]) q = np.array([1.4, 0.7]) # the query of "she" scores = K @ q # one dot product per earlier token weights = np.exp(scores - scores.max()) weights /= weights.sum() # softmax: positive, adds up to 1 out = weights @ V # the weighted average of the values for tok, s, w in zip(tokens, scores, weights): print(f"{tok:5s} score {s:+.2f} weight {w:.3f}") print("output:", out.round(3))

Токен сравнил свой запрос с четырьмя ключами, и все оценки вышли отрицательными: −3, −5, −6, −9. Что даст softmax?

Все веса будут равны нулю: подходящих токенов нет Веса всё равно сложатся в 1, и больше всего достанется токену с оценкой −3 Веса станут отрицательными, и токен «оттолкнёт» эти значения

Softmax сравнивает оценки только друг с другом: $e^{-3}$ намного больше, чем $e^{-9}$, а после деления на сумму доли всегда положительны и дают 1. Поэтому головам и нужен токен-«отстойник» вроде метки начала.

Зачем делить на √d

В виджете векторы двумерные. У Ростка запрос и ключ каждой головы состоят из 64 чисел, и скалярное произведение — это сумма 64 произведений. Чем больше слагаемых, тем сильнее разброс суммы.

Чем это плохо? Вспомните главу 3: если логиты большие, softmax становится почти «жёстким» — один вес близок к 1, остальные к 0. А у такого softmax почти нулевой градиент: если чуть-чуть подвинуть оценки, веса почти не изменятся. Обучение такого внимания буксует с первых же шагов. Лекарство простое: поделить оценки на $\sqrt d$, чтобы их разброс снова стал около единицы:

$$s_j = \frac{q \cdot k_j}{\sqrt{d}}.$$

Не подглядывать в будущее

Как модель учится? Мы даём ей кусок текста из $T$ токенов и просим на каждой позиции угадать следующий токен — все $T$ предсказаний сразу, а ошибка — среднее по всем позициям (глава 2). Это очень выгодно: один проход по тексту даёт $T$ обучающих примеров.

Но тут есть ловушка. Если токен на позиции $t$ может смотреть на позицию $t+1$, он видит ровно тот ответ, который должен угадать. Модель быстро научится списывать — и окажется бесполезной, когда начнёт писать текст сама: будущего тогда ещё нет, списывать неоткуда. Поэтому перед softmax всем оценкам «из будущего» ставят $-\infty$:

$$s_{tj} = -\infty \quad \text{при } j > t, \qquad e^{-\infty} = 0.$$

Такие позиции получают вес ровно 0, а оставшиеся веса строки, как и положено, дают в сумме 1. Это называется причинной маской (causal mask): каждый токен видит себя и прошлое, но не будущее. Соберите её сами:

Нижняя часть виджета показывает цену внимания. Каждый токен сравнивает себя с каждым предыдущим, так что сравнений примерно $T^2/2$. Удвоили длину текста — работы стало вчетверо больше. На полном контексте Ростка в 512 токенов одна голова считает 262 144 оценки (почти половину из них тут же закрывает маска), а все 48 голов модели — больше 12 миллионов за один проход. Именно поэтому длинный контекст у больших моделей — дорогое удовольствие.

Все токены сразу

До сих пор мы смотрели на один токен. Но считать их по одному расточительно: запросы, ключи и значения всех $T$ токенов можно сложить в матрицы $Q$, $K$ и $V$ (по строке на токен) и получить все оценки одним умножением. Элемент $(t, j)$ матрицы $QK^\top$ — это как раз $q_t \cdot k_j$. Вся глава умещается в одну формулу:

$$\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt d} + M\right) V,$$

где $M$ — маска: нули на диагонали и под ней, $-\infty$ над ней, а softmax берётся по каждой строке отдельно. Размеры: $Q$, $K$, $V$ — это $T \times d$, таблица оценок — $T \times T$, результат — снова $T \times d$, по вектору на токен.

Почему это важно? Из главы 5 вы помните, что видеокарты (и процессоры тоже) лучше всего умеют умножать большие матрицы. Внимание в матричной форме — это два умножения матриц и softmax, то есть ровно то, что железо делает быстрее всего. Во многом поэтому трансформеры и вытеснили рекуррентные сети, которые читали текст строго по одному слову.

Несколько голов

Один запрос — один вопрос. Но токену обычно нужно узнать сразу несколько вещей: кто тут главный герой, какое слово стояло прямо перед ним, где началось предложение, о каком предмете шла речь. А одна голова — это одно взвешенное среднее: смотреть в несколько мест с разными целями сразу она не может.

Решение — многоголовое внимание. Вектор токена у Ростка состоит из 384 чисел. Запрос, ключ и значение режем на 6 кусков по 64 числа, и каждая тройка кусков работает как отдельная маленькая голова — со своими оценками и своими весами. Шесть голов считаются параллельно, их выходы склеиваются обратно в вектор из 384 чисел, а последняя матрица $W_O$ перемешивает их между собой:

$$\operatorname{MultiHead}(x) = \big[\,o^{(1)}, o^{(2)}, \dots, o^{(6)}\,\big]\, W_O.$$

Параметров от этого не прибавляется: матрицы $W_Q$, $W_K$, $W_V$ остаются размером $384 \times 384$, мы лишь разбиваем их столбцы на группы. Вместе с $W_O$ это $4 \times 384^2 = 589\,824$ параметра на слой. Вот все 48 голов Ростка сразу:

Попробуйте найти три типа голов:

• «На соседа» — голова, которая почти всегда смотрит на предыдущий токен. На карте это яркая полоска прямо под диагональю. Такая голова сообщает каждому токену, кто стоял перед ним.
• «На начало» — голова, сливающая внимание в метку ◆. На карте — яркий левый столбец. Скорее всего, это та самая голова, которой часто нечего искать.
• «Индукция» — выберите пример «Выдуманное имя». Дракона зовут Zorbin Plax: такого имени Росток никогда не встречал, и в первый раз угадать его нельзя. Но во второй раз, дойдя до «Z», можно посмотреть, что шло после «Z» в прошлый раз, — и скопировать. Индукционная голова зажигает на карте клетки далеко слева от диагонали — там, где имя прозвучало впервые. Под виджетом видно, насколько лучше Росток угадывает кусочки текста, которые уже встречались.

Соберём всё вместе на numpy: матрицы, маску и головы. Функция ниже — полноценное многоголовое причинное внимание, только со случайными, не обученными весами; в конце проверим, что матричная версия даёт то же, что и расчёт руками для одного токена.

import numpy as np def softmax(x): x = x - x.max(axis=-1, keepdims=True) e = np.exp(x) return e / e.sum(axis=-1, keepdims=True) def attention(x, Wq, Wk, Wv, n_head): T, C = x.shape d = C // n_head q, k, v = x @ Wq, x @ Wk, x @ Wv # (T, C) each # cut the C columns into heads: (n_head, T, d) q, k, v = (a.reshape(T, n_head, d).transpose(1, 0, 2) for a in (q, k, v)) scores = q @ k.transpose(0, 2, 1) / np.sqrt(d) # (n_head, T, T) future = np.triu(np.ones((T, T), dtype=bool), k=1) # above the diagonal scores[:, future] = -np.inf # no peeking w = softmax(scores) # every row adds up to 1 y = w @ v # (n_head, T, d) return y.transpose(1, 0, 2).reshape(T, C), w # glue the heads back rng = np.random.default_rng(0) T, C, H = 5, 12, 3 x = rng.standard_normal((T, C)) Wq, Wk, Wv = (rng.standard_normal((C, C)) / np.sqrt(C) for _ in range(3)) y, w = attention(x, Wq, Wk, Wv, H) print("output:", y.shape) print("head 0, weights:\n", w[0].round(2)) # the same row by hand: token 3 in head 0 d = C // H q3 = (x[3] @ Wq)[:d] k0 = (x[:4] @ Wk)[:, :d] s = k0 @ q3 / np.sqrt(d) print("by hand matches:", np.allclose(np.exp(s) / np.exp(s).sum(), w[0, 3, :4]))

Посмотрите на верхний треугольник напечатанной матрицы: там нули — это маска. А каждая строка в сумме даёт 1.

Внимание в model.py

А вот как то же самое записано в настоящем Ростке — класс из того самого файла, по которому обучали модель, без единой правки:

class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y)

Пройдём по строчкам.

• self.qkv — это $W_Q$, $W_K$ и $W_V$, склеенные в одну матрицу $384 \times 1152$: одно большое умножение быстрее трёх маленьких. .split(C, dim=-1) режет результат обратно на три части по 384.
• B, T, C — размер батча (сколько текстов обрабатываем разом), длина текста и ширина вектора, 384. Тексты батча друг друга не видят: это просто параллельная работа.
• .view(B, T, n_head, head_dim).transpose(1, 2) — те же 6 кусков по 64, что и в numpy-версии: головы становятся ещё одним измерением вроде батча и дальше считаются параллельно.
• q_norm, k_norm — QK-norm из заметки выше.
• apply_rope — добавляет к запросам и ключам сведения о позиции. Это тема следующей главы, а зачем она нужна, станет ясно через минуту.
• F.scaled_dot_product_attention(..., is_causal=True) — вся формула $\operatorname{softmax}(QK^\top/\sqrt d + M)V$ одним вызовом. PyTorch сам делит на $\sqrt d$ и сам ставит маску, а на видеокарте выбирает быстрый алгоритм вроде FlashAttention.
• transpose + view склеивают головы обратно в 384 числа, а self.proj — это $W_O$.

Есть у этого механизма одна неожиданная слепота. Посмотрите ещё раз на формулу $o = \sum_j a_j v_j$ и ответьте на вопрос.

Возьмём только внимание из этой главы, без apply_rope. Перемешаем все токены перед текущим, а его самого оставим на месте. Изменится ли его выход?

Да, ведь порядок слов меняет смысл Нет: веса зависят только от содержания запроса и ключей, а от суммы порядок слагаемых не зависит Изменится, потому что маска закроет другие позиции

Оценки $q\cdot k_j$ считаются по содержимому векторов, а взвешенная сумма не зависит от порядка слагаемых. Маска тоже не спасает: набор прошлых токенов тот же. Для такого внимания «the dog bit the boy» и «the boy bit the dog» — один и тот же мешок слов. Поэтому трансформеру нужен отдельный способ сообщить позиции; в Ростке это RoPE — вращения из следующей главы.

Внимание как мягкий поиск по словарю

В Python есть словари: d[key] находит ровно одну запись с совпадающим ключом и возвращает её значение. Внимание — «мягкая» версия того же: вместо точного совпадения — степень похожести, вместо одной записи — смесь всех записей с весами по похожести. Будь softmax бесконечно резким, внимание превратилось бы в обычный поиск самого похожего ключа. Мягкость нужна для обучения: у жёсткого выбора нет градиента, а у взвешенной смеси он есть, и модель может понемногу подкручивать и запросы, и ключи.

Росток сейчас

Мы заменили окно из восьми токенов вниманием. Эта модель — уже настоящий маленький трансформер: эмбеддинги, один слой внимания с четырьмя головами и небольшая сеть после него (что она делает, расскажет следующая глава). Она видит до 256 токенов назад, и имя из начала истории больше не теряется за краем окна. На отложенном тексте она ошибается в среднем на 2,87 ната на токен, а сеть с окном из главы 6 — на 3,31. Переведём в перплексию (глава 2): $e^{3{,}31} \approx 27$, а $e^{2{,}87} \approx 18$ — вместо двадцати семи вариантов следующего токена модель в среднем колеблется между восемнадцатью.

Но слой внимания один, а значит, за один проход сведения могут перепрыгнуть от токена к токену только один раз: узнать «кто стоял перед кем» и тут же этим воспользоваться, как индукционная голова, модель не может. Мы проверили на драконе по имени Zorbin Plax: встретив «Z» во второй раз, gpt-1 ждёт продолжения «or» с вероятностью меньше процента. А позиции в ней уже есть — это та самая строчка apply_rope, — только как они устроены, мы пока не разобрали. В следующей главе соберём полный блок трансформера, разберёмся с позициями-вращениями и поставим блоки друг на друга.