Che cos'è l'attenzione alle query raggruppate (GQA)?

Ainslie et al hanno concepito la Grouped Query Attention come un'ottimizzazione della Multi-Head Attention (MHA), l'innovativo algoritmo di auto-attenzione introdotto nell'importante articolo del 2017 "Attention is All You Need" che ha introdotto le reti neurali dei trasformatori. Più specificamente, è stato proposto come un'applicazione più limitata e una generalizzazione della Multi-Query Attention (MQA), una precedente ottimizzazione dell'MHA.

Sebbene la Multi-Head Attention standard abbia rappresentato un balzo evolutivo in avanti nel machine learning, così come nell'elaborazione del linguaggio naturale (NLP) e nell'AI generativa, è estremamente esigente in termini di risorse di calcolo e larghezza di banda di memoria. Via via che gli LLM diventavano più grandi e sofisticati, questi requisiti di utilizzo della memoria sono diventati un ostacolo al progresso, specialmente per gli LLM decoder-only autoregressivi utilizzati per la generazione di testo, la sintesi e altre attività di AI generativa.

La ricerca successiva si è concentrata sulle tecniche per migliorare o semplificare la Multi-Head Attention. Alcune, come la Flash Attention e la Ring Attention, migliorano il modo in cui le GPU utilizzate per addestrare ed eseguire i modelli gestiscono i calcoli e lo storage. Altri, come GQA e MQA, hanno esplorato le modifiche al modo in cui le architetture trasformative elaborano i token.

La Grouped Query Attention punta a bilanciare i compromessi tra la Multi-Head Attention standard e la Multi-Query Attention. La prima aumenta al massimo l'accuratezza, a discapito di un maggiore sovraccarico della larghezza di banda della memoria e di una riduzione della velocità. La seconda invece aumenta al massimo la velocità e l'efficienza a discapito dell'accuratezza.

Per comprendere come l'attenzione alle query raggruppate ottimizza i modelli trasformatori, è fondamentale capire prima il funzionamento della Multi-Head Attention in generale. Sia la GQA che la MQA si limitano a perfezionare, piuttosto che sostituire, la metodologia di base dell'MHA.

La forza trainante degli LLM e di altri modelli che utilizzano l'architettura trasformativa è rappresentata dall'auto-attenzione, un framework per comprendere le relazioni tra ciascuno dei diversi token in una sequenza. L'auto-attenzione consente a un LLM di interpretare i dati di testo non solo attraverso definizioni di base statiche, ma anche tramite il contesto fornito da altre parole e frasi.

Negli LLM autoregressivi utilizzati per la generazione di testo, il meccanismo di attenzione aiuta il modello a prevedere il token successivo in una sequenza, determinando a quali token precedenti valga la pena "prestare attenzione" in quel momento. Alle informazioni provenienti dai token ritenuti più rilevanti viene assegnato un peso di attenzione maggiore, mentre alle informazioni provenienti dai token ritenuti irrilevanti viene assegnato un peso di attenzione che si avvicina a 0.

Il meccanismo di Multi-Head Attention alla base dei modelli trasformativi genera informazioni contestuali dettagliate calcolando l'auto-attenzione molte volte in parallelo e suddividendo i livelli di attenzione in più "head".

Gli autori di "Attention is All You Need" hanno articolato il suo meccanismo di attenzione utilizzando la terminologia di un database relazionale: query, chiave e valore. I database relazionali sono progettati per semplificare lo storage e il recupero dei dati rilevanti: assegnano un identificatore univoco ("chiave") a ciascun dato e ogni chiave è associata a un valore corrispondente. L'obiettivo di un database relazionale è quello di far corrispondere ogni query alla chiave giusta.

Per ogni token in una sequenza, la Multi-Head Attention richiede la creazione di tre vettori.

