Глава 8 из 17 40 мин
Трансформер
Проследим путь от эмбеддинга до следующего токена. Соберём блок Ростка, разберём нормировку и позиции, затем посчитаем все параметры модели.
В этой главе
- собрать блок трансформера из внимания, сети SwiGLU, остаточных связей и RMSNorm
- понять, как RoPE кодирует позиции поворотом и учитывает расстояние между токенами
- посчитать параметры Ростка и любой похожей модели
- заглянуть внутрь модели логит-линзой и увидеть, как предсказание созревает от слоя к слою
Мы уже умеем превращать текст в токены и векторы, обрабатывать их нейросетью и смешивать сведения с помощью внимания. Теперь соберём из этих частей трансформер Ростка. Будем двигаться по пути одного вектора: от эмбеддинга через повторяющиеся блоки к оценкам следующего токена. По дороге разберём оставшиеся детали model.py — нормировку, остаточные связи, SwiGLU и кодирование позиций.
Для сборки понадобятся три идеи: остаточные связи, которые добавляют результат слоя к его входу; нормировка, которая выравнивает масштаб входных чисел; и позиционные кодировки, позволяющие учитывать порядок токенов. Затем разберём SwiGLU — сеть, которая преобразует представление каждой позиции после внимания. Начнём со схемы всего блока.
Росток на одной схеме
Вот вся модель. Текст входит сверху, предсказание выходит снизу. Толстая линия слева — остаточный поток, пунктирная рамка — блок, который повторяется 8 раз, жёлтый пунктир справа — одна и та же матрица, работающая и на входе, и на выходе. Нажмите на любую деталь, и рядом со схемой появится её код — дословно из файла, по которому обучали Ростка, с номерами строк.
scripts/micro-llm/model.py. Переключатель сверху меняет только числа в конфигурации: у Ростка, у gpt-4 этой главы и у gpt-1 из прошлой код один и тот же.Обратите внимание на переключатель: gpt-1 из прошлой главы, gpt-4 из этой и полный Росток — это один и тот же код. Меняются лишь несколько чисел: ширина потока, число блоков и голов, размер сети. Вся разница между маленькой и большой моделью — в этих числах, данных и времени обучения. Дальше пройдём по схеме сверху вниз.
Остаточный поток
Главная линия схемы — вертикальная. Токен входит в модель вектором из 384 чисел и таким же вектором выходит. По дороге ни один блок этот вектор не заменяет — каждый лишь прибавляет к нему свою поправку:
$$x \leftarrow x + \operatorname{Attn}\big(\operatorname{Norm}(x)\big),$$ $$x \leftarrow x + \operatorname{FFN}\big(\operatorname{Norm}(x)\big).$$Представления, которые проходят через последовательные блоки с такими прибавлениями, называют остаточным потоком. Его можно представить общей рабочей записью: внимание и сеть читают текущее состояние и добавляют свои поправки. Но это не архив отдельных страниц. Все поправки складываются в одном векторе фиксированной длины, поэтому могут усиливать, ослаблять и даже отменять прежние компоненты. Гарантии сохранения всей исходной информации нет.
Почему не проще — $x \leftarrow f(x)$, как в сетях из главы 4? Причин две.
Градиенту нужна дорога. При обратном проходе (глава 4) производные слоёв перемножаются по цепочке. Если каждая из них хоть немного меньше единицы, через восемь блоков сигнал заметно ослабнет, а через сотню почти исчезнет. С прибавлением всё иначе:
$$\frac{\partial}{\partial x}\big(x + f(x)\big) = 1 + f'(x).$$Для вектора вместо единицы в формуле стоит единичная матрица. Эта прямая ветвь даёт градиенту дополнительный путь, по которому ему не нужно проходить через производную нелинейной части блока. Остаточные связи облегчают обучение глубоких сетей, хотя сами по себе не гарантируют, что градиенты нигде не затухнут и не вырастут. Они были ключевой частью ResNet (He и др., 2015) и затем вошли в трансформеры.
Небольшие поправки в начале. В инициализации выходные матрицы внимания и сети (proj и w2) получают веса в $\sqrt{2 \cdot 8}=4$ раза меньшего масштаба. Это уменьшает начальный вклад каждого блока в остаточный поток. Модель всё ещё учитывает контекст через внимание, но исходные эмбеддинги не сразу перекрываются суммой больших случайных поправок. По мере обучения блоки учатся менять представления полезным для предсказания образом.
RMSNorm: держать числа в узде
У постоянного дописывания есть обратная сторона: размер вектора в потоке растёт. Мы измерили на gpt-4 из этой главы: типичное число в потоке (среднеквадратичное, в среднем по отложенному тексту) сразу после эмбеддинга — около 0,18, после первого блока — уже около 10, а после четвёртого — около 35, в двести раз больше. А слою удобнее всего, когда вход всегда примерно одного масштаба, где бы в модели этот слой ни стоял.
Поэтому перед каждым слоем вектор нормируют. Росток использует RMSNorm — среднеквадратичную нормировку:
$$\operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2},$$ $$\hat x_i = \frac{x_i}{\sqrt{\operatorname{RMS}(x)^2 + \varepsilon}}\, g_i .$$Каждую координату делим на корень из среднего квадрата координат с добавкой $\varepsilon$. До обучаемого масштабирования среднеквадратичная величина результата близка к единице, если исходный вектор не слишком мал. Затем каждую координату умножаем на свой коэффициент $g_i$, который учится вместе с моделью. Так слой получает управляемый масштаб входа, сохраняя возможность усиливать разные координаты по-разному. Всего таких коэффициентов $d$.
Обратите внимание, где стоит нормировка: на входе в слой, а не в самом потоке. Нормируется копия, которую слой читает, а поток остаётся нетронутым — иначе мы бы снова «переписывали тетрадь». Такую схему называют pre-norm. В самом первом трансформере 2017 года нормировка стояла после сложения (post-norm), и глубокие модели с ней обучались капризно; Сюн и др. (2020) показали, почему pre-norm стабильнее. Ещё одна нормировка стоит в самом конце, перед выходным слоем.
Каждый токен думает сам: SwiGLU
Внимание — единственное место, где токены обмениваются сведениями. Но собрать сведения мало, их надо переработать. Для этого в каждом блоке после внимания стоит маленькая нейросеть из двух слоёв, почти такая же, как в главе 4, — сеть прямого распространения (feed-forward network). Она обрабатывает каждый токен отдельно, одними и теми же весами на всех позициях. Здесь живёт большая часть параметров: $3 \times 384 \times 1\,024 = 1\,179\,648$ на блок против $589\,824$ у внимания — две трети блока.
В первом трансформере сеть была такой: $\operatorname{ReLU}(x W_1)\,W_2$ — расширить вектор в 4 раза, обнулить отрицательное, сжать обратно. Росток, как и Llama, использует вариант SwiGLU с тремя матрицами:
$$\operatorname{FFN}(x) = \big(\operatorname{silu}(x W_1) \odot x W_3\big)\,W_2,$$ $$\operatorname{silu}(a) = \frac{a}{1 + e^{-a}}.$$Здесь $\odot$ — поэлементное умножение. Вектор расширяется до 1 024 чисел дважды, двумя разными матрицами. Первая копия проходит через SiLU — гладкую родственницу ReLU, — вторая остаётся как есть, и они перемножаются: одна копия служит воротами для другой. Нейрон может сказать «пропускаю, только если выполнены оба условия» или даже поменять знак сигнала — одиночной ReLU такое не под силу.
А откуда странное число 1 024? В классической сети скрытый слой в 4 раза шире потока, и матриц две: $2 \times 4d^2 = 8d^2$ параметров. У SwiGLU матриц три, поэтому ширину берут равной $\tfrac{2}{3}$ от обычной, чтобы параметров вышло столько же: $\tfrac23 \cdot 4 \cdot 384$ — ровно $1\,024$. Вот сеть Ростка целиком:
Где я? Позиции поворотом
В конце прошлой главы выяснилось, что внимание не различает порядок слов: для него «the dog bit the boy» и «the boy bit the dog» — один и тот же набор векторов. Порядок надо как-то сообщить.
Самое прямолинейное решение — завести ещё одну таблицу эмбеддингов, на этот раз для позиций, и прибавлять её строку к вектору токена. Так сделано в GPT-2 (2019). Но тогда модель знает только абсолютные номера и отдельно учит, что «слово двумя позициями раньше» на десятом месте и на трёхсотом — одно и то же отношение. А для языка важнее как раз относительное: насколько далеко стоит нужное слово.
RoPE (rotary position embedding; Су и др., 2021) ничего не прибавляет — он поворачивает. Запрос одной головы — это 64 числа; разобьём их на 32 пары, и каждая пара станет точкой на плоскости. Токен на позиции $m$ поворачивает свою $i$-ю пару на угол $m\,\theta_i$, где
$$\theta_i = 10000^{-2i/64}, \qquad i = 0, 1, \dots, 31.$$Ключи поворачиваются так же, каждый по своей позиции $n$. Значения не поворачивают.
Для фиксированных запроса и ключа поворот сохраняет длины, а их скалярное произведение зависит от разности углов поворота. Поэтому явная зависимость оценки от позиций входит через относительное смещение $m-n$. Это не значит, что внимание зависит только от расстояния: сами векторы несут содержание текста, а в последующих слоях — и результаты предыдущих вычислений. Сравним одну и ту же пару векторов на разных позициях.
У разных пар координат разный темп поворота. Первая пара делает оборот примерно за 6,3 позиции; последняя — примерно за 47 тысяч. Один поворот периодичен, а сочетание разных частот позволяет различать больше взаимных расположений. Это похоже на несколько часовых стрелок с разной скоростью. Но аналогия не обещает точного отсчёта на произвольной длине: способность модели использовать расстояния зависит также от обучения и длины доступного контекста.
В коде это две функции. Первая один раз готовит таблицы косинусов и синусов для всех позиций и частот. Вторая поворачивает пары $(x_i, x_{i+32})$: первая половина вектора — это «иксы» всех пар, вторая — «игреки».
Проверим главное свойство на numpy. Одни и те же запрос и ключ на позициях 5 и 2 и на позициях 105 и 102 — расстояние одно, значит, и оценка должна совпасть:
Запрос стоит на позиции 100, ключ — на позиции 90. Потом ту же пару векторов переносим на позиции 20 и 10. Что станет с оценкой внимания?
Повороты на $m\theta$ и $n\theta$ входят в скалярное произведение только через разность $(m-n)\theta$. Это и проверила ячейка выше: пары позиций (5, 2), (105, 102) и (400, 397) дали одно и то же число.
В нашем варианте RoPE частоты заданы заранее и не обучаются. Формулу поворота можно вычислить для любого номера позиции, но реализация Ростка заранее строит таблицы только для настроенного контекста. Она не получает неограниченную память автоматически. Даже если расширить таблицы, качество за пределами обученной длины нужно проверять отдельно. Именно различие между вычислимым поворотом и освоенным диапазоном важно помнить, обсуждая длинный контекст.
Блок целиком
Теперь у нас есть все части блока. Вот он. Весь его смысл умещается в двух строчках метода forward, а комментарии к ним пересказывают эту главу лучше любой схемы:
Блок трансформера — это два шага: сначала токены обмениваются сведениями (внимание), потом каждый обдумывает их сам (сеть SwiGLU). Оба шага обновляют вектор прибавлением поправки.
Восемь этажей и выход
Блоки ставятся друг на друга: у Ростка их восемь, одинаковых по устройству, но каждый со своими весами. Зачем больше одного? Потому что каждый следующий блок читает то, что записали предыдущие, и может на этом строить. В прошлой главе мы видели индукционную голову: ей нужно, чтобы кто-то раньше записал в каждый токен, какое слово стояло перед ним. Одному слою такое не под силу, а двум — уже да. Глубина позволяет строить цепочки таких шагов.
Проверим на деле. Ниже три модели читают один и тот же текст: однослойная gpt-1 из прошлой главы, четырёхблочная gpt-4 из этой и полный Росток.
В примере с повторяющимся «Zorbin Plax» четырёхблочная модель заметно повышает вероятность уже встречавшихся продолжений: около 98% для «or» и 93% для «bin». У однослойной модели вероятность «or» остаётся ниже процента. Это показывает разницу в поведении на данном примере; сами числа не доказывают, что результат обеспечила именно пара голов, рассмотренная ранее. У полного Ростка вероятность «or» при округлении доходит до 100%, а «bin» — до 98%. Округлённые 100% здесь не означают математически точную единицу.
После последнего блока поток ещё раз нормируется и попадает в выходной слой. Тот считает скалярное произведение вектора со строкой эмбеддинга каждого из 8 192 токенов — получаются логиты. И вот тонкость: в model.py это буквально та же матрица, что и на входе:
self.head.weight = self.embed.weight # weight tying
Это называется связанными весами (weight tying; Пресс и Вольф, 2017). Смысл простой: на входе строка таблицы говорит «вот что значит этот токен», на выходе та же строка спрашивает «похоже ли то, что модель хочет сказать, на этот токен?». Одна таблица учится с двух сторон, а модель экономит $8\,192 \times 384 = 3\,145\,728$ параметров — без этого трюка Ростку понадобилось бы 20,45 млн чисел вместо 17,31 млн.
Дальше всё знакомо: softmax превращает логиты в вероятности, а при обучении кросс-энтропия с настоящим следующим токеном даёт ошибку (глава 2). Маска loss_mask понадобится, когда мы будем учить Ростка разговаривать (глава 12).
Считаем параметры
Теперь можно разобрать Ростка по косточкам и пересчитать всё до последнего числа. Обозначим ширину потока $d = 384$, число блоков $L = 8$, ширину сети $d_{ff} = 1\,024$ и размер словаря $V = 8\,192$:
- эмбеддинг (он же выходной слой): $V d = 3\,145\,728$;
- внимание в одном блоке: $W_Q, W_K, W_V, W_O$ — это $4d^2 = 589\,824$, плюс $2 \times 64$ усилений QK-norm;
- сеть SwiGLU в одном блоке: $3 d\, d_{ff} = 1\,179\,648$;
- две нормировки в блоке: $2d = 768$, и одна финальная: $384$.
Это и есть 17,31 миллиона параметров Ростка. Посчитайте свою модель:
Пощёлкайте заготовки. У крошечной gpt-1 почти всё — эмбеддинг: 80% параметров уходит на словарь, и на «мышление» остаётся совсем немного. У Ростка на словарь приходится уже 18%, а у модели размером с GPT-2 small с её огромным словарём доля снова растёт. Видно и то, что цена блока растёт как квадрат ширины потока.
Ширину потока $d$ удвоили, $d_{ff}$ тоже удвоили (чтобы сохранить пропорцию), число блоков оставили. Во сколько раз вырастет число параметров внутри блоков?
И внимание ($4d^2$), и сеть ($3d \cdot d_{ff}$, где $d_{ff}$ растёт вместе с $d$) — квадратичны по ширине. Удвоение $d$ даёт $2^2 = 4$; только крошечные усиления нормировок растут вдвое. Эмбеддинг $Vd$ за пределами блоков тоже вырастет лишь вдвое.
Логит-линза: промежуточные прогнозы
Обычно выходной слой читает остаточный поток только после последнего блока. Для исследования можно применить ту же финальную нормировку и выходную матрицу к промежуточным состояниям. Получим распределения по словарю, с которыми удобно сравнивать окончательный прогноз. Этот приём называют логит-линзой (logit lens). Он не раскрывает внутренний «план» модели: мы намеренно читаем промежуточный вектор способом, для которого он отдельно не обучался.
На что посмотреть:
- Строка «эмбеддинг» часто повторяет входной токен. Здесь одна таблица служит и для чтения токенов, и для получения выходных оценок: собственная строка токена нередко даёт большое скалярное произведение с его вектором. Это не математическая гарантия — другая строка может дать ещё большее значение. На этом этапе блоки трансформера ещё не обработали контекст.
- Где впервые появляется правильный ответ? Рамка отмечает совпадение с настоящим следующим словом. Для лёгких продолжений ответ часто созревает в средних слоях, для трудных — только в последних.
- Как меняется распределение? Цвет показывает вероятность самого вероятного токена при чтении через линзу. На трёх заготовленных текстах её среднее значение было 77% после седьмого блока и 59% после восьмого. Это наблюдение о выбранных примерах, а не доказательство, что последний блок специально «убирает самоуверенность». Промежуточные состояния не обучались напрямую выдавать вероятности через эту линзу. Нажмите на клетку и сравните пять лучших вариантов.
Весь model.py
Вот файл полностью, от первой строки до последней. Теперь вам понятно в нём всё, кроме функции generate: выбору слов посвящена глава 11. Одно предупреждение: числа в Config — лишь значения по умолчанию, а настоящие размеры Ростка (поток из 384 чисел, 6 голов, сеть на 1 024) передаёт ему train.py.
Росток сейчас
Это уже Росток в миниатюре: тот же model.py, только поток шириной 192 числа вместо 384 и четыре блока вместо восьми. На отложенном тексте он ошибается в среднем на 2,26 ната на токен — против 2,87 у однослойной gpt-1 из прошлой главы. Разница кажется небольшой, но вспомните перплексию из главы 2: $e^{2{,}87} \approx 18$, а $e^{2{,}26} \approx 9{,}6$. Раньше модель в среднем колебалась между восемнадцатью вариантами следующего токена, теперь — между девятью с половиной.
Архитектура готова, и дальше мы её уже не меняем. Всё, что отделяет эту модель от настоящего Ростка, — размер, данные и время. Следующая часть курса — о них: откуда взять сотни миллионов слов, как их почистить и как три с половиной часа учить модель на ноутбуке.