DeepFusion: Accelerating MoE Training via Federated Knowledge Distillation from Heterogeneous Edge Devices

1. Il Problema

I recenti modelli di linguaggio di grandi dimensioni (LLM) basati su architetture Mixture-of-Experts (MoE) (es. Qwen-MoE, DeepSeek-MoE) stanno rivoluzionando l'IA generativa. Tuttavia, questi modelli richiedono quantità enormi e diversificate di dati di addestramento per attivare efficacemente tutti i loro "esperti" (sottoreti specializzate).
Le sfide principali identificate sono:

• Carenza di dati pubblici: Si prevede che i dati testuali pubblici di alta qualità si esauriranno tra il 2026 e il 2032.
• Limitazioni dell'Apprendimento Federato (FL) Tradizionale: L'FL standard (es. FedAvg) richiede che ogni dispositivo edge ospiti una copia locale del modello globale. Per i modelli MoE, che possono avere centinaia di miliardi di parametri, questo è impossibile per dispositivi con risorse limitate (smartphone, dispositivi embedded).
• Approcci Esistenti Impraticabili: I tentativi recenti di FL per MoE richiedono che i dispositivi eseguano modelli "esperto locali" ridotti (pruned). Anche se più piccoli, questi modelli mantengono l'architettura MoE complessa, risultando ancora troppo pesanti per dispositivi a risorse limitate (es. NVIDIA Jetson Nano).
• Eterogeneità e "View-Mismatch": Esiste un disallineamento fondamentale tra i modelli locali (teacher) e il modello globale MoE (student). Hanno architetture diverse, spazi latenti differenti e prospettive predittive non allineate, rendendo inefficaci le tecniche tradizionali di distillazione della conoscenza basate sull'allineamento diretto dei logit o delle feature.

2. Metodologia: DEEPFUSION

DEEPFUSION è un framework scalabile di addestramento federato MoE basato sulla distillazione della conoscenza federata. Permette a dispositivi eterogenei e vincolati di risorse di partecipare al training senza ospitare il modello MoE globale.

Il processo si articola in tre fasi principali:

A. Progettazione FL "One-Shot" (A Scatto Singolo)

A differenza dell'FL tradizionale che richiede molteplici round di comunicazione, DEEPFUSION adotta un approccio one-shot:

1. Ogni dispositivo edge configura e addestra autonomamente un LLM locale (su misura per le sue risorse e applicazioni) utilizzando i propri dati privati.
2. Una volta completato l'addestramento locale, il modello LLM ottimizzato viene caricato sul server centrale in un singolo round di comunicazione.
3. Il server esegue la distillazione della conoscenza per integrare le competenze locali nel modello MoE globale.

B. Clustering della Conoscenza Locale

Per gestire la scalabilità con migliaia di dispositivi:

• I modelli LLM caricati vengono raggruppati in K domini di conoscenza locali (basati sulla similarità semantica dei loro dati di addestramento, stimata tramite embedding a basso rango).
• Per ogni dominio, viene creato un modello proxy locale ($\bar{m}_i$) ottenendo la media dei pesi dei modelli clusterizzati. Questo riduce il carico computazionale sul server, distillando solo K modelli proxy invece di N modelli individuali.

C. Distillazione della Conoscenza Cross-Architettura (Il Cuore Innovativo)

Il problema principale è l'allineamento tra il modello proxy (teacher, architettura compatta) e il modello base MoE (student, architettura standard). Per risolvere il problema del "view-mismatch" (disallineamento delle prospettive predittive), il paper introduce il modulo View-Aligned Attention (VAA):

• Meccanismo VAA: Utilizza un meccanismo di attenzione multi-testa per integrare le rappresentazioni di feature a più stadi dal modello student (MoE base).
• Allineamento: Il modulo VAA permette al modello student di acquisire una prospettiva predittiva simile a quella del teacher, allineando le feature intermedie prima della distillazione.
• Funzione di Perdita: La distillazione combina tre termini: perdita di cross-entropy (sugli etichette dure), perdita di allineamento delle feature (MSE tra feature teacher e student allineate) e perdita di divergenza KL (tra i logit soft).

D. Tuning del Modello MoE Globale

I K modelli base MoE risultanti dalla distillazione vengono fusi nel modello MoE globale:

• Gli strati FFN (Feed-Forward Networks) di ciascun modello base diventano gli esperti specializzati del MoE globale.
• Gli strati comuni (Embedding, Self-Attention, Output) vengono inizializzati tramite la media dei pesi dei modelli base.
• Viene eseguita una fase di fine-tuning globale: gli esperti FFN vengono congelati (poiché già addestrati), mentre gli strati di routing (gating), self-attention e output vengono ottimizzati per coordinare gli esperti.

3. Contributi Chiave

1. Primo Framework Scalabile: È il primo sistema che permette a dispositivi edge eterogenei e vincolati di risorse di partecipare al training federato di modelli MoE di livello industriale tramite distillazione della conoscenza.
2. Modulo VAA (View-Aligned Attention): Risolve il problema critico del disallineamento delle prospettive predittive tra architetture eterogenee (teacher piccoli vs student MoE grandi), permettendo una distillazione efficace basata sulle feature intermedie.
3. Efficienza e Privacy: Elimina la necessità di caricare modelli MoE completi sui dispositivi edge e riduce drasticamente i costi di comunicazione grazie all'approccio one-shot.

4. Risultati Sperimentali

Il framework è stato valutato su modelli MoE di livello industriale (Qwen-MoE e DeepSeek-MoE) utilizzando dataset reali in ambito medico (QA a scelta multipla) e finanziario (QA a risposta aperta).

• Prestazioni: DEEPFUSION raggiunge prestazioni vicine all'addestramento centralizzato (DeepSpeed) e supera i baseline federati (FedJETS, FedKMT, OFA-KD).
  • Miglioramento della perplessità dei token (token perplexity) fino al 5,28% rispetto ai baseline.
  • Eccellente performance nel task di QA a risposta aperta, dove la capacità di ragionamento è cruciale.
• Riduzione dei Costi di Comunicazione: Rispetto ai metodi federati MoE esistenti, DEEPFUSION riduce i costi di comunicazione fino al 71%.
• Footprint di Memoria: I dispositivi edge richiedono 3,3x - 9,3x in meno di memoria rispetto ai metodi che richiedono modelli MoE locali (es. FedJETS), permettendo l'uso di hardware molto limitato (es. Jetson Nano, iPhone).

5. Significato e Impatto

DEEPFUSION rappresenta un passo avanti significativo per l'addestramento di LLM su larga scala in scenari reali:

• Superamento della Carenza di Dati: Sfrutta la vastità dei dati privati distribuiti sui dispositivi edge senza violare la privacy.
• Democratizzazione dell'IA: Permette a dispositivi con risorse limitate di contribuire allo sviluppo di modelli di intelligenza artificiale complessi, superando le barriere hardware.
• Robustezza Architetturale: La soluzione al problema del "view-mismatch" apre la strada a nuove tecniche di distillazione della conoscenza tra modelli con architetture radicalmente diverse, un problema fino ad oggi poco esplorato nel contesto federato.

In sintesi, DEEPFUSION dimostra che è possibile costruire modelli MoE globali ricchi di conoscenza aggregando l'intelligenza distribuita di milioni di dispositivi eterogenei, in modo efficiente, privato e scalabile.