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) з двома лінійними перетвореннями (після першого - активаційна функція 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. У наступному документі дані архітектури будуть розглянуті докладніше.

Література

Статті

Вихідні коди