Реализация трансформатора с нуля

7 удивительных вещей, которые вы могли не знать о Трансформере

Введение

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

Цель этого поста не в том, чтобы обсудить всю реализацию — для этого есть множество отличных ресурсов, — а в том, чтобы выделить семь вещей, которые я нашел особенно удивительными или познавательными и о которых вы, возможно, не знали. о. Я сделаю это конкретным, указав на определенные строки в моем коде, используя этого робота гиперссылок 🤖 (попробуйте!). Код должен быть простым для понимания: он хорошо документирован, автоматически тестируется и проверяется с помощью Github Actions.

Структура этого поста проста. Первые три пункта связаны с реализацией многоголового внимания; последние четыре относятся к другим компонентам. Я предполагаю, что вы концептуально знакомы с трансформером и многоголовым вниманием (если нет; отличной отправной точкой является Иллюстрированный трансформер), и начните с чего-то, что чрезвычайно помогло мне лучше понять механику многоголового внимания. Это также самый важный момент.

Table of Contents
Introduction
1. Multi-head attention is implemented with one weight matrix
2. The dimensionality of the key, query and value vectors is not a hyperparameter
3. Scaling in dot-product attention avoids extremely small gradients
4. The source embedding, target embedding AND pre-softmax linear share the same weight matrix
5. During inference we perform a decoder forward pass for every token; during training we perform a single forward pass for the entire sequence
6. Transformers can handle arbitrarily long sequences, in theory
7. Transformers are residual streams
Wrapping up

1. Мультиголовное внимание реализовано с одной весовой матрицей

Прежде чем мы углубимся в это; помните, что для каждого заголовка внимания нам нужен запрос, ключ и вектор значений для каждого входного токена. Затем мы вычисляем показатели внимания как softmax по масштабированному точечному произведению между одним запросом и всеми ключевыми векторами в предложении (🤖). Приведенное ниже уравнение вычисляет взвешенное по вниманию среднее по всем векторам значений для каждого входного маркера одновременно. Q – это матрица, в которой сложены qвекторы uery для всех входных токенов; K и V делают то же самое для векторов key и value.

Как же тогда мы можем получить эти векторы запросов, ключей и значений эффективно для всех токенов и головок? Оказывается, мы можем сделать это за один раз, используя единую матрицу весовW. Это отличается от трехпроекционных весовых матриц, которые можно ожидать после прочтения бумаги или Иллюстрированного трансформера. Давайте рассмотрим, как это работает.

Предположим, что наш ввод состоит из 4 токенов: ['привет', 'как', 'есть', 'вы'], а размер нашего встраивания равен 512. Пока игнорируем пакеты, пусть X будет матрицей 4x512, укладывающей вложения токенов в виде строк.

Пусть W будет весовой матрицей с 512 строками и 1536 столбцами. Теперь мы увеличим эту 512x1536 размерную весовую матрицу W (🤖), чтобы выяснить, зачем нам нужны 1536 размеры и как их умножить. с X приводит к матрице P (для pпроекций), которая содержит все необходимые нам векторы запроса, ключа и значения. (В коде 🤖 я называю эту матрицу qkv)

Умножение матриц за многоголовым вниманием

Каждый элемент в результирующей 4x1536 матрице P=X@W представляет собой сумму поэлементного произведения (другими словами: скалярного произведения) между вектором-строкой в X (вложение) и вектор-столбец вW (некоторые веса).

Чтобы напомнить об умножении матриц, на изображении ниже показано, как вычислить первый элемент простой матрицы 3x3 при умножении двух матриц 3x3. Та же стратегия применима к нашим большим матрицам X и W.

Таким образом, каждый элемент в iй строке P[i, :] в нашей проекционной матрице представляет собой линейную комбинацию iго токена встраивания X[i, :] и один из весовых столбцов в W. Это означает, что мы можем просто сложить больше столбцов в нашей весовой матрице W, чтобы создать более независимые линейные комбинации (скаляры) каждого токена, вложенного в X. Другими словами, каждый элемент в P представляет собой отдельный скаляр «представление» или «сводка» маркера, внедренного в X, взвешенного по столбцу в W. Это ключ к пониманию того, как восемь «головок» с векторами «запрос», «ключ» и «значение» прячутся в каждой из строк P.

