Two-Stage Regularization-Based Structured Pruning for LLMs

1. Il Problema

La diffusione e il deployment dei Grandi Modelli Linguistici (LLM) sono attualmente ostacolati dal loro enorme numero di parametri, che richiede risorse computazionali e di memoria significative.

• Limitazioni delle metodologie esistenti: Le tecniche di pruning strutturato (rimozione di parametri interi, come canali o layer) esistenti tendono a identificare i layer "meno importanti" basandosi su metriche statiche e rimuoverli direttamente.
• Conseguenze: Questo approccio diretto causa una perdita di conoscenza critica contenuta anche nei layer considerati meno importanti. Per recuperare le prestazioni, è necessario un estensivo retraining (addestramento di recupero), che comporta un elevato costo computazionale e di tempo.

2. Metodologia: TRSP

Gli autori propongono TRSP (Two-Stage Regularization-Based Structured Pruning), un metodo innovativo che separa il processo di selezione dei layer da quello di rimozione, introducendo una fase di regolarizzazione per trasferire la conoscenza prima del taglio.

Il processo si articola in quattro fasi principali:

1. Preparazione dei Dati: Viene selezionata una piccola porzione di dati (es. 128 istanze) da benchmark standard (come WikiText-2 o Alpaca) per guidare il processo di regolarizzazione.
2. Fase 1: Apprendimento Iterativo dei Pesi dei Layer (Regularizzazione 1):
   • A ogni layer del transformer viene assegnato un peso apprendibile iniziale ($S[i]$).
   • L'output di ogni layer viene moltiplicato per questo peso.
   • Viene ottimizzata una funzione di perdita che include la perdita di modellazione linguistica più una penalità di regolarizzazione basata sulla norma $\ell_1$ dei pesi dei layer.
   • Questo processo è iterativo: in ogni passo, il layer con il peso più basso viene identificato e "mascherato" (escluso temporaneamente), permettendo agli altri layer di adattarsi progressivamente.
3. Fase 2: Trasferimento della Conoscenza (Regularizzazione 2):
   • Una volta identificati i layer da rimuovere (insieme $P$), viene applicata una regolarizzazione aggiuntiva sulla differenza tra l'output e l'input di questi specifici layer ($X_{out} - X_{in}$).
   • L'obiettivo è forzare questi layer a comportarsi come funzioni identità (cioè $X_{out} \approx X_{in}$), spingendo la conoscenza critica a migrare verso i layer rimanenti.
   • Questo passaggio riduce drasticamente l'impatto della rimozione successiva.
4. Pruning (Rimozione):
   • Dopo la regolarizzazione, i layer con i pesi più bassi vengono rimossi definitivamente dal modello.
   • Vantaggio chiave: Non è necessario alcun retraining (fine-tuning) dopo la rimozione, poiché il modello è già stato "preparato" a funzionare senza quei layer.

3. Contributi Chiave

• Conservazione della Conoscenza: A differenza dei metodi che rimuovono i layer direttamente, TRSP ridistribuisce attivamente la conoscenza dai layer da eliminare a quelli preservati tramite regolarizzazione, mitigando la perdita di capacità del modello.
• Efficienza e Assenza di Retraining: Il metodo è "retraining-free" dopo il pruning. Questo elimina il costo computazionale significativo associato al recupero delle prestazioni tipico delle tecniche attuali.
• Costo Minimo: Richiede una quantità di dati estremamente ridotta (pochi centinaia di campioni) per la fase di regolarizzazione, rendendo il processo molto leggero.
• Prestazioni Superiori: Dimostra di superare i metodi di pruning strutturato a livello di layer più avanzati (come SLEB, ShortGPT, LaCo) sia nelle task di generazione che in quelle zero-shot.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli (Phi-2, OPT, LLaMA2, LLaMA3) con diversi rapporti di pruning (dal 25% al 60%).

• Prestazioni: TRSP ottiene la perplessità (PPL) più bassa e la migliore accuratezza media rispetto ai baselines.
  • Esempio: Su LLaMA2-7B con pruning del 25%, TRSP riduce la perplessità del 20% rispetto al secondo metodo migliore (ShortGPT) e mantiene un calo di accuratezza media di soli l'1% rispetto al modello denso.
• Accelerazione End-to-End: La rimozione di interi layer porta a un significativo aumento del throughput e una riduzione della latenza.
  • Su OPT-13B con pruning del 50%, si ottiene un aumento del throughput del 75% e una riduzione della latenza del 46%.
• Robustezza: Il metodo si dimostra robusto su diversi dataset di calibrazione e mantiene buone prestazioni anche con rapporti di pruning elevati (fino al 60%), dove altri metodi falliscono.
• Scalabilità: I risultati su Llama-2-70B confermano che il metodo scala efficacemente anche su modelli molto grandi.

5. Significato e Impatto

TRSP rappresenta un passo avanti significativo verso il deployment efficiente degli LLM.

• Superamento del compromesso Tradizionale: Risolve il dilemma tra compressione del modello e mantenimento delle prestazioni, eliminando la necessità di costosi cicli di ri-addestramento.
• Accessibilità: Rendendo possibile la creazione di modelli più piccoli e veloci senza richiedere risorse di calcolo massicce per il fine-tuning, TRSP democratizza l'accesso a LLM efficienti per applicazioni reali.
• Nuova Direzione di Ricerca: Introduce l'idea di utilizzare la regolarizzazione non solo per prevenire l'overfitting, ma come meccanismo attivo per la migrazione della conoscenza prima della distruzione strutturale del modello, offrendo una nuova prospettiva per le strategie di pruning.

In sintesi, TRSP offre una soluzione pratica ed efficace per comprimere gli LLM, garantendo che i modelli ridotti mantengano quasi tutte le capacità cognitive di quelli originali, pronti per l'uso immediato.