Transformer Explainer: живая GPT-2 в браузере — как пользоваться и перевод статьи

Transformer Explainer — интерактивная визуализация группы Polo Club of Data Science (Технологический институт Джорджии): в браузере работает настоящая модель GPT-2 (small), вы вводите свой текст и видите, как он превращается в токены, эмбеддинги, веса внимания и вероятности следующего токена. Под визуализацией на той же странице — поясняющая статья; ниже её полный перевод.

Лицензия и источник

Инструмент и текст статьи распространяются по лицензии MIT, © 2022 Polo Club of Data Science. Перевод сделан по исходному коду страницы в репозитории poloclub/transformer-explainer (файл src/components/article/Article.svelte). Рисунки статьи — это интерактивные фрагменты визуализации; здесь они описаны словами. Сайт инструмента: poloclub.github.io/transformer-explainer.

Как пользоваться вместе с курсом

1. Когда открывать. После урока 4.4.2 «Self-attention» — и ещё раз после 4.5.1 «Блок трансформера». Инструмент почти без текста: всё видно на схемах, а статья ниже объясняет, что вы видите.

2. Начните с готового примера («Data visualization empowers users to») и посмотрите на столбики вероятностей следующего токена справа.

3. Двигайте ползунок температуры — 0,2, 1, 2 — и смотрите, как распределение становится острее или площе (это наш урок 4.1.3 «Сэмплирование»).

4. Кликните по блоку внимания — раскроется расчёт Q, K, V и матрица внимания по головам. Наведите курсор на токен, чтобы увидеть, на какие токены он «смотрит». Заметьте треугольную маску: будущие токены закрыты.

5. Введите свой текст — по-английски (GPT-2 обучалась на английском; русский текст разобьётся на множество мелких токенов, и это само по себе полезный эксперимент к уроку 4.2.1).

Инструмент можно запустить и у себя: нужен Node.js 20+, затем git clone https://github.com/poloclub/transformer-explainer.git, npm install, npm run dev и адрес http://localhost:5173 в браузере.

Перевод статьи

Что такое трансформер?

Трансформер — архитектура нейросетей, которая коренным образом изменила подход к искусственному интеллекту. Впервые трансформер был представлен в знаковой статье «Attention is All You Need» в 2017 году и с тех пор стал основной архитектурой моделей глубокого обучения, на которой работают генерирующие текст модели вроде GPT от OpenAI, Llama от Meta и Gemini от Google. Помимо текста, трансформер применяют для генерации звука, распознавания изображений, предсказания структуры белков и даже в играх, что показывает его универсальность во множестве областей.

По сути, генерирующие текст трансформеры работают по принципу предсказания следующего токена: по текстовому промпту пользователя — какой токен (слово или часть слова) скорее всего последует за этим вводом? Главное нововведение и сила трансформеров — механизм самовнимания (self-attention), который позволяет обрабатывать последовательности целиком и улавливать дальние зависимости эффективнее прежних архитектур.

Семейство моделей GPT-2 — яркий пример генерирующих текст трансформеров. Transformer Explainer работает на модели GPT-2 (small) со 124 миллионами параметров. Это не самая новая и не самая мощная модель-трансформер, но у неё много общих архитектурных компонентов и принципов с нынешними передовыми моделями, поэтому она идеальна, чтобы начать разбираться в основах.

Архитектура трансформера

Любой генерирующий текст трансформер состоит из трёх ключевых компонентов:

• Эмбеддинг: входной текст делится на меньшие единицы — токены, которыми могут быть слова или части слов. Токены превращаются в числовые векторы — эмбеддинги, которые отражают смысл слов.
• Блок трансформера — основной строительный блок модели, который обрабатывает и преобразует входные данные. Каждый блок включает:
  • механизм внимания — центральный компонент блока трансформера. Он позволяет токенам «общаться» с другими токенами, улавливая контекст и связи между словами;
  • слой MLP (многослойный перцептрон) — сеть прямого распространения, которая работает с каждым токеном независимо. Если задача слоя внимания — передавать информацию между токенами, то задача MLP — уточнять представление каждого токена.
