источник), эта маска имеет нижнетреугольный вид.
Один набор $Q$, $K$ и $V$ может отражать только один вид зависимостей между токенами, и матрицы извлекают лишь ограниченный набор информации из входных представлений. Чтобы скомпенсировать эту неоптимальность, авторы архитектуры предложили подход с несколькими «головами» внимания (multi-head attention): по сути вместо одного слоя внимания мы применяем несколько параллельных с разными весами, а потом агрегируем результаты. Рисунок ниже показывает, как выглядит multi-head attention:
Подход к обработке последовательностей целиком через внимание позволяет избавиться от такого понятия, как скрытое состояние, обновляющееся рекуррентно: каждый токен может напрямую «прочитать» любую часть последовательности, наиболее полезную для предсказания. В частности, отсутствие рекуррентности означает, что мы можем применять слой ко всей последовательности одновременно, так как матричные умножения прекрасно параллелятся.
Однако стоит помнить о затратах памяти и времени: поскольку каждый элемент последовательности взаимодействует с каждым, легко показать, что сложность self-attention составляет $O(n^2)$ по длине последовательности, а простые реализации, формирующие полную матрицу внимания, будут расходовать ещё и $O(n^2)$ памяти. С оптимизацией вычислительной сложности внимания связано множество работ как инженерного, так и архитектурного плана: в частности, есть подходы, которые позволяют сократить время работы self-attention до линейного или существенно уменьшают константы за счёт учёта иерархии памяти GPU.
Например, на графиках ниже сравнивается время работы и потребление памяти Transformer со стандартным вниманием и с механизмом из статьи Longformer:
Вторая часть трансформерного блока называется feed-forward network (FFN) и представляет собой два обычных полносвязных слоя, применяемых независимо к каждому элементу входной последовательности. В последних архитектурах размер промежуточного представления (то есть выхода первого слоя) бывает весьма большим — в 4 раза больше выходов блока. Из-за этого вычислительной стоимостью FFN не стоит пренебрегать: несмотря на квадратичную асимптотику внимания, в больших моделях или на коротких последовательностях FFN может занимать существенно больше времени по сравнению с self-attention. В виде формулы применение FFN можно представить так:
$$ \text{FFN}(x)=\text{act}\left(xW_1+b_1\right)W_2+b_2, $$
Промежуточные активации $\text{act}$ в FFN бывают разными: начиналось всё с широко известной ReLU, но в какой-то момент сообщество перешло на GELU (Gaussian Error Linear Unit) с формулой $x\Phi(x)$, где $\Phi$ — функция распределения стандартной нормальной случайной величины.
Сравнение ReLU, ELU и GELU (взято из оригинальной статьи)
Скажем ещё пару слов о layer normalization: как было показано в ряде работ, их положение внутри residual-ветки довольно важно. В стандартной архитектуре используется формулировка PostLN, где нормализация применяется после остаточной связи. Однако такое применение нормализации оказывается довольно нестабильным при обучении моделей с большим числом слоёв: вместо этого предлагается использовать PreLN (справа на рисунке снизу), где нормализация применяется ко входу residual-ветки.
Внимательный читатель может заметить, что все операции внутри transformer-блока, строго говоря, инвариантны к порядку элементов в последовательности. Например, результат внимания зависит от скалярных произведений между эмбеддингами токенов, но расположение этих токенов внутри текста значения не имеет. Таким образом, итоговые представления каждого токена на выходе из модели будут одинаковыми вне зависимости от порядка слов, что вряд ли нас устроит. Как с этим справиться?
На помощь приходит такая вещь, как позиционные эмбеддинги. Это вспомогательные представления, которые прибавляются к обычным эмбеддингам токенов входной последовательности и позволяют слоям внимания различать одинаковые токены на разных местах. Исторически первым подходом были фиксированные эмбеддинги, однозначно кодирующие позицию тригонометрическими функциями (ниже $pos$ — номер позиции, $i$ — индекс элемента в векторе, кодирующем эту позицию, $d$ — размерность эмбеддинга):
С момента появления архитектуры Transformer, однако, появилось множество других способов кодировать позиции токенов. Например, можно просто сделать позиционные эмбеддинги обучаемыми наряду с эмбеддингами токенов. Иной подход — напрямую учесть тот факт, что нам важны не абсолютные позиции токенов, а расстояние между ними, и обучать относительные позиционные представления: подобный подход заметно улучшает качество на чувствительных к порядку слов задачах, а его более современные модификации регулярно используются в самых мощных моделях.
Иллюстрация из статьи про метод позиционного кодирования ALiBi: метод добавляет необучаемые константы к весам внимания в зависимости от расстояния между токенами ключа и значения
Несомненно, Transformer-модели не были бы так интересны, если бы практически все задачи NLP сейчас не решались бы с помощью этой архитектуры. Главными факторами, повлиявшими на бурный рост популярности идеи self-attention, послужили два семейства хорошо всем известных архитектур — BERT и GPT, которые в некотором роде являются энкодером и декодером трансформера, которые зажили своей жизнью.
Модель GPT (Generative Pretrained Transformer) хронологически появилась раньше и представляет собой обычную языковую модель, реализованную в виде последовательности слоев Transformer-декодера. В качестве задачи при обучении выступает обычное предсказание следующего токена (то есть многоклассовая классификация по словарю). Важно, что в качестве маски внимания как раз выступает нижнетреугольная матрица: в противном случае возникла бы утечка в данных из-за того, что токены из «прошлого» будут видеть «будущее». Полученную модель можно использовать для генерации текстов и всех задач, которые на это опираются; даже ChatGPT, обученная на специальных инструкциях, по своей сути незначительно отличается от базовой модели.
Иллюстрация задачи при обучении GPT. Источник
Как понятно из названия, модель Bidirectional Encoder Representations from Transformers (или BERT) отличается от GPT двунаправленностью внимания: это значит, что при обработке входной последовательности все токены могут использовать информацию друг о друге. Это делает такую архитектуру более удобной для задач, где нужно сделать предсказание относительно всего входа целиком без генерации, например, при классификации предложений или поиске пар похожих документов. Важно, что при этом BERT не учится генерировать тексты с нуля: одна из его задач при обучении — это masked language modeling (предсказание случайно замаскированных слов по оставшимся, изображено на рисунке ниже), а вторая — next sentence prediction (предсказание по паре текстовых фрагментов, следуют они друг за другом или нет).
Заметим, что самое ключевое отличие в моделях BERT и GPT (а не в задачах для обучения или применениях) можно свести к использованию разных видов внимания, изображенных на рисунке снизу.
Отличия между вниманием в BERT и GPT (рисунок взят из оригинальной статьи про BERT)
К сожалению, если вы просто напишете код Transformer-нейросети и попробуете сразу обучить что-то содержательное, используя привычные для других архитектур гиперпараметры, то вас с большой вероятностью постигнет неудача. Оптимизационный процесс для таких моделей зачастую требуется изменить, и недостаточное внимание к этому может повлечь за собой существенные потери в итоговом качестве или вообще привести к нестабильному обучению.
Первый момент, на который стоит обратить внимание, — размер батча для обучения. Практически все современные Transformer-модели обучаются на больших батчах, которые для самых больших языковых моделей могут достигать миллионов токенов. Разумеется, ни одна современная GPU не может обработать столько данных за один шаг: на помощь приходят распределенное обучение и чуть более универсальный трюк с аккумуляцией градиентов по микробатчам. Также в последних статьях зачастую прибегают к увеличению размера батча по ходу обучения: идея заключается в том, что на ранних этапах важнее быстрее совершить много шагов градиентного спуска, а на поздних становится важнее иметь точную оценку градиента.
Размер батча может играть большую роль даже для сравнительно маленьких моделей. Иллюстрация из статьи Training Tips for the Transformer Model.
Второй немаловажный фактор — выбор оптимизатора и расписания для learning rate. Обучить трансформер стандартным SGD, скорее всего, не выйдет: в оригинальной статье в качестве оптимизатора использовался Adam, и де-факто он остаётся стандартом до сих пор. Однако стоит заметить, что для больших размеров батча Adam порой работает плохо: из-за этого порой приходится прибегать к алгоритмам наподобие LAMB, нормализующим обновления весов для каждого слоя.
Разумеется, успех этого семейства архитектур на множестве текстовых задач не мог остаться незамеченным для исследователей в других доменах. Одним из наиболее ярких примеров областей, в которой Transformer-модели нашли новое приложение, несомненно, является компьютерное зрение. К примеру, архитектура ViT (Vision Transformer) в свое время побила рекорды качества по классификации изображений, задействуя идею self-attention для картинок, разделенных на множество «лоскутных» (patches) сегментов квадратной формы. Как пишут авторы статьи, идея использовать Transformer-архитектуру в зрении пришла к ним после наблюдения за успехами таких моделей в NLP: использование такого общего подхода, как self-attention, позволяет избежать необходимости явно закладывать в архитектуру особенности задачи (это ещё называют inductive bias) при достаточном времени обучения, числе параметров и размере выборки.
Также именно на трансформерах базируется генеративная часть DALL-E — модели, положившей начало активным исследованиям последних лет в генерации изображений по тексту. Концептуально DALL-E довольно проста: её можно рассматривать как авторегрессивную «языковую модель», генерирующую изображение по одному «визуальному токену» за шаг.
Применяют трансформеры и к обучению с подкреплением: ярким примером является работа Decision Transformer, в которой предлагают использовать авторегрессивное моделирование с использованием этой архитектуры для построения агента. Авторы показали, что такой же подход, который используют для генерации текстов, можно использовать для предсказания действий в динамической среде: как показано на рисунке ниже, модель последовательно принимает стандартные тройки из закодированных состояний, текущих действий и наград и в качестве ответа на каждом шаге выдаёт следующее действие.
Архитектура Decision Transformer (иллюстрация из оригинальной статьи).