Трансформер з нуля на PyTorch: увага, енкодер і декодер

Трансформери показали, що моделювання послідовностей можна повністю побудувати на увазі, без рекурентності й без згорток. Автори оригінальної статті про Transformer запропонували архітектуру енкодер-декодер на основі self-attention і показали, що вона чудово навчається паралельно та сильно виступає в машинному перекладі.

Коли трансформери трапляються лише через високорівневі бібліотеки, їх легко сприйняти як магію. Написати трансформер самому — один із найкращих способів побачити, що насправді відбувається всередині. Нижче — компактна версія ключових частин на PyTorch: спершу кодування позицій, потім обидва різновиди уваги й feed-forward блок, далі шари енкодера й наприкінці начерк декодера. Корисні супутники — сама стаття, Annotated Transformer від Harvard NLP і довідник із torch.nn.

Навіщо писати його вручну

Статтю написано ясно, але перетворення її на робочий код вимагає безлічі дрібних рішень: форми тензорів, маскування, нормалізація, residual-з'єднання і те, як модулі стикуються між собою. Annotated Transformer прославився тим, що перетворює статтю на порядкову навчальну реалізацію, і про архітектуру стає значно легше міркувати.

Власна реалізація також робить відчутними три визначальні ідеї:

1. Увага бере на себе роботу, яку раніше виконувала рекурентність, — пов'язувати токени між собою.
2. Позицію треба додавати явно, бо сама собою увага не відчуває порядку слів.
3. Блоки трансформера модульні, що чудово лягає на torch.nn.Module, ModuleList, лінійні шари, шари нормалізації та dropout.

Проста уявна картинка

У класичній моделі два стеки. Стек енкодера в кожному шарі застосовує self-attention і feed-forward мережу. Стек декодера робить те саме, а також маскований self-attention і увагу до виходу енкодера.

Кожен підшар обгорнуто в residual-з'єднання та layer normalization. Нормалізація працює за останніми вимірами ознак і може використовувати навчувані параметри масштабу та зсуву. Якщо дивитися так, блок виявляється значно менш загадковим, ніж звучить. Здебільшого це лінійні проєкції, увага, нормалізація й residual-шляхи.

Крок 1: ембединги та позиційне кодування

Трансформери працюють із векторами, тому ID токенів спершу перетворюються на щільні ембединги. Потім додається позиційне кодування, щоб модель могла відрізнити «cat chased dog» від «dog chased cat». У статті для цього використовували синусоїдальні кодування.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model)

        self.register_buffer("pe", pe)

    def forward(self, x):
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len]

Вхідний етап поєднує ембединг із масштабувальним множником:

class InputEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.scale = math.sqrt(d_model)

    def forward(self, x):
        return

Множення на sqrt(d_model) відповідає оригінальній архітектурі й тримає величини ембедингів у пропорції з рештою мережі.

Крок 2: scaled dot-product attention

Це ядро моделі. Для запитів Q, ключів K і значень V увага множить Q на транспоновану K, ділить на квадратний корінь із розмірності ключів, застосовує softmax і використовує отримані ваги, щоб змішати V.

Ділення на sqrt(d_k) важливе. Скалярні добутки зростають разом із розмірністю, і без масштабування softmax став би надто гострим.

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

Ця маленька функція дає кожному токену змогу динамічно вирішувати, які інші токени заслуговують на увагу. У self-attention запити, ключі й значення беруться з однієї послідовності. В увазі енкодер-декодер декодер звертається із запитами до виходів енкодера.

Крок 3: multi-head attention

Один великий патерн уваги замінюється кількома меншими. Входи проєктуються в кілька підпросторів, у кожному паралельно працює увага, а результати конкатенуються й проєктуються назад у розмірність моделі.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch_size, seq_len, d_model = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # (batch, heads, seq, d_k)

    def combine_heads(self, x):
        batch_size, heads, seq_len, d_k = x.size()
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, heads * d_k)

    def forward(self, q, k, v, mask=None):
        q = self.split_heads(self.w_q(q))
        k = self.split_heads(self.w_k(k))
        v = self.split_heads(self.w_v(v))

        attn_output, attn_weights = scaled_dot_product_attention(q, k, v, mask)
        output = self.combine_heads(attn_output)
        return self.w_o(output), attn_weights

