LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

Riepilogo Tecnico: LEAP: Potatura Adattiva End-to-End Apprendibile dei Modelli Linguistici di Grande Dimensione

Enunciato del Problema
Il deployment dei Modelli Linguistici di Grande Dimensione (LLM) è sempre più ostacolato da vincoli di memoria e calcolo. Sebbene la sparsità dei pesi offra una soluzione, il panorama è diviso in regimi strutturati, semi-strutturati (ad es. 2:4) e non strutturati. I recenti avanzamenti nei kernel GPU (ad es. FlashLLM, SpInfer) e nell'hardware dataflow hanno spostato il collo di bottiglia dall'esecuzione della sparsità non strutturata alla sua induzione con una perdita minima di accuratezza.

I metodi attuali all'avanguardia per la potatura non strutturata, come Wanda, SparseGPT, Thanos e ADMM, si basano su surrogati livello per livello derivati dal principio del Cerebro Ottimale (OBS). Questi metodi minimizzano gli errori di ricostruzione locali piuttosto che la perdita del modello end-to-end. Di conseguenza, accumulano errori locali che si moltiplicano nelle reti profonde, portando a una significativa degradazione dell'accuratezza, in particolare sotto sparsità aggressive (50–60%).

Al contrario, i metodi di maschera apprendibile end-to-end come MaskLLM e PATCH ottimizzano con successo le maschere direttamente contro la perdita di modellazione linguistica, ma sono limitati a pattern semi-strutturati. La loro parametrizzazione si basa su una distribuzione categorica su pattern validi all'interno di un gruppo. Sebbene trattabile per gruppi piccoli (ad es. $\binom{4}{2}=6$ pattern per una sparsità 2:4), questo approccio diventa combinatoriamente intrattabile per la sparsità non strutturata. Per una riga di larghezza $d=4096$ a una sparsità del 50%, il numero di maschere valide è $\binom{4096}{2048} \approx 10^{1229}$, rendendo impossibile l'archiviazione e l'indicizzazione dei logit dei pattern.

Metodologia: LEAP
Gli autori introducono LEAP (Potatura Adattiva End-to-End Apprendibile), un framework progettato per rendere trattabile l'apprendimento di maschere non strutturate end-to-end sostituendo la parametrizzazione intrattabile "categorica-sui-pattern" con un approccio per-peso.

• Relassamento Bernoulli per Peso: Invece di assegnare un logit a ogni pattern valido, LEAP assegna un logit apprendibile a ogni singolo peso. Per una matrice di pesi $W \in \mathbb{R}^{m \times n}$, viene introdotta una matrice di logit parallela $P \in \mathbb{R}^{m \times n}$. La maschera stocastica $M$ è generata tramite il relassamento Gumbel-sigmoid:
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  dove $g$ è rumore Gumbel, $\sigma$ è la funzione sigmoidale, $\alpha$ è un fattore di scala e $\tau$ è un parametro di temperatura. Questa formulazione scala linearmente con il numero di pesi ($O(mn)$) piuttosto che con il numero di pattern validi.
• Strategia di Ottimizzazione:
  • Pesi Congelati: I pesi preaddestrati $W$ vengono mantenuti congelati. Solo i logit della maschera $P$ vengono addestrati. Questo preserva la calibrazione del modello preaddestrato e semplifica il deployment.
  • Inizializzazione: $P$ è inizializzato utilizzando una maschera Wanda one-shot (impostando le voci selezionate a $+s$ e le altre a $-s$) per fornire un punto di partenza ragionevole, evitando problemi di avvio a freddo.
  • Programmi: Un programma di scala ($\alpha$) aumenta per spingere la sigmoidale verso valori binari, mentre un programma di temperatura ($\tau$) decade per affinare la maschera.
  • Funzione di Perdita: L'obiettivo combina la perdita di modellazione linguistica su un flusso di calibrazione con due termini di regolarizzazione:
    1. Regolarizzatore di Sparsità Globale: Impone la densità target globalmente su tutti i livelli, permettendo ai singoli livelli di regolare la propria densità in base all'importanza end-to-end.
    2. Termine Consapevole della Magnitudine: Sposta l'ottimizzazione verso il mantenimento di pesi di maggiore magnitudine per prevenire minimi degeneri in cui pesi critici vengono eliminati.
