Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space

1. Il Problema

I modelli Transformer basati sull'attenzione multi-testa (MHA) soffrono di due limiti fondamentali quando si tratta di contesti lunghi:

• Complessità Computazionale Quadratica: Il calcolo dell'attenzione scala quadraticamente rispetto alla lunghezza della sequenza ($S^2$) e alla dimensione del modello, rendendo costosi l'addestramento e la fase di prefill (elaborazione dell'input).
• Crescita Lineare della KV-Cache: Durante la generazione autoregressiva, la cache Chiave-Valore (KV) cresce linearmente con la lunghezza della sequenza e la dimensione nascosta. Questo crea un collo di bottiglia nella larghezza di banda della memoria, limitando l'efficienza hardware e rendendo costoso servire modelli con contesti estesi (es. modelli di ragionamento).

Le soluzioni esistenti presentano compromessi:

• GQA (Grouped Query Attention): Riduce la dimensione della KV-cache condividendo le teste di Key e Value, migliorando la velocità di decodifica, ma non riduce il costo computazionale (FLOPs) per l'addestramento o il prefill.
• MLA (Multi-Latent Attention): Comprime le chiavi e i valori in uno spazio latente condiviso per ridurre la cache. Tuttavia, richiede proiezioni "up" (up-projections) prima dell'attenzione, il che significa che non riduce i FLOPs durante l'addestramento e il prefill rispetto a MHA, e introduce complessità nell'integrazione con le posizioni (RoPE).

2. Metodologia: CCA e CCGQA

Gli autori propongono Compressed Convolutional Attention (CCA), un metodo che esegue l'intera operazione di attenzione all'interno di uno spazio latente compresso, combinando compressione dei parametri e riduzione dei calcoli.

A. Compressed Convolutional Attention (CCA)

Il flusso di lavoro di CCA è il seguente:

1. Down-Projection: Query ($Q$), Key ($K$) e Value ($V$) vengono proiettati linearmente in uno spazio latente compresso di dimensione $\tilde{e} = E/C$ (dove $C$ è il fattore di compressione).
2. Operazioni Convolutionali: A differenza di MLA, CCA applica convoluzioni sia lungo la dimensione della sequenza che lungo quella dei canali alle latenti di $Q$ e $K$. Questo aggiunge espressività e migliora il trasferimento di informazioni.
3. QK-Mean: Viene aggiunta una media dei valori di $Q$ e $K$ (pre e post-convoluzione) come connessione residua. Questo aiuta a condividere informazioni tra $Q$ e $K$ e aumenta la sparsità della diagonale dell'attenzione.
4. Value-Shift: Per le latenti di $V$, viene introdotta un'operazione di "shift" che combina valori derivati dall'input corrente e dall'input precedente della sequenza. Questo introduce un bias induttivo utile.
5. Attenzione Compressa: L'operazione di softmax attention viene eseguita interamente nello spazio latente compresso.
6. Up-Projection: L'output viene proiettato di nuovo alla dimensione originale del residuo.

Vantaggi chiave:

• Elimina la necessità di matrici di up-projection separate per $K$ e $V$ (necessarie in MLA), riducendo i parametri.
• Integra RoPE (Rotary Positional Embeddings) direttamente nello spazio latente senza bisogno di teste separate.
• Riduce i FLOPs di attenzione di un fattore $C$.

B. Compressed Convolutional Grouped Query Attention (CCGQA)

Gli autori combinano CCA con GQA. Poiché la compressione dei parametri e la condivisione dei parametri sono strategie ortogonali, CCGQA applica la condivisione delle teste di Key e Value (tipica di GQA) all'interno dello spazio latente già compresso.

• Permette di decouplare i fattori di compressione per Query ($C_1$) e per Key/Value ($C_2$).
• Permette di raggiungere compressioni della KV-cache ancora più elevate (es. 8x) senza penalità di qualità.

3. Contributi Chiave

1. Riduzione Unica di Parametri, Cache e FLOPs: CCA è il primo metodo che riduce simultaneamente i parametri, la dimensione della KV-cache e il costo computazionale (FLOPs) di un fattore di compressione desiderato.
2. Pareto Frontier Ottimale: CCGQA crea un fronte di Pareto fluido che permette di bilanciare i vincoli di memoria (bandwidth) e di calcolo (compute) senza sacrificare la qualità del modello.
3. Integrazione Elegante di RoPE: A differenza di MLA, CCA gestisce le posizioni direttamente nello spazio latente, evitando la complessità di cache RoPE separate.
4. Efficienza Pratica: L'implementazione di un kernel fuso su GPU (H100) dimostra che le operazioni aggiuntive (convoluzioni) hanno un overhead trascurabile rispetto ai benefici.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli densi (1B parametri) e MoE (Mixture of Experts, 350M attivi/1.5B totali) addestrati sul dataset Zyda2.

• Qualità del Modello:
  • CCA supera costantemente sia GQA che MLA a parità di compressione della KV-cache su modelli densi e MoE.
  • CCGQA supera MHA standard pur avendo una KV-cache compressa e meno FLOPs.
  • Su modelli MoE, CCGQA con 8x compressione della KV-cache ottiene prestazioni pari a MHA standard, ma con metà della cache di GQA/MLA.
• Efficienza Computazionale:
  • Prefill: Su GPU H100, CCA riduce la latenza di prefill di circa 1.7x rispetto a MHA a una lunghezza di sequenza di 16k.
  • Training/Backward: L'addestramento è accelerato di circa 1.3x.
  • FLOPs: CCA riduce i FLOPs di addestramento e prefill di un fattore $C$ (es. 4x o 8x), a differenza di GQA e MLA che non riducono i FLOPs di prefill.
• Ablation Study: Le analisi mostrano che le convoluzioni sequenziali sono il componente principale del miglioramento delle prestazioni, mentre QK-Mean e Value-Shift forniscono guadagni aggiuntivi significativi, specialmente nei modelli MoE.

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo nell'efficienza dei Transformer per contesti lunghi:

• Scalabilità: Permette di elaborare sequenze molto più lunghe (es. 16k+) con lo stesso budget computazionale, rendendo fattibili modelli di ragionamento estesi.
• Versatilità: Essendo ortogonale alle tecniche di parallelismo (Tensor Parallelism, Context Parallelism), CCA/CCGQA si integra facilmente nelle infrastrutture esistenti senza richiedere modifiche architetturali complesse come quelle di MLA.
• Futuro: Dimostra che esiste una ridondanza significativa nei metodi di attenzione classici e che operazioni aggiuntive ben scelte (come le convoluzioni) nello spazio latente possono compensare la perdita di espressività dovuta alla compressione dimensionale.

In sintesi, CCA e CCGQA offrono una soluzione "best-of-both-worlds", superando i limiti di memoria e calcolo delle tecniche attuali senza compromettere la qualità del modello, rendendo i Transformer più efficienti sia per l'addestramento che per l'inferenza.