5. Αρχιτεκτονική LLM

Reading time: 18 minutes

Αρχιτεκτονική LLM

tip

Ο στόχος αυτής της πέμπτης φάσης είναι πολύ απλός: Αναπτύξτε την αρχιτεκτονική του πλήρους LLM. Συνδυάστε τα πάντα, εφαρμόστε όλα τα επίπεδα και δημιουργήστε όλες τις λειτουργίες για να παράγετε κείμενο ή να μετατρέπετε κείμενο σε IDs και αντίστροφα.

Αυτή η αρχιτεκτονική θα χρησιμοποιηθεί τόσο για την εκπαίδευση όσο και για την πρόβλεψη κειμένου μετά την εκπαίδευση.

Παράδειγμα αρχιτεκτονικής LLM από https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

Μια υψηλού επιπέδου αναπαράσταση μπορεί να παρατηρηθεί σε:

https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31

  1. Είσοδος (Κωδικοποιημένο Κείμενο): Η διαδικασία ξεκινά με κωδικοποιημένο κείμενο, το οποίο μετατρέπεται σε αριθμητικές αναπαραστάσεις.
  2. Επίπεδο Ενσωμάτωσης Κωδικών και Επίπεδο Θέσης: Το κωδικοποιημένο κείμενο περνά μέσα από ένα επίπεδο ενσωμάτωσης κωδικών και ένα επίπεδο ενσωμάτωσης θέσης, το οποίο καταγράφει τη θέση των κωδικών σε μια ακολουθία, κρίσιμο για την κατανόηση της σειράς των λέξεων.
  3. Μπλοκ Transformer: Το μοντέλο περιέχει 12 μπλοκ transformer, το καθένα με πολλαπλά επίπεδα. Αυτά τα μπλοκ επαναλαμβάνουν την εξής ακολουθία:
  • Masked Multi-Head Attention: Επιτρέπει στο μοντέλο να εστιάζει σε διάφορα μέρη του εισερχόμενου κειμένου ταυτόχρονα.
  • Κανονικοποίηση Επίπεδου: Ένα βήμα κανονικοποίησης για τη σταθεροποίηση και τη βελτίωση της εκπαίδευσης.
  • Επίπεδο Feed Forward: Υπεύθυνο για την επεξεργασία των πληροφοριών από το επίπεδο προσοχής και την πρόβλεψη του επόμενου κωδικού.
  • Επίπεδα Dropout: Αυτά τα επίπεδα αποτρέπουν την υπερβολική προσαρμογή απορρίπτοντας τυχαία μονάδες κατά τη διάρκεια της εκπαίδευσης.
  1. Τελικό Επίπεδο Έξοδου: Το μοντέλο εξάγει έναν τε tensor διαστάσεων 4x50,257, όπου 50,257 αντιπροσωπεύει το μέγεθος του λεξιλογίου. Κάθε γραμμή σε αυτόν τον tensor αντιστοιχεί σε ένα διάνυσμα που το μοντέλο χρησιμοποιεί για να προβλέψει την επόμενη λέξη στην ακολουθία.
  2. Στόχος: Ο στόχος είναι να ληφθούν αυτές οι ενσωματώσεις και να μετατραπούν ξανά σε κείμενο. Συγκεκριμένα, η τελευταία γραμμή της εξόδου χρησιμοποιείται για να παραγάγει την επόμενη λέξη, που αναπαρίσταται ως "forward" σε αυτό το διάγραμμα.

Αναπαράσταση Κώδικα

python
import torch
import torch.nn as nn
import tiktoken

class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

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

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

# Shortcut connection for feed forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

return x


class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(
cfg["emb_dim"], cfg["vocab_size"], bias=False
)

def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
out = model(batch)
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

Συνάρτηση Ενεργοποίησης GELU

python
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

Σκοπός και Λειτουργικότητα

  • GELU (Gaussian Error Linear Unit): Μια συνάρτηση ενεργοποίησης που εισάγει μη γραμμικότητα στο μοντέλο.
  • Ομαλή Ενεργοποίηση: Σε αντίθεση με το ReLU, το οποίο μηδενίζει τις αρνητικές εισόδους, το GELU χαρτογραφεί ομαλά τις εισόδους στις εξόδους, επιτρέποντας μικρές, μη μηδενικές τιμές για αρνητικές εισόδους.
  • Μαθηματικός Ορισμός:

note

Ο στόχος της χρήσης αυτής της συνάρτησης μετά από γραμμικά επίπεδα μέσα στο επίπεδο FeedForward είναι να αλλάξει τα γραμμικά δεδομένα σε μη γραμμικά, επιτρέποντας στο μοντέλο να μάθει πολύπλοκες, μη γραμμικές σχέσεις.

Δίκτυο Νευρώνων FeedForward

Σχήματα έχουν προστεθεί ως σχόλια για να κατανοηθούν καλύτερα τα σχήματα των μητρών:

python
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[1](x) # x shape remains: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[2](x) # x shape: (batch_size, seq_len, emb_dim)
return x  # Output shape: (batch_size, seq_len, emb_dim)

Σκοπός και Λειτουργικότητα

  • Δίκτυο FeedForward κατά Θέση: Εφαρμόζει ένα δίκτυο πλήρως συνδεδεμένο δύο επιπέδων σε κάθε θέση ξεχωριστά και ομοιόμορφα.
  • Λεπτομέρειες Επιπέδου:
  • Πρώτο Γραμμικό Επίπεδο: Επεκτείνει τη διάσταση από emb_dim σε 4 * emb_dim.
  • Ενεργοποίηση GELU: Εφαρμόζει μη γραμμικότητα.
  • Δεύτερο Γραμμικό Επίπεδο: Μειώνει τη διάσταση πίσω σε emb_dim.

note

Όπως μπορείτε να δείτε, το δίκτυο Feed Forward χρησιμοποιεί 3 επίπεδα. Το πρώτο είναι ένα γραμμικό επίπεδο που θα πολλαπλασιάσει τις διαστάσεις κατά 4 χρησιμοποιώντας γραμμικά βάρη (παράμετροι προς εκπαίδευση μέσα στο μοντέλο). Στη συνέχεια, η συνάρτηση GELU χρησιμοποιείται σε όλες αυτές τις διαστάσεις για να εφαρμόσει μη γραμμικές παραλλαγές ώστε να συλλάβει πλουσιότερες αναπαραστάσεις και τελικά ένα άλλο γραμμικό επίπεδο χρησιμοποιείται για να επιστρέψει στο αρχικό μέγεθος των διαστάσεων.

Μηχανισμός Πολυκεφαλικής Προσοχής

Αυτό έχει ήδη εξηγηθεί σε προηγούμενη ενότητα.

Σκοπός και Λειτουργικότητα

  • Πολυκεφαλική Αυτοπροσοχή: Επιτρέπει στο μοντέλο να εστιάζει σε διαφορετικές θέσεις μέσα στην είσοδο κατά την κωδικοποίηση ενός token.
  • Κύρια Στοιχεία:
  • Ερωτήσεις, Κλειδιά, Τιμές: Γραμμικές προβολές της εισόδου, που χρησιμοποιούνται για τον υπολογισμό των σκορ προσοχής.
  • Κεφάλια: Πολλαπλοί μηχανισμοί προσοχής που εκτελούνται παράλληλα (num_heads), καθένας με μειωμένη διάσταση (head_dim).
  • Σκορ Προσοχής: Υπολογίζονται ως το εσωτερικό γινόμενο των ερωτήσεων και των κλειδιών, κλιμακωμένα και μάσκες.
  • Μάσκα: Εφαρμόζεται μια αιτιώδης μάσκα για να αποτραπεί το μοντέλο από το να εστιάζει σε μελλοντικά tokens (σημαντικό για αυτοπαραγωγικά μοντέλα όπως το GPT).
  • Βάρη Προσοχής: Softmax των μάσκων και κλιμακωμένων σκορ προσοχής.
  • Διάνυσμα Πλαισίου: Ζυγισμένο άθροισμα των τιμών, σύμφωνα με τα βάρη προσοχής.
  • Προβολή Εξόδου: Γραμμικό επίπεδο για να συνδυάσει τις εξόδους όλων των κεφαλιών.

note

Ο στόχος αυτού του δικτύου είναι να βρει τις σχέσεις μεταξύ των tokens στο ίδιο πλαίσιο. Επιπλέον, τα tokens χωρίζονται σε διαφορετικά κεφάλια προκειμένου να αποτραπεί η υπερβολική προσαρμογή, αν και οι τελικές σχέσεις που βρίσκονται ανά κεφάλι συνδυάζονται στο τέλος αυτού του δικτύου.

