L'obiettivo di questa quinta fase è molto semplice: Sviluppare l'architettura del LLM completo. Metti tutto insieme, applica tutti i livelli e crea tutte le funzioni per generare testo o trasformare testo in ID e viceversa.
Questa architettura sarà utilizzata sia per l'addestramento che per la previsione del testo dopo che è stato addestrato.
Una rappresentazione ad alto livello può essere osservata in:
Input (Testo Tokenizzato): Il processo inizia con testo tokenizzato, che viene convertito in rappresentazioni numeriche.
Layer di Embedding dei Token e Layer di Embedding Posizionale: Il testo tokenizzato passa attraverso un layer di embedding dei token e un layer di embedding posizionale, che cattura la posizione dei token in una sequenza, fondamentale per comprendere l'ordine delle parole.
Blocchi Transformer: Il modello contiene 12 blocchi transformer, ciascuno con più livelli. Questi blocchi ripetono la seguente sequenza:
Attenzione Multi-Testa Mascherata: Consente al modello di concentrarsi su diverse parti del testo di input contemporaneamente.
Normalizzazione del Livello: Un passaggio di normalizzazione per stabilizzare e migliorare l'addestramento.
Layer Feed Forward: Responsabile dell'elaborazione delle informazioni dal layer di attenzione e della formulazione di previsioni sul token successivo.
Layer di Dropout: Questi layer prevengono l'overfitting eliminando casualmente unità durante l'addestramento.
Layer di Output Finale: Il modello produce un tensore di dimensione 4x50,257, dove 50,257 rappresenta la dimensione del vocabolario. Ogni riga in questo tensore corrisponde a un vettore che il modello utilizza per prevedere la prossima parola nella sequenza.
Obiettivo: L'obiettivo è prendere questi embedding e convertirli di nuovo in testo. In particolare, l'ultima riga dell'output viene utilizzata per generare la prossima parola, rappresentata come "avanti" in questo diagramma.
Rappresentazione del Codice
import torchimport torch.nn as nnimport tiktokenclassGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.5* x * (1+ torch.tanh(torch.sqrt(torch.tensor(2.0/ torch.pi)) *(x +0.044715* torch.pow(x, 3))))classFeedForward(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"]),)defforward(self,x):return self.layers(x)classMultiHeadAttention(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_outself.num_heads = num_headsself.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dimself.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 outputsself.dropout = nn.Dropout(dropout)self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))defforward(self,x):b, num_tokens, d_in = x.shapekeys = 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 maskattn_scores = queries @ keys.transpose(2, 3)# Dot product for each head# Original mask truncated to the number of tokens and converted to booleanmask_bool = self.mask.bool()[:num_tokens,:num_tokens]# Use the mask to fill attention scoresattn_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_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec)# optional projectionreturn context_vecclassLayerNorm(nn.Module):def__init__(self,emb_dim):super().__init__()self.eps =1e-5self.scale = nn.Parameter(torch.ones(emb_dim))self.shift = nn.Parameter(torch.zeros(emb_dim))defforward(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.shiftclassTransformerBlock(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"])defforward(self,x):# Shortcut connection for attention blockshortcut = xx = 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 blockshortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_shortcut(x)x = x + shortcut # Add the original input backreturn xclassGPTModel(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 _ inrange(cfg["n_layers"])])self.final_norm =LayerNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)defforward(self,in_idx):batch_size, seq_len = in_idx.shapetok_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 logitsGPT_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)
Funzione di Attivazione GELU
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classGELU(nn.Module):def__init__(self):super().__init__()defforward(self,x):return0.5* x * (1+ torch.tanh(torch.sqrt(torch.tensor(2.0/ torch.pi)) *(x +0.044715* torch.pow(x, 3))))
Scopo e Funzionalità
GELU (Gaussian Error Linear Unit): Una funzione di attivazione che introduce non linearità nel modello.
Attivazione Liscia: A differenza di ReLU, che annulla gli input negativi, GELU mappa dolcemente gli input sugli output, consentendo valori piccoli e non nulli per gli input negativi.
Definizione Matematica:
L'obiettivo dell'uso di questa funzione dopo i livelli lineari all'interno del livello FeedForward è cambiare i dati lineari in dati non lineari per consentire al modello di apprendere relazioni complesse e non lineari.
Rete Neurale FeedForward
I formati sono stati aggiunti come commenti per comprendere meglio le forme delle matrici:
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classFeedForward(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"]),)defforward(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)
Scopo e Funzionalità
Rete FeedForward a livello di posizione: Applica una rete completamente connessa a due strati a ciascuna posizione separatamente e in modo identico.
Dettagli del Livello:
Primo Livello Lineare: Espande la dimensionalità da emb_dim a 4 * emb_dim.
Attivazione GELU: Applica non-linearità.
Secondo Livello Lineare: Riduce la dimensionalità di nuovo a emb_dim.
Come puoi vedere, la rete Feed Forward utilizza 3 strati. Il primo è uno strato lineare che moltiplicherà le dimensioni per 4 utilizzando pesi lineari (parametri da addestrare all'interno del modello). Poi, la funzione GELU è utilizzata in tutte quelle dimensioni per applicare variazioni non lineari per catturare rappresentazioni più ricche e infine un altro strato lineare è utilizzato per tornare alla dimensione originale.
Meccanismo di Attenzione Multi-Testa
Questo è già stato spiegato in una sezione precedente.
Scopo e Funzionalità
Auto-Attenzione Multi-Testa: Consente al modello di concentrarsi su diverse posizioni all'interno della sequenza di input durante la codifica di un token.
Componenti Chiave:
Query, Chiavi, Valori: Proiezioni lineari dell'input, utilizzate per calcolare i punteggi di attenzione.
Teste: Molteplici meccanismi di attenzione che funzionano in parallelo (num_heads), ciascuno con una dimensione ridotta (head_dim).
Punteggi di Attenzione: Calcolati come il prodotto scalare di query e chiavi, scalati e mascherati.
Mascheramento: Una maschera causale è applicata per impedire al modello di prestare attenzione ai token futuri (importante per modelli autoregressivi come GPT).
Pesi di Attenzione: Softmax dei punteggi di attenzione mascherati e scalati.
Vettore di Contesto: Somma pesata dei valori, secondo i pesi di attenzione.
Proiezione di Uscita: Strato lineare per combinare le uscite di tutte le teste.
L'obiettivo di questa rete è trovare le relazioni tra i token nello stesso contesto. Inoltre, i token sono divisi in diverse teste per prevenire l'overfitting anche se le relazioni finali trovate per testa sono combinate alla fine di questa rete.
Inoltre, durante l'addestramento viene applicata una maschera causale in modo che i token successivi non vengano presi in considerazione quando si cercano le relazioni specifiche a un token e viene applicato anche un dropout per prevenire l'overfitting.
Normalizzazione del Livello
# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04classLayerNorm(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))defforward(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
Scopo e Funzionalità
Layer Normalization: Una tecnica utilizzata per normalizzare gli input attraverso le caratteristiche (dimensioni di embedding) per ciascun esempio individuale in un batch.
Componenti:
eps: Una piccola costante (1e-5) aggiunta alla varianza per prevenire la divisione per zero durante la normalizzazione.
scale e shift: Parametri apprendibili (nn.Parameter) che consentono al modello di scalare e spostare l'output normalizzato. Sono inizializzati a uno e zero, rispettivamente.
Processo di Normalizzazione:
Calcola Media (mean): Calcola la media dell'input x attraverso la dimensione di embedding (dim=-1), mantenendo la dimensione per il broadcasting (keepdim=True).
Calcola Varianza (var): Calcola la varianza di x attraverso la dimensione di embedding, mantenendo anche la dimensione. Il parametro unbiased=False garantisce che la varianza venga calcolata utilizzando l'estimatore biased (dividendo per N invece di N-1), che è appropriato quando si normalizza sulle caratteristiche piuttosto che sui campioni.
Normalizza (norm_x): Sottrae la media da x e divide per la radice quadrata della varianza più eps.
Scala e Sposta: Applica i parametri apprendibili scale e shift all'output normalizzato.
L'obiettivo è garantire una media di 0 con una varianza di 1 attraverso tutte le dimensioni dello stesso token. L'obiettivo di questo è stabilizzare l'addestramento delle reti neurali profonde riducendo il cambiamento interno della covariata, che si riferisce al cambiamento nella distribuzione delle attivazioni della rete a causa dell'aggiornamento dei parametri durante l'addestramento.
Blocco Transformer
Sono state aggiunte forme come commenti per comprendere meglio le forme delle matrici:
Composizione dei Livelli: Combina attenzione multi-testa, rete feedforward, normalizzazione dei livelli e connessioni residue.
Normalizzazione dei Livelli: Applicata prima dei livelli di attenzione e feedforward per un addestramento stabile.
Connessioni Residue (Scorciatoie): Aggiungono l'input di un livello alla sua uscita per migliorare il flusso del gradiente e abilitare l'addestramento di reti profonde.
Dropout: Applicato dopo i livelli di attenzione e feedforward per la regolarizzazione.
Funzionalità Passo-Passo
Primo Percorso Residuo (Auto-Attenzione):
Input (shortcut): Salva l'input originale per la connessione residua.
Dropout (drop_shortcut): Applica dropout per la regolarizzazione.
Aggiungi Residuo (x + shortcut): Combina con l'input originale.
Secondo Percorso Residuo (FeedForward):
Input (shortcut): Salva l'input aggiornato per la prossima connessione residua.
Layer Norm (norm2): Normalizza l'input.
FeedForward Network (ff): Applica la trasformazione feedforward.
Dropout (drop_shortcut): Applica dropout.
Aggiungi Residuo (x + shortcut): Combina con l'input dal primo percorso residuo.
Il blocco transformer raggruppa tutte le reti insieme e applica alcune normalizzazioni e dropout per migliorare la stabilità e i risultati dell'addestramento.
Nota come i dropout siano effettuati dopo l'uso di ciascuna rete mentre la normalizzazione è applicata prima.
Inoltre, utilizza anche scorciatoie che consistono nell'aggiungere l'uscita di una rete con il suo input. Questo aiuta a prevenire il problema del gradiente che svanisce assicurando che i livelli iniziali contribuiscano "tanto" quanto quelli finali.
GPTModel
Le forme sono state aggiunte come commenti per comprendere meglio le forme delle matrici:
Token Embeddings (tok_emb): Converte gli indici dei token in embedding. Come promemoria, questi sono i pesi dati a ciascuna dimensione di ciascun token nel vocabolario.
Positional Embeddings (pos_emb): Aggiunge informazioni posizionali agli embedding per catturare l'ordine dei token. Come promemoria, questi sono i pesi dati ai token in base alla loro posizione nel testo.
Dropout (drop_emb): Applicato agli embedding per la regolarizzazione.
Blocchi Transformer (trf_blocks): Stack di n_layers blocchi transformer per elaborare gli embedding.
Normalizzazione Finale (final_norm): Normalizzazione dello strato prima dello strato di output.
Strato di Output (out_head): Proietta gli stati nascosti finali alla dimensione del vocabolario per produrre logit per la previsione.
L'obiettivo di questa classe è utilizzare tutte le altre reti menzionate per prevedere il prossimo token in una sequenza, fondamentale per compiti come la generazione di testo.
Nota come utilizzerà tanti blocchi transformer quanto indicato e che ogni blocco transformer utilizza una rete di attenzione multi-testa, una rete feed forward e diverse normalizzazioni. Quindi, se vengono utilizzati 12 blocchi transformer, moltiplica questo per 12.
Inoltre, uno strato di normalizzazione è aggiunto prima dell'output e uno strato lineare finale è applicato alla fine per ottenere i risultati con le dimensioni appropriate. Nota come ogni vettore finale abbia la dimensione del vocabolario utilizzato. Questo perché sta cercando di ottenere una probabilità per ogni possibile token all'interno del vocabolario.
Numero di Parametri da addestrare
Avendo definito la struttura GPT, è possibile scoprire il numero di parametri da addestrare:
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
Calcolo Passo-Passo
1. Strati di Embedding: Token Embedding & Position Embedding
Avendo un modello che predice il prossimo token come quello precedente, è sufficiente prendere i valori dell'ultimo token dall'output (poiché saranno quelli del token previsto), che saranno un valore per voce nel vocabolario e poi utilizzare la funzione softmax per normalizzare le dimensioni in probabilità che sommano a 1 e poi ottenere l'indice della voce più grande, che sarà l'indice della parola all'interno del vocabolario.
defgenerate_text_simple(model,idx,max_new_tokens,context_size):# idx is (batch, n_tokens) array of indices in the current contextfor _ inrange(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 contextidx_cond = idx[:,-context_size:]# Get the predictionswith 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 probabilitiesprobas = torch.softmax(logits, dim=-1)# (batch, vocab_size)# Get the idx of the vocab entry with the highest probability valueidx_next = torch.argmax(probas, dim=-1, keepdim=True)# (batch, 1)# Append sampled index to the running sequenceidx = torch.cat((idx, idx_next), dim=1)# (batch, n_tokens+1)return idxstart_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 dropoutout =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]))