• Вероятности на выходе: последние слои — линейный и softmax — превращают обработанные эмбеддинги в вероятности, позволяя модели предсказывать следующий токен последовательности.

Эмбеддинг

Допустим, вы хотите сгенерировать текст трансформером. Вы вводите промпт, например: «Data visualization empowers users to» («Визуализация данных даёт пользователям возможность…»). Этот ввод нужно превратить в формат, который модель может понять и обработать. Здесь и нужен эмбеддинг: он превращает текст в числовое представление, с которым модель может работать. Чтобы превратить промпт в эмбеддинг, нужно 1) токенизировать ввод, 2) получить эмбеддинги токенов, 3) добавить информацию о позициях и, наконец, 4) сложить кодирования токенов и позиций, чтобы получить итоговый эмбеддинг. Посмотрим, как делается каждый шаг.

Рисунок 1: раскрытый слой эмбеддинга — как промпт превращается в векторное представление: (1) токенизация, (2) эмбеддинг токенов, (3) позиционное кодирование, (4) итоговый эмбеддинг.

Шаг 1. Токенизация

Токенизация — это разбиение входного текста на меньшие, более удобные части — токены. Токеном может быть слово или часть слова. Слова «Data» и «visualization» соответствуют отдельным токенам, а слово «empowers» делится на два токена. Полный словарь токенов определяется до обучения модели: в словаре GPT-2 50 257 уникальных токенов. Разбив входной текст на токены с собственными номерами (ID), мы можем получить их векторные представления из эмбеддингов.

Шаг 2. Эмбеддинг токенов

GPT-2 (small) представляет каждый токен словаря 768-мерным вектором; размерность вектора зависит от модели. Эти векторы эмбеддингов хранятся в матрице размера (50 257, 768) — это примерно 39 миллионов параметров! Такая большая матрица позволяет модели придавать каждому токену смысл: токены со схожим употреблением или значением в языке оказываются близко друг к другу в этом многомерном пространстве, а непохожие — дальше.

Проверка переводчика: $50\,257 \cdot 768 = 38\,597\,376 \approx 39$ млн — почти треть всех 124 млн параметров GPT-2 small.

Шаг 3. Позиционное кодирование

Слой эмбеддинга также кодирует информацию о позиции каждого токена во входном промпте. Разные модели используют разные способы позиционного кодирования. GPT-2 обучает собственную матрицу позиционного кодирования с нуля, прямо в процессе обучения.

Другой подход к позиционному кодированию (в исходнике статьи этот фрагмент закомментирован, но переведём его тоже). Другие модели — исходный трансформер и BERT — используют для позиционного кодирования синусоидальные функции. Такое синусоидальное кодирование детерминировано и устроено так, чтобы отражать и абсолютную, и относительную позицию каждого токена. Каждой позиции в последовательности сопоставляется уникальное математическое представление из комбинации синусов и косинусов. Для данной позиции синус задаёт чётные измерения вектора позиционного кодирования, а косинус — нечётные. Периодичность гарантирует, что у каждой позиции согласованное кодирование, не зависящее от окружающего контекста. Формулы:

$$PE_{(pos,\,2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right), \qquad PE_{(pos,\,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d}}\right),$$

где $pos$ — позиция токена, $i$ — номер пары измерений, $d$ — размерность эмбеддинга. (Формулы восстановлены переводчиком по статье «Attention is All You Need»; в исходнике статьи на их месте заглушка.) Подробнее — урок 4.5.2 «Позиционное кодирование».

Шаг 4. Итоговый эмбеддинг

Наконец, мы складываем кодирования токенов и позиций и получаем итоговое представление-эмбеддинг. Это объединённое представление отражает и смысл токенов, и их позицию во входной последовательности.