Система модулів PyTorch тут чудово пасує, бо кожна навчувана проєкція — це просто лінійний шар усередині nn.Module.

Крок 4: feed-forward мережа та layer norm

Після уваги кожна позиція проходить через ту саму feed-forward мережу. В оригінальній архітектурі це два лінійні шари з нелінійністю між ними. Вихід приєднується до residual-шляху й нормалізується.

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        return self.net(x)

LayerNorm у PyTorch нормалізує за останніми вимірами ознак і може навчати поелементні масштаб і зсув — саме так нормалізацію використовують у стеках трансформера. Коли ці частини готові, шар енкодера збирається легко:

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff, dropout)

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

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))

        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        return x

Крок 5: стек енкодера

Оригінальна модель складає в стос кілька однакових шарів енкодера. У PyTorch це зазвичай роблять через nn.ModuleList, який зберігає підмодулі й водночас реєструє їх для оптимізації та збереження чекпоінтів.

class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super().__init__()
        self.embed = InputEmbedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, mask=None):
        x = self.embed(src)
        x = self.pos(x)
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(x, mask)

        return x

На цьому етапі вже є робочий енкодер для задач на кшталт класифікації чи навчання представлень послідовностей. Повній моделі енкодер-декодер далі потрібні стек декодера й фінальна проєкція на словник — таку схему описує стаття, і її ж повторює реалізація Harvard.

Крок 6: декодер

Повний шар декодера вводить дві додаткові ідеї:

1. Маскований self-attention, який не дає токену дивитися на наступні токени під час навчання.
2. Cross-attention, за якого декодер звертає увагу на виходи енкодера.

У грубому псевдокоді один шар декодера працює так:

# Pseudocode structure
x = norm1(x + masked_self_attention(x, x, x))
x = norm2(x + cross_attention(x, encoder_output, encoder_output))
x = norm3(x + feed_forward(x))

Коли енкодер зрозумілий, декодер — це здебільшого те саме ще раз, плюс ще один блок уваги й каузальна маска.

Перевіряйте форми на ходу

Звичка, що заощаджує години, — проганяти випадковий батч через кожну частину одразу, щойно її написано. Подайте тензор форми (batch, sequence, model dimension) у позиційне кодування — форма не повинна змінитися. Пропустіть його через multi-head attention — вихід теж має зберегти цю форму, а ваги уваги повернуться у формі (batch, heads, sequence, sequence). Feed-forward блок усередині розширює останній вимір і повертає його до розміру моделі, тож форма виходу знову збігається з входом.

Якщо якась із цих перевірок не проходить, проблема майже завжди в тому, як голови розділяються чи об'єднуються. Упіймати її на цьому етапі значно дешевше, ніж виявити після довгого навчання, яке тихо видає сміття.

Ще один корисний тест — навмисно крихітна модель. Двох шарів, невеликої розмірності моделі та іграшкового датасету, наприклад копіювання короткої послідовності, досить, щоб побачити, чи справді падає loss. Якщо модель не може навчитися копіювати, десь баг, а якщо може — отже, найважче вона вже робить правильно.

Часті баги

Коли написаний із нуля трансформер відмовляється навчатися, причина зазвичай одна з цих:

• форми тензорів ламаються під час розділення чи об'єднання голів
• помилки broadcasting маски, особливо в self-attention декодера
• забули позиційне кодування, і увага не бачить порядку
• нормалізація стоїть не там, де у варіанті, з яким порівнюють
• шари зберігаються у звичайному списку Python замість ModuleList, тож так і не реєструються

Архітектура елегантна. Проблеми зазвичай ховаються в «сантехніці».

Що винести

Жоден компонент не складний до неможливості. Сила — у комбінації, яка лишається простою, добре паралелиться й багато виражає. Інгредієнти — навчувані ембединги, позиційні сигнали, self-attention, позиційно незалежні feed-forward шари, residual-обходи й нормалізація. Цей рецепт став основою сучасних моделей послідовностей.

Щоб по-справжньому його зрозуміти, вийдіть за межі діаграм. Зберіть модель, друкуйте форми тензорів, спостерігайте, як зміщуються оцінки уваги, і навмисно ламайте маски, перш ніж їх лагодити. Після цього модель перестає здаватися магією й починає виглядати як інженерія.