Do Value Vectors in Deep Layers Need Context from the Residual Stream?

Sintesi Tecnica: Bank of Values (BoV)

Definizione del Problema
L'architettura transformer, l'ossatura dei moderni Large Language Models (LLM), si basa su strati di attenzione in cui i vettori query, key e value vengono calcolati a partire dallo stream residuo. Sebbene lo stream residuo accumuli efficacemente il contesto per le query e le key, la necessità di contesto per i vettori value negli strati profondi rimane scarsamente esaminata. L'attenzione standard calcola i vettori value ($V = x_i W_V$) utilizzando lo stream residuo ($x_i$), che mescola l'informazione originale del token con il contesto accumulato. Recenti lavori che tentano di iniettare informazioni sui token prive di contesto nei vettori value (ad esempio, Value Residual Learning, SVFormer) hanno mostrato risultati contrastanti, spesso fallendo nel isolare il contributo specifico delle componenti prive di contesto o applicandoli in modo inefficiente in tutti gli strati. Il problema centrale affrontato è se gli strati profondi richiedano realmente vettori value dipendenti dal contesto, o se beneficerebbero maggiormente preservando l'informazione originale e priva di contesto del token senza l'overhead computazionale della proiezione dello stream residuo.

Metodologia
Gli autori propongono un'indagine sistematica sulla composizione dei vettori value, seguita da un nuovo design architettonico chiamato Bank of Values (BoV).

1. Studio di Ablazione Sistematica:
   Utilizzando un modello da 12 strati e 135 milioni di parametri, gli autori hanno condotto esperimenti su quattro assi per determinare la composizione ottimale dei vettori value negli strati profondi:

   • Additivo vs. Sostitutivo: Confronto tra la sostituzione dell'intero vettore value standard con una fonte priva di contesto rispetto alla loro combinazione lineare.
   • Coefficienti: Test di coefficienti apprendibili (non vincolati) contro coefficienti fissi.
   • Fonte di Informazione Priva di Contesto: Confronto tra una proiezione specifica per strato dell'embedding iniziale ($x_0 W_V$) e un vettore value condiviso dal primo strato ($V_1$).
   • Strati Target: Valutazione dell'impatto dell'applicazione di queste modifiche a tutti gli strati, a ogni strato alternato o solo all'ultimo terzo degli strati.

   Risultati Chiave dall'Ablazione:

   • Gli strati profondi preferiscono fortemente componenti prive di contesto ($\tilde{V}$) rispetto a componenti dipendenti dal contesto ($V$). Nelle varianti additive, i coefficienti appresi per $\tilde{V}$ erano significativamente più grandi (spesso >10) rispetto a quelli per $V$ (spesso spinti verso lo 0).
   • La sostituzione di $V$ con $\tilde{V}$ produce guadagni significativi di performance solo nell'ultimo terzo degli strati; l'applicazione di questo approccio agli strati superficiali degrada le prestazioni.
   • Una volta che gli strati profondi hanno accesso a una componente priva di contesto di alta qualità, l'aggiunta della componente dipendente dal contesto fornisce un beneficio trascurabile.
   • I vettori privi di contesto specifici per strato ($x_0 W_V$) superano i vettori condivisi ($V_1$), poiché permettono a ogni strato di proiettare il token in uno spazio unico.

2. Architettura Bank of Values (BoV):
   Basandosi sulla scoperta che gli strati profondi beneficiano di vettori value privi di contesto, deterministici e indipendenti dalla sequenza di input, gli autori propongono BoV.

   • Meccanismo: Nell'ultimo terzo degli strati, la standard proiezione densa della matrice value ($x_i W_V$) è sostituita da una tabella di lookup sparsa ($E_V$).
   • Implementazione: Per un token con ID $i$, il vettore value viene recuperato direttamente come $v_i = \gamma_v E_V[i]$, dove $E_V \in \mathbb{R}^{|V| \times d}$ è una tabella di parametri apprendibili che memorizza i vettori value pre-calcolati per ogni token del vocabolario, e $\gamma_v$ è un coefficiente di scala apprendibile.
   • Efficienza di Inferenza: A differenza dell'attenzione standard, che memorizza (cache) i vettori value per tutti i token precedenti (richiedendo memoria $O(Td)$ per strato), BoV non effettua il caching dei valori. Invece, esegue un'operazione di gather $O(Td)$ dalla tabella di parametri statica $E_V$ durante il forward pass. La tabella $E_V$ può essere spostata nella memoria host, con solo le voci richieste pre-caricate (prefetched).

Contributi Chiave

• Insight: Il paper dimostra che i vettori value negli strati profondi dei transformer beneficiano principalmente di informazioni originali del token, prive di contesto, e che la componente dipendente dal contesto derivata dallo stream residuo è in gran parte ridondante in questi strati.
• Architettura: Introduzione di Bank of Values (BoV), un metodo che sostituisce il calcolo dei vettori value negli strati profondi con un lookup statico da una tabella specifica per token.
• Efficienza: BoV elimina la necessità di calcolare e memorizzare i vettori value per gli strati profondi, riducendo l'uso di memoria per l'inferenza a lungo contesto e abbassando i FLOPs per token.

Risultati
Gli autori hanno valutato BoV su modelli da 135M e 780M di parametri, addestrati sotto budget di calcolo controllati ($1.50 \times 10^{18}$ e $3.91 \times 10^{19}$ FLOPs, rispettivamente).

• Validation Loss: BoV ha ottenuto una validation loss inferiore rispetto ai baseline di attenzione standard su una partizione di validazione del dataset ClimbMix a entrambe le scale di modello.
• Benchmark Performance: Alla scala di 780M, BoV ha migliorato il punteggio medio attraverso 21 benchmark DCLM CORE rispetto al baseline di attenzione standard.
• Confronto: BoV ha eguagliato o superato i metodi esistenti che iniettano informazioni sui token nei vettori value (ad esempio, Value Residual Learning) utilizzando meno memoria e meno FLOPs. Nello specifico, ha superato i metodi che aggiungono informazioni sul token sopra il calcolo dello stream residuo, i quali incorrono comunque nel costo di calcolo e caching dei valori dipendenti dal contesto.

Significatività
Il paper sostiene che BoV offra un'alternativa più efficiente alla standard attention per gli strati profondi, sfruttando l'intuizione secondo cui questi strati non richiedono contesto per calcolare i vettori value. Rendendo i vettori value parametri del modello sparsi anziché calcoli dinamici, BoV riduce l'impronta di memoria per l'inferenza a lungo contesto e il costo computazionale senza sacrificare (e in alcuni casi migliorando) le prestazioni del modello. Ciò sfida il paradigma predefinito di calcolare tutti i vettori di attenzione dallo stream residuo e suggerisce una via verso architetture transformer più efficienti.