MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

Riepilogo Tecnico: MaskPro – Apprendimento Probabilistico a Spazio Lineare per la Sparsità Rigida (N:M) su LLM

Enunciato del Problema
Il rapido scalamento dei Large Language Models (LLM) ha creato un collo di bottiglia significativo nell'efficienza dell'inferenza, rendendo necessarie tecniche di compressione efficaci. La sparsità semi-strutturata, in particolare il pattern (N:M) in cui $N$ elementi sono mantenuti su ogni $M$ pesi consecutivi, offre una soluzione hardware-friendly per l'accelerazione e la riduzione della memoria. Tuttavia, identificare la maschera (N:M) ottimale è un formidabile problema di ottimizzazione combinatoria. Gli approcci esistenti rientrano in due categorie con limitazioni distinte:

1. Metodi basati su regole: Questi impiegano ricerche greedy, livello per livello, utilizzando metriche create a mano (ad esempio, minimizzare l'errore di ricostruzione livello per livello). Sebbene computazionalmente economici, soffrono di bias intrinseci e di scostamenti considerevoli rispetto alle prestazioni del compito end-to-end.
2. Metodi basati sull'apprendimento (ad esempio, MaskLLM): Questi ottimizzano le maschere direttamente tramite discesa del gradiente su una perdita di generazione. Sebbene efficaci, comportano costi di addestramento proibitivi. Nello specifico, MaskLLM richiede di memorizzare i logit per tutte le maschere possibili all'interno di un gruppo di $M$ pesi, portando a una complessità di memoria di $O(\binom{M}{N} \frac{d}{M})$. Per $M$ grandi, questo scala esponenzialmente, rendendolo non fattibile per modelli di grandi dimensioni. Inoltre, il vasto spazio combinatorio induce un'alta varianza nei gradienti della politica, causando instabilità nell'addestramento.

Metodologia: MaskPro
Gli autori propongono MaskPro, un nuovo framework probabilistico a spazio lineare progettato per apprendere la sparsità (N:M) in modo efficiente e stabile.

1. Rilassamento Probabilistico a Spazio Lineare:
   Invece di definire una distribuzione categorica su tutte le $\binom{M}{N}$ maschere possibili per un gruppo di $M$ pesi, MaskPro riformula il problema. Stabilisce una distribuzione categorica sui $M$ elementi individuali all'interno di ciascun gruppo. La maschera (N:M) viene quindi generata eseguendo un campionamento N-volte senza reimmissione da questa distribuzione.

   • Base Teorica: Gli autori dimostrano che qualsiasi maschera (N:M) può essere rappresentata come una somma probabilistica di $N$ vettori di base distinti.
   • Efficienza di Memoria: Questa riformulazione riduce il numero di parametri apprendibili (logit) da $O(\binom{M}{N} \frac{d}{M})$ a un valore lineare $O(d)$, dove $d$ è il numero totale di pesi. Questo elimina l'esplosione della memoria associata ai precedenti metodi basati sull'apprendimento.

2. Stimatore di Gradiente della Politica Raffinato (PGE):
   Per ottimizzare i logit, MaskPro utilizza gradienti della politica. Gli autori identificano che i gradienti della politica standard che utilizzano valori di perdita grezzi sono instabili negli LLM a causa dell'alta varianza causata dalle differenze tra i minibatch (dove la variazione della perdita dovuta al campionamento dei dati spesso supera di gran lunga la variazione della perdita dovuta al campionamento della maschera).

   • Residui di Perdita: Invece di utilizzare la perdita grezza $f(m \odot w, \xi)$, MaskPro aggiorna i logit basandosi sul residuo di perdita rispetto a una maschera di baseline iniziale $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. Questo isola l'impatto del cambiamento della maschera.
   • Tracciatore a Media Mobile: Per stabilizzare ulteriormente l'addestramento e prevenire esplosioni numeriche, viene introdotto un tracciatore a media mobile $\delta$ per livellare la baseline. La regola di aggiornamento diventa:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • Garanzie Teoriche: Il documento fornisce dimostrazioni che mostrano come questi stimatori raffinati rimangano stimatori non distorti del vero gradiente della politica. Inoltre, sotto la condizione che la perdita della maschera campionata sia inferiore al doppio della perdita di baseline (una condizione facilmente soddisfatta nella pratica), la varianza dello stimatore raffinato è strettamente inferiore a quella del gradiente della politica vanilla.

Contributi Chiave

1. Framework a Spazio Lineare: MaskPro riduce il requisito di memoria per l'apprendimento della sparsità (N:M) da esponenziale $O(\binom{M}{N} \frac{d}{M})$ a lineare $O(d)$ modellando il processo di selezione come un campionamento N-volte senza reimmissione da una distribuzione categorica su $M$ elementi.
2. Gradiente della Politica Potenziato: Gli autori introducono un meccanismo di aggiornamento a varianza ridotta utilizzando residui di perdita e un tracciatore a media mobile, affrontando l'instabilità dell'addestramento causata dall'alta varianza dei gradienti della politica in grandi spazi combinatori.
3. Analisi Teorica: Vengono fornite dimostrazioni complete riguardanti la non distorsione degli stimatori e le proprietà di riduzione della varianza delle regole di aggiornamento proposte.

Risultati Sperimentali
Sono stati condotti esperimenti estesi su molteplici LLM (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) su varie attività downstream (ad esempio, HellaSwag, RACE, PIQA) e pattern di sparsità ((2:4), (4:8), (8:16)).

• Prestazioni: MaskPro supera costantemente i metodi basati su regole all'avanguardia (SparseGPT, Wanda, Pruner-Zero) e raggiunge prestazioni comparabili a MaskLLM, che richiede molta memoria, spesso con un miglioramento medio superiore al 2% in termini di accuratezza.
• Scalabilità della Memoria: Mentre MaskLLM incontra errori Out-of-Memory (OOM) su 8 GPU A100 per la sparsità (4:8) su modelli da 7B, MaskPro mantiene un utilizzo stabile della memoria (~36-37 GB) indipendentemente dal pattern di sparsità.
• Efficienza dei Dati: MaskPro dimostra una notevole robustezza al volume dei dati. Raggiunge prestazioni stabili anche quando addestrato con un solo campione di dati, whereas MaskLLM richiede centinaia di migliaia di campioni per convergere.
• Efficienza di Addestramento: Il metodo richiede significativamente meno campioni di addestramento e meno tempo di addestramento rispetto a MaskLLM, mantenendo al contempo un utilizzo della memoria comparabile ai metodi basati su regole.

Significato e Affermazioni
Il documento afferma che MaskPro fornisce un percorso praticabile per il dispiegamento locale, pratico e scalabile degli LLM, risolvendo i colli di bottiglia di memoria e stabilità dell'apprendimento della sparsità semi-strutturata basato sull'apprendimento. Gli autori sottolineano che il loro approccio colma il divario tra le alte prestazioni dei metodi basati sull'apprendimento e l'efficienza dei metodi basati su regole. Riducendo l'impronta di memoria a uno spazio lineare e stabilizzando il processo di addestramento attraverso gradienti della politica raffinati, MaskPro permette l'apprendimento di maschere (N:M) ottimali su modelli su larga scala senza il sovraccarico computazionale proibitivo precedentemente associato a tali compiti. Il metodo è presentato come un framework generalizzabile per l'apprendimento probabilistico delle maschere, robusto alla scarsità di dati e alla scala del modello.