Künstliche Neuronale Netze im Überblick 11: Implementierung eines Transformers

Neuronale Netze sind der Motor vieler Anwendungen in KI und GenAI. Diese Artikelserie gibt einen Einblick in die einzelnen Elemente. Der elfte und letzte Teil implementiert einen Transformer.

Der Transformer ist eine neuronale Architektur, die vollständig auf Aufmerksamkeitsmechanismen basiert und ohne Rekursion und Faltung auskommt, um Sequenzen parallel zu verarbeiten. Seine wichtigste Innovation ist die skalierte Skalarprodukt-Attention-Unterschicht, die gleichzeitig die Beziehungen zwischen allen Positionen in der Eingabe berechnet. Ein Transformer-Encoder stapelt mehrere Schichten von Multi-Head-Self-Attention- und positionsbezogenen Feed-Forward-Netzwerken, die jeweils in Restverbindungen und Schichtnormalisierung eingebettet sind. Der Decoder fügt maskierte Selbstaufmerksamkeit und Encoder-Decoder-Aufmerksamkeit hinzu, um eine autoregressive Generierung zu ermöglichen.

Wir beginnen mit der Formalisierung der skalierten Skalarprodukt-Aufmerksamkeit. Bei gegebenen Abfrage-, Schlüssel- und Wertematrizen Q, K und V mit den Formen (batch_size, num_heads, seq_len, d_k) berechnen wir Rohwerte, indem wir das Punktprodukt von Q mit der Transponierten von K bilden. Anschließend skalieren wir diese Werte mit √d_k, um zu verhindern, dass Extremwerte zu verschwindenden Gradienten führen, wenden Softmax an, um Aufmerksamkeitsgewichte zu erhalten, und multiplizieren mit V, um die beachtete Ausgabe zu erhalten:

Attention(Q, K, V) = softmax( (Q · Kᵀ) / √d_k ) · V

