ML: Трансформер

Введение

В этом документе мы продолжаем обсуждать механизм внимания. Многоголовое внимание было введено в статье "Attention is All You Need" (2017) для архитектуры трансформера (transformer). Это был вариант энкодер-декодера для задачи перевода, который не использовал рекуррентных слоёв. Вместо этого последовательности векторов слов пропускались через несколько слоёв с маскированным вниманием. В результате появилась возможность распараллеливания, что существенно ускорило обучение (по сравнению с рекуррентными сетями).

В дальнейшем различные части трансформера были использованы в таких моделях как GPT (2018) и BERT (2018), которым посвящён следующий документ.

Общая архитектура

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

Декодер в режиме принудительного обучения (teacher forcing) на вход получает эту память и слова предложения-перевода. На выходе учится предсказывать этот же перевод сдвинутый влево на одно слово. В режиме тестирования (или "честного" обучения) на декодер сначала подают служебное слово "<BOS>" (begin of sentence) и ожидают на выходе слово "кот". Затем на вход подаётся "<BOS> кот", а на выходе получают "кот сидит" и т.д., пока декодер не выдаст "<EOS>" (end of sentence).

Энкодер и декодер состоят из стопки однотипных блоков, которые осуществляют "глубокое" преобразование входных тензоров.

Энкодер Трансформера

Рассмотрим подробнее энкодер трансформера (рисунок справа). На его вход подаётся тензор формы (N,B,E), где N - число слов во входной последовательности, B - число одновременно обрабатываемых примеров (батч) и E - размерность их векторов эмбединга. Этот тензор пропускается через функцию многоголового само-внимания (Multi-Head Attention): три стрелки на рисунке - это совпадающие запросы, ключи и значения. Результат складывается с исходным тензором и нормируется (см. ниже). Получившийся тензор поступает в полносвязный слой (Feed Forward) c двумя линейными преобразованиями (после первого - активационная функция ReLU): $$ \text{FFN}(\mathbf{x}) = \max(0,~\mathbf{x}\cdot\mathbf{W}_1+\mathbf{b}_1)\cdot\mathbf{W}_2+\mathbf{b}_2. $$ Выход этого слоя снова суммируется с его входом и нормируется. Подобные вычисления повторяются несколько раз (на рисунке множитель L× означает L таких слоёв-блоков с различными параметрами). На выходе последнего блока получается тензор исходной формы (N,B,E), который описывает слова последовательности с учётом контекста всего текста.

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

☝ Нормализация борется с ситуацией, кода веса нейрона "загоняют" его выход в очень большие или очень маленькие значения, что замедляет процесс обучения. Для устранения этого эффекта, из выходов нейронов (до или после активационной функции) вычитается среднее значение и результат делят на стандартное отклонение (корень из дисперсии). Существует два метода нормализции нейронов скрытых слоёв: batch (2015) и layer (2016) normalization. В первом методе усреднения проиводятся по примерам батча, а во втором - по всем нейронам данного слоя. Оба метода ускоряют обучение, однако второй проще, т.к. одинаковым образом работает при обучении и тестировании и не зависит от размера батча. В Трансформере усреднение проводится во всем компонентам вектора эмбединга независимо для каждого входа (слова).

Энкодер в PyTorch

В PyTorch энкодер трансформера строится в два этапа. Сперва определяется TransformerEncoderLayer, а затем с его помощью создаётся собственно энкодер TransformerEncoder:

✒ nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1,activation='relu')

Параметры: d_model = E – размерность входа (вектора одного токена), nhead = H – число голов (должно быть делителем параметра d_model), dim_feedforward – размерность полносвязной сети. Функция активации activation используется в полносвязном слое Feed Forward, а dropout задаёт долю элементов матрицы которые случайным образом делаются нулевыми (борьба с переобучением на этапе тренировки). Слой дропаут стоит сразу после функции softmax в nn.MultiheadAttention.

✒ nn.TransformerEncoder(encoder_layer, num_layers, norm=None)

Параметр encoder_layer является экземпляром класса TransformerEncoderLayer, а num_layers = L задаёт число последовательных блоков, подобных тому, что приведен на рисунке выше.

Приведём пример создания энкодера трансформера: 
N, B, E, H = 10, 32, 512, 8    # число слов, размер батча, размерность эмбединга, голов

encoder_layer       = nn.TransformerEncoderLayer(d_model=E, nhead=H)
transformer_encoder = nn.TransformerEncoder     (encoder_layer, num_layers=1)

src = torch.rand(N, B, E)
out = transformer_encoder(src)   # out.shape == src.shape
Список параметров для одного слоя имеет вид (в именах опущен префикс layers.0.): 
self_attn.in_proj_weight  : 786432  (1536, 512)  # (3*E, E)  Wq, Wk, Wv
self_attn.in_proj_bias    :   1536  (1536,)      # (3*E,)    Bq, Bk, Bv
self_attn.out_proj.weight : 262144  (512, 512)   # (E,E)     Wo
self_attn.out_proj.bias   :    512  (512,)       # (E,)      Bo

linear1.weight            :1048576  (2048, 512)  # (dim_feedforward, E)  W1
linear1.bias              :   2048  (2048,)      # (dim_feedforward,)    B1
norm1.weight              :    512  (512,)       # (E,)
norm1.bias                :    512  (512,)       # (E,)

