HEAPr: Hessian-based Efficient Atomic Expert Pruning in Output Space

Il paper introduce HEAPr, un nuovo algoritmo di pruning che scompone gli esperti nei modelli MoE in "esperti atomici" indivisibili e utilizza informazioni di secondo ordine semplificate nell'output per ottenere una compressione quasi senza perdita con una complessità computazionale ridotta, superando le prestazioni dei metodi di pruning esistenti.

L'essenziale

Immagina di avere un gigantesco esercito di specialisti (chiamati "Esperti") all'interno di un'intelligenza artificiale. Questi esperti lavorano insieme per rispondere alle tue domande.

In passato, questi eserciti erano "densi": tutti lavoravano su ogni domanda, rendendo il sistema lentissimo e costosissimo da ospitare. Poi, è arrivata l'idea degli MoE (Mixture of Experts): invece di far lavorare tutti, ne scegliamo solo pochi (i più adatti) per ogni domanda. È come se, invece di chiamare un intero reparto ospedaliero per una semplice influenza, chiamassi solo il medico di turno giusto.

Il problema? Anche se ne usi solo pochi per rispondere, devi ospitare tutti gli esperti nella memoria del computer. È come avere un hotel con 10.000 stanze, ma ne usi solo 10 ogni giorno. Il costo per mantenere l'hotel aperto è enorme, anche se le stanze sono vuote.

La soluzione: HEAPr (Il "Potatore Atomico")

Gli scienziati hanno provato a tagliare interi esperti (come se licenziassimo un intero reparto), ma spesso questo rovinava la qualità delle risposte. È come se licenziassimo un intero dipartimento di un'azienda: potresti risparmiare spazio, ma perdi competenze preziose.

Gli autori di questo paper hanno inventato HEAPr, un metodo intelligente per fare un "taglio chirurgico" molto più preciso. Ecco come funziona, spiegato con analogie semplici:

1. Scomporre l'Esperto in "Atomi"

Immagina che ogni "Esperto" non sia un blocco unico e indivisibile, ma sia composto da piccoli mattoncini (chiamati "Esperti Atomici").

• Metodo vecchio: Se un esperto era inutile, lo buttavi via tutto intero.
• Metodo HEAPr: Guardi dentro l'esperto e vedi che è fatto di 100 mattoncini. Forse 20 di questi mattoncini non servono a nulla, mentre gli altri 80 sono fondamentali. HEAPr ti permette di rimuovere solo quei 20 mattoncini inutili, lasciando intatto il resto dell'esperto.

2. Come decide cosa tagliare? (La "Sfera di Cristallo" Matematica)

Come fa il sistema a sapere quali mattoncini sono inutili senza provare e sbagliare?
HEAPr usa una palla di cristallo matematica (basata sulla teoria del "Cervello Ottimale" o Optimal Brain Surgeon).
Invece di guardare solo quanto un mattoncino è "grande" o "attivo", HEAPr calcola quanto danni farebbe rimuoverlo.

• Se rimuovi un mattoncino e il danno è quasi zero, via!
• Se rimuoverlo farebbe crollare la qualità della risposta, lo tieni.