In PyTorch lässt sich die wie folgt umsetzen:

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Berechnet die skalierte Skalarprodukt-Aufmerksamkeit.
    
    Q, K, V haben die Form (batch_size, num_heads, seq_len, d_k).
    Mask, falls angegeben, wird zu den Bewertungen hinzugefügt, um die Aufmerksamkeit auf bestimmte Positionen zu verhindern.
    """
    d_k = Q.size(-1)
    # Berechne die rohen Aufmerksamkeitswerte.
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    # Wende die Maske an (z. B. um zu verhindern, dass nachfolgende Token im Decoder Beachtung finden).
    if mask is not None:
        scores = scores + mask
    # Normalisieren, um Aufmerksamkeitsgewichte zu erhalten.
    attn_weights = F.softmax(scores, dim=-1)
    # Berechne die gewichtete Summe der Werte.
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

In dieser Funktion extrahieren wir die Dimension d_k aus Q, berechnen die Skalarprodukte, skalieren sie und fügen optional vor dem Softmax eine Maske hinzu. Die Maske enthält große negative Werte (−∞) an unzulässigen Positionen, sodass diese Positionen nach dem Softmax das Gewicht Null erhalten.

Multi-Head-Attention erweitert diesen Ansatz, indem es dem Modell ermöglicht, Informationen aus mehreren Darstellungsunterräumen gemeinsam zu berücksichtigen. Zunächst projizieren wir den Eingabetensor X der Form (batch_size, seq_len, d_model) mithilfe gelernter linearer Schichten auf Abfragen, Schlüssel und Werte. Anschließend teilen wir jede dieser Projektionen entlang der Merkmalsdimension in num_heads separate Köpfe auf, wenden die Skalarprodukt-Aufmerksamkeit parallel auf jeden Kopf an, verknüpfen die Ergebnisse und projizieren sie zurück auf das ursprüngliche d_model:

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model muss durch num_heads teilbar sein"
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        # Lineare Projektionen für Abfragen, Schlüssel, Werte und die endgültige Ausgabe.
        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 forward(self, X, mask=None):
        batch_size, seq_len, _ = X.size()
        # Projektionen der Eingaben auf Q, K, V.
        Q = self.W_q(X)
        K = self.W_k(X)
        V = self.W_v(X)
        # Umformen und transponieren, um Köpfe zu trennen.
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        # Skalarprodukt-Aufmerksamkeit anwenden.
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)
        # Köpfe verknüpfen und zurück auf d_model projizieren.
        concat = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        output = self.W_o(concat)
        return output

Da der Transformer keine integrierte Vorstellung von Reihenfolge hat, werden den Token-Einbettungen Positionscodierungen hinzugefügt, um dem Modell Informationen über die Position jedes Elements in der Sequenz zu liefern. Der ursprüngliche Transformer verwendet sinusförmige Codierungen, die wie folgt definiert sind:

P[pos, 2i ] = sin( pos / (10000^(2i/d_model)) )

P[pos, 2i+1 ] = cos( pos / (10000^(2i/d_model)) )

für pos in [0, L−1] und i in [0, d_model/2−1].

Wir implementieren dies wie folgt:

import torch
import math

def get_sinusoidal_positional_encoding(L, d_model):
    # Erstelle einen Tensor der Form (L, d_model).
    P = torch.zeros(L, d_model)
    position = torch.arange(0, L).unsqueeze(1).float()
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
    # Sinus auf gerade Indizes anwenden
    P[:, 0::2] = torch.sin(position * div_term)
    # Cosinus auf ungerade Indizes anwenden.
    P[:, 1::2] = torch.cos(position * div_term)
    return P

Jede Encoder-Schicht besteht aus einer Selbstaufmerksamkeits-Unterschicht, gefolgt von einem positionsbezogenen Feedforward-Netzwerk. Beide Unterschichten sind in Restverbindungen eingeschlossen. Darauf folgen eine Schichtnormalisierung und ein Dropout. Das Feedforward-Netzwerk hat die Form:

FFN(x) = ReLU(x·W₁ + b₁)·W₂ + b₂

und wird unabhängig auf jede Position angewendet. Wir erstellen eine Encoder-Schicht in PyTorch:

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Selbstaufmerksamkeit mit Restverbindung und Normalisierung.
        attn_out = self.self_attn(x, mask)
        x = x + self.dropout1(attn_out)
        x = self.norm1(x)
        # Feed-Forward mit Restverbindung und Normalisierung.
        ffn_out = self.ffn(x)
        x = x + self.dropout2(ffn_out)
        x = self.norm2(x)
        return x

Um den vollständigen Encoder zu erstellen, stapeln wir N solcher Schichten und wenden Positionscodierungen auf die Eingabe an:

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout):
        super(TransformerEncoder, self).__init__()
        self.pos_encoder = get_sinusoidal_positional_encoding
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, src, src_mask=None):
        # src: (batch_size, seq_len, d_model)
        seq_len = src.size(1)
        # Positionskodierung hinzufügen.
        pos_enc = self.pos_encoder(seq_len, src.size(2)).to(src.device)
        x = src + pos_enc.unsqueeze(0)
        # Durchlaufen jeder Encoder-Schicht.
        for layer in self.layers:
            x = layer(x, src_mask)
        return self.norm(x)

Die Implementierung eines Decoders folgt dem gleichen Muster, umfasst jedoch eine maskierte Selbstaufmerksamkeits-Unterschicht, um die Berücksichtigung nachfolgender Positionen zu verhindern, sowie eine Encoder-Decoder-Aufmerksamkeits-Unterschicht, die die Ausgabe des Encoders berücksichtigt. Eine abschließende lineare Softmax-Schicht ordnet die Decoder-Ausgabe den Wahrscheinlichkeiten im Zielvokabular zu.

Durch die Codierung jeder Komponente – von der Skalarprodukt-Aufmerksamkeit über Multi-Head-Aufmerksamkeit, Positionskodierung, Feed-Forward-Netzwerke bis hin zu Encoder-Schichten – erhalten Sie Einblicke in den Informationsfluss durch den Transformer. Auf dieser Grundlage lässt sich das Modell leicht an Aufgaben wie maschinelle Übersetzung, Textzusammenfassung oder sogar Bildgenerierung anpassen oder erweitern.

Vollständiger Encoder-Decoder-Transformer

Um Sequenz-zu-Sequenz-Aufgaben wie maschinelle Übersetzung oder Zusammenfassung durchzuführen, müssen wir den Encoder zu einem gepaarten Decoder erweitern, der jeweils ein Token generiert und dabei die Ausgabe des Encoders berücksichtigt. Ein vollständiger Transformer umfasst somit Token-Einbettungen, Positionskodierungen, einen Stapel von Encoder-Schichten, einen Stapel von Decoder-Schichten und eine abschließende lineare Projektion in das Zielvokabular.

Nachfolgend finden Sie eine schrittweise Implementierung in PyTorch, wobei zu jeder Zeile eine Erläuterung erfolgt.

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Berechnet die skalierte Skalarprodukt-Aufmerksamkeit.
    
    Q, K, V sind Form (batch_size, num_heads, seq_len, d_k).
    Mask, falls angegeben, enthält -inf unzulässige Positionen.
    """
    d_k = Q.size(-1)
    
    # Berechne die rohen Aufmerksamkeitswerte durch Matrixmultiplikation der Abfragen mit den Schlüsseln.
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # Wenn eine Maske angegeben ist, füge sie hinzu (Positionen mit -inf bleiben nach Softmax Null).
    if mask is not None:
        scores = scores + mask
    # Normalisiere die Werte zu Wahrscheinlichkeiten.
    attn_weights = F.softmax(scores, dim=-1)
    # Multipliziere die Wahrscheinlichkeiten mit den Werten, um die beachteten Ausgaben zu erhalten.
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        # Sicherstellen, dass d_model gleichmäßig durch die Anzahl der Köpfe teilbar ist.
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        # Lineare Projektionen für Abfragen, Schlüssel, Werte.
        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)
        # Endgültige lineare Projektion nach Verkettung aller Köpfe.
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        # Projektion der Eingabetensoren in Q, K, V.
        Q = self.W_q(query)
        K = self.W_k(key)
        V = self.W_v(value)
        # In (batch, heads, seq_len, d_k) umformen und transponieren.
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        # Skalierte Skalarprodukt-Aufmerksamkeit pro Kopf anwenden.
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)
        # Köpfe verknüpfen: zurück transponieren und Kopfdimension zusammenführen.
        concat = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        # Endgültige lineare Projektion.
        output = self.W_o(concat)
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # Erstelle einmal sinusförmige Positionskodierungen.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        
        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)
        
        # Batch-Dimension hinzufügen und als Puffer registrieren, damit sie sich mit dem Modell mitbewegt.
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        # x hat die Form (batch_size, seq_len, d_model).
        # Die Positionskodierungen bis zur Eingabelänge hinzufügen.
        x = x + self.pe[:, :x.size(1)]
        return x

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(TransformerEncoderLayer, self).__init__()
        # Selbstaufmerksamkeits-Unterschicht.
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        # Positionsbezogenes Feedforward-Netzwerk.
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        # Layer-Normalisierungsmodule.
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        # Dropout zur Regularisierung.
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        # Selbstaufmerksamkeit anwenden, dann addieren und normieren.
        attn_out = self.self_attn(x, x, x, src_mask)
        x = self.norm1(x + self.dropout1(attn_out))
        # Feedforward-Netzwerk anwenden, dann addieren und normieren.
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_out))
        return x

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(TransformerDecoderLayer, self).__init__()
        # Maskierte Selbstaufmerksamkeit für Zielsequenz.
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        # Encoder-Decoder-Aufmerksamkeit für die Quelle.
        self.src_attn = MultiHeadAttention(d_model, num_heads)
        # Feed-Forward-Netzwerk.
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        
        # Layer-Normen und Dropouts für jede Unterschicht.
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, memory, src_mask=None, tgt_mask=None):
        # Maskierte Selbstaufmerksamkeit auf dem Decodereingang.
        self_attn_out = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout1(self_attn_out))
        # Encoder-Decoder-Aufmerksamkeit über Encoder-Ausgaben.
        src_attn_out = self.src_attn(x, memory, memory, src_mask)
        x = self.norm2(x + self.dropout2(src_attn_out))
        # Feed-forward und add & norm.
        ffn_out = self.ffn(x)
        x = self.norm3(x + self.dropout3(ffn_out))
        return x

