QKV Projections Require a Fraction of Their Memory

Dit paper introduceert PAMM, een nieuwe tensorcompressietechniek die het geheugengebruik van de QKV-projecties in LLM's tot wel 512 keer verlaagt zonder de prestaties te schaden, en volledig compatibel is met bestaande efficiënte aandachtmechanismen.

De kern

🧠 De "Grote Vergeten" Slurper in AI's Geheugen

Stel je voor dat je een gigantische bibliotheek bouwt (een Groot Taalmodel of LLM, zoals die van onschatbare waarde zijn voor AI). Om deze bibliotheek te bouwen, moet je miljoenen boeken lezen en onthouden. Maar er is een probleem: de bibliotheek is zo groot dat de ruimte (het computergeheugen) op is voordat je klaar bent.

Meestal denken onderzoekers dat het probleem ligt bij het lezen van de boeken (de berekeningen). Maar dit nieuwe onderzoek, getiteld "QKV PROJECTIONS REQUIRE A FRACTION OF THEIR MEMORY", wijst op een heel ander probleem: het onthouden van de boeken die je net hebt gelezen.

🏗️ De Bouwplaats: De "Q, K, V" Werknemers

In de bouw van een AI (de 'Transformer'-architectuur) zijn er drie speciale werknemers die constant aan het werk zijn:

1. Q (Query): Vraagt: "Wat zoek ik?"
2. K (Key): Vraagt: "Wat heb ik?"
3. V (Value): Vraagt: "Wat is de inhoud?"

Elke keer als de AI een zin leest, moeten deze drie werknemers een enorme stapel papier (data) van elke letter en elk woord opslaan in hun geheugen, zodat ze later terug kunnen kijken om de zin te begrijpen. Dit papier stapelt zich op tot een enorme berg die 20% van je geheugen opvult, terwijl het eigenlijk alleen maar tijdelijk nodig is.

🎒 De Oplossing: PAMM (De Slimme Koffer)

De auteurs van dit paper, Malik Khalaf en zijn team, hebben een nieuwe methode bedacht genaamd PAMM (Point-Approximate Matrix Multiplication).

Stel je voor dat je op reis gaat en je moet al je kleding meenemen.

• De oude manier: Je pakt je hele koffer, vult hem tot de rand met exact dezelfde T-shirts, broeken en sokken, en sleept die zware koffer mee.
• De PAMM-methode: Je merkt dat je 500 T-shirts hebt, maar ze zijn allemaal bijna hetzelfde. In plaats van 500 shirts mee te nemen, neem je één perfect shirt mee (de "generateur") en een lijstje met aantekeningen: "Shirt 1 is dit shirt, maar 10% lichter. Shirt 2 is dit shirt, maar 5% donkerder."

In plaats van 500 shirts te dragen, draag je nu één shirt en een klein briefje. Je hebt 99% minder gewicht, maar je kunt nog steeds precies hetzelfde kledingstuk "reconstrueren" als je erom vraagt.

🚀 Hoe werkt het in de praktijk?

1. Kiezen van een "Voorbeeld": De AI kijkt naar de enorme hoeveelheid data die ze moet onthouden. Ze kiest een paar kleine, representatieve voorbeelden (bijvoorbeeld 1 op de 512).
2. De "Rekentruc": De rest van de data wordt niet opgeslagen. In plaats daarvan wordt er gezegd: "Die andere 511 stukjes data lijken zo veel op dat ene voorbeeld, dat we ze gewoon als een versie daarvan kunnen behandelen."
3. Het Resultaat: De AI hoeft niet meer die enorme berg papier op te slaan. Ze slaat alleen de voorbeelden en de kleine lijstjes op.

📉 Wat levert dit op?

• Ruimtebesparing: Ze besparen tot 512 keer minder geheugen voor deze specifieke taken. Het is alsof je van een vrachtwagen vol met zand overstapt op een fiets.
• Geen kwaliteitsverlies: Het gekke is: de AI wordt er niet slimmer of dommer van. De "reconstructie" is zo goed, dat de AI precies even goed presteert als met de volle geheugen.
• Snelheid: Het kost bijna geen extra tijd om deze truc uit te voeren. Het is alsof je niet meer hoeft te sjouwen, maar wel even snel kunt rennen.

