Riepilogo Tecnico: Upcycling Continuo di LLM tramite Sparsità per Bank a Predittore-Gated

Problema

Il costo computazionale dei Large Language Models (LLM), in particolare dei Transformer decoder-only, è sempre più dominato dalle dimensioni dei parametri e dalla computazione per token richiesta dalle reti Feed-Forward (FFN) dense. Nelle moderne architetture SwigLU, la FFN espande lo stato nascosto verso una dimensione intermedia elevata, lo elabora e lo proietta nuovamente. Durante l'addestramento e l'inferenza a lungo contesto, questa computazione densa viene applicata a tutti i token, anche se solo un sottoinsieme di canali intermedi è critico per uno specifico token.

Le soluzioni esistenti presentano dei limiti:

• Mixture-of-Experts (MoE): Riduce la computazione indirizzando i token verso pool di esperti specifici, ma introduce pool di parametri separati, requisiti complessi di bilanciamento del carico e overhead di comunicazione all-to-all.
• Sparsità Post-hoc: I metodi che derivano maschere sparse dalle attivazioni native delle FFN arrivano spesso troppo tardi per saltare le computazioni dominanti di gate e up-projection, fallendo nel fornire un'accelerazione a livello hardware.
• Sparsità Contestuale: Sebbene lavori precedenti dimostrino che la computazione degli LLM dipenda dall'input, i metodi senza addestramento possono rivelarsi fragili in compiti di generazione complessi.

Il paper affronta la sfida di convertire un LLM denso pre-addestrato in un modello a canali sparsi attraverso l'addestramento continuo (upcycling) senza una degradazione significativa delle prestazioni, con l'obiettivo di creare un percorso sparso orientato all'hardware dove il segnale di routing è disponibile prima della computazione principale della FFN.

Metodologia

L'approccio proposto, Predictor-Gated Bank-Wise Sparsity, trasforma un modello Qwen2.5-8B denso in una variante sparsa attraverso una specifica ricetta di addestramento continuo.

1. Architettura: Predictor-Gated Sparse SwigLU

Inveve di utilizzare gli output del gate nativo per determinare la sparsità, gli autori introducono un predittore a basso rango ($P_l$) che opera sullo stato nascosto del token ($x_{l,t}$) prima della computazione della FFN.

• Predittore: Un modulo leggero a due strati ($R^d \to R^{448} \to R^{d_i}$) che mappa lo stato nascosto ai logit sui canali intermedi della FFN.
• Meccanismo di Routing: Il predittore genera i logit che vengono utilizzati per generare un segnale di routing. Questo sostituisce lo standard gate denso $\sigma(u_{l,t})$ di SwigLU.
• Regola Bank-wise Top-k: La dimensione intermedia ($d_i = 14336$) è partizionata in bank fissi di dimensione 64. Per ogni bank, viene applicata una regola top-k rigida ($k=16$). Ciò mantiene 16 canali attivi per ogni bank da 64 canali, ottenendo una sparsità di 4x nella larghezza dell'attivazione intermedia della FFN ($d_s = 3584$).
• Esecuzione:
  • Forward Pass: Viene applicata una maschera binaria rigida. I canali inattivi sono esattamente zero.
  • Backward Pass: La selezione top-k non differenziabile è gestita tramite un gradiente surrogato soft (gradiente locale simile a una sigmoide) attorno alla soglia di selezione, permettendo alla perdita di linguaggio (language modeling loss) di aggiornare i parametri del predittore.

2. Ricetta di Addestramento (Dense-to-Sparse Continual Training)

Il processo converte un modello Qwen2.5-8B addestrato su un contesto di 8K in un modello sparso a 32K attraverso quattro fasi distinte:

1. Fase Dense 8K: Addestramento da zero utilizzando l'architettura standard Qwen2.5-8B.
2. Fase Dense 32K: Ripresa dal checkpoint 8K, estensione del contesto a 32K e continuazione dell'addestramento denso per circa 3T di token.
3. Warm-up del Predittore: Introduzione del predittore a basso rango come ramo ausiliario per imitare il segnale del gate nativo. Questo ramo non influenza la perdita principale.
4. Predictor-Gated Dense FFN: Il routing del gate della FFN principale passa attraverso il predittore (usando la sigmoide) in modo che il predittore influenzi la perdita, ma senza applicare ancora la selezione top-k rigida.
5. Predictor-Gated Top-k Sparse FFN: Attivazione della selezione bank-wise top-k per imporre la sparsità di 4x.

Nota: Il ramo sparso viene introdotto tardi nella fase 32K (dopo circa 3T di token) per riflettere uno scenario pratico di upcycling in cui un modello denso viene convertito in sparso.

3. Implicazioni Hardware