• Un vettore di query, Q, rappresenta le informazioni che il token sta "cercando". Per esempio, il vettore di query per un sostantivo potrebbe rappresentare una ricerca di aggettivi che lo descrivono.
  
  
• Un vettore di chiave, K, rappresenta le informazioni contenute nel token. I punteggi di allineamento, che rappresentano la rilevanza del vettore di chiave di ciascun token rispetto al vettore di query di ciascuno degli altri token, vengono utilizzati per calcolare i pesi di attenzione.
  
  
• Un vettore di valore, V, rappresenta le informazioni contestuali che verranno aggiornate dai contributi ponderati in base all'attenzione dei vettori chiave di altri token.

Le interazioni matematiche tra questi tre vettori, mediate dal meccanismo di attenzione, sono il modo in cui un modello regola la sua comprensione specifica del contesto di ciascun token.
 

Per generare ciascuno di questi tre vettori per un dato token, il modello inizia con l'incorporamento del vettore originale di quel token: una codifica numerica in cui ogni dimensione del vettore corrisponde a qualche elemento astratto del significato semantico del token. Il numero di dimensioni in questi vettori è un iperparametro predefinito.

I vettori Q, K e V per ogni token vengono generati facendo passare l'incorporamento del token originale attraverso un livello lineare che precede il primo livello di attenzione. Questo strato lineare è suddiviso in tre matrici univoche di pesi del modello: W_Q, W_K e W_V. I valori di peso specifici in esso contenuti vengono appresi attraverso un pre-addestramento auto-supervisionato su un enorme set di dati di esempi di testo.

Moltiplicando l'incorporamento del vettore originale del token per W_Q, W_K e W_V si ottengono rispettivamente il vettore di query, il vettore di chiave e il vettore di valore corrispondenti. Il numero di dimensioni d che ogni vettore contiene è determinato dalla dimensione di ogni matrice di peso. Q e K avranno lo stesso numero di dimensioni, d_k.

Questi tre vettori vengono quindi passati al livello di attenzione.

Nel livello di attenzione, i vettori Q, K e V vengono utilizzati per calcolare un punteggio di allineamento tra ciascun token in ciascuna posizione di una sequenza. Questi punteggi di allineamento vengono poi normalizzati in pesi di attenzione utilizzando una funzione softmax.

Per ogni token x in una sequenza, i punteggi di allineamento vengono calcolati attraverso il calcolo del "dot product, o prodotto scalare, del vettore di query Q_x di quel token con il vettore di chiave K di ciascuno degli altri token: in altre parole, moltiplicandoli insieme. Se una relazione significativa tra due token si riflette in somiglianze tra i rispettivi vettori, moltiplicandoli insieme si otterrà un valore elevato. Se i due vettori non sono allineati, moltiplicandoli insieme si otterrà un valore piccolo o negativo. La maggior parte dei modelli trasformativi utilizza una variante chiamata Scaled Dot Product Attention in cui QK viene scalato, ovvero moltiplicato, per $\frac{1}{\sqrt{d_k}}$  per migliorare la stabilità dell'addestramento.

Questi punteggi di allineamento delle chiavi di query vengono poi immessi in una funzione softmax. Softmax normalizza tutti gli input a un valore compreso tra 0 e 1 in modo che la somma sia 1. Gli output della funzione softmax sono i pesi di attenzione, ognuno dei quali rappresenta la quota (su 1) dell'attenzione del token x da prestare a ciascuno degli altri token. Se il peso di attenzione di un token è vicino a 0, verrà ignorato. Un peso di attenzione pari a 1 indica che un token riceve l'intera attenzione di xe tutti gli altri saranno ignorati.

Per finire, il vettore di valore per ogni token viene moltiplicato per il suo peso di attenzione. Questi contributi ponderati per l'attenzione di ogni token precedente vengono calcolati come media e aggiunti all'incorporamento vettoriale originale per il token x. Con questo, l'incorporamento del token xviene ora aggiornato per riflettere il contesto fornito dagli altri token nella sequenza che lo riguardano.

L'incorporazione vettoriale aggiornata viene poi inviata a un altro livello lineare, con la propria matrice W_Z, dove il vettore aggiornato al contesto viene normalizzato di nuovo a un numero coerente di dimensioni e poi inviato al livello di attenzione successivo. Ogni livello di attenzione progressivo acquisisce una maggiore sfumatura contestuale.
 