🤝 Samenvatting in één zin

Dit onderzoek laat zien dat we de geheugenvraag van AI's drastisch kunnen verkleinen door niet alles letterlijk op te slaan, maar slimme "steekproeven" te nemen en de rest te benaderen, net als het inpakken van een koffer met slechts één voorbeeldkledingstuk en een notitieblok.

Dit maakt het mogelijk om veel grotere en slimmere AI-modellen te bouwen op hardware die we vandaag al hebben, zonder dat we duurdere computers hoeven te kopen.

Technische samenvatting

Titel: QKV-projecties vereisen slechts een fractie van hun geheugen

Auteurs: Malik Khalaf, Yara Shamshoum, Nitzan Hodos, Yuval Sieradzki, Assaf Schuster (Technion, Israël)

---

1. Het Probleem

Hoewel er veel onderzoek is gedaan naar het verbeteren van de reken- en geheugenefficiëntie van Large Language Models (LLM's), richt de meeste aandacht zich op het benaderen van de scaled dot-product attention zelf (bijv. FlashAttention). Een cruciaal maar vaak over het hoofd gezien component is het geheugengebruik van de lineaire projectielagen die de input $x$ omzetten in de Query ($Q$), Key ($K$) en Value ($V$) tensoren.

• Geheugennood: Tijdens het trainen moeten de activeringen van de input $x$ worden opgeslagen voor de backpropagatie. Deze activeringen kunnen tot 20% van het totale piekgeheugen van de GPU beslaan binnen de attention-blokken.
• Schalingsprobleem: In tegenstelling tot parameters of optimizer-staten (die vaak worden gecomprimeerd), schalen activaties lineair met de batchgrootte en de sequentielengte.
• Redundantie: De auteurs identificeren dat er aanzienlijke redundantie bestaat in de sequentiemengte (token-embeddings). Tokens in een batch vertonen vaak vergelijkbare patronen, padding of lokale contextuele gelijkenis, waardoor de rang van de activatiematrix veel lager is dan de totale hoeveelheid tokens.

2. Methodologie: Point-Approximate Matrix Multiplication (PAMM)

De auteurs stellen PAMM voor, een nieuwe techniek voor tensorcompressie die specifiek is ontworpen om de geheugenfootprint van de $Q, K, V$ projecties te elimineren tijdens het trainen, zonder de modelprestaties significant te beïnvloeden.

Kernprincipes:
In plaats van de volledige input-matrix $X \in \mathbb{R}^{b \times n}$ (waarbij $b$ het aantal tokens is en $n$ de hidden dimension) op te slaan, comprimeert PAMM deze naar een kleine set van representatieve punten.

Het Algoritme (in twee fasen):

1. Compressie (Forward Pass):

   • Er wordt een kleine set van $k$ "generator"-punten $C \in \mathbb{R}^{k \times n}$ geselecteerd door willekeurig $k$ rijen uit $X$ te bemonsteren (zonder teruglegging).
   • Voor elke rij $A_i$ in de originele matrix wordt de dichtstbijzijnde lijn bepaald die wordt opgespannen door een van de generators $C_j$.
   • De rij $A_i$ wordt vervangen door een geschaalde versie van de generator: $\tilde{A}_i = \alpha_i C_{f(i)}$.
   • Er worden twee auxiliary vectoren opgeslagen:
     • $f$: Een toewijzingsfunctie die aangeeft welke generator voor elke rij wordt gebruikt.
     • $\alpha$: Schalingsfactoren (cosine similarity gebaseerd).
   • Een tolerantieparameter $\epsilon$ bepaalt of een rij wordt behouden of als nulvector wordt behandeld (in de praktijk werkt $\epsilon = \infty$ het beste, wat betekent dat alle rijen worden toegewezen).

2. Benaderende Matrixvermenigvuldiging (Backward Pass):

   • Tijdens backpropagatie moet de gradient $\nabla W = X^\top \cdot \nabla Z$ worden berekend.
   • In plaats van de volledige $X$ te reconstrueren, gebruikt PAMM de gecomprimeerde vorm:
     $$\tilde{O} = \sum_{j=1}^{k} C_j^\top \cdot \left( \sum_{i: f(i)=j} \alpha_i B_i \right)$$
   • Hierbij wordt eerst de matrix $B$ (de gradient) samengevoegd tot een kleinere matrix $\tilde{B} \in \mathbb{R}^{k \times m}$, waarna de vermenigvuldiging $C^\top \tilde{B}$ wordt uitgevoerd.
   • Dit vermindert de rekencost van $O(bnm)$ naar $O(kn(b+m))$, waarbij $k \ll b$.

Theoretische Garanties:
De auteurs bewijzen dat als het aantal tokens $b$ toeneemt, het aantal benodigde generators $k$ slechts logarithmisch hoeft te groeien om de data-distributie goed te dekken, dankzij de hoge dichtheid van de data in de sequentielengte.

3. Belangrijkste Bijdragen

• Nieuwe Compressie-richting: In plaats van te focussen op de hidden dimension (zoals bij low-rank adapters), exploiteert PAMM redundantie in de sequentielengte.
• Extreme Compressie: PAMM kan de activeringen van $Q, K, V$ comprimeren met een factor van tot wel 512x (ratio $r = 1/512$).
• Composabiliteit: De methode is volledig compatibel met bestaande efficiënte technieken zoals FlashAttention, Gradient Checkpointing en Low-Rank Adapters (LoRA).
• Eenvoud: Het algoritme is eenvoudig te implementeren en vereist geen complexe clustering-algoritmen; willekeurige bemonstering bleek voldoende.

4. Resultaten

De auteurs hebben PAMM getest op verschillende modellen (LLaMA van 60M tot 7B, RoBERTa-base, en Pixtral-12B) voor zowel pretraining als finetuning.

• Geheugenreductie:
  • PAMM reduceert het geheugengebruik voor $Q, K, V$ activeringen met >97% (tot op 1/512e van de originele grootte).
  • Bij LLaMA-1B daalt het geheugen van 3 GB naar slechts 24 MB voor deze lagen.
• Modelprestaties (Perplexiteit):
  • Bij extreme compressie ($r=1/512$) blijft de perplexiteit gelijk aan of zelfs beter dan de baseline zonder compressie.
  • Bij sommige grotere modellen verbeterde de prestatie zelfs licht, wat suggereert dat de redundantie in de input het trainingsproces kan hinderen en dat PAMM hier als een regularisator werkt.
• Doorvoer (Throughput):
  • De overhead door de compressie en decompressie is verwaarloosbaar, vooral bij grotere modellen (< 2.7% vertraging voor LLaMA-1B).
  • De rekentijd neemt toe door de extra berekeningen, maar dit wordt gecompenseerd door de verminderde geheugentransfers en de schaalbaarheid.
• Vergelijking met andere methoden:
  • PAMM presteert aanzienlijk beter dan bestaande methoden zoals CompAct (Gaussian random projections) en Uniform-CRS (willekeurige rij/kolom sampling). Deze alternatieven leiden tot significante degradatie in perplexiteit bij hoge compressie, terwijl PAMM dit niet doet.
  • De parameter $\epsilon = \infty$ (geen neighborhood constraint) bleek optimaal, wat betekent dat het toewijzen van elke token aan een generator, zelfs als de afstand groot is, werkt vanwege de clusterachtige aard van attention-activaties.

5. Betekenis en Conclusie

Dit paper introduceert een doorbraak in het trainen van LLM's met beperkte GPU-geheugenresources.

• Praktische Impact: Het maakt het mogelijk om grotere batchgroottes of langere sequenties te trainen op bestaande hardware, of om zeer grote modellen te trainen op kleinere clusters.
• Paradigmaverschuiving: Het toont aan dat de "memory wall" in transformer-training grotendeels veroorzaakt wordt door de opslag van input-activaties voor de lineaire projecties, en dat deze redundantie effectief kan worden benut zonder de kwaliteit van het model te offeren.
• Toekomst: De methode is een plug-in oplossing die direct kan worden toegepast in bestaande frameworks en kan worden gecombineerd met andere optimalisaties zoals FlashAttention en LoRA.

Samenvattend biedt PAMM een elegante en uiterst efficiënte oplossing voor het geheugenvraagstuk bij het trainen van moderne taalmodellen, waarbij de kosten van het opslaan van $Q, K, V$ activeringen effectief worden "uitgewist".