class Transformer(nn.Module):
    def __init__(self,
                 src_vocab_size,
                 tgt_vocab_size,
                 d_model=512,
                 num_heads=8,
                 d_ff=2048,
                 num_encoder_layers=6,
                 num_decoder_layers=6,
                 dropout=0.1):
        super(Transformer, self).__init__()
        # Token-Einbettung für Quelle und Ziel.
        self.src_embed = nn.Sequential(
            nn.Embedding(src_vocab_size, d_model),
            PositionalEncoding(d_model)
        )
        
        self.tgt_embed = nn.Sequential(
            nn.Embedding(tgt_vocab_size, d_model),
            PositionalEncoding(d_model)
        )
        # Gestapelte Encoder- und Decoder-Schichten.
        self.encoder_layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_encoder_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            TransformerDecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_decoder_layers)
        ])
        # Endgültige lineare Projektion auf die Vokabulargröße.
        self.generator = nn.Linear(d_model, tgt_vocab_size)
        self.d_model = d_model

    def encode(self, src, src_mask=None):
        # Positionskodierung einbetten und hinzufügen.
        x = self.src_embed(src) * math.sqrt(self.d_model)
        # Durch jede Encoder-Schicht hindurchlaufen.
        for layer in self.encoder_layers:
            x = layer(x, src_mask)
        return x

    def decode(self, tgt, memory, src_mask=None, tgt_mask=None):
        # Ziel einbetten und Positionskodierung hinzufügen.
        x = self.tgt_embed(tgt) * math.sqrt(self.d_model)
        # Durchlaufen jeder Decoderschicht.
        for layer in self.decoder_layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return x

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        # Encoder-Ausgabe berechnen.
        memory = self.encode(src, src_mask)
        # Decoder-Ausgabe unter Berücksichtigung des Encoder-Speichers berechnen.
        output = self.decode(tgt, memory, src_mask, tgt_mask)
        # Auf Vokabular-Logits projizieren.
        return self.generator(output)