Блок трансформера

Основная обработка в трансформере происходит в блоке трансформера, который состоит из многоголового самовнимания и слоя многослойного перцептрона. Большинство моделей состоят из нескольких таких блоков, поставленных друг за другом. Представления токенов меняются от слоя к слою, от первого блока до последнего, и модель шаг за шагом строит сложное понимание каждого токена. Такой послойный подход приводит к представлениям входа всё более высокого порядка. Модель GPT-2 (small), которую мы разбираем, состоит из 12 таких блоков.

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

Механизм самовнимания позволяет модели улавливать связи между токенами последовательности, так что представление каждого токена зависит от остальных. Несколько голов внимания позволяют модели смотреть на эти связи с разных сторон; например, одна голова может улавливать близкие синтаксические связи, а другая — более широкий смысловой контекст. Дальше мы шаг за шагом разберём, как вычисляется многоголовое самовнимание.

Шаг 1. Матрицы запросов, ключей и значений.

Рисунок 2: вычисление матриц Query, Key и Value из исходного эмбеддинга.

Вектор эмбеддинга каждого токена превращается в три вектора: запрос (Query, Q), ключ (Key, K) и значение (Value, V). Эти векторы получаются умножением входной матрицы эмбеддингов на обучаемые матрицы весов для Q, K и V. Аналогия с веб-поиском поможет интуитивно понять эти матрицы:

• Запрос (Q) — текст, который вы вводите в строку поисковика. Это токен, о котором вы хотите «узнать больше».
• Ключ (K) — заголовок каждой веб-страницы в окне результатов поиска. Он представляет возможные токены, на которые запрос может обратить внимание.
• Значение (V) — само содержимое показанных веб-страниц. Сопоставив поисковый запрос (Query) с подходящими результатами (Key), мы хотим получить содержимое (Value) самых релевантных страниц.

С помощью этих значений QKV модель вычисляет оценки внимания (attention scores), которые определяют, сколько внимания получает каждый токен при генерации предсказаний.

Шаг 2. Разделение на головы. Векторы запроса, ключа и значения делятся на несколько голов — у GPT-2 (small) на 12 голов. Каждая голова независимо обрабатывает свой сегмент эмбеддингов, улавливая разные синтаксические и смысловые связи. Такое устройство позволяет параллельно выучивать разнообразные языковые признаки и усиливает выразительность модели.

Проверка переводчика: $768 / 12 = 64$ — размер векторов Q, K, V в одной голове.

Шаг 3. Маскированное самовнимание. В каждой голове выполняется маскированное самовнимание. Этот механизм позволяет модели генерировать последовательности, сосредотачиваясь на релевантных частях входа и не имея доступа к будущим токенам.

Рисунок 3: вычисление маскированного самовнимания по матрицам Query, Key и Value.

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

Пояснение переводчика: всё вместе — это формула $\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$, где масштаб $\sqrt{d_k} = \sqrt{64} = 8$, а маска $M$ содержит $-\infty$ выше диагонали.

Шаг 4. Выход и склейка. Модель умножает маскированные оценки самовнимания на матрицу значений (Value) и получает итоговый выход механизма самовнимания. У GPT-2 12 голов самовнимания, каждая улавливает свои связи между токенами. Выходы этих голов склеиваются (конкатенируются) и проходят через линейную проекцию.

MLP: многослойный перцептрон

Рисунок 4: слой MLP проецирует представления после самовнимания в пространство большей размерности, чтобы повысить выразительность модели.

После того как несколько голов самовнимания уловили разнообразные связи между входными токенами, склеенные выходы проходят через слой многослойного перцептрона (MLP), чтобы повысить выразительность модели. Блок MLP состоит из двух линейных преобразований с функцией активации GELU между ними.