L'uso delle medie dei contributi ponderati per l'attenzione di altri token anziché tenere conto di ogni singolo elemento del contesto ponderato in base all'attenzione è matematicamente efficiente, tuttavia comporta una perdita di dettagli.

Per compensare, le reti trasformative dividono l'incorporamento del token di input originale in hpezzi di dimensioni uguali. Allo stesso modo, dividono W_Q, W_K e W_V in h sottoinsiemi chiamati rispettivamente head di query, head di chiave e head di valore. Ognuna di queste head riceve una parte dell'incorporazione del token originale. I vettori prodotti da ciascuna di queste triplette parallele di head (query, chiave e valore) vengono inseriti in una head di attenzione corrispondente .Per finire, gli output di questi h circuiti paralleli vengono nuovamente concatenati insieme per aggiornare il token.

Durante l'addestramento, ogni circuito apprende pesi distinti che catturano un aspetto separato dei significati semantici. Questo, a sua volta, aiuta il modello a elaborare i diversi modi in cui le implicazioni di una parola possono essere influenzate dal contesto di altre parole che la circondano.

Lo svantaggio della Multi-Head Attention standard non è tanto la presenza di qualche difetto cruciale, quanto piuttosto la mancanza di qualsiasi ottimizzazione. L'MHA è stato il primo algoritmo di questo tipo e rappresenta l'esecuzione più complessa del meccanismo generale di calcolo dell'attenzione.

La maggior parte dell'inefficienza dell'MHA deriva dall'abbondanza di calcoli e parametri del modello. Nell'MHA standard, ogni head (query, chiave e valore) in ogni blocco di attenzione ha la propria matrice di pesi. Ad esempio, un modello con otto head di attenzione ad ogni livello di attenzione, molto meno della maggior parte dei moderni LLM, richiederebbe 24 matrici di peso uniche per le sole head Q, K e V del livello. Ciò comporta un numero enorme di calcoli intermedi a ogni livello.

Una conseguenza di questa configurazione è che è costosa dal punto di vista computazionale. I requisiti di calcolo per l'MHA scalano con proporzione quadratica rispetto alla lunghezza della sequenza: raddoppiare il numero di token in una sequenza di input richiede il quadruplo della complessità. Ciò impone rigorosi limiti pratici alla dimensione delle finestre di contesto.

L'MHA mette anche a dura prova la memoria del sistema. Le GPU non hanno molta memoria integrata per memorizzare gli output dell'enorme quantità di calcoli intermedi che devono essere richiamati in ogni fase di elaborazione successiva. Questi risultati intermedi vengono invece memorizzati nella memoria ad elevata larghezza di banda (HBM), che non si trova sul chip della GPU stessa. Questo comporta una piccola latenza ogni volta che le chiavi e i valori devono essere letti dalla memoria. Quando i modelli trasformativi hanno iniziato a scalare fino a molti miliardi di parametri, il tempo e i calcoli necessari per addestrare ed eseguire l'inferenza sono diventati un rallentamento per le prestazioni del modello.

Ulteriori progressi hanno richiesto metodi per ridurre il numero di passaggi computazionali senza ridurre la capacità dei trasformatori di apprendere e riprodurre modelli linguistici intricatamente complessi; l'introduzione dell'MQA e, successivamente, del GQA è avvenuta in questo contesto..

La Multi-Query Attention (MQA) è un meccanismo di attenzione più efficiente dal punto di vista computazionale che semplifica la Multi-Head Attention per ridurre l'utilizzo della memoria e i calcoli intermedi. Invece di addestrare un'unica head di chiave e una head di valore per ogni head di attenzione, l'MQA utilizza una singola head di chiave e una singola head di valore a ogni livello. Pertanto, i vettori di chiave e i vettori di valore vengono calcolati una sola volta; questo singolo set di vettori di chiave e valore viene quindi condiviso tra tutte le head di attenzione.