Выявление головок внимания и векторов запросов, ключей и значений

Мы можем разложить столбцы 1536, которые мы выбрали для W (и в итоге получить количество столбцов в P), в 1536 = 8 * 3 * 64. Теперь мы обнаружили восемь головок, по три 64-мерных вектора в каждой строке в P! Каждый такой вектор или кусок состоит из 64 различных взвешенных линейных комбинаций вложенных токенов, и мы выбираем интерпретировать их определенным образом. Вы можете увидеть визуальное представление P и то, как его разложить, на изображении ниже. Декомпозиция также происходит в коде (🤖).

Для нескольких предложений в партии просто представьте третье измерение «за» P, которое превращает 2D-матрицу в 3D-матрицу.

Кодер-декодер внимание

Для кодировщика-декодера это немного сложнее. Напомним, что внимание кодировщика-декодера позволяет каждому декодеру следить за вложениями, выдаваемыми самым верхним кодировщиком.

Для внимания кодировщика-декодера нам нужны векторы запросов для встраивания маркера декодера и векторы ключа и значения для самого верхнего кодировщика встраивания токенов. Вот почему мы разделяем W на две — матрицу 512x512 и 512x1024 (🤖) — и выполняем две отдельные проекции: одну для получения векторов ключа и значения из вложений кодировщика (🤖), а другую для получения векторов запроса для вложений декодера (🤖).

Наконец, обратите внимание, что нам нужна вторая весовая матрица (🤖) во внимании с несколькими головками, чтобы смешать векторы значений из каждой головы и получить одно контекстное вложение для каждого токена (🤖).

2. Размерность векторов ключа, запроса и значения не является гиперпараметром.

Я никогда особо не задумывался об этом, но всегда предполагал, что размерность векторов запроса, ключа и значения является гиперпараметром. Как оказалось, он динамически устанавливается равным количеству измерений встраивания, деленному на количество головок: qkv_dim = embed_dim/num_heads = 512/8 = 64(🤖).

Это похоже на выбор дизайна Vaswani et al. чтобы количество параметров в многоголовом внимании оставалось постоянным, независимо от количества выбранных головок. Хотя можно было бы ожидать, что количество параметров будет расти с увеличением количества головок, на самом деле происходит следующее: размерность векторов запроса, ключа и значения уменьшается.

Если мы посмотрим на рисунок выше, на котором показано R=X@W, и представим себе внимание одной головы, это станет ясно. Количество элементов в X, W и R остается таким же, как и с восемью головками, но способ интерпретации элементов в Rменяется. С одной головкой у нас есть только одна проекция запроса, ключа и значения для каждого встраивания токена (строка в P), и они будут охватывать одну треть каждой строки: 512 элементов — столько же, сколько размер встраивания.

Какой тогда смысл в нескольких головах, спросите вы? Ну, Васвани и др. утверждают, что это позволяет головам захватывать разные подпространства представления. Например; одна голова может отслеживать синтаксические отношения, а другая больше фокусируется на позиционной информации. Есть довольно много работ, в которых исследуется, действительно ли это происходит на практике, например. в переводе. На самом деле, я сам несколько лет назад поработал над этим в резюме.

3. Масштабированиеточечного произведения внимания позволяет избежать очень малых градиентов

Как и в предыдущем пункте, я никогда не задумывался о том, почему мы делим логиты внимания на некоторую константу (🤖), но на самом деле это довольно просто.

Напомним, что каждый логит является результатом скалярного произведения (то есть суммы поэлементного произведения) между запросом и ключевым вектором. Таким образом, большее количество измерений qkv_dim приводит к большему количеству продуктов в этой сумме, что приводит к более высокой дисперсии логитов внимания. Как видно из приведенных ниже примеров, преобразование softmax на логитах с высокой дисперсией приводит к чрезвычайно малым выходным вероятностям и, следовательно, к маленьким градиентам.

4. Исходное встраивание, целевое встраивание и pre-softmax linear используют одну и ту же матрицу весов.

Теперь мы отойдем от многоголового внимания и погрузимся в «привязку веса» — распространенную практику в моделях последовательностей. Я нахожу это довольно интересным, потому что встраивание весовых матриц на самом деле компенсирует огромное количество параметров по сравнению с остальной частью модели. Учитывая словарь из 30 тысяч токенов и размер внедрения 512, эта матрица содержит 15,3 миллиона параметров!

