3BASiL: An Algorithmic Framework for Sparse plus Low-Rank Compression of LLMs

Deze paper introduceert 3BASiL-TM, een efficiënte post-training methode die een nieuwe 3-Block ADMM-algoritme en transformer-matching gebruikt om Large Language Models succesvol te comprimeren via een som van sparse en low-rank matrices, wat aanzienlijk betere prestaties en snellere compressie oplevert dan bestaande methoden.

De kern

Stel je voor dat je een enorme, superintelligente robot hebt die alles kan lezen, schrijven en begrijpen. Dit is een Large Language Model (LLM), zoals de modellen die nu overal worden gebruikt. Het probleem is dat deze robot zo groot en zwaar is dat hij alleen in een gigantisch datacenter past, met duizenden dure computers. Je kunt hem niet op je telefoon of laptop zetten.

De onderzoekers van deze paper (3BASiL) hebben een slimme manier bedacht om deze robot te "verkleinen" zonder dat hij zijn intelligentie verliest. Ze noemen hun methode 3BASiL.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De Zware Koffer

Stel je voor dat de kennis van de robot opgeslagen is in een enorme koffer vol met duizenden zware stenen. Je wilt deze koffer vervoeren, maar hij is te zwaar.

• De oude manier: Mensen probeerden gewoon willekeurig stenen weg te gooien (verwijderen) of ze in te pakken in heel kleine dozen (kwantisatie). Maar vaak bleef de koffer nog steeds te zwaar, of de robot werd "dom" omdat hij belangrijke stenen kwijtraakte.
• De nieuwe aanpak: In plaats van stenen weg te gooien, proberen we de koffer te herschikken. We zeggen: "Laten we de inhoud splitsen in twee delen: een lege, lichte structuur (waar we veel weg kunnen laten) en een kleine, krachtige kern (die de echte magie bevat)."

2. De Oplossing: 3BASiL (De Slimme Sorteerder)

De auteurs hebben een algoritme bedacht dat de koffer in tweeën deelt:

1. De "Spaarne" (Sparse): Dit is het grootste deel van de koffer, maar het is grotendeels leeg. Het is als een raamkozijn: er zit veel ruimte tussen de balken, maar de structuur staat er nog steeds. Omdat het leeg is, kost het weinig ruimte om op te slaan.
2. De "Kern" (Low-Rank): Dit is een klein, compact blokje dat de belangrijkste informatie bevat. Het is als een magneet die de essentie van de stenen vasthoudt.

Hoe doen ze dit?
Ze gebruiken een wiskundige techniek genaamd ADMM.

• De Analogie: Stel je voor dat je een zware muur moet afbreken en herbouwen in een lichter huis. Je doet dit niet in één keer. Je werkt in drie stappen die elkaar afwisselen (vandaar "3-Block"):
  1. Je kijkt naar de lege plekken en maakt ze nog leger.
  2. Je kijkt naar de rest en maakt een compacte kern van de overgebleven stenen.
  3. Je kijkt of de nieuwe muur nog steeds lijkt op de oude, en past het een beetje aan.
  4. Je herhaalt dit tot het perfect is.

Het mooie van 3BASiL is dat ze dit gelijktijdig doen. Veel oude methoden deden eerst het leegmaken en daarna het compact maken, wat vaak leidde tot fouten. 3BASiL doet het in één keer, zoals een meester-architect die beide ideeën tegelijk in zijn hoofd heeft.

3. De Fijnafstelling: TM (De "Matchmaker")

Nadat ze de robot hebben verkleind, is hij nog niet perfect. Hij doet het goed op kleine stukjes, maar als je hem een heel verhaal laat lezen, loopt hij vast.

