L'attenzione lineare è tutto ciò di cui hai bisogno. Auto-attenzione a una frazione del… | di Sam Maddrell-Mander | Giugno 2024 | Intelligenza-Artificiale

Questo è il genere di cose che chiunque abbia passato molto tempo a lavorare con i trasformatori e l'attenzione verso se stessi avrà sentito centinaia di volte. È assolutamente vero, lo abbiamo sperimentato tutti mentre provi ad aumentare la dimensione del contesto del tuo modello, tutto improvvisamente si ferma. Ma poi allo stesso tempo, praticamente ogni settimana a quanto pare, c'è un nuovo modello all'avanguardia con una nuova durata del contesto da record. (Gemini ha una lunghezza del contesto di 2 milioni di token!)

Esistono molti metodi sofisticati come RingAttention che rendono possibile l'addestramento di lunghezze di contesto incredibilmente lunghe in sistemi distribuiti di grandi dimensioni, ma ciò che mi interessa oggi è una domanda più semplice.

Fino a che punto possiamo arrivare con la sola attenzione lineare?

Questo sarà un po' un tour a tappe, ma abbi pazienza mentre toccheremo alcuni punti chiave prima di approfondire i risultati.

Possiamo sostanzialmente riassumere il meccanismo tradizionale dell’attenzione con due punti chiave:

• Innanzitutto, la tipica espressione di attenzione softmax prende il prodotto della query e delle matrici chiave, lo normalizza per la stabilità, quindi prende il softmax (in termini di riga) per ottenere i punteggi di attenzione tra ciascun elemento della sequenza.
• In secondo luogo, la complessità temporale è dominata dai prodotti scalari N² e quello all'interno del softmax è il fattore limitante. È lì che calcoliamo i punteggi di attenzione.

Ciò è espresso nella forma tradizionale come:

Si scopre che se chiediamo ai nostri amici matematici possiamo pensare a questo in modo leggermente diverso. Il softmax può essere pensato come uno dei tanti modi per descrivere la distribuzione di probabilità che mette in relazione i token tra loro. Possiamo usare qualsiasi misura di somiglianza che ci piace (il prodotto scalare è uno dei più semplici) e finché lo normalizziamo, va bene.

È un po' sciatto dirlo È attenzione, poiché in realtà è solo l'attenzione che conosciamo e amiamo quando la funzione di somiglianza è l'esponenziale del prodotto scalare di query e chiavi (riportate di seguito) come troviamo nel softmax. Ma è qui che diventa interessante, se invece di usare questa espressione, cosa succederebbe se potessimo approssimarla?

Possiamo supporre che esista una mappa delle caratteristiche “phi†che ci dà un risultato quasi equivale a prendere l'esponenziale del prodotto scalare. E, cosa fondamentale, scrivere l'espressione in questo modo ci permette di giocare con l'ordine delle operazioni di moltiplicazione di matrici.

Nel carta propongono l'Unità Lineare Esponenziale (ELU) come mappa delle caratteristiche a causa di una serie di proprietà utili:

1. Per valori superiori a 0 l'ELU(x) fornisce un risultato lineare, che pur non essendo uguale all'esponenziale preserva l'ordine relativo tra i punteggi.
2. Per valori inferiori o uguali a 0 il termine esponenziale preserva la natura continua della funzione e garantisce che i gradienti non svaniscano.

Non spenderemo troppo tempo su questo qui, ma questo è abbastanza ben verificato empiricamente come una buona approssimazione della funzione softmax.

Ciò che questo ci consente di fare è cambiare l’ordine delle operazioni. Possiamo prendere prima il prodotto della nostra mappa delle caratteristiche di K con V per creare un blocco KV, quindi il prodotto con Q. Il prodotto quadrato supera la dimensione della dimensione del modello anziché la lunghezza della sequenza.

Mettere tutto insieme nell’espressione di attenzione lineare ci dà:

Dove dobbiamo calcolare i termini tra parentesi solo una volta per riga di query.

(Se vuoi approfondire come si inserisce il mascheramento casuale e come vengono calcolati i gradienti, dai un'occhiata al documento. Oppure guarda questo spazio per un blog futuro.)

Il caso matematico è forte, ma personalmente finché non vedo alcuni benchmark rimango sempre un po' sospettoso.

Iniziamo esaminando i frammenti di codice per descrivere ciascuno di questi termini. L'attenzione di Softmax sembrerà molto familiare, non stiamo facendo nulla di stravagante qui.

class TraditionalAttention(nn.Module):
def __init__(self, d_k):
super(TraditionalAttention, self).__init__()
self.d_k = d_k
def forward(self, Q, K, V):
Z = torch.sqrt(torch.tensor(self.d_k, device=Q.device, dtype=torch.float32))
scores = torch.matmul(Q, K.transpose(-2, -1)) / Z
attention_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attention_weights, V)
return output

Quindi, per l'attenzione lineare, iniziamo ottenendo le matrici Query, Key e Value, quindi applichiamo la mappatura delle funzionalità ELU(x) a Query e Keys. Quindi utilizziamo la notazione einsum per eseguire le moltiplicazioni.

class LinearAttention(nn.Module):
def __init__(self):
super(LinearAttention, self).__init__()
self.eps = 1e-6
def elu_feature_map(self, x):
return F.elu(x) + 1
def forward(self, Q, K, V):
Q = self.elu_feature_map(Q)
K = self.elu_feature_map(K)
KV = torch.einsum("nsd,nsd->ns", K, V)
# Compute the normalizer
Z = 1/(torch.einsum("nld,nd->nl", Q, K.sum(dim=1))+self.eps)
# Finally compute and return the new values
V = torch.einsum("nld,ns,nl->nd", Q, KV, Z)
return V.contiguous()

Vederlo scritto nel codice va benissimo, ma cosa significa effettivamente sperimentalmente? Di quale incremento prestazionale stiamo parlando? Può essere difficile apprezzare il grado di accelerazione passando da un collo di bottiglia quadratico a uno lineare, quindi ho eseguito il seguente esperimento.

Prenderemo un singolo livello di attenzione, con una dimensione del modello d_k fissa di 64, e valuteremo il tempo impiegato per un passaggio in avanti di un set di sequenze di dimensioni batch di 32 dimensioni. L'unica variabile da modificare sarà la lunghezza della sequenza, da 128 a 6000 (la lunghezza del contesto GPT-3 come riferimento se 2048). Ogni esecuzione viene eseguita 100 volte per ottenere una media e una deviazione standard e gli esperimenti vengono eseguiti utilizzando una GPU Nvidia T4.

Per un esperimento così semplice i risultati sono piuttosto sorprendenti.

I risultati mostrano che anche per un giocattolo incredibilmente piccolo otteniamo una velocità fino a 60 volte.

Discussione

Ci sono alcuni ovvi suggerimenti qui:

1. Il vantaggio dell'attenzione lineare è enorme: in termini di velocità, un rendimento più elevato è sempre una buona cosa. Oppure in termini di requisiti di memoria per elaborare lunghe sequenze. In ambienti con poca memoria questo potrebbe essere un grande vantaggio.
2. Il grafico del rapporto ha una piega sorprendente: ci porta a sospettare che qui si stia verificando un'ulteriore ottimizzazione di livello inferiore, il che significa che il rapporto atteso non si materializza del tutto. Dobbiamo quindi prendere questo risultato con le pinze.

Per completezza, inoltre, non confondere questo come detto “L'attenzione lineare è 60 volte più veloce per i modelli piccoli”. In realtà i livelli feed-forward rappresentano spesso una parte più grande dei parametri in un Transformer e anche la codifica/decodifica è spesso un componente di dimensione limitante. Ma in questo problema ben definito, piuttosto impressionante!