[Перевод] Знакомство с трансформерами. Часть 207.12.2021 15:47

Публикуем вторую часть материала о трансформерах. В первой части речь шла о теоретических основах трансформеров, были показаны примеры их реализации с использованием PyTorch. Здесь поговорим о том, какое место слои внутреннего внимания занимают в нейросетевых архитектурах, и о том, как создают трансформеры, ориентированные на решение различных задач.

Разработка трансформеров

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

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

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

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

К созданию базовых блоков трансформеров подходят по-разному, тут имеются некоторые вариации, но большинство таких блоков структурированы примерно так, как показано ниже.

Базовый блок трансформера

Получается, что в этом блоке, последовательно, входные данные проходят через следующие слои:

Слой внутреннего внимания.
Слой нормализации.
Слой прямого распространения (по одному многослойному перцептрону на каждый вектор).
Ещё один слой нормализации.

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

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

Вот как выглядит блок трансформера в PyTorch:

class TransformerBlock(nn.Module):
  def __init__(self, k, heads):
    super().__init__()

    self.attention = SelfAttention(k, heads=heads)

    self.norm1 = nn.LayerNorm(k)
    self.norm2 = nn.LayerNorm(k)

    self.ff = nn.Sequential(
      nn.Linear(k, 4 * k),
      nn.ReLU(),
      nn.Linear(4 * k, k))

  def forward(self, x):
    attended = self.attention(x)
    x = self.norm1(attended + x)
    
    fedforward = self.ff(x)
    return self.norm2(fedforward + x)

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

Трансформеры и решение задач классификации

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

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

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

Выход: результат классификации

Чаще всего классификаторы последовательностей создают из слоёв, выполняющих преобразование последовательности в последовательность. Делают это, применяя операцию GAP (Global Average Pooling, глобальный усредняющий пуллинг) к итоговой выходной последовательности и отображая результат на классификационный вектор, обрабатываемый с помощью функции Softmax.

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

Вход: использование позиционных эмбеддингов или кодировок

Мы уже говорили о принципах, на которых основаны эмбеддинг-слои. Именно их мы и будем использовать для представления слов.

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

Для этого можно применить довольно простое решение: создаётся ещё один вектор той же длины, представляющий позицию слова в текущем предложении, который добавляется к эмбеддингу слов. Сделать это можно двумя способами.

Первый — это использование позиционных эмбеддингов. Мы просто включаем в состав эмбеддинга позиции — так же, как делали со словами. Например, мы создавали векторы эмбеддинга $v_{cat}$ и $v_{susan}$ , а теперь будем создавать ещё и векторы $v_{12}$ и $v_{25}$ . И так — вплоть до значений, соответствующих ожидаемой длине последовательности. Минус этого подхода заключается в том, что в процессе обучения модели нам нужно будет обработать последовательности всех длин, встречающихся в наборе, иначе модель не научится тому, как воспринимать соответствующие позиционные эмбеддинги. А сильные стороны этого решения в том, что работает оно очень хорошо, и в том, что его легко реализовать.

Второй — это использование позиционных кодировок. Работают они так же, как и позиционные эмбеддинги, за исключением того, что модели не предлагается, в процессе обучения, обрабатывать позиционные вектора. Мы всего лишь выбираем некую функцию f: N→ R^k , применяемую для сопоставления позиций с векторами, содержащими реальные значения, и позволяем сети разобраться с тем, как интерпретировать эти кодировки. Плюс тут в том, что, если функция выбрана удачно, у сети должна появиться возможность работать с последовательностями, которые длиннее, чем те, которые предлагались ей в процессе обучения (модель вряд ли будет показывать на них хорошие результаты, но, как минимум, можно будет её на таких последовательностях проверить). А минус этого подхода в том, что выбор кодировочной функции влияет на сложный гиперпараметр, что немного усложняет реализацию системы.

Мы, ради простоты, воспользуемся в нашей реализации трансформера-классификатора позиционными эмбеддингами.

Реализация трансформера-классификатора в PyTorch

Вот полноценный трансформер, предназначенный для классификации текстов, реализованный средствами PyTorch:

class Transformer(nn.Module):
    def __init__(self, k, heads, depth, seq_length, num_tokens, num_classes):
        super().__init__()

        self.num_tokens = num_tokens
        self.token_emb = nn.Embedding(num_tokens, k)
        self.pos_emb = nn.Embedding(seq_length, k)

        # Последовательность блоков трансформера, на которую 
        # возлагается обязанность решения сложных задач
        tblocks = []
        for i in range(depth):
            tblocks.append(TransformerBlock(k=k, heads=heads))
        self.tblocks = nn.Sequential(*tblocks)

        # Настраиваем соответствие итоговой выходной последовательности ненормализованным значениям, получаемым для различных классов
        self.toprobs = nn.Linear(k, num_classes)

    def forward(self, x):
        """
        :param x: A (b, t) тензор целочисленных значений, представляющий 
                  слова (в некоем заранее заданном словаре).
        :return: A (b, c) тензор логарифмических вероятностей по 
                 классам (где c - это количество классов).
        """
        # генерируем эмбеддинги токенов
        tokens = self.token_emb(x)
        b, t, k = tokens.size()

        # генерируем позиционные эмбеддинги
        positions = torch.arange(t)
        positions = self.pos_emb(positions)[None, :, :].expand(b, t, k)
        
        x = tokens + positions
        x = self.tblocks(x)
        
        # Выполняем операцию усредняющего пуллинга по t измерениям и проецируем 
        # на вероятности, соответствующие классам
        x = self.toprobs(x.mean(dim=1))
        return F.log_softmax(x, dim=1)

