Capire la Multi-Head Attention nei Transformer

Introduzione al meccanismo di self-attention

Il meccanismo di attention, o più specificamente la self-attention, è stato introdotto per la prima volta nel famoso paper  “Attention is All You Need“ di Vaswani et al., nel 2017. 

Questo meccanismo è ciò che alimenta il modello Transformers, che a sua volta alimenta alcuni dei modelli più potenti che abbiamo, come LLM e VLM. 

Quindi sì, è piuttosto importante!

Ma ecco il punto: molte persone (e anch'io ero tra queste!) credono che l'attenzione sia stata inventata nel 2017.

Non è così.

Il meccanismo di attention originale è stato introdotto per la prima volta nel 2015, quando Bahdanau et al. lo introdussero nel contesto della traduzione automatica. La loro idea era di permettere al decoder di concentrarsi su parti diverse della frase di input a ogni passo di decodifica, invece di fare affidamento su un singolo vettore di contesto fisso.

Quel piccolo cambiamento fece una grande differenza — e piantò il seme di ciò che sarebbe diventato la self-attention e, alla fine, la multi-head attention.

A questo punto potresti avere una domanda:

“Che cosa significa concretamente attention?”

Ti introduco il concetto con due frasi semplici:

• “The cat sat on the mat.” → Il nostro modello potrebbe focalizzare l'attenzione da sat sia su cat sia su mat.
• “She drank her tea because she was thirsty.” → L'attenzione da drank punta sia a tea sia a thirsty, collegando causa e azione.

Da questi esempi possiamo iniziare a capire cosa cerca di fare l'attenzione. È tutta una questione di focus: aiutare il modello a decidere quali altre parole sono rilevanti per quella che sta elaborando, spesso in base al significato semantico o alle relazioni grammaticali.

Una cosa importante da notare è che tutti i token vengono utilizzati nel calcolo dei pesi e dei punteggi di attention (ne parleremo più avanti), non solo alcuni token specifici, come si vede nel diagramma sopra. Questo significa anche che ci sarà un'attenzione più forte tra alcune parole rispetto ad altre, come mostrato dallo spessore delle frecce.

È questo che rende i transformer così bravi a comprendere il linguaggio: usano l'attenzione per mettere in relazione in modo dinamico ogni parola con tutte le altre.

Passi intermedi

Ora abbiamo chiarito cos'è l'attenzione, ma non sappiamo ancora come avviene l'elaborazione interna.

Qui entrano in gioco Query, Key e Value (comunemente abbreviati in Q, K e V). Per il resto dell'articolo, resteremo sul nostro esempio “the cat sat on the mat”. 

Per prima cosa, quando una frase viene data in input al Transformer, viene immediatamente tokenizzata da un tokenizer; un tokenizer famoso è il Byte-Pair Encoding. 

Possiamo pensare alla tokenizzazione come alla suddivisione di un testo in unità più piccole. Poiché tutte le parole della nostra frase sono piuttosto brevi, in questo esempio possiamo assumere che ogni unità corrisponda a una parola. Visivamente, la frase dopo il passaggio attraverso il tokenizer dovrebbe apparire così:

Nota che anche la punteggiatura (il punto) è stata convertita in un Token. Di nuovo, per semplicità, assumiamo che il punto sia un token a sé stante. 

Dopo la tokenizzazione, usiamo un vocabolario del tokenizer per convertire ogni token nel rispettivo Token ID. Il vocabolario del tokenizer si può immaginare come un database/tabella che contiene tutti i token che il modello può riconoscere e i rispettivi Token ID.

Successivamente, ognuno di questi token ID univoci viene mappato al rispettivo embedding tramite una matrice di embedding. È importante notare che tutti questi token embedding saranno vettori della stessa dimensione; nel nostro esempio, assumiamo 512.

L'ultimo passaggio è l'aggiunta dei positional embedding, così che il modello abbia l'informazione di posizione di ciascun token. I due metodi più comuni sono:

1. Codifica sinusoidale (usata nel Transformer originale)
2. Positional embedding appresi (usati in BERT, GPT)

La meccanica di base di questi positional embedding è oltre lo scopo di questo tutorial. Per ora ci basta sapere che abbiamo ottenuto i nostri vettori di input tramite questa formula:

Vettore di input = Token Embedding + Positional Embedding

Query, key e value

Siamo finalmente pronti a introdurre i vettori Q, K, V. Per cominciare, nota che per ogni parola in input creiamo tre vettori:

• Query (Q) – ciò che la parola sta cercando
• Key (K) – ciò che la parola offre alle altre
• Value (V) – la vera informazione che la parola contiene

Queste sono trasformazioni apprendibili e sono anche le rispettive matrici. Matematicamente, queste trasformazioni si scrivono così:

Da questo, ogni token può essere rappresentato da tre nuovi vettori. Ho creato il diagramma qui sotto per aiutare a visualizzare il processo:

Calcolo dello scaled dot-product 

Nel paper “Attention is all you need” è stata fornita una formula matematica che mostra come i nuovi vettori vengano utilizzati per calcolare gli Attention Weights. Vediamola passo per passo, assumendo che stiamo calcolando gli Attention Weights quando siamo su “sat”:

Per iniziare, usiamo il vettore Query di “sat” per confrontarlo con i vettori Key di ogni altra parola — compresa se stessa. Il risultato è un insieme di punteggi di similarità (scalari) che ci dicono quanto ogni parola sia rilevante per “sat”. Più formalmente sono noti come Attention Scores.

Questi punteggi vengono poi scalati tramite sqrt(d_k) per evitare che diventino troppo grandi e per aiutare a stabilizzare i gradienti durante l'addestramento tramite backpropagation. Questi punteggi scalati vengono quindi passati a una softmax, trasformandoli in attention weights. 

La softmax ci dà una probabilità che dice al modello quanto peso assegnare a ciascuna parola e quindi i pesi di attention per un particolare token sommeranno a 1. 

Infine, facciamo una somma pesata dei vettori value usando questi attention weights. 

Otteniamo così un nuovo vettore che rappresenta “sat” nel contesto, arricchito dalle informazioni di tutti i token, ma soprattutto di cat, mat e di altri token rilevanti.

Questo processo avviene per ogni token della frase. Ogni token presta attenzione a ogni altro token, e anche a se stesso. Per questo la chiamiamo self-attention.

Nota anche che abbiamo calcolato il nuovo vettore arricchito per un singolo token, ma in pratica lo faremmo per molti, se non per tutti, i token in una volta sola, poiché la moltiplicazione di matrici è estremamente efficiente. 

Possiamo farlo perché l'architettura è completamente parallelizzabile. A differenza delle recurrent neural network (RNN) che elaborano in sequenza, questo meccanismo consente ai Transformer di elaborare tutti i token contemporaneamente, rendendoli veloci e potenti.

Per aiutarti a visualizzare, guarda il diagramma sopra. Si vede che la matrice di input X viene trasformata nelle rispettive matrici Query, Key e Value, quindi viene applicata la formula della scaled dot-product attention, ottenendo i vettori di embedding arricchiti di tutti i token del nostro input - A.

Implementare la scaled dot-product attention in PyTorch

Ora scriviamo in codice quanto appena discusso in PyTorch. Sorprendentemente, è piuttosto semplice:

import torchimport torch.nn.functional as Fdef scaled_dot_product_attention(Q, K, V):    d_k = Q.size(-1)  # the dimension of the key vectors    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))    attn_weights = F.softmax(scores, dim=-1)    output = torch.matmul(attn_weights, V)    return output, attn_weightsOfferto da

Implementando una semplice funzione, possiamo replicare quanto sopra! Ecco i punti più importanti:

• Q, K e V sono tensori di forma [batch_size, seq_len, dim]
• L'output è un nuovo insieme di vettori in cui ogni token è rappresentato nel contesto dell'intera sequenza
• Gli Attention Scores sono le similarità grezze non normalizzate
• Gli Attention Weights sono normalizzati (tramite softmax) per focalizzare e sommano a 1

Tipi diversi di attention

C'è però un dettaglio importante che non ho ancora detto. Vedi, non esiste un solo tipo di attention nel Transformer: in realtà ce ne sono tre:

1. Encoder self-attention
2. Decoder self-attention
3. Encoder–decoder (cross) attention

Ognuna svolge un ruolo leggermente diverso, quindi passiamole rapidamente in rassegna.

Encoder self-attention

Questo è ciò che abbiamo esplorato finora, in cui nell'encoder ogni token presta attenzione a tutti gli altri token nella stessa sequenza di input, come nell'esempio “the cat sat on the mat.”. 

La cosa importante da notare qui è che non ci sono restrizioni, poiché l'intero input è noto in anticipo, permettendo al modello di costruire rappresentazioni contestuali profonde di ogni token, capendo ciò che viene prima e ciò che viene dopo. 

Nota anche che questo avviene nell'encoder. 

Decoder self-attention

Qui le cose si fanno leggermente più difficili. Poiché ora siamo nel decoder, non vogliamo che ogni token veda il futuro, perché sarebbe barare durante la generazione di testo.

Prendiamo un altro esempio: “The dog barked very loudly.” Quando il modello sta generando “barked” (assumendo per semplicità che ogni parola sia un token), non dovrebbe sapere che il token successivo è “very”, altrimenti il modello non imparerebbe a generare nuovi token. 

Per risolvere questo problema, applichiamo una causal mask (a volte chiamata anche maschera triangolare) che fa sì che ogni token possa prestare attenzione solo a se stesso e ai token precedenti, mai a quelli successivi.

Dal diagramma sopra, possiamo vedere che i punteggi di attenzione verso il futuro vengono rimossi tramite mascheramento, consentendo al modello di imparare a generare token nuovi, accurati e sensati.

Encoder-decoder (cross) attention

Una volta che il decoder ha iniziato a generare, ha comunque bisogno delle informazioni dall'input originale (cioè dalla cosa a cui sta rispondendo). Qui entra in gioco la cross attention.

Qui, il decoder presta attenzione all'output dell'encoder. Quindi, mentre la self-attention del decoder è mascherata, la cross-attention non lo è, e quindi il decoder può guardare a tutti i token dell’encoder contemporaneamente.

Quando stavo imparando questo concetto, mi piaceva molto l'esempio della traduzione. Immagina di tradurre “Bonjour tout le monde” in inglese: il decoder dovrebbe poter guardare tutte le parole francesi mentre decide ciascuna parola inglese, giusto? 

È esattamente ciò che fa la cross attention. Tuttavia, non approfondiremo oltre la cross attention qui, perché non è necessaria per lo scopo di questo articolo.

Perché serve la Multi-Head Attention

Ora capiamo come funziona la self-attention. Ma potresti avere questa domanda:

Se abbiamo già l'attenzione, perché ci servono più teste di attenzione?

Per questo, vorrei presentare l'analogia della torcia, che mi ha aiutato a costruire l'intuizione quando ho iniziato a imparare.

Pensa a una singola testa di attenzione come a una torcia che cerca di illuminare le parti importanti di una frase. Ma una frase avrà molteplici tipi diversi di relazioni da comprendere: soggetto-verbo, oggetto-pronome, aggettivo-sostantivo, dipendenze a lungo raggio, ecc.

Non sarebbe meglio poter catturare più di queste relazioni invece di una sola ogni volta?

La soluzione, quindi, è dare al modello più torce, ciascuna focalizzata su parti diverse della frase. Questa è la multi-head attention.

Invece di calcolare un unico set di Q, K, V ed eseguire un singolo calcolo di attention, dividiamo gli embedding in parti più piccole ed eseguiamo più teste indipendenti di attention in parallelo. 

Come si vede nel diagramma sopra, l'input embedding (per un singolo token in questo caso, per semplicità) è stato suddiviso equamente tra le n teste diverse. 

Nel nostro esempio, supponiamo che la dimensione fosse 512 e il numero di teste 8; quindi, ogni testa di attention elaborerà in parallelo (512/8=64) dimensioni, cercando di catturare le diverse relazioni. 

Per rafforzare ulteriormente l'intuizione su quali tipi di relazioni potrebbero (dico potrebbero perché non comprendiamo ancora completamente quale relazione guardi ciascuna testa) elaborare queste teste, facciamo un esempio semplice:

Consideriamo questa frase:

“The quick brown fox jumps over the lazy dog.”

• Una testa potrebbe concentrarsi su coppie aggettivo–sostantivo: “quick → fox”, “lazy → dog”
• Un'altra testa potrebbe focalizzarsi su soggetto–verbo: “fox → jumps”
• E un'altra testa potrebbe catturare relazioni a lunga distanza: “jumps → over → dog”

La cosa principale da notare è che tutte queste teste girano contemporaneamente e i loro output vengono concatenati e proiettati nuovamente alla dimensione originale dell'embedding (nel nostro caso 512). 

Visualizzare più teste tramite codice

Per consolidare completamente la comprensione della Multi-Head Attention, usiamo la libreria HuggingFace per vedere come funziona il meccanismo di Multi-Head Attention. 

!pip install transformers bertviz torchOfferto da

Per prima cosa, esegui la cella di codice sopra per scaricare tutti i pacchetti rilevanti.

from transformers import BertTokenizer, BertModelfrom bertviz import head_viewimport torchmodel_name = 'bert-base-uncased' # We will be using the Bert modelmodel = BertModel.from_pretrained(model_name, output_attentions=True)tokenizer = BertTokenizer.from_pretrained(model_name)model.eval() # As we are not training, we have set the model to evaluation mode so no updates are made# This is our input sentencesentence = "The quick brown fox jumps over the lazy dog."# Using the tokenizer to convert the sentence into tokensinputs = tokenizer(sentence, return_tensors='pt')input_ids = inputs['input_ids']# Forward pass with attentionswith torch.no_grad():    outputs = model(**inputs)    attentions = outputs.attentions  # Tuple of (num_layers, batch, num_heads, seq_len, seq_len)# Here we decode the tokens tokens = tokenizer.convert_ids_to_tokens(input_ids[0])head_view(attentions, tokens) # VisualizationOfferto da

L'output ottenuto è visibile qui sotto. Nota come ci siano 12 colori diversi, il che significa che in questo modello ci sono 12 teste. Anche lo spessore delle linee è correlato alla magnitudine dei punteggi di attention tra i token. 

Implementare la Multi-Head Attention: best practice

In questa sezione, scriveremo il pseudocodice del meccanismo di Multi-Head Attention. Supponiamo innanzitutto che il nostro input: X abbia forma [batch, seq_len, d_model]. 

Poi possiamo procedere con il pseudocodice: 

Q = X @ W_Q # Calculating the Query valuesK = X @ W_K # Calculating the Key valuesV = X @ W_V# Calculating the Value valuesQ_i, K_i, V_i = split_into_heads(Q, K, V, num_heads) # Now we split the Q,K,V into multiple different heads (by evenly splitting the dimensions)# Here we compute scaled dot-product attention for each headscores = (Q_i @ K_i.T) / sqrt(head_dim)scores += maskA_i = softmax(scores)A_i = dropout(A_i)Z_i = A_i @ V_iZ = concat_heads(Z_i) # We have to concatenate all of the Heads# This is our final linear projection matrixoutput = Z @ W_Ooutput = dropout(output)Offerto da

Tutto questo dovrebbe esserti ormai familiare, tranne forse la matrice finale di proiezione lineare. 

Rivediamola velocemente. 

Dopo che tutte le singole teste di attention hanno terminato il loro lavoro, i loro output vengono concatenati in un unico vettore. Ma questo output combinato deve ancora essere rimappato alla dimensione originale del modello ed è qui che entra in gioco la matrice finale di proiezione lineare. Mappa semplicemente l'output concatenato alla stessa dimensione dell'embedding di input, così può essere passato allo strato successivo del Transformer.

Parliamo anche di alcune ottimizzazioni e best practice:

1. Parallelizzazione: poiché ogni testa è indipendente, possiamo calcolarle tutte in parallelo.
2. Moltiplicazione di matrici: mettiamo tutto in batch usando operazioni di algebra lineare (ad es. matmul), che sono altamente ottimizzate.
3. FlashAttention (Avanzato): è una tecnica più recente che utilizza trick a basso livello in CUDA per ridurre l'overhead di memoria e velocizzare l'attenzione, soprattutto su sequenze lunghe.

In termini di best practice, quando implementiamo il livello di attention, ci sono 3 iperparametri principali da considerare:

1. num_heads: Quante teste di attention eseguire in parallelo, assicurandoci che dividano equamente d_model.
2. head_dim: Dimensione di ciascuna testa.
3. dropout_rate: Per prevenire l'overfitting.

Conclusione

Abbiamo coperto molto in questo articolo. Siamo partiti da un'idea semplice: l'attenzione riguarda il focus. Poi abbiamo esplorato come i Transformer usano questo focus calcolando come ogni parola si relaziona a ogni altra parola usando Query, Key e Value.

Abbiamo visto come la self-attention (e brevemente altre forme di Attention come la Causal e la Cross Attention) fornisca a ciascun token una rappresentazione ricca e consapevole del contesto. Poi siamo andati oltre con la multi-head attention, dove invece di fare affidamento su una sola "torcia", ne forniamo al modello più di una, ciascuna capace di catturare schemi diversi nella frase, tutti in una volta.

È questo che dà ai Transformer il loro vero potere. Non si limitano a leggere; comprendono la struttura sottostante, il significato e le relazioni su intere sequenze. 

In questo articolo abbiamo coperto un aspetto dei Transformer, ma c'è ancora molto da imparare, quindi dai un'occhiata a questo corso completo e dettagliato: Transformer Models with PyTorch Course. 

Se vuoi passare subito al lato applicativo, ti consiglio anche il progetto Developing Multi-Input Models For OCR o il corso Working with HuggingFace. Se preferisci gli articoli, dai un'occhiata a questa Guida completa alla costruzione di un modello Transformer con PyTorch e alla mia altra guida sui Vision Transformer.