Technische Samenvatting: Efficiënte LLM's met AMP: Pruning van Attention Heads en MLP

Probleemstelling

Large Language Models (LLM's) hebben menselijk niveau van prestaties bereikt in cognitieve taken, maar lijden onder uitgebreide parameteraantallen, wat leidt tot hoge computationele kosten en trage inferentie. Deze factoren hinderen de implementatie in omgevingen met beperkte middelen. Hoewel pruning een veelbelovende strategie is om de modelgrootte te verminderen, kampen bestaande methoden met aanzienlijke trade-offs:

• Structured Pruning: Verwijdert volledige componenten (bijv. lagen, heads), maar vereist vaak computationeel dure fine-tuning of leidt tot prestatieverlies als kritieke componenten worden verwijderd.
• Semi-Structured Pruning: Vereist gespecialiseerde hardware (bijv. specifieke 2:4 GPU-sparsity patronen) om versnellingen te bereiken, wat de implementatie op consumentenhardware beperkt.
• Unstructured Pruning: Creëert onregelmatige ijle matrices die de inferentie op algemene hardware niet versnellen zonder gespecialiseerde software-optimalisaties.
• Algemene Beperkingen: Veel bestaande technieken vertrouwen op loss-gebaseerde metrieken, Taylor-expansies of complexe heuristieken die computationeel duur, inflexibel of afhankelijk van gespecialiseerde hardware zijn. Er is behoefte aan een hardware-agnostische, efficiënte methode die LLM's comprimeert met minimaal prestatieverlies en tastbare versnellingen in inferentie.

Methodologie: AMP (Attention Heads and MLP Pruning)

De auteurs stellen AMP voor, een structured pruning-methode die minder kritieke attention heads binnen Multi-Head Attention (MHA) en neuronen binnen Multilayer Perceptrons (MLP) verwijdert. In tegenstelling tot methoden die uitsluitend vertrouwen op gewicht-magnitudes of veranderingen in de loss, evalueert AMP de structurele belangrijkheid door inputdata op gewichten te projecteren om de bijdrage van elke component aan de residual stream te beoordelen.

Kernmechanisme

De methode werkt volgens het principe dat informatie door de residual stream stroomt. AMP berekent een belangrijkheidsscore voor elke component op basis van de magnitude van de activatiebijdrage aan de uiteindelijke output van de laag.

1. MHA Component:

   • De MHA-output wordt gedecomposeerd in een som van bijdragen van individuele attention heads ($O = \sum h_n W_n$).
   • De belangrijkheid van de $n$-de head ($I_n$) wordt gedefinieerd als de $\ell_1$-norm van het product van de output van de head ($h_n$) en de bijbehorende output projectie matrix ($W_n$):
     $$I_n = \|h_n W_n\|_1$$
   • Deze aanpak houdt rekening met de niet-lineaire interacties tussen gewichten en inputs, in tegen tegenstelling tot eenvoudige weight-norm pruning.

2. MLP Component:

   • De methode richt zich op de SwiGLU-mechanisme die wordt gebruikt in modellen zoals LLaMA, specifiek de Up en Gate projecties.
   • Omdat SwiGLU elementgewijze vermenigvuldiging gebruikt, moeten neuronen in paren worden gepruned uit beide projecties om de architecturale integriteit te behouden.
   • De belangrijkheid van het $m$-de paar neuronen ($I_m$) wordt berekend als het gemiddelde van de $\ell_1$-norm van het elementgewijze product van de Gate en Up projectie outputs over een steekproef van input tokens:
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. Pruning Proces:

   • Een kleine kalibratieset (bijv. 50 willekeurige samples van Alpaca) wordt gebruikt om gemiddelde activatiewaarden te extraheren.
   • Componenten (heads en neuronen) worden gerangschikt op belangrijkheid over alle lagen.
   • Om een doelstelling compressieratio $c$ te bereiken, worden de onderste $c\%$ attention heads en de onderste $c\%$ MLP-neuronen uniform over alle lagen verwijderd.
   • Post-Training: Het geprunede model ondergaat een lichtgewicht fine-tuning proces met behulp van Low-Rank Adaptation (LoRA) om prestaties te herstellen.

Belangrijkste Bijdragen

• Nieuwe Belangrijkheidsmetriek: AMP introduceert een data-gedreven metriek die inputactivaties op gewichten projecteert om de belangrijkheid van componenten te evalueren, wat verschilt van traditionele weight-only of loss-gebaseerde heuristieken.
• Hardware-Agnostische Versnelling: De methode bereikt praktische inferentie-versnellingen (tot 1.25$\times$) zonder gespecialiseerde hardware of sparse tensor cores te vereisen, waardoor het geschikt is voor consumenten-GPU's.
• Efficiëntie: Het beslissingsproces voor pruning duurt slechts enkele minuten op een consumenten-GPU, en de daaropvolgende fine-tuning vereist verwaarloosbare middelen vergeleken met volledige hertraining.
• Validatie van Coherentie: De auteurs stellen een "coherence check" experiment voor en voeren dit uit, waarbij zij aantonen dat het verwijderen van de meest belangrijke componenten geïdentificeerd door AMP catastrofale prestatiedalingen veroorzaakt, wat de metriek valideert in het onderscheiden van kritieke versus niet-kritieke structuren.
• Flexibiliteit: De methode is toepasbaar op verschillende LLM-families (LLaMA, Phi) en staat flexibele compressieratio's toe zonder gebonden te zijn aan vaste sparsity-patronen (bijv. 2:4).

Experimentele Resultaten

De auteurs evalueerden AMP op LLaMA 7B, LLaMA-2 7B, Phi-1.5 en Phi-2 modellen met behulp van zero-shot benchmarks (WinoGrande, HellaSwag, ARC-e/c, PIQA) en perplexiteit (WikiText2).

• Prestaties:
  • LLaMA 7B: Bij een pruning ratio van 20% behaalde AMP een gemiddelde nauwkeurigheid van 64,52%, waarmee het state-of-the-art methoden zoals LLM-Pruner (59,71%) en Shortened LLaMA (63,46%) overtrof.
  • LLaMA-2 7B: Bij een pruning ratio van 30% behaalde AMP een gemiddelde nauwkeurigheid van 61,02%, waarmee het SliceGPT (51,50%) en LLM Surgeon (59,03%) overtrof.
  • Verbetering: AMP overtrof bestaande structured pruning technieken met tot wel 1,49 procentpunt in gemiddelde taaknauwkeurigheid.
• Inferentie Snelheid:
  • AMP behaalde een 1.25$\times$ versnelling voor LLaMA 7B en 1.19$\times$ voor LLaMA-2 7B.
  • Dit presteert beter dan Shortened LLaMA (1.13$\times$) en evenaart semi-gestructureerde methoden zoals Wanda (1.24$\times$) zonder gespecialiseerde hardware te vereisen.
• Perplexity: Fine-tuned geprunede modellen vertoonden aanzienlijk lagere perplexiteit op WikiText2 vergeleken met niet-fine-tuned versies, waarbij het gat groter werd bij hogere pruning ratio's.
• Ablatie Studies:
  • Coherence Check: Het omkeren van de pruning-volgorde (het verwijderen van de meest belangrijke componenten) resulteerde in nauwkeurigheidsdalingen van >25 procentpunten, wat de validiteit van de metriek bevestigt.
  • Component Isolatie: Alleen MHA heads of alleen MLP neuronen prunen resulteerde in aanzienlijk slechtere prestaties dan het gelijktijdig prunen van beide, wat wijst op de noodzaak van een gebalanceerde aanpak.

Betekenis en Claims

Het artikel claimt dat AMP de kritieke trade-offs in modelcompressie aanpakt door een methode aan te bieden die:

1. Effectief is: Het overtreft huidige state-of-the-art structured pruning methoden in nauwkeurigheid terwijl het hoge compressieratio's behoudt.
2. Efficiënt is: Het vereist minimale computationele middelen voor zowel de pruning-beslissing (minuten) als de prestatieherstel (lichtgewicht LoRA fine-tuning).
3. Praktisch is: Het maakt tastbare inferentie-versnellingen mogelijk op standaard hardware, wat implementatie in omgevingen met beperkte middelen vergemakkelijkt.
4. Duurzaam is: In lijn met Green AI principes, vermindert de methode de operationele kosten met ~17,39% en de CO2-emissies met ~19,97% door verminderde inferentie latentie en parameteraantallen.

De auteurs benadrukken dat hun methode succesvol niet-essentiële structuren identificeert en verwijdert, waardoor de implementatie van hoog geoptimaliseerde LLM's mogelijk is zonder de noodzaak van gespecialiseerde hardware-infrastructuur. Ze benadrukken ook de cruciale rol van post-training (fine-tuning) bij het herstellen van de modelprestaties na pruning.