Sintesi Tecnica: Pruning Guidato dall'Attribuzione e Massimizzazione della Copertura per la Compressione di MoE Strutturali

1. Definizione del Problema

I modelli Mixture-of-Experts (MoE) offrono un'elevata capacità di parametri con una computazione gestibile, attivando solo un sottoinsieme di esperti per ogni token. Tuttavia, l'implementazione di questi modelli rimane costosa a causa dei sostanziali footprint di memoria e dell'overhead di inferenza. Sebbene il pruning strutturale sia una soluzione promettente, i metodi esistenti affrontano limitazioni significative:

• Granularità Grossolana: Gli approcci precedenti operano a livello di esperto, rimuovendo interi esperti o classificandoli in base all'importanza. Ciò non riesce a catturare la ridondanza interna negli esperti "ad alta importanza", dove l'informazione è spesso concentrata in pochissimi canali.
• Stima Inaccurata dell'Importanza: Le metriche esistenti per l'importanza dell'esperto (es. pesi del router, utilizzo dei token, statistiche grezze di attivazione/gradiente) sono spesso disallineate con il reale contributo di un esperto alla perdita finale. Le statistiche del router riflettono la frequenza di selezione piuttosto che l'utilità, e le statistiche grezze presentano bias dipendenti dallo strato che rendono inaffidabile il confronto tra strati diversi.
• Allocazione Inefficiente del Budget: Allocare i budget di pruning basandosi su classifiche grossolane degli esperti porta a uno spreco di capacità su strutture a basso contributo e non riesce a ottimizzare il compromesso tra compressione e accuratezza nei MoE su larga scala con centinaia di esperti.

2. Metodologia

Gli autori propongono un framework di pruning strutturale su misura per i MoE, composto da tre componenti principali: Attribution-Guided Loss Approximation (ALA), Coverage-Maximized Budget Allocation (CBA) e Alignment-Aware Redistribution (AAR).

2.1 Attribution-Guided Loss Approximation (ALA)

Per affrontare l'alto costo dell'ablatore per singolo esperto (che scala linearmente con il numero di esperti), gli autori derivano un proxy efficiente per l'importanza dell'esperto utilizzando un'espansione di Taylor del primo ordine.

• Derivazione: La variazione della perdita indotta dalla rimozione di un esperto $e$ è approssimata come $\Delta L^{(e)} \approx -(\frac{\partial L}{\partial z_{\ell,e}})^\top z_{\ell,e}$, dove $z_{\ell,e}$ è l'output dell'esperto e $\frac{\partial L}{\partial z_{\ell,e}}$ è il gradiente rispetto a tale output.
• Efficienza: Ciò consente la stima dell'importanza per tutti gli esperti in uno strato tramite un singolo passaggio backward, riducendo il tempo di calibrazione di 14–26$\times$ rispetto all'ablazione esaustiva.

2.2 Coverage-Maximized Budget Allocation (CBA)

Motivati dall'osservazione che l'informazione negli esperti MoE è altamente concentrata in una piccola frazione di canali, gli autori riformulano l'obiettivo di pruning. Invece di allocare un rapporto di pruning fisso per esperto, l'obiettivo è massimizzare la copertura cumulativa del punteggio dei canali sotto un budget globale.

• Meccanismo: I canali all'interno di ogni esperto sono ordinati per punteggio (derivato da ALA). L'algoritmo cerca un rapporto di copertura target $\rho$ per ogni esperto tale che il numero totale di canali mantenuti soddisfi il budget di pruning globale massimizzando la somma dei punteggi dei canali mantenuti.
• Beneficio: Questo approccio dà priorità al mantenimento dei canali ad alto contributo e pruna aggressivamente le code a basso punteggio, evitando lo spreco di capacità causato dai metodi uniformi o basati su rapporti.

2.3 Alignment-Aware Redistribution (AAR)

Per garantire la compatibilità con la quantizzazione a basso bit e i kernel efficienti per l'hardware (es. GEMM), le dimensioni dei canali devono essere multipli di una dimensione di blocco $a$ (es. 64 o 128).

• Processo: Dopo che CBA determina il numero ottimale di canali, AAR esegue un allineamento verso il basso al multiplo più vicino di $a$. I canali "residui" rilasciati da questo arrotondamento vengono ridistribuiti agli esperti con i più grandi residui frazionari (usando la regola del più grande resto di Hamilton).
• Obiettivo: Ciò assicura che il modello prunato soddisfi i vincoli del kernel senza deviazioni significative dall'allocazione ottimale basata sulla copertura, prevenendo il degrado delle prestazioni dovuto al padding o a accessi inefficienti alla memoria.

