Einen Transformer von Grund auf in PyTorch bauen: Attention, Encoder und Decoder

Transformer haben gezeigt, dass sich Sequenzmodellierung vollständig auf Attention aufbauen lässt, ohne Rekurrenz und ohne Faltungen. Die Autoren des ursprünglichen Transformer-Papers schlugen eine Encoder-Decoder-Architektur auf Basis von Self-Attention vor und zeigten, dass sie sich sehr gut parallel trainieren lässt und bei maschineller Übersetzung stark abschneidet.

Wer Transformern nur über High-Level-Bibliotheken begegnet, hält sie leicht für Magie. Selbst einen zu schreiben, ist einer der besten Wege, um zu sehen, was im Inneren tatsächlich passiert. Es folgt eine kompakte PyTorch-Version der Kernteile: zuerst die Kodierung von Positionen, dann beide Arten von Attention und der Feed-forward-Block, danach die Encoder-Schichten und zum Schluss eine Skizze des Decoders. Hilfreiche Begleiter sind das Paper selbst, der Annotated Transformer von Harvard NLP und die torch.nn-Referenz.

Warum von Hand schreiben

Das Paper ist klar, doch es in lauffähigen Code zu übersetzen, erfordert viele kleine Entscheidungen: Tensor-Shapes, Masking, Normalisierung, Residual-Verbindungen und wie die Module zusammenpassen. Der Annotated Transformer wurde berühmt, weil er das Paper in eine zeilenweise kommentierte Lehrimplementierung verwandelt, was das Nachdenken über die Architektur stark erleichtert.

Der Eigenbau macht außerdem drei prägende Ideen greifbar:

1. Attention übernimmt die Aufgabe, die früher die Rekurrenz erledigte, nämlich Tokens zueinander in Beziehung zu setzen.
2. Position muss explizit eingespeist werden, denn Attention allein hat kein Gespür für die Wortreihenfolge.
3. Transformer-Blöcke sind modular, was sehr gut zu torch.nn.Module, ModuleList, linearen Schichten, Normalisierungsschichten und Dropout passt.

Ein einfaches mentales Bild

Das klassische Modell hat zwei Stacks. Der Encoder-Stack wendet in jeder Schicht Self-Attention und ein Feed-forward-Netz an. Der Decoder-Stack tut das ebenfalls, dazu kommen maskierte Self-Attention und Attention über die Ausgabe des Encoders.

Jede Unterschicht ist in eine Residual-Verbindung und Layer Normalization eingebettet. Die Normalisierung wirkt über die letzten Merkmalsdimensionen und kann lernbare Skalierungs- und Verschiebungsparameter nutzen. So betrachtet ist ein Block weniger geheimnisvoll, als er klingt. Er besteht hauptsächlich aus linearen Projektionen, Attention, Normalisierung und Residual-Pfaden.

Schritt 1: Embeddings und Positional Encoding

Transformer arbeiten mit Vektoren, also werden Token-IDs zuerst zu dichten Embeddings. Dann wird ein Positional Encoding addiert, damit das Modell „cat chased dog“ von „dog chased cat“ unterscheiden kann. Das Paper verwendete dafür sinusförmige Kodierungen.

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]

Die Eingabestufe kombiniert ein Embedding mit einem Skalierungsfaktor:

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

Die Multiplikation mit sqrt(d_model) folgt dem ursprünglichen Design und hält die Größenordnung der Embeddings im Verhältnis zum Rest des Netzes.

Schritt 2: Scaled Dot-Product Attention

Das ist der Kern des Modells. Für Queries Q, Keys K und Values V multipliziert Attention Q mit der Transponierten von K, teilt durch die Quadratwurzel der Key-Dimension, wendet Softmax an und mischt V mit den resultierenden Gewichten.

Die Division durch sqrt(d_k) ist wichtig. Skalarprodukte wachsen mit der Dimension, und ohne Skalierung würde Softmax viel zu spitz.

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

Diese kleine Funktion lässt jedes Token dynamisch entscheiden, welche anderen Tokens Aufmerksamkeit verdienen. Bei Self-Attention stammen Queries, Keys und Values aus derselben Sequenz. Bei Encoder-Decoder-Attention richtet der Decoder seine Queries an die Ausgaben des Encoders.

Schritt 3: Multi-Head Attention

Ein großes Attention-Muster wird durch mehrere kleinere ersetzt. Die Eingaben werden in mehrere Unterräume projiziert, in jedem läuft parallel Attention, und die Ergebnisse werden verkettet und zurück auf die Modelldimension projiziert.

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

Das Modulsystem von PyTorch passt hier gut, denn jede gelernte Projektion ist einfach eine lineare Schicht in einem nn.Module.

Schritt 4: Feed-forward-Netz und Layer Norm

