Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space

Probleemstelling

De huidige Transformer-architecturen, die gebaseerd zijn op Multi-Headed Attention (MHA), kampen met twee fundamentele beperkingen bij het verwerken van lange contexten:

1. Rekencomplexiteit: De aandacht-mechanisme heeft een kwadratische complexiteit ($O(S^2)$) ten opzichte van de sequentielengte ($S$) en een lineaire complexiteit ten opzichte van de verborgen dimensie. Dit maakt training en "prefill" (verwerking van input) duur.
2. KV-Cache Overhead: Tijdens autoregressieve generatie groeit de Key-Value (KV) cache lineair met de sequentielengte. Dit veroorzaakt enorme kosten voor geheugenbandbreedte, wat de inferentiesnelheid beperkt, vooral bij lange contexten.

Bestaande oplossingen zoals Grouped Query Attention (GQA) en Multi-Latent Attention (MLA) proberen deze problemen aan te pakken, maar hebben tekortkomingen:

• GQA deelt parameters voor keys en values om de cache te verkleinen, maar vermindert de rekenkosten (FLOPs) voor training en prefill niet significant.
• MLA projecteert keys en values naar een kleinere, gedeelde latente ruimte om de cache te comprimeren. Dit vermindert echter de rekenkosten niet (en verhoogt ze zelfs door up-projections) en introduceert complexiteit bij het integreren van positiespecifieke encoding (zoals RoPE), wat vaak aparte heads vereist.

Methodologie: Compressed Convolutional Attention (CCA)

De auteurs introduceren Compressed Convolutional Attention (CCA), een methode die de volledige attentie-operatie uitvoert binnen een gecomprimeerde latente ruimte. Dit leidt tot een drastische reductie in parameters, KV-cache en FLOPs.

Kerncomponenten van CCA:

1. Compressie in Latente Ruimte: In plaats van keys, queries en values (QKV) te projecteren naar de volledige dimensie en ze daar te comprimeren (zoals bij MLA), projecteert CCA de input direct naar een gecomprimeerde latente ruimte met dimensie $\tilde{e} = E/C$ (waarbij $C$ de compressiefactor is). De volledige attentie-berekening (QK-product en gewogen som) vindt plaats in deze kleine ruimte.
2. Convolutionele Mixing: Om de expressiviteit van de compressie te compenseren, worden convoluties toegepast op de gecomprimeerde query ($\tilde{q}$) en key ($\tilde{k}$) vectoren. Dit omvat:
   • Sequentie- en kanaal-mixing: Convoluties over de sequentielengte en binnen de kanalen van een head.
   • QK-Mean: Een operatie die het gemiddelde van de pre- en post-convolutie waarden van Q en K toevoegt aan de vectoren. Dit fungeert als een skip-connection en helpt informatie te delen, wat de diagonalisatie van de attentiematrix verbetert.
3. Value-Shift: Voor de value-vectoren wordt een unieke aanpak gebruikt waarbij elke head twee soorten values ontvangt: één van de huidige input en één van de vorige sequentiestap. Dit introduceert een inductieve bias die de modelstabiliteit en prestaties verbetert.
4. Integratie van RoPE: Omdat de attentie volledig in de latente ruimte plaatsvindt, kan de Rotary Positional Embedding (RoPE) naadloos worden toegepast zonder de noodzaak van aparte "RoPE-heads" of complexe up-projections, zoals bij MLA het geval is.

Compressed Convolutional Grouped Query Attention (CCGQA):
De auteurs combineren CCA met het principe van GQA. Binnen de gecomprimeerde ruimte worden keys en values gedeeld over groepen van heads. Dit creëert een extra compressie-laag zonder kwaliteitsverlies, waardoor gebruikers kunnen schakelen tussen reken- en geheugenbeperkingen.

Belangrijkste Bijdragen

• Nieuwe Architectuur: CCA is de eerste methode die attentie volledig in een gecomprimeerde latente ruimte uitvoert voor zowel training als inferentie, wat leidt tot een directe reductie in FLOPs met een factor $C$.
• Orthogonaliteit: Het artikel toont aan dat parameter-deling (GQA) en parameter-compressie (CCA) orthogonaal zijn en effectief kunnen worden gecombineerd (CCGQA) om een Pareto-grens te creëren tussen rekenkosten en geheugenbandbreedte.
• Kernel-Implementatie: De auteurs hebben een gefuseerde kernel voor H100 GPU's ontwikkeld die convoluties en Flash Attention combineert, wat de theoretische snelheidswinst in de praktijk realiseert.

Resultaten

Experimenten werden uitgevoerd op zowel dichte (dense) modellen als Mixture-of-Experts (MoE) modellen, vergeleken met MHA, GQA en MLA.

1. Prestaties:

   • CCA en CCGQA presteren consistent beter dan GQA en MLA bij gelijke KV-cache-compressie.
   • In MoE-modellen behaalde CCGQA met 8x KV-cache compressie (ten opzichte van MHA) geen prestatieverlies vergeleken met standaard MHA, terwijl het aanzienlijk minder parameters en FLOPs gebruikte.
   • CCA overtrof MLA met 16x minder decode-FLOPs en CCGQA met 8x minder FLOPs, terwijl de kwaliteit gelijk bleef of beter was.

2. Snelheid en Latentie:

   • Prefill: Op H100 GPU's reduceerde de CCA-kernel de prefill-latentie met ongeveer 1,7x bij een sequentielengte van 16k tokens ten opzichte van MHA.
   • Backward Pass: De training (backward pass) was ongeveer 1,3x sneller.
   • Schaling: Omdat de kwadratische termen ($S^2$) worden verkleind met de compressiefactor, neemt de snelheidswinst toe naarmate de sequentielengte groeit.

3. Efficiëntie:

   • CCA vermindert de parameters en de KV-cache gelijktijdig met de compressiefactor.
   • In vergelijking met MLA, die extra parameters nodig heeft voor up-projections, is CCA zuiniger in parameters.

Betekenis en Toekomstperspectief

Deze paper biedt een doorbraak in de efficiëntie van Transformer-modellen voor lange contexten. Door de rekenkosten en het geheugenverbruik gelijktijdig te verlagen, maakt CCA het haalbaar om modellen te trainen en te serveren met veel langere contextvensters zonder de hardware-eisen exponentieel te laten stijgen.

• Flexibiliteit: De methode is zeer flexibel en kan worden aangepast aan verschillende hardware-beperkingen (rekenbound vs. geheugenbound).
• Compatibiliteit: Het werkt goed met bestaande parallelisme-strategieën (zoals Tensor Parallelism) en vereist geen ingewikkelde communicatie-overhead.
• Toekomst: De auteurs suggereren dat CCA orthogonaal is tot methoden die sequentiecompressie toepassen (zoals NSA of MoBA), wat de weg vrijmaakt voor verdere optimalisaties door beide technieken te combineren.

Kortom, CCA en CCGQA stellen een nieuwe standaard voor efficiënte attentie-mechanismen die de "quadratische bottleneck" van Transformers effectief verzachten door slimme compressie en convolutie binnen de latente ruimte.