Представьте, что у вас есть три таких матрицы: одна сопоставляет индексы исходных токенов с вложениями, другая сопоставляет целевые токены с вложениями, а третья сопоставляет каждое из самых верхних контекстуализированных вложений токенов декодера с логитами по целевому словарю ( линейный слой до softmax). Ага; у нас остается 46 миллионов параметров.

В коде вы можете видеть, что я инициализирую один слой встраивания в основном классе преобразователя (🤖), который я использую в качестве встраивания кодировщика (�), встраивания декодера (🤖) и веса преобразования декодера до softmax (🤖).

5. Во время вывода мы выполняем прямой проход декодера для каждого токена; на тренировке выполняем один проход вперед на всю последовательность

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

Вывод

Допустим, мы находимся в режиме вывода, в котором авторегрессивно (один за другим) прогнозируем целевые токены. Преобразователь всегда выводит распределение по словарю для каждого токена в пакете. Следующий токен прогнозируется на основе выходного распределения последнего индекса токена в пакете (🤖). Это означает, что мы фактически отбрасываем все выходные распределения для всех предыдущих индексов.

Обучение и принуждение учителей

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

Обратите внимание, что каждый прогноз токена основан на предыдущих токенах достоверности, а не на ранее предсказанных токенах! Обратите также внимание, что этот единственный прямой проход эквивалентен авторегрессионному декодированию с использованием в качестве входных данных только токенов истинности и игнорированием предыдущих прогнозов (!), но гораздо более эффективным. Мы используем внимание маску, чтобы ограничить модуль внутреннего внимания декодера вниманием к будущим токенам (меткам) и мошенничеством.

Я думаю, полезно понимать, что этот способ обучения, называемый принуждением учителя, применяется не только к моделям перевода, но и к наиболее популярным предварительно обученным моделям авторегрессионного языка, таким как GPT-3.

6. Теоретически трансформаторы могут обрабатывать произвольно длинные последовательности…

…на практике, однако, внимание с несколькими головками имеет требования к вычислительным ресурсам и памяти, которые ограничивают длину последовательности примерно 512 токенами. Такие модели, как BERT, на самом деле налагают жесткие ограничения на длину входной последовательности, потому что они используют изученные вложения вместо кодирования синусоиды. Эти изученные позиционные вложенияподобны встраиваниям токенов и аналогичным образом работают только для предопределенного набора позиций до некоторого числа (например, 512 для BERT).

7. Трансформаторы – остаточные потоки

До финальной точки. Мне нравится думать о преобразователе как о нескольких «остаточных потоках». Это похоже на то, как LSTM поддерживает горизонтальный «поток памяти» слева направо при обработке новых токенов один за другим и регулирует поток информации с помощью шлюзов.

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

Следствием этого остаточного потока является то, что количество измерений промежуточных представлений токенов должно быть одинаковым на всех (под)уровнях, поскольку остаточные соединения добавляют два вектора (пример в энкодере: 🤖). Вдобавок ко всему, поскольку слои кодировщика (и аналогичного декодера) накладываются друг на друга, их выходная форма должна совпадать с их входной формой.

Завершение…

Спасибо за прочтение этого поста! Дайте мне знать, если вам понравилось, есть вопросы или вы заметили ошибку. Вы можете написать мне в Twitter или подключиться в LinkedIn. Ознакомьтесь с другими моими сообщениями в блоге на jorisbaan.nl/posts.

Несмотря на то, что я написал код, чтобы его было легко понять (он хорошо документирован, модульно протестирован и проверен), пожалуйста, используйте на практике официальную реализацию PyTorch 😁.

Спасибо Дэвиду Стэпу за идею реализовать трансформер с нуля, Деннису Улмеру и Элизе Бассигнане за отзывы об этом посте, Лукасу де Хаасу за сеанс поиска ошибок и Розали Габбельс за создание визуальных эффектов. Некоторые замечательные ресурсы, которые я искал для вдохновения, — это учебник и реализация преобразователя PyTorch; Трансформер Филиппа Липпе учебник; Annotated Transformer Александра Раша и Illustrated Transformer Джея Аламмара.