Επιπλέον, κατά την εκπαίδευση εφαρμόζεται μια αιτιώδης μάσκα ώστε τα μελλοντικά tokens να μην λαμβάνονται υπόψη κατά την αναζήτηση συγκεκριμένων σχέσεων με ένα token και εφαρμόζεται επίσης κάποια dropout για να αποτραπεί η υπερβολική προσαρμογή.

Κανονικοποίηση Επίπεδου

python
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5 # Prevent division by zero during normalization.
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

Σκοπός και Λειτουργικότητα

  • Layer Normalization: Μια τεχνική που χρησιμοποιείται για να κανονικοποιήσει τις εισόδους σε όλη τη διάρκεια των χαρακτηριστικών (embedding dimensions) για κάθε μεμονωμένο παράδειγμα σε μια παρτίδα.
  • Συστατικά:
  • eps: Ένας μικρός σταθερός αριθμός (1e-5) που προστίθεται στη διακύμανση για να αποτραπεί η διαίρεση με το μηδέν κατά τη διάρκεια της κανονικοποίησης.
  • scale και shift: Μαθητές παράμετροι (nn.Parameter) που επιτρέπουν στο μοντέλο να κλιμακώνει και να μετατοπίζει την κανονικοποιημένη έξοδο. Αρχικοποιούνται σε ένα και μηδέν, αντίστοιχα.
  • Διαδικασία Κανονικοποίησης:
  • Υπολογισμός Μέσου (mean): Υπολογίζει τον μέσο όρο της εισόδου x σε όλη τη διάρκεια της διάστασης embedding (dim=-1), διατηρώντας τη διάσταση για broadcasting (keepdim=True).
  • Υπολογισμός Διακύμανσης (var): Υπολογίζει τη διακύμανση του x σε όλη τη διάρκεια της διάστασης embedding, διατηρώντας επίσης τη διάσταση. Η παράμετρος unbiased=False διασφαλίζει ότι η διακύμανση υπολογίζεται χρησιμοποιώντας τον μεροληπτικό εκτιμητή (διαιρώντας με N αντί για N-1), που είναι κατάλληλο όταν κανονικοποιούμε σε χαρακτηριστικά αντί για δείγματα.
  • Κανονικοποίηση (norm_x): Αφαιρεί τον μέσο όρο από το x και διαιρεί με την τετραγωνική ρίζα της διακύμανσης συν eps.
  • Κλίμακα και Μετατόπιση: Εφαρμόζει τις μαθητές παραμέτρους scale και shift στην κανονικοποιημένη έξοδο.

note

Ο στόχος είναι να διασφαλιστεί ένας μέσος όρος 0 με διακύμανση 1 σε όλες τις διαστάσεις του ίδιου token. Ο στόχος αυτού είναι να σταθεροποιήσει την εκπαίδευση βαθιών νευρωνικών δικτύων μειώνοντας την εσωτερική μετατόπιση παραλλαγών, η οποία αναφέρεται στην αλλαγή της κατανομής των ενεργοποιήσεων του δικτύου λόγω της ενημέρωσης των παραμέτρων κατά τη διάρκεια της εκπαίδευσης.

Transformer Block

Οι σχήματα έχουν προστεθεί ως σχόλια για να κατανοήσουμε καλύτερα τα σχήματα των μητρών:

python
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"]
)
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for attention block
shortcut = x  # shape: (batch_size, seq_len, emb_dim)
x = self.norm1(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.att(x)    # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for feedforward block
shortcut = x       # shape: (batch_size, seq_len, emb_dim)
x = self.norm2(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.ff(x)     # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

return x  # Output shape: (batch_size, seq_len, emb_dim)

Σκοπός και Λειτουργικότητα

  • Σύνθεση Στρωμάτων: Συνδυάζει multi-head attention, feedforward network, layer normalization και residual connections.
  • Layer Normalization: Εφαρμόζεται πριν από τα στρώματα προσοχής και feedforward για σταθερή εκπαίδευση.
  • Residual Connections (Συντομεύσεις): Προσθέτει την είσοδο ενός στρώματος στην έξοδό του για να βελτιώσει τη ροή του gradient και να επιτρέψει την εκπαίδευση βαθιών δικτύων.
  • Dropout: Εφαρμόζεται μετά από τα στρώματα προσοχής και feedforward για κανονικοποίηση.

Λειτουργικότητα Βήμα-Βήμα

  1. Πρώτη Διαδρομή Residual (Self-Attention):
  • Είσοδος (shortcut): Αποθηκεύστε την αρχική είσοδο για τη σύνδεση residual.
  • Layer Norm (norm1): Κανονικοποιήστε την είσοδο.
  • Multi-Head Attention (att): Εφαρμόστε self-attention.
  • Dropout (drop_shortcut): Εφαρμόστε dropout για κανονικοποίηση.
  • Προσθήκη Residual (x + shortcut): Συνδυάστε με την αρχική είσοδο.
  1. Δεύτερη Διαδρομή Residual (FeedForward):
  • Είσοδος (shortcut): Αποθηκεύστε την ενημερωμένη είσοδο για την επόμενη σύνδεση residual.
  • Layer Norm (norm2): Κανονικοποιήστε την είσοδο.
  • FeedForward Network (ff): Εφαρμόστε τη μετασχηματιστική διαδικασία feedforward.
  • Dropout (drop_shortcut): Εφαρμόστε dropout.
  • Προσθήκη Residual (x + shortcut): Συνδυάστε με την είσοδο από την πρώτη διαδρομή residual.

note

Το μπλοκ transformer ομαδοποιεί όλα τα δίκτυα μαζί και εφαρμόζει κάποια κανονικοποίηση και dropouts για να βελτιώσει τη σταθερότητα και τα αποτελέσματα της εκπαίδευσης.
Σημειώστε πώς γίνονται τα dropouts μετά τη χρήση κάθε δικτύου ενώ η κανονικοποίηση εφαρμόζεται πριν.

Επιπλέον, χρησιμοποιεί επίσης συντομεύσεις που συνίστανται στο να προσθέτουμε την έξοδο ενός δικτύου με την είσοδό του. Αυτό βοηθά στην πρόληψη του προβλήματος της εξαφάνισης του gradient διασφαλίζοντας ότι τα αρχικά στρώματα συμβάλλουν "τόσο όσο" και τα τελευταία.

GPTModel

Οι σχήματα έχουν προστεθεί ως σχόλια για να κατανοήσουμε καλύτερα τα σχήματα των μητρών:

python
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
# shape: (vocab_size, emb_dim)

self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
# shape: (context_length, emb_dim)

self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
# Stack of TransformerBlocks

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
# shape: (emb_dim, vocab_size)

def forward(self, in_idx):
# in_idx shape: (batch_size, seq_len)
batch_size, seq_len = in_idx.shape

# Token embeddings
tok_embeds = self.tok_emb(in_idx)
# shape: (batch_size, seq_len, emb_dim)

# Positional embeddings
pos_indices = torch.arange(seq_len, device=in_idx.device)
# shape: (seq_len,)
pos_embeds = self.pos_emb(pos_indices)
# shape: (seq_len, emb_dim)

# Add token and positional embeddings
x = tok_embeds + pos_embeds  # Broadcasting over batch dimension
# x shape: (batch_size, seq_len, emb_dim)

x = self.drop_emb(x)  # Dropout applied
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.trf_blocks(x)  # Pass through Transformer blocks
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.final_norm(x)  # Final LayerNorm
# x shape remains: (batch_size, seq_len, emb_dim)

logits = self.out_head(x)  # Project to vocabulary size
# logits shape: (batch_size, seq_len, vocab_size)

return logits  # Output shape: (batch_size, seq_len, vocab_size)

Σκοπός και Λειτουργικότητα

  • Embedding Layers:
  • Token Embeddings (tok_emb): Μετατρέπει τους δείκτες των tokens σε embeddings. Ως υπενθύμιση, αυτά είναι τα βάρη που δίνονται σε κάθε διάσταση κάθε token στο λεξιλόγιο.
  • Positional Embeddings (pos_emb): Προσθέτει πληροφορίες θέσης στα embeddings για να καταγράψει τη σειρά των tokens. Ως υπενθύμιση, αυτά είναι τα βάρη που δίνονται στα tokens σύμφωνα με τη θέση τους στο κείμενο.
  • Dropout (drop_emb): Εφαρμόζεται στα embeddings για κανονικοποίηση.
  • Transformer Blocks (trf_blocks): Στοίβα n_layers transformer blocks για την επεξεργασία των embeddings.
  • Final Normalization (final_norm): Κανονικοποίηση επιπέδου πριν από το επίπεδο εξόδου.
  • Output Layer (out_head): Προβάλλει τις τελικές κρυφές καταστάσεις στο μέγεθος του λεξιλογίου για να παραγάγει logits για πρόβλεψη.

note

Ο στόχος αυτής της κλάσης είναι να χρησιμοποιήσει όλα τα άλλα αναφερόμενα δίκτυα για να προβλέψει το επόμενο token σε μια ακολουθία, το οποίο είναι θεμελιώδες για εργασίες όπως η γεννήτρια κειμένου.

Σημειώστε πώς θα χρησιμοποιήσει τόσα πολλά transformer blocks όσο υποδεικνύεται και ότι κάθε transformer block χρησιμοποιεί ένα multi-head attestation net, ένα feed forward net και αρκετές κανονικοποιήσεις. Έτσι, αν χρησιμοποιηθούν 12 transformer blocks, πολλαπλασιάστε αυτό με 12.

Επιπλέον, ένα κανονικοποιητικό επίπεδο προστίθεται πριν από την έξοδο και ένα τελικό γραμμικό επίπεδο εφαρμόζεται στο τέλος για να αποκτήσει τα αποτελέσματα με τις κατάλληλες διαστάσεις. Σημειώστε πώς κάθε τελικός διανύσματος έχει το μέγεθος του χρησιμοποιούμενου λεξιλογίου. Αυτό συμβαίνει επειδή προσπαθεί να αποκτήσει μια πιθανότητα ανά πιθανό token μέσα στο λεξιλόγιο.

Αριθμός Παραμέτρων προς εκπαίδευση

Έχοντας καθορίσει τη δομή του GPT, είναι δυνατόν να βρείτε τον αριθμό των παραμέτρων προς εκπαίδευση:

python
GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

model = GPTModel(GPT_CONFIG_124M)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")
# Total number of parameters: 163,009,536

Βήμα-Βήμα Υπολογισμός

1. Ενσωματωμένα Επίπεδα: Ενσωμάτωση Token & Ενσωμάτωση Θέσης

  • Επίπεδο: nn.Embedding(vocab_size, emb_dim)
  • Παράμετροι: vocab_size * emb_dim
python
token_embedding_params = 50257 * 768 = 38,597,376
  • Επίπεδο: nn.Embedding(context_length, emb_dim)
  • Παράμετροι: context_length * emb_dim
python
position_embedding_params = 1024 * 768 = 786,432

Συνολικές Παράμετροι Ενσωμάτωσης

python
embedding_params = token_embedding_params + position_embedding_params
embedding_params = 38,597,376 + 786,432 = 39,383,808

2. Μπλοκ Μετασχηματιστή

Υπάρχουν 12 μπλοκ μετασχηματιστή, οπότε θα υπολογίσουμε τις παραμέτρους για ένα μπλοκ και στη συνέχεια θα πολλαπλασιάσουμε με το 12.

Παράμετροι ανά Μπλοκ Μετασχηματιστή

α. Πολυκεφαλική Προσοχή

  • Συστατικά:

  • Γραμμικό Επίπεδο Ερώτησης (W_query): nn.Linear(emb_dim, emb_dim, bias=False)

  • Γραμμικό Επίπεδο Κλειδιού (W_key): nn.Linear(emb_dim, emb_dim, bias=False)

  • Γραμμικό Επίπεδο Τιμής (W_value): nn.Linear(emb_dim, emb_dim, bias=False)

  • Έξοδος Πρόβλεψης (out_proj): nn.Linear(emb_dim, emb_dim)

  • Υπολογισμοί:

  • Κάθε από τα W_query, W_key, W_value:

python
qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824

Δεδομένου ότι υπάρχουν τρία τέτοια επίπεδα:

python
total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472
  • Έξοδος Πρόβλεψης (out_proj):
python
out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592
  • Συνολικές Παράμετροι Πολυκεφαλικής Προσοχής:
python
mha_params = total_qkv_params + out_proj_params
mha_params = 1,769,472 + 590,592 = 2,360,064

β. Δίκτυο Τροφοδοσίας

  • Συστατικά:

  • Πρώτο Γραμμικό Επίπεδο: nn.Linear(emb_dim, 4 * emb_dim)

  • Δεύτερο Γραμμικό Επίπεδο: nn.Linear(4 * emb_dim, emb_dim)

  • Υπολογισμοί:

  • Πρώτο Γραμμικό Επίπεδο:

python
ff_first_layer_params = (emb_dim * 4 * emb_dim) + (4 * emb_dim)
ff_first_layer_params = (768 * 3072) + 3072 = 2,359,296 + 3,072 = 2,362,368
  • Δεύτερο Γραμμικό Επίπεδο:
python
ff_second_layer_params = (4 * emb_dim * emb_dim) + emb_dim
ff_second_layer_params = (3072 * 768) + 768 = 2,359,296 + 768 = 2,360,064
  • Συνολικές Παράμετροι Δικτύου Τροφοδοσίας:
python
ff_params = ff_first_layer_params + ff_second_layer_params
ff_params = 2,362,368 + 2,360,064 = 4,722,432

γ. Κανονικοποιήσεις Επίπεδου

  • Συστατικά:
  • Δύο περιπτώσεις LayerNorm ανά μπλοκ.
  • Κάθε LayerNorm έχει 2 * emb_dim παραμέτρους (κλίμακα και μετατόπιση).
  • Υπολογισμοί:
python
layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072

δ. Συνολικές Παράμετροι ανά Μπλοκ Μετασχηματιστή

python
pythonCopy codeparams_per_block = mha_params + ff_params + layer_norm_params_per_block
params_per_block = 2,360,064 + 4,722,432 + 3,072 = 7,085,568

Συνολικοί Παράμετροι για Όλα τα Μπλοκ Μετασχηματιστή

python
pythonCopy codetotal_transformer_blocks_params = params_per_block * n_layers
total_transformer_blocks_params = 7,085,568 * 12 = 85,026,816

3. Τελικά Στρώματα

a. Τελική Κανονικοποίηση Στρώματος

  • Παράμετροι: 2 * emb_dim (κλίμακα και μετατόπιση)
python
pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

β. Στρώμα Έξοδου Πρόβλεψης (out_head)

  • Στρώμα: nn.Linear(emb_dim, vocab_size, bias=False)
  • Παράμετροι: emb_dim * vocab_size
python
pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376

4. Συνοψίζοντας Όλες τις Παραμέτρους

python
pythonCopy codetotal_params = (
embedding_params +
total_transformer_blocks_params +
final_layer_norm_params +
output_projection_params
)
total_params = (
39,383,808 +
85,026,816 +
1,536 +
38,597,376
)
total_params = 163,009,536

Δημιουργία Κειμένου

Έχοντας ένα μοντέλο που προβλέπει το επόμενο token όπως το προηγούμενο, χρειάζεται απλώς να πάρουμε τις τελευταίες τιμές token από την έξοδο (καθώς θα είναι αυτές του προβλεπόμενου token), οι οποίες θα είναι μια τιμή ανά είσοδο στο λεξιλόγιο και στη συνέχεια να χρησιμοποιήσουμε τη συνάρτηση softmax για να κανονικοποιήσουμε τις διαστάσεις σε πιθανότητες που αθροίζουν σε 1 και στη συνέχεια να πάρουμε τον δείκτη της μεγαλύτερης εισόδου, ο οποίος θα είναι ο δείκτης της λέξης μέσα στο λεξιλόγιο.

Κώδικας από https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

python
def generate_text_simple(model, idx, max_new_tokens, context_size):
# idx is (batch, n_tokens) array of indices in the current context
for _ in range(max_new_tokens):

# Crop current context if it exceeds the supported context size
# E.g., if LLM supports only 5 tokens, and the context size is 10
# then only the last 5 tokens are used as context
idx_cond = idx[:, -context_size:]

# Get the predictions
with torch.no_grad():
logits = model(idx_cond)

# Focus only on the last time step
# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
logits = logits[:, -1, :]

# Apply softmax to get probabilities
probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

# Get the idx of the vocab entry with the highest probability value
idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

# Append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

return idx


start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)

encoded_tensor = torch.tensor(encoded).unsqueeze(0)
print("encoded_tensor.shape:", encoded_tensor.shape)

model.eval() # disable dropout

out = generate_text_simple(
model=model,
idx=encoded_tensor,
max_new_tokens=6,
context_size=GPT_CONFIG_124M["context_length"]
)

print("Output:", out)
print("Output length:", len(out[0]))

Αναφορές