def generate_square_subsequent_mask(sz):
    """
    Erstellt eine Maske für kausale Aufmerksamkeit, sodass Position i nur
    auf Positionen ≤ i achten kann. Masken-Einträge sind 0, wo erlaubt, und
    -inf, wo nicht erlaubt.
    """
    mask = torch.triu(torch.full((sz, sz), float('-inf')), diagonal=1)
    return mask

# Anwendungsbeispiel mit Dummy-Daten:
# Vokabulargrößen und Sequenzlängen definieren.
src_vocab_size, tgt_vocab_size = 10000, 10000
batch_size, src_len, tgt_len = 2, 20, 22

# Transformer instanziieren.
model = Transformer(src_vocab_size, tgt_vocab_size)

# Beispiel für Quell- und Zieltoken-Indizes.
src = torch.randint(0, src_vocab_size, (batch_size, src_len))
tgt = torch.randint(0, tgt_vocab_size, (batch_size, tgt_len))

# Keine Füllmaske für dieses Beispiel.
src_mask = None
# Kausale Maske für den Decoder.
tgt_mask = generate_square_subsequent_mask(tgt_len)

# Der Vorwärtsdurchlauf liefert Logits der Form (batch_size, tgt_len, tgt_vocab_size).
logits = model(src, tgt, src_mask, tgt_mask)

In dieser Implementierung wendet jede Encoder-Schicht Multi-Head-Selbstaufmerksamkeit und ein positionsbezogenes Feedforward-Netzwerk an, jeweils mit Restverbindungen und Schichtnormalisierung. Jede Decoder-Schicht fügt einen maskierten Selbstaufmerksamkeitsschritt hinzu, um das Einsehen zukünftiger Token zu verhindern, sowie eine zusätzliche Encoder-Decoder-Aufmerksamkeit, die es dem Decoder ermöglicht, sich auf relevante Teile der Quellsequenz zu konzentrieren. Mit Sinuskurven erstellte Positionscodierungen fügen dem Modell Reihenfolgeinformationen hinzu, und die letzte lineare Schicht projiziert die Decoder-Ausgaben auf rohe Token-Scores.

Auf dieser Grundlage können Sie das Modell anhand von gelabelten (gepaarten) Textdaten trainieren, indem Sie einen geeigneten Verlust definieren (beispielsweise die Kreuzentropie zwischen den vorhergesagten Logits und den tatsächlichen Token-Indizes) und einen der zuvor beschriebenen Optimierer verwenden.

… und noch mehr

Diese Blogserie hat einen Rundflug über das große Terrain der Künstlichen Neuronalen Netze (KNNs) geboten. Dabei haben wir noch gar nicht alle Anwendungen betrachtet, in denen KNNs eine Rolle spielen. Hier sei exemplarisch das Beispiel Reinforcement Learning erwähnt, etwa DQN (Deep-Q Learning).

Durch ihre Struktur erweisen sich KNNs als adäquates Mittel, um statistisches Pattern-Matching durchzuführen, also um bestimmte Muster in Eingangsdaten aufzuspüren. Diese Fähigkeit kommt speziell bei Transformer-Architekturen zum Tragen, die Prompts verarbeiten, um daraus Rückmeldungen zu liefern. Zusätzlich findet dort häufig Reinforcement Learning statt, um die Large Language Models zu bestimmten gewünschten Abläufen zu „überreden“, etwa zum sogenannten Reasoning.

Heutige KNNs sind an biologische neuronale Netze (BNNs) angelehnt, sind diesen gegenüber aber gewaltig eingeschränkt. In Zukunft könnten Wissenschaftler versuchen, sich bei KNNs noch deutlicher ihren biologischen Vorbildern anzunähern. Insgesamt dürften Künstliche Neuronale Netze in der Zukunft Generative KI eine fundamentale und zentrale Rolle spielen.

(rme)