• Efficienza dell'Addestramento: Il metodo viene addestrato per circa 2.000 iterazioni su un piccolo flusso di testo generale di calibrazione (SlimPajama).

Contributi Chiave

1. Identificazione dell'Ostacolo Combinatorio: Il documento identifica formalmente che le parametrizzazioni "categoriche-sui-pattern" (utilizzate da MaskLLM/PATCH) non possono scalare alla sparsità non strutturata a causa della crescita esponenziale dei pattern validi.
2. Riformulazione Trattabile: Gli autori propongono una riformulazione Bernoulli per-peso tramite Gumbel-sigmoid che abilita l'apprendimento differenziabile end-to-end per maschere non strutturate alla scala degli LLM.
3. Framework LEAP: Un framework di potatura leggero che opera su pesi congelati, richiedendo solo un piccolo dataset di calibrazione e circa 2.000 passaggi di addestramento.
4. Validazione Empirica: Valutazione estesa su cinque famiglie di LLM (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) che vanno da 0,5B a 8B parametri a sparsità del 50% e 60%.

Risultati
LEAP è stato valutato rispetto a baseline livello per livello (Wanda, SparseGPT, Thanos, ADMM) e alla baseline apprendibile semi-strutturata MaskLLM.

• Miglioramenti dell'Accuratezza: Su dieci impostazioni (modello, sparsità), LEAP ha migliorato l'accuratezza zero-shot media su sei task di +2,59 punti rispetto ad ADMM, la migliore baseline livello per livello.
• Miglioramenti Specifici: Il guadagno più grande è stato di +5,40 punti su LLaMA-3.2 1B a una sparsità del 60% (50,39 contro 44,99).
• Perplessità: LEAP ha raggiunto costantemente una perplessità (PPL) inferiore rispetto alle baseline. Ad esempio, su LLaMA-3.1 8B a una sparsità del 50%, LEAP ha raggiunto 7,66 PPL rispetto ai 9,12 di ADMM.
• Allocazione della Sparsità: Gli esperimenti hanno mostrato che, sotto ottimizzazione end-to-end con un budget globale, LEAP converge a un'allocazione di densità per blocco quasi uniforme, in contrasto con le variazioni non uniformi osservate nei metodi semi-strutturati.
• Costo Computazionale: Sebbene più costoso dei metodi one-shot (richiedendo ore su H100 rispetto a minuti), LEAP è significativamente più economico di MaskLLM (ad es. 276 ore GPU per LLaMA-3.1 8B contro migliaia di ore A100 per MaskLLM su scale comparabili).

Significato e Affermazioni
Il documento afferma che LEAP risolve l'ostacolo fondamentale che impedisce l'apprendimento di maschere end-to-end nella sparsità non strutturata. Disaccoppiando l'apprendimento della maschera dagli aggiornamenti dei pesi, offre una via pratica per comprimere i modelli fondazionali senza sacrificare la calibrazione dei pesi preaddestrati.

Gli autori sottolineano che LEAP produce maschere nel regime non strutturato del 50%–60%, direttamente compatibili con i kernel accelerati esistenti (ad es. FlashLLM, SpInfer, MACKO). Il significato non risiede nella creazione di nuovi kernel di inferenza, ma nell'abilitare il deployment di questi kernel su modelli potati di alta qualità, precedentemente irraggiungibili tramite metodi livello per livello. Il lavoro suggerisce che la sparsità non strutturata, quando indotta tramite apprendimento end-to-end, può offrire un compromesso favorevole tra compressione e accuratezza per il deployment su hardware commerciale.

Limitazioni e Ambito

• Pesi Congelati: L'ambito corrente mantiene i pesi congelati. Sebbene l'ottimizzazione congiunta di pesi e maschere sia indicata come un'estensione naturale, non è implementata in questo lavoro.
• Allocazione Uniforme: L'osservazione di una densità per blocco uniforme è presentata come un risultato specifico del regime per i modelli e i livelli di sparsità studiati, non come una regola generale contro l'allocazione non uniforme.
• Sicurezza: Gli autori avvertono che i modelli compressi possono esibire modalità di guasto diverse su input a coda lunga e raccomandano una rivalutazione per deployment critici per la sicurezza.