Attention Sinks

Un sink di attenzione è un token che riceve una quota elevata dei pesi di attenzione in molte teste e livelli, anche se il suo contenuto è irrilevante per la richiesta; spesso si tratta del primo token della sequenza. Guangxuan Xiao e colleghi di MIT, Meta AI e Carnegie Mellon University hanno dato un nome al fenomeno in Efficient Streaming Language Models with Attention Sinks, pubblicato a ICLR 2024. Hanno mostrato che i modelli linguistici di grandi dimensioni (LLM) come Llama 2, MPT, Falcon e Pythia si degradano quando una cache KV scorrevole elimina i primi token e che mantenerne quattro consente di ripristinare una generazione stabile.

Perché si formano i sink#

La self-attention trasforma i punteggi di ogni query in pesi mediante una softmax, e la loro somma è sempre uno. Quando nessun token nel contesto è rilevante per una testa, il peso deve comunque essere assegnato a qualcosa. In un modello linguistico causale, i primi token sono visibili a tutti quelli successivi, perciò l'addestramento li trasforma nella destinazione predefinita per il peso di cui una testa non ha bisogno. Nell'analisi StreamingLLM di Llama 2, ogni livello oltre i due inferiori concentrava l'attenzione sui token iniziali in tutte le teste. Un logit sink aggiunto alla softmax mostra l'effetto:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Senza sink, la testa distribuisce il suo peso totale su sei valori irrilevanti e restituisce un output con norma 0,83. Con un logit sink pari a 3, il sink riceve il 77% del peso e l'output scende a 0,19, così che la testa non fa quasi nulla per questa query. I modelli open-weight gpt-oss di OpenAI integrano questo meccanismo con un logit sink appreso per ogni testa di attenzione, offrendo a ciascuna un punto in cui assegnare il peso residuo senza contribuire all'output.

Generazione in streaming con StreamingLLM#

Un modello che genera testo all'infinito supera qualsiasi finestra di contesto e non può mantenere ogni token nella cache KV; la semplice attenzione a finestra, che conserva solo i token più recenti, non funziona più appena vengono rimossi i primi token. StreamingLLM mantiene permanentemente le chiavi e i valori dei quattro token iniziali accanto a una finestra scorrevole di quelli recenti e assegna le posizioni in base all'ordine nella cache anziché nel testo originale, riapplicando a ogni passaggio gli embedding posizionali rotativi (RoPE):

Questo ha mantenuto stabile la modellazione linguistica per oltre 4 milioni di token e ha raggiunto una velocità fino a 22,2 volte superiore rispetto a una finestra scorrevole che ricalcola la cache per ogni nuovo token. La finestra continua a limitare ciò che il modello può ricordare: i token eliminati dalla parte centrale sono persi per sempre, quindi un chatbot in streaming o un agente di AI ricorda solo gli scambi più recenti. Gli autori hanno anche preaddestrato modelli con un token segnaposto dedicato e apprendibile all'inizio di ogni campione, che ha assunto autonomamente il ruolo del sink.

Sink di attenzione e termini correlati#

• Attenzione a finestra scorrevole: Limita ogni token ai vicini recenti. Applicata in inferenza a un modello addestrato con attenzione completa, non funziona più quando i sink escono dalla finestra.
• Degrado del contesto: Calo di accuratezza al crescere di un prompt, anche quando rientra ancora nella finestra di contesto.

Sink nei modelli di visione#

I Transformer visivi (ViT) mostrano un comportamento correlato. In Vision Transformers Need Registers, ricercatori di Meta e Inria hanno individuato token con norme insolitamente elevate in patch di sfondo con poche informazioni nei ViT supervisionati e auto-supervisionati, riutilizzati per calcoli interni. Aggiungere alcuni token register all'input ha assorbito quel ruolo, rimosso gli artefatti dalle mappe delle caratteristiche e dalle mappe di attenzione e migliorato la predizione densa.

SAM 2 di Meta, eseguibile con il pacchetto Ultralytics, gestisce la memoria video secondo uno schema simile, che mantiene un riferimento iniziale e una finestra: per ogni nuovo fotogramma, la sua attenzione alla memoria legge le voci della banca di memoria relative ai fotogrammi indicati dall'utente, conservate per l'intero video, oltre a quelle dei 6 fotogrammi più recenti, così che il costo dell'attenzione per fotogramma rimanga costante indipendentemente dalla durata del video.