3. Contributi Chiave

1. Osservazione della Ridondanza Interna: Il documento dimostra che l'informazione MoE è altamente concentrata in una piccola frazione di canali, rendendo le metriche di importanza a livello di esperto troppo grossolane per un pruning strutturale efficace.
2. Nuovo Obiettivo di Pruning: Gli autori introducono la copertura del punteggio dei canali come obiettivo di pruning, riformulando l'allocazione della capacità per massimizzare la copertura sotto un budget globale. È il primo lavoro che ottimizza esplicitamente la copertura del punteggio invece di semplici rapporti di pruning.
3. Framework Efficiente: Propongono un framework scalabile che combina:
   • ALA: Un'approssimazione basata sull'attribuzione per la stima dell'importanza dell'esperto che riduce le ore di GPU di 20$\times$ rispetto all'ablazione.
   • CBA: Un algoritmo basato sulla ricerca binaria per allocare i budget che massimizzano la copertura del punteggio.
   • AAR: Una strategia di ridistribuzione per soddisfare i vincoli dei kernel a basso bit.
4. Validazione Empirica: Esperimenti estesi sulle architetture DeepSeek e Qwen MoE mostrano che il metodo raggiunge oltre 5$\times$ di compressione con una perdita di accuratezza trascurabile sui compiti di conoscenza generale e mantiene forti prestazioni nei benchmark di ragionamento.

4. Risultati Sperimentali

Il framework è stato valutato su DeepSeek-MoE-16B, DeepSeek-V2-Lite, Qwen1.5-MoE-A2.7B e Qwen3-30B-A3B.

• Conoscenza Generale: Sotto un rapporto di pruning del 25% combinato con quantizzazione a 4 bit ($P25\% Q4b$), il metodo ha ottenuto una riduzione dello storage superiore a 5$\times$. Su Qwen3-30B-A3B, l'impronta di memoria è stata ridotta di 5,27$\times$ (da 61,06 GB a 11,64 GB) con un calo di accuratezza media di al massimo l'1% sui benchmark generali (es. MMLU, ARC).
• Compiti di Ragionamento: Sui benchmark di ragionamento (GSM8K, MATH500), i modelli compressi si sono costantemente avvicinati o hanno superato i rispettivi originali. Ad esempio, Qwen3-30B-A3B con il 50% di pruning ha mantenuto un punteggio di 94,5 su MATH500 (rispetto al baseline di 92,8).
• Efficienza: Il metodo ha superato i baseline state-of-the-art (inclusi Wanda, MoNE, EAC-MoE e PuzzleMoE) attraverso diversi benchmark. Ha inoltre dimostrato prestazioni robuste sotto diverse strategie di routing (top-1 a top-6) e corpora di calibrazione.
• Throughput: La ridistribuzione consapevole dell'allineamento ha migliorato significativamente il throughput di inferenza. Per Qwen3-30B-A3B, le dimensioni allineate (128) hanno prodotto 14,21 token/s rispetto ai 10,23 token/s per le dimensioni non allineate sotto il 50% di sparsità.

5. Significato e Rivendicazioni

Il documento sostiene di fornire una via pratica verso l'implementazione efficiente dei MoE affrontando il disallineamento tra l'importanza a livello di esperto e la ridondanza interna dei canali.

• Scalabilità: Sostituendo l'ablazione esaustiva con l'approssimazione basata sull'attribuzione, il metodo scala a MoE con centinaia di esperti, un regime in cui i precedenti metodi basati sulla perdita erano computazionalmente proibitivi.
• Ottimizzazione Fine-Grained: Il passaggio dal pruning a livello di esperto alla massimizzazione della copertura a livello di canale permette un'allocazione più precisa della capacità, preservando l'informazione critica e rimuovendo il rumore.
• Compatibilità Hardware: L'integrazione della ridistribuzione consapevole dell'allineamento assicura che i modelli compressi siano non solo più piccoli, ma anche compatibili con la quantizzazione a basso bit e i kernel ad alto throughput, affrontando un collo di bottiglia critico nell'implementazione reale.

Gli autori pongono questo lavoro come un passo fondamentale nel muovere la compressione MoE dalle decisioni grossolane a livello di esperto all'ottimizzazione fine a livello di canale, consentendo l'implementazione di modelli MoE su larga scala con requisiti di memoria e computazione significativamente ridotti.