Внимание

«У Мии была новая кукла. Она брала её с собой в парк, в магазин и к бабушке. Вечером Мия уложила спать…» Вероятное продолжение — «куклу». Чтобы связать его с началом рассказа, нужно сохранить сведения, которые появились раньше последних нескольких слов. Именно здесь короткое окно нашей сети начинает мешать.

Наша сеть получает только восемь последних токенов. Сведения за этой границей не попадают на вход, как бы хорошо мы ни обучили её веса. Теперь введём внимание (attention): механизм, который смешивает информацию из разных позиций доступного контекста. Это важная часть трансформеров, в том числе Ростка. Его расчёт состоит из знакомых операций — скалярных произведений, 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$ миллиона весов в одном только первом слое, причём большинство из них почти ничему не научится: сигнал до них будет доходить редко.

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

Нам нужно обрабатывать позиции общими весами и выбирать, какую информацию из доступного контекста смешать для очередного прогноза. Внимание позволяет это сделать, не выделяя отдельную матрицу весов каждому месту в окне. Само ограничение контекста при этом остаётся: полный Росток видит до 512 токенов, а модель в конце этой главы — до 256.

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

Прежде чем выводить формулы, посмотрим на готовый результат. У полного Ростка 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: побеждает наименее плохой ключ. Внимание не умеет «никуда не смотреть».

Этот опыт помогает понять одно из возможных объяснений большого веса у начальной метки. Softmax всегда распределяет всю массу между доступными позициями; у некоторых моделей часть массы устойчиво достаётся первым токенам. Такое поведение называют attention sink. Но по одной карте нельзя заключить, что голова «ничего не делает» или что значение метки безобидно: результат зависит ещё от векторов значений и последующих преобразований. Карта показывает коэффициенты смешивания, а не полное объяснение ответа модели.

Вот та же арифметика на 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 и зафиксируем запрос. Если одинаково переставить доступные пары ключ–значение, изменится ли результат их взвешенного смешивания?

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

При неизменном запросе перестановка пар ключ–значение лишь переставляет слагаемые взвешенной суммы; результат остаётся прежним. Это свойство отдельной операции с зафиксированными векторами. У целой многослойной причинной сети векторы уже могут зависеть от предшествующего текста, поэтому вывод о её полной нечувствительности к порядку был бы неверен. 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, — только как они устроены, мы пока не разобрали. В следующей главе соберём полный блок трансформера, разберёмся с позициями-вращениями и поставим блоки друг на друга.