linear2.weight            :1048576  (512, 2048)  # (E, dim_feedforward)  W2
linear2.bias              :    512  (512,)       # (E,)                  B2
norm2.weight              :    512  (512,)       # (E,)
norm2.bias                :    512  (512,)       # (E,)

total                     :3152384
PyTorch хранит матрицы линеных преобразований в транспонированном виде: $\text{line}(\mathbf{x})=\mathbf{x}\cdot \mathbf{W}^T+\mathbf{b}$. Поэтому в полносвязном слое происходят умножения: (*,E) @ (E, FF) @ (FF, E) = (*,E), где FF = dim_feedforward.
В Vaswani A., et al., (2017), как и выше, были использованы значения E = 512, FF = 2048, поэтому в модуле Feed Forward размерности векторов эмбединга сначала увеличиваются в 4 раза, а потом возвращаются к первоначальному значению.

Кодирование номеров слов

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

В исходной статье (2017) выбирались достаточно специфические периодические функции следующего вида: $$ \text{PosEmb}(\text{pos},~2i) = \sin(\text{pos}/10000^{2i/E}),~~~~~~~ \text{PosEmb}(\text{pos},~2i+1) = \cos(\text{pos}/10000^{2i/E}), $$ где pos - номер слова в предложении, а $i$ - номер компоненты вектора эмбединга. Получившиеся $E$-мерные векторы эмбединга складывались с векторами эмбединга слов.

В дальнейшем (GPT, BERT) использововались обучаемые векторы кодирования положения слов. Для этого, по-мимо эмбединга слов словаря, вводится отдельный (также $E$-мерный) эмбединг положения (для каждого положения pos слова в предложении свой вектор). Векторы слова и положения, как и выше, складываются, а затем поступают в энкодинг трансформера.

Декодер

Добавим теперь к энкодеру декодер, получив полную архитектуру Трансформера. На вход декодера подаются слова целевого предложения перевода. Эти слова векторизуются (с отличным от энкодера эмбедингом) и к ним добавляются векторы номера позиции слова (positional encoding).

Затем векторы проходят блок само-внимания (как в энкодере), для уточнения контекстного смысла векторов. В отличии от энкодера, это самовнимание с маской (Masked Multi-Head Attention), для того, чтобы декодер не заглядывал в "ответ" (подробнее см. ниже). Выход блока самовнимания суммируется с его входом и нормируется.

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

Завершает блок декодера полносвязная сеть (Feed Forward) из двух слоёв (как в энкодере). Таких последовательных блоков декодер имеет несколько (их число L× обычно совпадает с числом блоков энкодера). Естественно, параметры для обучения у блоков отличаются.

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

Трансформер в PyTorch

Трансформер можно собрать из энкодера и декодера (для которого есть свой класс nn.TransformerDecoder, аналогичный nn.TransformerEncoder). Впрочем, можно сразу воспользоваться классом nn.Transformer:

✒ nn.Transformer
(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation='relu', custom_encoder=None, custom_decoder=None)

Смысл параметров понятен из их названий. Функция прямого распространения:

✒ nn.Transformer.forward
(src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)

кроме исходного (src) и целевого (tgt) тензоров содержит маски, играющие важную роль в процессе обучения.

Маски

Для удобства работы с последовательностями переменной длинны (при формировании батчей), гиперпараметры $N$ и $M$ полагаются достаточно большими и более короткие предложения "добиваются" (padding) специальным токеном <PAD> с выделенным индексом (обычно 0). Например, пусть $B, N=1, 10$ (один пример и максимум десять слов в исходном предложении). Тогда для примера из начала документа, последовательность, поступающая в энкодер, имеет вид: 

The cat sits on the mat . <PAD> <PAD> <PAD>
Так как слова <PAD> необходимо игнорировать, в энкодер (и далее в функцию само-внимания) передаётся не только тензор $(N,B,E)$, но и логическая маска src_key_padding_mask: $(B,N)$, в которой значениями True отмечаются "забитые" слова (для каждого примера B). Например, для предложения про кота эта маска имеет вид: 
torch.tensor([[False, False, False, False, False, False, False, True, True, True]])
Маска используется в функции само-внимания для исключения "забитых" слов. Технически это делается замещением элементов матрицы $\mathbf{Q}\cdot\mathbf{K}: ~(N,M)$ большими отрицательными числами -inf в колонках ключей для слов <PAD>. После прохождения через софтмакс соответствующие этим ключам веса будут равны нулю (см предыдущий документ).

В декодере есть два блока внимания (само-внимания и внимание на исходной последовательности на выходе энкодера). Кроме этого, декодер не должен "заглядывать вперёд". Очередное сгенерённое слово $w_i$ в блоке само-внимания может использовать только предыдущие слова $w_1,...,w_{i-1}$. Поэтому требуется три маски: 

def get_tgt_mask(size):
    m = torch.from_numpy(np.triu(np.ones( (size, size) ), k=1).astype('uint8'))
    m = m.float().masked_fill(m == 1, float('-inf')).masked_fill(m == 0, float(0.0))
    return m

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

Литература

Статьи

Исходники