Dense2MoE: Pushing the Pareto Frontier of On-Device LLMs via Unified Pruning and Upcycling

Riepilogo Tecnico: Dense2MoE

Enunciato del Problema
Il deployment di Large Language Models (LLM) su dispositivi edge con risorse limitate (ad esempio, guida autonoma, edge AI) è gravemente ostacolato dal "muro della memoria per l'inferenza". Secondo la teoria hardware Roofline, gli LLM densi sono limitati dalla memoria durante il decoding autoregressivo, il che significa che ridurre semplicemente i FLOPs teorici è insufficiente per migliorare la latenza. Le soluzioni esistenti affrontano un compromesso fondamentale:

1. Potatura Strutturata: Metodi come UIDL riducono la latenza eliminando strati ridondanti, ma sacrificano inevitabilmente l'accuratezza del modello rimuovendo i Perceptron a Strati Multipli (MLP), che fungono da principali repository di conoscenza degli LLM.
2. Upcycling MoE: La conversione di modelli densi in architetture a Miscela di Esperti (MoE) aumenta la capacità senza aumentare linearmente il calcolo dinamico. Tuttavia, l'upcycling standard soffre spesso di "collasso della rappresentazione" a causa di un'inizializzazione omogenea degli esperti (duplicazione degli MLP nativi), richiedendo un pre-addestramento continuo massiccio per indurre la specializzazione. Inoltre, i recenti metodi di riutilizzo sparso si concentrano spesso sulla ridondanza intra-strato, fallendo nel addressing delle sovrapposizioni strutturali inter-strato a livello macro che causano ridondanza nell'attenzione e sovraccarichi nel caricamento della KV-cache.

Metodologia: Dense2MoE
Gli autori propongono Dense2MoE, un framework unificato che integra la potatura strutturata degli strati con l'upcycling MoE tramite una nuova tecnica chiamata Layer-Fusion Upcycling (LF-UC). Il metodo opera in tre fasi:

1. Analisi della Similarità Inter-Strato: Il framework identifica la ridondanza strutturale analizzando gli strati adiacenti per alta similarità coseno sia negli output del decoder che negli input degli MLP. Crucialmente, impone un vincolo rigoroso di corrispondenza della norma delle feature per garantire che il riutilizzo degli MLP ridondanti non introduca anomalie di attivazione o shift di scala.
2. Layer-Fusion Upcycling (LF-UC): Invece di scartare completamente gli strati ridondanti, Dense2MoE pota fisicamente solo i moduli di attenzione (MHA) di questi strati per eliminare le operazioni ad alta intensità di larghezza di banda della memoria. Gli MLP pre-addestrati di questi strati potati vengono preservati e riutilizzati come esperti eterogenei all'interno degli strati base mantenuti.
   • Questo crea un pool di esperti dove gli "esperti base" sono MLP nativi duplicati e gli "esperti supplementari" sono gli MLP preservati dagli strati potati.
   • Un router dinamico a livello di token instrada selettivamente i token verso questi esperti o li bypassa tramite percorsi residui.
   • Questo design sfrutta l'eterogeneità strutturale inter-strato per mitigare il collasso del routing "Winner-Takes-All" comune nell'upcycling omogeneo.
3. Pre-Addestramento Continuo: L'architettura modificata subisce una fase di pre-addestramento continuo leggera (225 miliardi di token) per ottimizzare il meccanismo di routing e allineare gli esperti fusi. Questo budget consuma meno dell'1% del calcolo richiesto per addestrare un modello MoE da zero.

Per determinare la profondità di potatura ottimale, gli autori impiegano una metrica di ricompensa consapevole dell'hardware, ispirata a MnasNet, che mappa i vincoli di potatura sulla frontiera di Pareto tra latenza e accuratezza per hardware target specifici (ad esempio, NVIDIA Jetson Thor-U).

Contributi Chiave

• Framework Unificato: Dense2MoE integra senza soluzione di continuità la potatura strutturata degli strati e l'upcycling MoE, risolvendo il compromesso tra efficienza di inferenza e accuratezza del modello.
• Nuova Architettura (LF-UC): Potando i colli di bottiglia dell'attenzione mentre ricicla gli MLP come esperti eterogenei, il metodo riduce il consumo di larghezza di banda della memoria senza distruggere la capacità rappresentativa.
• Ottimalità di Pareto all'Avanguardia: Il framework avanza la frontiera di Pareto per la latenza di inferenza su dispositivo rispetto all'accuratezza del modello, superando baseline dense, metodi di compressione all'avanguardia e approcci standard di upcycling.
• Scalabilità Robusta: Il metodo è validato su diverse scale di modello (da 0.5B a 7B) e famiglie architettoniche distinte (Qwen2.5 e LLaMA2).

Risultati Sperimentali
Valutato su 9 benchmark diversificati (inclusi MMLU, ARC, GSM8K e HumanEval) utilizzando modelli Qwen2.5 e LLaMA2:

• Prestazioni vs Latenza: La configurazione ottimale di Dense2MoE (mantenendo 19 strati da un seed Qwen2.5-0.5B a 24 strati) ha raggiunto una latenza di inferenza fisica di 271.0 ms sul SoC Thor-U (in calo rispetto ai 307.5 ms del modello seed) ottenendo al contempo un punteggio medio sui benchmark di 47.31, superando significativamente la baseline densa (37.53).
• Confronto con le Baseline: A differenza dei metodi di potatura tradizionali (ad esempio, UIDL, LLM-Streamline) e degli approcci di riutilizzo sparso (ad esempio, ToMoE, LLaMA-MoE), che hanno sofferto di un grave collasso della capacità nei compiti di ragionamento e coding, Dense2MoE ha mantenuto una robusta ritenzione rappresentativa. Ha superato la baseline densa di oltre 9.78 punti percentuali nel punteggio medio attivando solo 0.42B di parametri.
• Scalabilità: Quando applicato a modelli più grandi (Qwen2.5-1.5B e LLaMA2-7B), Dense2MoE ha costantemente superato le controparti dense rispettivamente di 3.04 e 3.54 punti percentuali, dopo il fine-tuning sullo stesso dataset da 225 miliardi di token.

Significato
Il paper afferma che Dense2MoE stabilisce una ricetta pratica e di standard industriale per il riutilizzo di LLM densi open-source in architetture su dispositivo altamente efficienti. Sfruttando sistematicamente la ridondanza strutturale inter-strato, consente l'ottimizzazione congiunta della latenza di inferenza e dell'accuratezza del modello. Il lavoro dimostra che è possibile rompere il muro della memoria per l'inferenza e preservare capacità di ragionamento complesse senza i costi proibitivi dell'addestramento di modelli MoE da zero o il degrado dell'accuratezza associato alla potatura aggressiva. Questo fornisce un continuum flessibile di modelli ottimali per deployment edge critici per la latenza.