3. Il trucco per non impazzire (L'Efficienza)

Calcolare questo "danno potenziale" per ogni singolo mattoncino di un modello gigante sarebbe come cercare di contare ogni granello di sabbia di un deserto: richiederebbe anni di tempo e computer enormi.
Gli autori hanno trovato un trucco geniale: invece di analizzare i mattoncini mentre sono "dentro" (nei parametri), li analizzano mentre escono (nell'output).

• Analogia: Invece di smontare ogni singolo ingranaggio di un orologio per vedere se è rotto, guardi l'orologio mentre segna l'ora. Se l'ora è corretta, non devi smontare tutto.
  Questo trucco riduce la complessità da "impossibile" a "facile", permettendo al sistema di fare il calcolo in pochissimo tempo (solo due passaggi in avanti e uno indietro).

I Risultati: Magia Pratica

Cosa succede quando applicano HEAPr?

• Taglio senza dolore: Possono rimuovere fino al 20-25% dei "mattoncini" (o anche di più) senza che l'intelligenza artificiale perda quasi nulla della sua capacità di ragionare. È come togliere il grasso in eccesso da un atleta senza togliere i muscoli.
• Risparmio reale: Il modello diventa più leggero e veloce (riducendo i calcoli necessari del 20%), rendendo possibile l'uso di questi modelli potenti su computer più piccoli o meno costosi.
• Test su giganti: Hanno provato questo metodo su modelli famosi come DeepSeek e Qwen, ottenendo risultati migliori rispetto a tutti i metodi precedenti.

In sintesi

HEAPr è come un giardiniere esperto che non taglia l'intero albero (l'esperto) perché sembra troppo grande, ma usa un paio di forbici microscopiche per rimuovere solo le foglie secche e i rami morti (gli esperti atomici inutili). Il risultato è un albero più sano, più leggero e che continua a dare frutti perfetti, occupando meno spazio nel giardino.

Sommario tecnico

1. Il Problema

Le architetture Mixture-of-Experts (MoE) nei Large Language Models (LLM) offrono prestazioni eccezionali e costi di inferenza ridotti rispetto ai modelli densi, attivando solo un sottoinsieme di parametri per ogni token. Tuttavia, presentano un collo di bottiglia critico: l'alto requisito di memoria. Anche se solo una frazione dei parametri è attiva durante l'inferenza (es. DeepSeek-V3 attiva 37B parametri su un totale di 671B), l'intero set di parametri deve risiedere nella memoria GPU, rendendo il deployment proibitivo su dispositivi con risorse limitate.

Le tecniche di pruning esistenti affrontano questo problema con due approcci principali, entrambi con limiti significativi:

• Pruning a livello di Expert (coarse-grained): Rimuove interi esperti. Sebbene riduca la memoria, spesso causa una significativa degradazione della precisione perché elimina blocchi funzionali completi.
• Pruning a livello di peso/struttura fine: Preserva meglio l'accuratezza ma non offre accelerazioni hardware dirette o riduzioni di memoria sufficienti.

Esistono metodi di fusione (merging) o decomposizione degli esperti, ma questi spesso soffrono di instabilità, conflitti parametrici o richiedono operazioni computazionalmente costose e ri-addestramento.

2. Metodologia: HEAPr

Il paper introduce HEAPr, un algoritmo di pruning che opera a un livello di granularità più fine: l'Atomic Expert.

Concetto di Atomic Expert

Invece di trattare un esperto come un blocco monolitico, HEAPr scompone ogni esperto in unità indivisibili chiamate "atomic experts". Formalmente, per un esperto $E_i$, le colonne delle matrici $W_{up}$ e $W_{gate}$ e la riga corrispondente di $W_{down}$ sono raggruppate insieme. L'output di un intero esperto è la somma degli output di questi atomic experts. Prunare a questo livello permette di rimuovere solo le unità che contribuiscono poco, isolando l'effetto del pruning senza interferire con le componenti rimanenti.

Stima dell'Importanza basata su Hessian (OBS)

Per determinare quali atomic experts rimuovere, HEAPr si basa sulla teoria dell'Optimal Brain Surgeon (OBS). L'obiettivo è minimizzare l'aumento della funzione di perdita ($\Delta \mathcal{L}$) quando un parametro viene rimosso.
La sfida principale è il costo computazionale e di memoria per calcolare la matrice Hessiana ($H$), che è proibitivo per modelli moderni.

HEAPr risolve questo problema attraverso due ottimizzazioni chiave:

1. Decomposizione in Atomic Experts: Dimostra che i parametri di diversi atomic experts all'interno dello stesso esperto sono disaccoppiati (le derivate incrociate sono zero). Questo riduce la complessità spaziale della stima dell'Hessiana da $O((3d_{model} \cdot d_{inter})^2)$ a $O((3d_{model})^2 \cdot d_{inter})$.
2. Trasformazione nello Spazio di Output: Invece di vincolare i parametri a zero (che richiederebbe di invertire l'Hessiana completa), HEAPr riformula il problema vincolando l'output dell'atomic expert a zero per un dato token.
   • Sfruttando l'equivalenza tra la Matrice di Informazione di Fisher e l'Hessiana attesa per reti ben convergenti, e applicando uno sviluppo di Taylor alla funzione dell'atomic expert, l'importanza viene calcolata nello spazio di output.
   • Poiché tutti gli atomic experts di uno stesso esperto condividono lo stesso gradiente della perdita rispetto all'output dell'esperto, è possibile calcolare una singola matrice di covarianza dei gradienti per esperto.
   • La complessità spaziale finale scende a $O(d_{model}^2)$, rendendo il calcolo fattibile.

Algoritmo

HEAPr richiede solo due forward pass e un backward pass su un piccolo set di calibrazione:

1. Stima della Covarianza dei Gradienti: Calcolo del gradiente condiviso per ogni esperto e costruzione della matrice di covarianza.
2. Calcolo dell'Importanza: Per ogni atomic expert, si calcola il punteggio di importanza $s_k$ come $s_k \approx \frac{1}{2} e_k(x)^\top \bar{G} e_k(x)$, dove $e_k(x)$ è l'output e $\bar{G}$ è la covarianza dei gradienti.
3. Ranking Globale: Gli atomic experts sono classificati globalmente su tutto il modello (non solo per layer) e i peggiori vengono rimossi.

3. Contributi Chiave

• Nuova Granularità: Introduzione del concetto di "Atomic Expert" per il pruning MoE, permettendo una rimozione più flessibile e precisa rispetto al pruning di interi esperti.
• Efficienza Computazionale: Sviluppo di uno schema di approssimazione del secondo ordine che riduce la complessità spaziale da $O(d^4)$ a $O(d^2)$, trasformando l'analisi dallo spazio dei parametri a quello degli output.
• Algoritmo Pratico: Un metodo che non richiede ri-addestramento (fine-tuning) e richiede solo un passaggio di calibrazione minimo (2 forward, 1 backward).
• Prestazioni Superiori: Validazione su modelli MoE di stato dell'arte (DeepSeek, Qwen) che dimostra come HEAPr superi i metodi esistenti (NAEE, MoE-I2, Sub-MoE, D2-MoE, ecc.).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli come DeepSeekMoE-16B-Base, Qwen1.5-MoE-A2.7B-Chat, Qwen2-57B-A14B e Qwen3-30B-A3B su sette task zero-shot (es. ARC, HellaSwag, MATHQA).

• Compressione "Quasi Lossless":
  • A un tasso di pruning del 20-25%, HEAPr mantiene prestazioni quasi identiche al modello originale su DeepSeekMoE-16B e Qwen1.5-MoE.
  • Su Qwen2-57B-A14B, mantiene le prestazioni quasi intatte anche a un 40% di pruning.
  • Su Qwen3-30B-A3B, la precisione media scende di soli 0.03 punti a un pruning del 25%.
• Riduzione FLOPs: Il metodo riduce i FLOPs di circa il 20% mantenendo l'accuratezza, superando i metodi di pruning a livello di esperto che spesso non offrono accelerazioni reali.
• Confronto con SotM: HEAPr supera metodi concorrenti come CAMERA-P (che usa un approccio euristico basato sull'energia di decodifica) e tecniche di fusione/decomposizione, specialmente nella capacità di effettuare un ranking globale coerente.
• Robustezza: Le prestazioni sono stabili indipendentemente dal set di dati di calibrazione (WikiText-2 vs C4) e migliorano con la dimensione del set di calibrazione.

5. Significato e Impatto

HEAPr rappresenta un passo avanti significativo nella compressione dei modelli MoE. Risolvendo il dilemma tra riduzione della memoria/accelerazione e mantenimento dell'accuratezza, permette di:

1. Rendere fattibile il deployment di modelli MoE su hardware con risorse limitate (es. GPU consumer o edge devices) riducendo drasticamente l'occupazione di memoria VRAM.
2. Superare i limiti del pruning strutturale dimostrando che la decomposizione fine degli esperti, guidata da informazioni di secondo ordine efficienti, è superiore ai metodi basati su clustering o fusione.
3. Fornire una nuova prospettiva teorica sulla ridondanza interna degli esperti MoE, suggerendo che una grande parte della capacità del modello risiede in unità atomiche specifiche che possono essere rimosse senza compromettere le prestazioni globali.

In sintesi, HEAPr offre un metodo scalabile, efficiente e ad alte prestazioni per comprimere i moderni LLM MoE, aprendo la strada a un'adozione più ampia di queste architetture in scenari reali vincolati dalle risorse.