Трансформер с нуля на 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-обходы и нормализация. Этот рецепт стал основой современных моделей последовательностей.

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