Sliding Window Attention

L'attenzione a finestra scorrevole (SWA) limita ogni token a considerare un numero fisso di token vicini: in un modello linguistico causale, gli w token più recenti, anziché l'intera sequenza come nell'autoattenzione completa. Il calcolo dell'attenzione cresce quindi linearmente con la lunghezza della sequenza anziché quadraticamente, e un modello generativo deve conservare nella cache KV solo chiavi e valori degli ultimi w token. Nel 2020 Longformer ha fatto della finestra scorrevole il fulcro della propria attenzione; nel 2023 Mistral 7B di Mistral AI l'ha introdotta in un modello linguistico di grandi dimensioni (LLM) aperto e molto usato, con una finestra di 4.096 token.

Come cresce la finestra con la profondità#

Un livello di SWA vede solo w token, ma i livelli Transformer sono impilati. Al livello 2, un token legge stati nascosti che riassumono già w token del livello 1, quindi l'informazione può viaggiare per circa altre w − 1 posizioni a ogni livello, proprio come le convoluzioni impilate ampliano il campo recettivo di una rete neurale convoluzionale (CNN). In PyTorch, la finestra è una maschera booleana di attenzione; moltiplicando la maschera per sé stessa si traccia la distanza percorsa dall'informazione:

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

Per 16 token, la finestra calcola 58 punteggi per livello invece di 136; la portata dell'ultimo token cresce di 3 token per livello: da 4 a 7, 10 e 13. La stessa maschera può essere passata come attn_mask a scaled_dot_product_attention di PyTorch. Per Mistral 7B, 32 livelli con una finestra di 4.096 token offrono un intervallo di attenzione teorico di circa 131K token, anche se le informazioni più lontane devono sopravvivere a molti passaggi prima di arrivare.

Mistral 7B memorizza inoltre chiavi e valori in un buffer circolare di w slot, scrivendo la posizione i nello slot i mod w e sovrascrivendo la voce più vecchia. Su una sequenza di 32K token, questo ha ridotto di otto volte la memoria della cache senza compromettere la qualità; la dimensione fissa della cache mantiene costante la memoria, indipendentemente dalla durata della generazione.

Alternanza di livelli locali e globali#

Le finestre pure non possono copiare in un solo passaggio un'informazione molto distante, quindi i modelli recenti combinano livelli SWA con livelli globali completi. Gemma 2 di Google li alternava uno a uno con una finestra di 4.096 token, mentre Gemma 3 usa cinque livelli locali per ogni livello globale, con una finestra di appena 1.024 token, mantenendo ridotta la cache KV in una finestra di contesto fino a 128K token. Solo i livelli globali memorizzano chiavi e valori per l'intero contesto, quindi cinque livelli su sei contengono appena 1.024 token ciascuno.

I Vision Transformer usano la stessa combinazione per le immagini ad alta risoluzione, in cui ogni patch è un token. Il SAM 3 di Meta, eseguibile nel pacchetto Ultralytics per la segmentazione delle immagini, applica l'attenzione all'interno di finestre non sovrapposte di patch 24 × 24 in 28 dei 32 blocchi del suo encoder d'immagine e globalmente sull'intera griglia di patch 72 × 72 di un'immagine da 1.008 pixel nei blocchi 7, 15, 23 e 31.

SWA e termini correlati#

• Swin Transformer: Divide l'immagine in finestre fisse e non sovrapposte, condivise da tutti i token al loro interno, e sposta la griglia tra i livelli in modo che le finestre adiacenti scambino informazioni.
• Attenzione sparsa: La famiglia più ampia di schemi fissi. SWA ne è il membro più semplice; altri aggiungono finestre dilatate, token globali o collegamenti casuali.
• Pozzi di attenzione: I primi token di una sequenza, che assorbono il peso di attenzione in eccesso. Applicare una finestra in inferenza a un modello addestrato con attenzione completa e rimuovere questi token interrompe la generazione; StreamingLLM li mantiene nella cache accanto alla finestra.
• Attenzione deformabile: Ogni query presta attenzione a pochi punti di campionamento di cui apprende gli offset, perciò il vicinato dipende dal contenuto.
• Attenzione Flash: Un kernel GPU per l'attenzione esatta. Le sue implementazioni accettano una dimensione della finestra, quindi i modelli SWA saltano del tutto i blocchi che ne restano fuori.