Riepilogo Tecnico: LLM Efficienti con AMP: Pruning delle Teste di Attenzione e dei MLP

Problema

I Large Language Models (LLM) hanno raggiunto prestazioni di livello umano in compiti cognitivi, ma soffrono di un numero esteso di parametri, il che comporta elevati costi computazionali e una lenta inferenza. Questi fattori ostacolano l'implementazione in ambienti con risorse limitate. Sebbene il pruning sia una strategia promettente per ridurre le dimensioni del modello, i metodi esistenti affrontano significativi compromessi:

• Pruning Strutturato: Rimuove interi componenti (ad es., layer, teste) ma spesso richiede un fine-tuning computazionalmente costoso o porta a un degrado delle prestazioni se vengono rimossi componenti critici.
• Pruning Semi-Strutturato: Richiede hardware specializzato (ad es., specifici pattern di sparsità GPU come il 2:4) per ottenere accelerazioni, limitando l'implementazione su hardware di consumo.
• Pruning Non Strutturato: Crea matrici sparse irregolari che non riescono ad accelerare l'inferenza su hardware general-purpose senza ottimizzazioni software specializzate.
• Limitazioni Generali: Molte tecniche esistenti si affidano a metriche basate sulla perdita (loss), espansioni di Taylor o euristiche complesse che sono computazionalmente costose, poco flessibili o richiedono hardware specializzato. È necessario un metodo hardware-agnostic ed efficiente per comprimere gli LLM con una perdita minima di prestazioni e accelerazioni tangibili dell'inferenza.

Metodologia: AMP (Attention Heads and MLP Pruning)

Gli autori propongono AMP, un metodo di pruning strutturato che rimuove le teste di attenzione meno critiche all'interno della Multi-Head Attention (MHA) e i neuroni all'interno dei Multilayer Perceptrons (MLP). A differenza dei metodi che si basano esclusivamente sulla magnitudo dei pesi o sui cambiamenti della perdita, AMP valuta l'importanza strutturale proiettando i dati di input sui pesi per valutare il contributo di ogni componente allo stream residuo.

Meccanismo Core

Il metodo opera sul principio secondo cui l'informazione fluisce attraverso lo stream residuo. AMP calcola un punteggio di importanza per ogni componente basandosi sulla magnitudo del contributo di attivazione al risultato finale del layer.

1. Componente MHA:

   • L'output della MHA è decomposto in una somma di contributi da singole teste di attenzione ($O = \sum h_n W_n$).
   • L'importanza della $n$-esima testa ($I_n$) è definita come la norma $\ell_1$ del prodotto dell'output della testa ($h_n$) e della sua corrispondente matrice di proiezione dell'output ($W_n$):
     $$I_n = \|h_n W_n\|_1$$
   • Questo approccio tiene conto delle interazioni non lineari tra pesi e input, a differenza del semplice pruning della norma dei pesi.

2. Componente MLP:

   • Il metodo prende di mira il meccanismo SwigLU utilizzato in modelli come LLaMA, specificamente le proiezioni Up e Gate.
   • Poiché SwigLU utilizza la moltiplicazione elemento per elemento, i neuroni devono essere rimossi in coppie da entrambe le proiezioni per mantenere l'integrità architettonica.
   • L'importanza della $m$-esima coppia di neuroni ($I_m$) è calcolata come la media della norma $\ell_1$ del prodotto elemento per elemento delle proiezioni Gate e Up attraverso un campione di token di input:
     $$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. Processo di Pruning:

   • Un piccolo set di calibrazione (ad es., 50 campioni casuali da Alpaca) viene utilizzato per estrarre i valori medi di attivazione.
   • I componenti (teste e neuroni) vengono classificati per importanza attraverso tutti i layer.
   • Per raggiungere un rapporto di compressione target $c$, l'ultimo $c\%$ delle teste di attenzione e l'ultimo $c\%$ dei neuroni MLP vengono rimossi uniformemente in tutti i layer.
   • Post-Training: Il modello sottoposto a pruning subisce un processo di fine-tuning leggero utilizzando la Low-Rank Adaptation (LoRA) per recuperare le prestazioni.

Contributi Chiave