Первое линейное преобразование увеличивает размерность входа в четыре раза — с 768 до 3072. Этот шаг расширения позволяет модели спроецировать представления токенов в пространство большей размерности, где можно уловить более богатые и сложные закономерности, не видные в исходной размерности.

Второе линейное преобразование затем уменьшает размерность обратно до исходных 768. Этот шаг сжатия возвращает представления к удобному размеру, сохраняя полезные нелинейные преобразования, внесённые на шаге расширения.

В отличие от механизма самовнимания, который объединяет информацию по токенам, MLP обрабатывает токены независимо и просто переводит представление каждого токена из одного пространства в другое, обогащая общую выразительность модели.

Вероятности на выходе

После того как вход прошёл через все блоки трансформера, выход подаётся в последний линейный слой, чтобы подготовить его к предсказанию токена. Этот слой проецирует итоговые представления в пространство размерности 50 257, где у каждого токена словаря есть соответствующее значение — логит (logit). Следующим словом может быть любой токен, и этот шаг позволяет просто упорядочить токены по тому, насколько вероятно, что они окажутся следующим словом. Затем мы применяем функцию softmax, чтобы превратить логиты в распределение вероятностей, дающее в сумме единицу. Это позволит выбрать следующий токен с учётом его вероятности.

Рисунок 5: каждому токену словаря сопоставляется вероятность по выходным логитам модели. Эти вероятности определяют, насколько вероятно, что токен окажется следующим словом последовательности.

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

• Температура = 1: деление логитов на единицу никак не влияет на выход softmax.
• Температура меньше 1: низкая температура делает модель увереннее и детерминированнее — распределение вероятностей становится острее, а выходы предсказуемее.
• Температура больше 1: высокая температура делает распределение мягче, и в сгенерированном тексте больше случайности — то, что некоторые называют «креативностью» модели.

Кроме того, процесс выбора можно дополнительно настроить параметрами top-k и top-p:

• top-k sampling: оставляет в кандидатах только $k$ токенов с наибольшими вероятностями, отбрасывая менее вероятные варианты.
• top-p sampling: рассматривает наименьший набор токенов, суммарная вероятность которых превышает порог $p$, — так в выборе участвуют только самые вероятные токены, но разнообразие сохраняется.

Настраивая температуру, top-k и top-p, можно найти баланс между детерминированными и разнообразными выходами и подстроить поведение модели под свои задачи.

Числовой пример переводчика: пусть логиты трёх токенов $(2, 1, 0)$. При $T = 1$ softmax даёт примерно $(0{,}67;\ 0{,}24;\ 0{,}09)$; при $T = 0{,}5$ логиты становятся $(4, 2, 0)$ и вероятности — $(0{,}87;\ 0{,}12;\ 0{,}02)$; при $T = 2$ — логиты $(1;\ 0{,}5;\ 0)$ и вероятности $(0{,}51;\ 0{,}31;\ 0{,}19)$.

Вспомогательные элементы архитектуры

Есть несколько вспомогательных элементов архитектуры, которые улучшают работу трансформеров. Они важны для общего качества модели, но не так важны для понимания основных идей архитектуры. Нормализация слоя (Layer Normalization), dropout и остаточные связи (residual connections) — ключевые компоненты трансформеров, особенно на этапе обучения. Нормализация слоя стабилизирует обучение и помогает модели быстрее сходиться. Dropout предотвращает переобучение, случайно отключая нейроны. Остаточные связи позволяют градиентам проходить прямо через сеть и помогают избежать проблемы затухающего градиента.

Нормализация слоя

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

Dropout

Dropout — метод регуляризации, который предотвращает переобучение нейросетей, во время обучения случайно обнуляя часть весов модели. Это побуждает модель выучивать более устойчивые признаки и снижает зависимость от отдельных нейронов, помогая сети лучше обобщать на новые, невиданные данные. При инференсе dropout отключается. По сути, это значит, что мы используем ансамбль обученных подсетей, что улучшает качество модели.