Questa semplificazione riduce notevolmente il numero di proiezioni lineari che il modello deve calcolare e memorizzare nella memoria ad elevata larghezza di banda. Secondo il documento del 2019 che ha introdotto l'MQA, l'MQA consente di avere uno storage di coppie chiave-valore (o cache KV) 10-100 volte più piccola e un'inferenza del decoder 12 volte più veloce. Il ridotto utilizzo della memoria di MQA accelera inoltre notevolmente l'addestramento supportando una maggiore dimensione del batch.

Nonostante i suoi benefici, l'MQA presenta diversi aspetti negativi inevitabili.

• Degrado delle prestazioni: non sorprende che la riduzione del numero di parametri unici e addestrabili del modello riduca la capacità di conoscenza e sfumatura del modello. L'MQA comporta un calo significativo della precisione rispetto all'MHA standard, rendendolo inadatto a determinate situazioni e casi d'uso.
  
  
• Deve essere addestrato da zero: un modello addestrato con MHA standard non può essere semplicemente adattato all'MQA, bensì deve essere addestrato con l'MQA da zero. Ciò significa che l'MQA non può essere utilizzata per ottimizzare i modelli esistenti e comporta un considerevole costo opportunità nella sperimentazione dell'MQA per nuovi modelli.
  
  
• Ridondanze nel parallelismo tensoriale: uno dei principali benefici dell'addestramento di modelli trasformativi su GPU è la capacità di eseguire più operazioni tensoriali complesse in parallelo. I valori K e V devono essere presenti su ogni nodo del cluster GPU che esegue queste operazioni, il che significa che in pratica devono essere replicati per ogni nodo. Questo non è un uso ottimale delle risorse di calcolo, nonostante sia ancora più efficiente dell'MHA standard.

La Grouped Query Attention è una formulazione più generale e flessibile della Multi-Query Attention che divide le head di query in più gruppi che condividono un set di chiavi e valori, anziché condividere un set di chiavi e valori in tutte le head di query.

In seguito alla pubblicazione di «GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" nel maggio 2023, molti LLM hanno adottato rapidamente la GQA. Ad esempio, Meta ha adottato per la prima volta la GQA per i suoi modelli Llama 2 nel luglio 2023 e ha mantenuto la GQA nei modelli Llama 3 introdotti nel 2024. Mistral AI ha utilizzato la GQA nel modello Mistral 7B rilasciato nel settembre 2023. Allo stesso modo, i modelli Granite 3.0 di IBM utilizzano la GQA per un'inferenza veloce.

In teoria, la GQA può essere vista come una generalizzazione dello spettro tra l'MHA standard e l'MQA completo. La GQA con lo stesso numero di gruppi di head di chiave-valore delle head di attenzione è l'equivalente dell'MHA standard; la GQA con un gruppo di head è l'equivalente dell'MQA.

Nella pratica, la GQA implica quasi sempre un approccio intermedio, in cui il numero di gruppi è di per sé un importante iperparametro.

La Grouped Query Attention offre diversi vantaggi che hanno portato alla sua adozione relativamente diffusa per i principali LLM.

• Utilizzo efficiente della GPU: la distribuzione di coppie chiave-valore della GQA utilizza al meglio il parallelismo tensoriale, riducendo la quantità di calcolo "sprecata" replicando valori ridondanti.
  
  
• Compromesso efficace: la GQA offre un compromesso ideale tra velocità di inferenza del decoder e precisione delle prestazioni, in quanto è pressoché accurata quanto l'MHA pur essendo veloce quasi quanto l'MQA.
  
  
• Riduzione del sovraccarico della larghezza di banda della memoria: come l'MQA, la GQA riduce significativamente il numero di calcoli intermedi da eseguire, memorizzare e recuperare al momento dell'inferenza.
  
  
• Addestramento flessibile: a differenza dell'MQA, la Group Query Attention non richiede che i modelli siano addestrati da zero utilizzando questo approccio. I modelli pre-addestrati utilizzando l'MHA standard possono essere adattati all'uso della GQA attraverso un processo di messa a punto chiamato "uptraining".