Nach der Attention durchläuft jede Position dasselbe Feed-forward-Netz. Das ursprüngliche Design verwendet zwei lineare Schichten mit einer Nichtlinearität dazwischen. Die Ausgabe kommt auf den Residual-Pfad und wird normalisiert.

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 in PyTorch normalisiert über die hinteren Merkmalsdimensionen und kann elementweise Skalierungs- und Bias-Werte lernen – genau so wird Normalisierung in Transformer-Stacks eingesetzt. Mit diesen Bausteinen lässt sich eine Encoder-Schicht leicht zusammensetzen:

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

Schritt 5: Der Encoder-Stack

Das ursprüngliche Modell stapelt mehrere identische Encoder-Schichten. In PyTorch geschieht das meist mit nn.ModuleList, das Untermodule aufnimmt und sie zugleich für Optimierung und Checkpointing registriert.

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

An diesem Punkt steht bereits ein brauchbarer Encoder für Aufgaben wie Klassifikation oder das Lernen von Sequenzrepräsentationen. Das vollständige Encoder-Decoder-Modell braucht dann noch einen Decoder-Stack und eine abschließende Projektion auf das Vokabular – eine Anordnung, die das Paper beschreibt und die Harvard-Implementierung nachbildet.

Schritt 6: Der Decoder

Eine vollständige Decoder-Schicht bringt zwei zusätzliche Ideen mit:

1. Maskierte Self-Attention, die ein Token im Training daran hindert, auf spätere Tokens zu schauen.
2. Cross-Attention, bei der der Decoder auf die Ausgaben des Encoders achtet.

In grobem Pseudocode funktioniert eine Decoder-Schicht so:

# 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))

Wenn der Encoder erst einmal verstanden ist, ist der Decoder größtenteils dasselbe noch einmal, mit einem weiteren Attention-Block und einer kausalen Maske.

Shapes unterwegs prüfen

Eine Gewohnheit, die Stunden spart: Schicken Sie einen zufälligen Batch durch jedes Teil, sobald es geschrieben ist. Geben Sie einen Tensor der Form (batch, sequence, model dimension) in das Positional Encoding, und die Form sollte sich nicht ändern. Schicken Sie ihn durch Multi-Head Attention, und auch die Ausgabe sollte diese Form behalten, während die Attention-Gewichte als (batch, heads, sequence, sequence) zurückkommen. Der Feed-forward-Block verbreitert die letzte Dimension intern und bringt sie wieder auf die Modellgröße, sodass die Ausgabeform erneut der Eingabe entspricht.

Schlägt eine dieser Prüfungen fehl, liegt das Problem fast immer daran, wie Heads aufgeteilt oder zusammengeführt werden. Es in diesem Stadium zu finden, ist viel billiger, als es nach einem langen Trainingslauf zu entdecken, der still und leise Müll produziert.

Ein weiterer nützlicher Test ist ein absichtlich winziges Modell. Zwei Schichten, eine kleine Modelldimension und ein Spielzeugdatensatz wie das Kopieren einer kurzen Sequenz reichen, um zu sehen, ob der Loss tatsächlich sinkt. Ein Modell, das das Kopieren nicht lernt, hat irgendwo einen Bug, und ein Modell, das es kann, macht den schwierigen Teil bereits richtig.

Häufige Bugs

Wenn ein selbst gebauter Transformer nicht trainieren will, liegt es meist an einem dieser Punkte:

• Tensor-Shapes, die beim Aufteilen oder Zusammenführen der Heads schiefgehen
• Broadcasting-Fehler bei Masken, besonders in der Self-Attention des Decoders
• vergessenes Positional Encoding, wodurch Attention blind für die Reihenfolge wird
• Normalisierung an anderer Stelle als in der Variante, mit der verglichen wird
• Schichten in einer gewöhnlichen Python-Liste statt in einer ModuleList, sodass sie nie registriert werden

Die Architektur ist elegant. Probleme verstecken sich meist in den Rohrleitungen.

Was hängen bleiben sollte

Keine einzelne Komponente ist überwältigend kompliziert. Die Stärke liegt in der Kombination, die einfach bleibt, gut parallelisiert und viel ausdrücken kann. Die Zutaten sind gelernte Embeddings, Positionssignale, Self-Attention, positionsweise Feed-forward-Schichten, Residual-Abkürzungen und Normalisierung. Dieses Rezept wurde zur Grundlage moderner Sequenzmodelle.

Um es wirklich zu verstehen, gehen Sie über die Diagramme hinaus. Bauen Sie eins, geben Sie die Tensor-Shapes aus, beobachten Sie, wie sich die Attention-Scores verschieben, und machen Sie die Masken absichtlich kaputt, bevor Sie sie reparieren. Danach fühlt sich das Modell nicht mehr wie Magie an, sondern wie Ingenieurarbeit.