Примечание переводчика: точнее, dropout обнуляет не веса, а активации (выходы нейронов), — именно так он реализован в nn.Dropout и в самой GPT-2.

Остаточные связи

Остаточные связи впервые появились в модели ResNet в 2015 году. Это архитектурное новшество совершило революцию в глубоком обучении, позволив обучать очень глубокие нейросети. По сути, остаточные связи — это «обходные пути», минующие один или несколько слоёв: вход слоя прибавляется к его выходу. Это помогает смягчить проблему затухающего градиента и упрощает обучение глубоких сетей из многих блоков трансформера, поставленных друг на друга. В GPT-2 остаточные связи используются в каждом блоке трансформера дважды: один раз перед MLP и один раз после него, — чтобы градиенты легче проходили и ранние слои получали достаточные обновления при обратном распространении.

Примечание переводчика: точнее, одна остаточная связь обходит подблок внимания, а вторая — подблок MLP (урок 3.4.2 «Остаточные связи и dropout»).

Интерактивные возможности

Transformer Explainer сделан интерактивным и позволяет исследовать внутреннее устройство трансформера. Вот с какими интерактивными возможностями можно поиграть:

• Введите свою текстовую последовательность, чтобы увидеть, как модель её обрабатывает и предсказывает следующее слово. Изучите веса внимания, промежуточные вычисления и то, как вычисляются итоговые вероятности на выходе.
• Двигайте ползунок температуры, чтобы управлять случайностью предсказаний модели. Посмотрите, как сделать выход модели детерминированнее или «креативнее», меняя значение температуры.
• Выберите способ выбора top-k или top-p, чтобы настроить выбор токенов при инференсе. Попробуйте разные значения и посмотрите, как меняется распределение вероятностей и как это влияет на предсказания модели.
• Взаимодействуйте с картами внимания, чтобы увидеть, как модель сосредотачивается на разных токенах входной последовательности. Наводите курсор на токены, чтобы подсветить их веса внимания, и исследуйте, как модель улавливает контекст и связи между словами.

Как устроен Transformer Explainer?

В Transformer Explainer прямо в браузере работает живая модель GPT-2 (small). Модель получена из реализации GPT на PyTorch в проекте nanoGPT Андрея Карпатого и сконвертирована в ONNX Runtime для бесшовного выполнения в браузере. Интерфейс написан на JavaScript с фронтенд-фреймворком Svelte и библиотекой D3.js для динамических визуализаций. Числовые значения обновляются вживую вслед за вводом пользователя.

Кто создал Transformer Explainer?

Transformer Explainer создали Эри Чо, Грейс К. Ким, Александр Карпеков, Алек Хелблинг, Джей Ван, Сонмин Ли, Бенджамин Гувер и Поло Чау в Технологическом институте Джорджии.

Что дальше

• Объёмная 3D-визуализация маленькой GPT шаг за шагом — LLM Visualization Брендана Бикрофта.
• Та же GPT-2, но собранная своими руками, — урок 4.6.2 «nanoGPT на PyTorch».
• У авторов есть и другие «объяснялки»: Diffusion Explainer (как Stable Diffusion превращает текст в картинку), CNN Explainer и GAN Lab — всё на poloclub.github.io.

Словарик

EN	RU
next-token prediction	предсказание следующего токена
token, vocabulary, token ID	токен, словарь, номер токена
token embedding, positional encoding	эмбеддинг токена, позиционное кодирование
multi-head self-attention, attention head	многоголовое самовнимание, голова внимания
query, key, value	запрос, ключ, значение
attention score, masked (causal) attention	оценка внимания, маскированное (каузальное) внимание
MLP, GELU	многослойный перцептрон, GELU
logit, softmax	логит, softmax
temperature, top-k, top-p sampling	температура, выбор top-k, top-p
layer normalization, dropout, residual connection	нормализация слоя, dropout, остаточная связь