• De Analogie: Stel je voor dat je een kopie maakt van een beroemd schilderij. De details zijn goed, maar de sfeer is net iets anders.
• De onderzoekers hebben een extra stap toegevoegd genaamd Transformer Matching (TM). Ze laten de verkleinde robot kijken naar de originele, zware robot en zeggen: "Kijk, als de grote robot dit woord zegt, wat zou jij dan moeten zeggen?"
• Ze passen de verkleinde robot een beetje aan (zonder de hele robot opnieuw te trainen) zodat hij precies hetzelfde gedrag vertoont als de grote versie. Dit is als het geven van een laatste polijstbeurt zodat de kopie ononderscheidbaar is van het origineel.

Waarom is dit geweldig?

1. Snelheid: Het duurt veel minder tijd om deze verkleinde robot te maken dan met oude methoden. Het is alsof je van handmatig stenen slepen overschakelt op een kraan.
2. Kwaliteit: De verkleinde robot is bijna net zo slim als de grote versie. De onderzoekers laten zien dat de "slimme" robot (3BASiL) veel minder fouten maakt dan andere verkleinde robots.
3. Toekomst: Omdat de robot nu lichter is, kun je hem straks op je eigen telefoon of laptop draaien, zonder dat je een datacenter nodig hebt.

Kortom:
3BASiL is een slimme manier om een gigantische, zware AI te verkleinen door hem op te splitsen in een "lege structuur" en een "kleine kern", en hem daarna een laatste duwtje te geven zodat hij precies doet wat de grote versie doet. Hierdoor kunnen we slimme AI overal mee naartoe nemen.

Technische samenvatting

1. Probleemstelling

Grote Taalmodellen (LLMs) bieden uitstekende prestaties, maar hun enorme reken- en geheugenvereisten maken implementatie op resource-beperkte apparaten of in real-time toepassingen moeilijk. Bestaande compressietechnieken zoals pruning (verwijdering van gewichten) en quantisatie zijn effectief, maar leiden vaak tot prestatieverlies.

Een veelbelovende richting is de Sparse plus Low-Rank (S + LR) decompositie, waarbij de oorspronkelijke gewichten $W$ worden benaderd als de som van een schaarse matrix $S$ en een laag-rang matrix $L$ ($W \approx S + L$). Dit combineert de voordelen van pruning (efficiëntie) met laag-rang aanpassing (flexibiliteit). Echter, bestaande methoden voor S + LR decompositie lijden onder twee belangrijke tekortkomingen:

1. Suboptimale optimalisatie: Ze gebruiken vaak "alternating minimization" (afwisselend optimaliseren van $S$ en $L$), wat geen gegarandeerde convergentie biedt en kan leiden tot een slecht gezamenlijk optimum.
2. Lagen-gewijze beperking: Bestaande methoden optimaliseren per laag onafhankelijk, wat cumulatieve fouten introduceert in de transformer-architectuur.

2. Methodologie

De auteurs introduceren 3BASiL-TM, een efficiënte "one-shot" post-training methode die bestaat uit twee hoofdstappen:

A. 3BASiL: 3-Block ADMM voor Laag-gewijze Decompositie

In plaats van afwisselend te optimaliseren, gebruiken de auteurs een 3-Block Alternating Direction Method of Multipliers (ADMM) framework. Dit lost het optimalisatieprobleem op om de reconstructiefout te minimaliseren:
$$\min_{S,L} \frac{1}{2} \|XW - X(S + L)\|_F^2 + \frac{\lambda}{2} \|W - (S + L)\|_F^2$$
onder de voorwaarden dat $S$ schaars is en $L$ een maximale rang $r$ heeft.

Het innovatieve aspect is de introductie van een hulpvariabele $D$ (een kopie van $S$), waardoor het probleem wordt opgesplitst in drie blokken:

1. S-blok: Update van de schaarse component.
2. L-blok: Update van de laag-rang component.
3. D-blok: Projectie op de schaarsheidsconstraint.

Voordelen:

• Gegarandeerde convergentie: De auteurs bewijzen theoretisch dat de methode convergeert als de penalty-parameter $\rho_t$ voldoende snel toeneemt.
• Gesloten-formule oplossingen: De updates voor $S$ en $L$ kunnen efficiënt worden berekend via matrixinversies en SVD (Singular Value Decomposition), wat sneller is dan gradient-based benaderingen die in andere methoden worden gebruikt.
• Efficiëntie: Door vooraf berekende eigenwaarde-decomposities en randomised-SVD te gebruiken, is de algoritme-complexiteit $O(N^3)$, wat zeer schaalbaar is voor LLMs.

B. Transformer-Matching (TM): Globale Verfijning

Na de laag-gewijze decompositie wordt een tweede stap toegepast: Transformer-Matching (TM).

• Doel: De auteurs merken op dat laag-gewijze optimalisatie imperfecte proxies zijn voor de werkelijke eind-tot-eind verliesfunctie. TM corrigeert dit door de output van een hele transformer-blok (bestaande uit meerdere lagen) af te stemmen op de output van het oorspronkelijke dichte model.
• Methode: Een gradient-based optimalisatie (Adam) wordt uitgevoerd om zowel de schaarse ($S$) als de laag-rang ($L$) componenten gezamenlijk over alle lagen in een transformer-blok te verfijnen.
• Universeel: Deze stap is onafhankelijk van de onderliggende decompositiemethode en kan worden toegepast op elke S + LR (of zelfs puur schaarse) compressie om de kwaliteit te verbeteren.
• Efficiëntie: Het vereist weinig extra geheugen en runtime, omdat het slechts één transformer-blok per keer verwerkt en gebruikmaakt van een kleine calibratiedataset.

3. Belangrijkste Bijdragen

1. 3BASiL (3-Block ADMM): Een nieuw algoritme dat S + LR decompositie behandelt binnen een unified optimalisatieraamwerk met theoretische convergentiegaranties, in tegenstelling tot de gebruikelijke alternatieve minimalisatie.
2. Transformer-Matching (TM): Een universele, geheugenefficiënte verfijningstechniek die de kwaliteit van de compressie aanzienlijk verbetert door transformer-niveau fouten te minimaliseren.
3. State-of-the-Art Prestaties: De combinatie 3BASiL-TM bereikt nieuwe records in compressie-efficiëntie en prestatiebehoud.

4. Experimentele Resultaten

De methode is getest op diverse LLaMA-modellen (van 1B tot 8B parameters) en een OPT-30B model.

• Perplexiteit: Onder een configuratie van (2:4 Sparse + 64 LR) voor LLaMA-8B, verkleint 3BASiL-TM de perplexiteitskloof met het dichte model met meer dan 30% in vergelijking met bestaande SOTA-methoden (zoals OATS en HASSLE-free).
• Zero-Shot Taken: De methode behaalt de hoogste scores op diverse zero-shot benchmarks (zoals PIQA, ARC, HellaSwag) na compressie.
• Snelheid:
  • De compressie is >2.5x sneller dan de huidige SOTA S + LR methode op een A100 GPU.
  • In vergelijking met de HASSLE-free-ALPS baseline is er een snelheidswinst van >7x bij het comprimeren van een Llama3.2-3B model.
• LoRA Fine-tuning: De gegenereerde laag-rang componenten dienen als een "slimme initialisatie" voor LoRA fine-tuning. Zelfs na fine-tuning behoudt 3BASiL-TM een significant voordeel ten opzichte van concurrenten.

5. Betekenis en Conclusie

3BASiL-TM is een doorbraak in het veld van modelcompressie voor LLMs. Het adresseert fundamentele beperkingen in bestaande S + LR methoden door:

1. Een robuustere wiskundige basis te bieden (ADMM met convergentiegarantie) voor het gezamenlijk optimaliseren van schaarse en laag-rang componenten.
2. De "stapel" van fouten die optreedt bij laag-gewijze compressie te doorbreken via transformer-matching.

Dit maakt het mogelijk om LLMs extreem efficiënt te comprimeren (voor snellere inferentie en minder geheugengebruik) zonder de kwaliteit van het model significant te verliezen. De code is open source beschikbaar, wat de reproduceerbaarheid en adoptie in de gemeenschap bevordert.