DECO: Sparse Mixture-of-Experts with Dense-Comparable Performance on End-Side Devices

Riepilogo Tecnico: DECO – Mixture-of-Experts Sparsi con Prestazioni Comparabili a Modelli Densi su Dispositivi Edge

Enunciazione del Problema
Sebbene le architetture Mixture-of-Experts (MoE) scalino la capacità del modello senza aumentare proporzionalmente il calcolo, tradizionalmente soffrono di ingenti footprint totali dei parametri. Ciò crea significativi colli di bottiglia nello storage e nell'accesso alla memoria, ostacolando il deployment efficiente su dispositivi edge, che richiedono elevate prestazioni, bassi costi computazionali e un minimo overhead di storage. La ricerca MoE esistente spesso dà priorità alle prestazioni e al basso calcolo, ma trascura l'obiettivo dell'"overhead di storage". Inoltre, molti modelli MoE ottengono i loro risultati addestrandosi su un numero sostanzialmente maggiore di token rispetto alle baseline dense, sotto lo stesso budget computazionale, anziché corrispondere all'esatto conteggio dei token. Il lavoro pone la domanda: Un modello MoE sparso può ottenere prestazioni comparabili a un modello denso dato lo stesso budget totale di parametri e lo stesso numero di token di addestramento?

Metodologia
Gli autori propongono DECO (DEnse-COmparable), un'architettura MoE sparso progettata per corrispondere alle prestazioni dei Transformer densi sotto vincoli identici. L'architettura si fonda su tre pilastri progettuali fondamentali:

1. Progettazione del Router:

   • Routing basato su ReLU: A differenza del routing TopK convenzionale non differenziabile, DECO impiega un routing basato su ReLU. Ciò consente rapporti di attivazione flessibili dipendenti dal token e una regolarità nell'apprendimento.
   • Scalabilità del Router per Esperto Apprendibile: Per affrontare gli squilibri nella scala dell'output tra esperti condivisi e instradati e per accomodare l'eterogeneità degli esperti, DECO sostituisce i fattori di scalatura scalari fissi con fattori di scalatura vettoriali apprendibili ($\alpha$). Ciò calibra il contributo dei singoli esperti instradati.

2. Progettazione degli Esperti:

   • Attivazione NormSiLU: Gli autori identificano che accoppiare il routing basato su ReLU con esperti attivati da SiLU standard porta a rapporti di attivazione in rapida crescita e a magnitudini di output in diminuzione. Per risolvere ciò, introducono NormSiLU, che applica una normalizzazione a due stadi prima dell'operatore SiLU:
     • Normalizzazione della media inter-esperto: Centra i pesi di proiezione in alto degli esperti intorno allo zero.
     • Normalizzazione RMS intra-esperto: Mantiene magnitudini di attivazione coerenti.
       Ciò stabilizza le tendenze di attivazione, riduce la necessità di una regolarizzazione aggressiva della sparsità e garantisce un utilizzo robusto dei parametri.
   • Esperti MLP Non-Gated: Contrariamente alla preferenza standard per gli MLP gated (SwiGLU), DECO utilizza esperti MLP non-gated. I risultati empirici mostrano che gli esperti non-gated esibiscono tendenze di rapporto di attivazione più stabili quando accoppiati con routing basato su ReLU, mentre le varianti gated soffrono di dinamiche di attivazione instabili e in rapida crescita che richiedono una pesante regolarizzazione.

3. Regolarizzazione Adattiva della Sparsità:

   • DECO impiega una regolarizzazione adattiva della sparsità basata sull'entropia del router. A differenza dei coefficienti statici, la forza della regolarizzazione ($\lambda$) viene scalata dinamicamente in base al livello di sparsità corrente rispetto all'obiettivo, garantendo un rapporto di attivazione stabile centrato precisamente al livello desiderato.

Risultati Chiave
Gli esperimenti sono stati condotti su quattro scale di parametri (da 0,11B a 1,18B) con un conteggio fisso di token di addestramento (circa 40 volte il conteggio dei parametri) e un rapporto di attivazione degli esperti instradati del 20%.

• Comparabilità con Modelli Densi: DECO raggiunge la parità di prestazioni con le baseline Transformer dense (che utilizzano FFN SwiGLU) in termini di Perplexity (PPL) sul set di validazione C4 e accuratezza media sui benchmark di ragionamento di senso comune, nonostante attivi solo il 20% degli esperti.
• Superiorità rispetto alle Baseline MoE: DECO supera le consolidate baseline MoE (TopP, DeepSeek-V3, ReMoE, BlockFFN) della stessa scala e rapporto di attivazione.
• Approfondimenti dalle Ablazioni:
  • La scalabilità apprendibile per esperto migliora significativamente le prestazioni rispetto alla scalatura fissa o scalare, affrontando l'eterogeneità osservata nelle norme di output degli esperti.
  • NormSiLU è critico; la rimozione dei suoi passaggi di normalizzazione porta a rapporti di attivazione in rapida crescita e a output SiLU in diminuzione, degradando le prestazioni.
  • Gli esperti non-gated sono superiori agli esperti gated specificamente nel contesto del routing basato su ReLU.
• Accelerazione Hardware: È stato implementato un kernel di accelerazione specializzato che sfrutta CUTLASS e tensor core. Su hardware reale (NVIDIA RTX 4090 e Jetson AGX), DECO ha ottenuto un speedup di 3,00× rispetto all'inferenza densa standard sfruttando l'attivazione sparsa per ridurre l'overhead di accesso alla memoria.

Significato e Affermazioni
Il lavoro afferma che DECO realizza con successo il "triangolo ideale" del deployment MoE su dispositivi edge: elevate prestazioni, basso costo computazionale e piccolo overhead di storage. La sua importanza risiede nel dimostrare che i modelli MoE sparsi possono corrispondere alle prestazioni dense senza richiedere un eccesso di token di addestramento o footprint di parametri massicci.

Gli autori attribuiscono questo successo a una revisione fondamentale della progettazione MoE che allinea il meccanismo di routing (basato su ReLU), l'attivazione degli esperti (NormSiLU) e la struttura degli esperti (non-gated) per creare una dinamica di addestramento stabile. Notano che la comparabilità con i modelli densi è condizionata, dipendendo da fattori quali il rapporto di attivazione, la granularità degli esperti e l'eterogeneità dei dati. Il lavoro suggerisce che dataset ad alta entropia e multi-dominio sono particolarmente adatti a questa architettura sparsa, poiché il modello può attivare efficacemente sottoinsiemi di parametri specializzati per diversi compiti.

Limitazioni
Gli autori riconoscono di non aver ancora valutato DECO nelle fasi di fine-tuning supervisionato (SFT) o apprendimento per rinforzo (RL), dove le architetture MoE potrebbero affrontare sfide uniche come l'instabilità del router. Inoltre, le leggi di scalatura per la soglia del rapporto di attivazione richiesta per la comparabilità con i modelli densi su dimensioni di modello più grandi e distribuzioni di dati diverse richiedono ulteriori indagini.