Этот трансформер, при глубине 6 и при максимальной длине последовательности 512, достигает точности классификации текстов около 85%, что сопоставимо с результатами RNN-моделей (моделей, основанных на рекуррентных нейронных сетях, Recurrent Neural Network), при условии, что обучение трансформера идёт гораздо быстрее. Для того чтобы довести результаты трансформера до уровня, приближающегося к человеческому, нужно обучать гораздо более глубокую модель на гораздо большем объёме данных. Позже мы поговорим о том, как это сделать.

Трансформеры, генерирующие тексты

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

Общая схема трансформера, генерирующего тексты

Если бы мы применяли RNN, то сейчас у нас уже было бы всё, что нужно, так как такие сети не могут заглядывать в «будущее» входных последовательностей: выход i зависит только от входов от 0 до i. А в случае с трансформерами выход зависит от всей входной последовательности, в результате предсказание следующего символа становится до смешного простым: достаточно взять его из входной последовательности.

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

Маскирование матрицы механизма внутреннего внимания позволяет обеспечить то, что элементы будут реагировать только на входные элементы, которые идут в последовательности до них. Обратите внимание на то, что символ умножения применяется здесь не в обычном для него смысле: мы, на самом деле, устанавливаем замаскированные элементы (белые квадраты) в ∞.

Так как нам нужно, чтобы соответствующие элементы, после применения функции Softmax, стали бы нулями, мы устанавливаем их в ∞. Вот как это выглядит в PyTorch:

dot = torch.bmm(queries, keys.transpose(1, 2))

indices = torch.triu_indices(t, t, offset=1)
dot[:, indices[0], indices[1]] = float('-inf')

dot = F.softmax(dot, dim=2)

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

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

Мы обучаем модель на последовательностях длиной 256, используя систему из 12 блоков трансформера и 256 измерений эмбеддингов. После примерно 24 часов обучения модели на видеокарте RTX 2080Ti (около 170 тысяч пакетов длиной 32) мы позволили модели сгенерировать текст на основе начального фрагмента длиной 256 символов. Для генерирования каждого символа мы передавали модели 256 предыдущих символов, и смотрели на то, что она предложит в качестве следующего символа (последний выходной вектор). При выборе символа использовался метод «температурного сэмплирования» с показателем 0,5. После выбора очередного символа мы переходили к следующему символу.

Модель выдала следующий текст (полужирным выделен начальный фрагмент):

1228X Human & Rousseau. Because many of his stories were originally published in long-forgotten magazines and journals, there are a number of [[anthology|anthologies]] by different collators each containing a different selection. His original books have been considered an anthologie in the [[Middle Ages]], and were likely to be one of the most common in the [[Indian Ocean]] in the [[1st century]]. As a result of his death, the Bible was recognised as a counter-attack by the [[Gospel of Matthew]] (1177–1133), and the [[Saxony|Saxons]] of the [[Isle of Matthew]] (1100–1138), the third was a topic of the [[Saxony|Saxon]] throne, and the [[Roman Empire|Roman]] troops of [[Antiochia]] (1145–1148). The [[Roman Empire|Romans]] resigned in [[1148]] and [[1148]] began to collapse. The [[Saxony|Saxons]] of the [[Battle of Valasander]] reported the y

Обратите внимание на то, что здесь правильно применяются синтаксические конструкции, используемые в Википедии для оформления ссылок, на то, что в ссылках приведены адекватные тексты. А самое важное то, что этот текст обладает хоть какой-то тематической согласованностью. А именно, сгенерированный текст придерживается тем, связанных с Библией и Римской империей, в нём, в разных местах, используются смежные термины. Конечно, нашей модели далеко до более продвинутых систем, вроде GPT-2, но даже здесь видны очевидные преимущества трансформеров перед схожими RNN-моделями: более быстрое обучение (сопоставимую RNN-модель пришлось бы обучать много дней) и лучшая долговременная согласованность.

Если вам интересно — что это за «Battle of Valasander», то знайте, что эту битву, похоже, сеть выдумала сама.

Модель, в описываемом состоянии, обеспечила уровень сжатия, соответствующий 1,343 бита на байт на проверочном наборе, что не так уж и сильно отличается от передовых 0,93 бита на байт, достигнутых моделью GPT-2 (подробнее о ней мы поговорим далее).

Особенности проектирования трансформеров

Для того чтобы понять то, почему трансформеры делают именно такими, полезно разобраться с базовыми факторами, влияющими на процесс их проектирования. Главная цель трансформеров заключалась в решении проблем архитектуры, которая до их появления считалась самой современной. Речь идёт о RNN. Обычно — это LSTM (Long Short-Term Memory, длинная цепь элементов краткосрочной памяти) или GRU (Gated Recurrent Units, управляемые рекуррентные блоки). Вот схема развёрнутой рекуррентной нейронной сети.

Схема рекуррентной нейронной сети

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

Одномерная свёртка

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

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

Остальные особенности архитектуры трансформеров основаны, по большей части, на единственном соображении. Это — глубина сети. В основе архитектур большинства трансформеров лежит желание обучить большую «стопку» блоков трансформера. Заметьте, например, что в трансформере есть всего два места, где производятся нелинейные преобразования данных. Это — место, где в блоке внутреннего внимания применяется функция Softmax, и место, где применяется ReLU (Rectified Linear Unit, выпрямленная линейная функция активации) в слое прямого распространения. Вся остальная модель полностью состоит из линейных трансформаций, что не мешает применению градиента.

Полагаю, что слой нормализации тоже отличается нелинейностью, но это — та самая нелинейность, которая, на самом деле, способствует стабильности градиента при его обратном распространении по сети.

Продолжение следует…

О, а приходите к нам работать?