• Nuova Metrica di Importanza: AMP introduce una metrica basata sui dati che proietta le attivazioni di input sui pesi per valutare l'importanza del componente, differenziandosi dalle tradizionali euristiche basate solo sui pesi o sulla perdita.
• Accelerazione Hardware-Agnostic: Il metodo ottiene accelerazioni pratiche dell'inferenza (fino a 1.25$\times$) senza richiedere hardware specializzato o tensor core sparsi, rendendolo adatto a GPU di consumo.
• Efficienza: Il processo decisionale del pruning richiede solo pochi minuti su una GPU di consumo, e il successivo fine-tuning richiede risorse trascurabili rispetto al retraining completo.
• Validazione della Coerenza: Gli autori propongono ed eseguono un esperimento di "coherence check", dimostrando che rimuovere i componenti più importanti identificati da AMP causa crolli catastrofici delle prestazioni, validando la capacità della metrica di distinguere tra strutture critiche e non critiche.
• Flessibilità: Il metodo è applicabile a diverse famiglie di LLM (LLaMA, Phi) e consente rapporti di compressione flessibili senza essere vincolato a pattern di sparsità fissi (ad es., 2:4).

Risultati Sperimentali

Gli autori hanno valutato AMP su modelli LLaMA 7B, LLaMA-2 7B, Phi-1.5 e Phi-2 utilizzando benchmark zero-shot (WinoGrande, HellaSwag, ARC-e/c, PIQA) e perplexity (WikiText2).

• Prestazioni:
  • LLaMA 7B: Con un rapporto di pruning del 20%, AMP ha raggiunto un'accuratezza media del 64.52%, superando i metodi allo stato dell'arte come LLM-Pruner (59.71%) e Shortened LLaMA (63.46%).
  • LLaMA-2 7B: Con un rapporto di pruning del 30%, AMP ha raggiunto il 61.02% di accuratezza media, superando SliceGPT (51.50%) e LLM Surgeon (59.03%).
  • Miglioramento: AMP ha superato le tecniche di pruning strutturato esistenti fino a 1.49 punti percentuali in termini di accuratezza media del task.
• Velocità di Inferenza:
  • AMP ha ottenuto un'accelerazione di 1.25$\times$ per LLaMA 7B e di 1.19$\times$ per LLaMA-2 7B.
  • Questo supera Shortened LLaMA (1.13$\times$) e compete con i metodi semi-strutturati come Wanda (1.24$\times$) senza richiedere hardware specializzato.
• Perplexity: I modelli sottoposti a fine-tuning dopo il pruning hanno mostrato una perplexity significativamente inferiore su WikiText2 rispetto alle versioni non sottoposte a fine-tuning, con un divario che si amplia a rapporti di pruning più elevati.
• Studi di Ablazione:
  • Coherence Check: Invertendo l'ordine del pruning (rimuovendo i componenti più importanti), si è verificato un calo dell'accuratezza superiore ai 25 punti percentuali, confermando la validità della metrica.
  • Isolamento dei Componenti: Il pruning solo delle teste MHA o solo dei neuroni MLP ha portato a prestazioni significativamente peggiori rispetto al pruning simultaneo di entrambi, indicando la necessità di un approccio bilanciato.

Significato e Rivendicazioni

L'articolo sostiene che AMP affronti i critici compromessi nella compressione dei modelli offrendo un metodo che è:

1. Efficace: Supera gli attuali metodi di pruning strutturato allo stato dell'arte in termini di accuratezza mantenendo alti i tassi di compressione.
2. Efficiente: Richiede risorse computazionali minime sia per la decisione di pruning (minuti) che per il recupero delle prestazioni (leggero fine-tuning tramite LoRA).
3. Pratico: Consente accelerazioni tangibili dell'inferenza su hardware standard, facilitando l'implementazione in contesti con risorse limitate.
4. Sostenibile: In linea con i principi della Green AI, il metodo riduce i costi operativi di circa il 17.39% e le emissioni di carbonio di circa il 19.97% grazie alla riduzione della latenza di inferenza e del numero di parametri.

Gli autori sottolineano che il loro metodo identifica con successo e rimuove le strutture non essenziali, consentendo l'implementazione di LLM altamente ottimizzati senza la necessità di infrastrutture hardware specializzate. Evidenziano inoltre il ruolo critico del post-training (fine-tuning) nel ripristinare le prestazioni del modello dopo il pruning.