L'addestramento è eseguito su GPU come simulazione funzionale. Sebbene l'attuale implementazione su GPU materializzi le proiezioni dense di gate e up prima del masking (limitando l'immediato guadagno in termini di tempo di esecuzione reale), il principio algoritmico permette a hardware dedicato per la sparsità di saltare le computazioni in tutte e tre le fasi della FFN (proiezioni di gate, up e down) per i canali inattivi, riducendo teoricamente la larghezza intermedia effettiva di 4x.

Contributi Chiave

1. Pattern di Canale Predictor-Gated Sparsity: Una nuova architettura in cui un predittore a basso rango genera i logit di routing dei canali prima della computazione della FFN, combinata con una regola top-k rigida per bank. Ciò consente di saltare i canali inattivi di gate, up e down, a patto di utilizzare kernel specializzati o hardware dedicato.
2. Ricetta di Addestramento Continuo Dense-to-Sparse: Un workflow pratico per convertire un modello denso Qwen2.5-8B (addestrato a 8K, continuato a 32K) in un modello FFN sparso a 4x con predittore-gated con una degradazione delle prestazioni limitata.
3. Lezioni di Addestramento e Modalità di Fallimento:
   • Bilanciamento dei Canali: Il paper dimostra che i bias di bilanciamento degli esperti in stile MoE non si trasferiscono pulitamente al routing dei canali per bank. Forzare esplicitamente il bilanciamento dei canali (tramite aggiornamenti online) è stato trovato dannoso o neutro rispetto al non applicare alcun bilanciamento.
   • Fallimento del Lungo Contesto Locale per Layer: L'valutazione su RULER-CWE ha rivelato un brusco calo delle prestazioni (cliff) tra 12K e 16K nel modello sparso, che era invisibile nei benchmark a contesto breve. Gli autori hanno identificato questo come un fallimento locale del layer (specificamente nel 7° layer) e hanno proposto un algoritmo di riparazione a singolo layer (intervento di fallback denso) che recupera sostanzialmente le prestazioni nell'intervallo di lunghezza interessato.

Risultati

Il modello sparso è stato valutato contro due baseline: un Dense Companion (completamente denso, stesso budget) e una baseline Dense-4x (sparsità 4x applicata post-hoc in base all'ampiezza dell'attivazione).

• Prossimità delle Prestazioni: In un ampio set di benchmark (Factual QA, Code, Reasoning, Math, Composite), il modello sparso con predittore-gated rimane costantemente più vicino al dense companion rispetto alla baseline naive Dense-4x.
• Guadagni Specifici:
  • Il modello sparso ha eguagliato o leggermente superato la baseline densa su MBPP, MultiPL-E e MMLU.
  • Ha ridotto sostanzialmente la degradazione osservata nella baseline Dense-4x su Chinese SimpleQA, MBPP+, MBXP e KOR-Bench.
• Prestazioni a Lungo Contesto:
  • Sui benchmark RULER (4K, 8K, 16K), il modello sparso ha performato in modo comparabile al modello denso.
  • A 32K, il modello sparso (80.4) ha superato la baseline Dense-4x (73.5) e ha leggermente superato il dense companion (79.2).
  • Analisi RULER-CWE: Il modello sparso ha mostrato un "cliff" (precipizio) nel recall tra 12K e 16K. L'applicazione del proposto intervento di riparazione a singolo layer (Layer 7) ha migliorato il recall medio nell'intervallo 8.7K–12.2K da 0.856 a 0.947 e nell'intervallo 12.8K–16.3K da 0.583 a 0.802.

Significato e Rivendicazioni

Il paper sostiene che l'upcycling sparso non è privo di perdite (lossless), ma è una via percorribile per costruire modelli sparsi orientati all'hardware partendo da checkpoint densi.

• Approccio Centrato sull'Addestramento: A differenza della compressione post-hoc, questo metodo tratta la sparsità dei canali come parte integrante del percorso di linguaggio, ottimizzata durante l'addestramento continuo.
• Allineamento Hardware: Posizionando il predittore di routing prima della computazione della FFN e utilizzando una maschera rigida, il metodo si allinea ai requisiti per l'hardware sparso dedicato per ottenere una vera accelerazione, piuttosto che una semplice sparsità teorica.
• Valore Diagnostico: Il lavoro evidenzia come la valutazione a lungo contesto sia critica per identificare modalità di fallimento (come il cliff locale del layer) che i benchmark a contesto breve mancano. Fornisce inoltre un risultato negativo riguardo alla trasferibilità delle strategie di bilanciamento MoE al routing a livello di canale nei contesti di addestramento continuo.

Gli autori concludono che, sebbene il modello sparso non sia un sostituto perfetto del modello denso in tutti gli scenari, offre una ricetta pratica per l'upcycling di modelli densi in varianti sparse efficienti con scambi di prestazioni gestibili.