Teacher-Guided Routing for Sparse Vision Mixture-of-Experts

1. Il Problema: Instabilità nell'Addestramento di MoE Sparsi

Le architetture Mixture-of-Experts (MoE) sparse sono diventate fondamentali per scalare i modelli di deep learning senza aumentare proporzionalmente il costo computazionale, attivando solo un sottoinsieme di "esperti" per ogni input. Tuttavia, l'addestramento di questi modelli, specialmente nel dominio della visione artificiale (Vision MoE o VMoE), presenta sfide critiche:

• Blocco del Gradiente (Gradient Blocking): Nel routing standard, il router riceve segnali di gradiente informativi solo dagli esperti selezionati durante il passaggio in avanti. Le route non selezionate non ricevono feedback, lasciando il router con informazioni limitate e altamente localizzate.
• Fluttuazione del Routing (Routing Fluctuation): A causa della mancanza di supervisione globale, il router fatica a imparare punteggi di selezione stabili, specialmente nelle fasi iniziali. Questo porta a dinamiche instabili dove l'assegnazione degli esperti per uno stesso input cambia frequentemente durante l'addestramento.
• Conseguenze: Tale instabilità impedisce una specializzazione stabile degli esperti, porta a un sottoutilizzo di alcuni esperti e complica l'ottimizzazione, poiché lo stesso campione può aggiornare esperti diversi in iterazioni successive. Le perdite di bilanciamento del carico (load-balancing losses) esistenti non sono sufficienti a risolvere questo problema.

2. Metodologia: TGR-MoE

Gli autori propongono TGR-MoE (Teacher-Guided Routing), un framework che stabilizza l'apprendimento del router MoE utilizzando la supervisione derivata da un modello denso pre-addestrato (Teacher).

Architettura e Funzionamento

1. Costruzione del Router Teacher:

   • Viene utilizzato un modello denso pre-addestrato (es. DeiT-III) come "Teacher". Il backbone del Teacher è congelato.
   • Viene aggiunto un router Teacher leggero (auxiliary router) che prende le rappresentazioni intermedie del Teacher e produce distribuzioni di assegnazione degli esperti.
   • Il router Teacher viene ottimizzato (solo durante l'addestramento) con una perdita di bilanciamento del carico ($L_{load}$) e una perdita di entropia ($L_{ent}$) per garantire distribuzioni stabili, bilanciate e confidenti.

2. Addestramento Guidato dello Studente:

   • Il modello MoE studente apprende imitando la distribuzione di routing del Teacher.
   • Viene introdotta una perdita di distillazione ($L_{distill}$) basata sulla divergenza di Kullback-Leibler (KL) tra la distribuzione del router studente ($p$) e quella del router teacher ($p_t$).
   • L'obiettivo totale dello studente è: $L_{student} = L_{task} + \lambda_{distill} \cdot L_{distill}$.

3. Vantaggi Chiave:

   • Prior Globale: Il router Teacher fornisce un "prior" informativo globale che supera il problema del blocco del gradiente, guidando il router anche verso route non selezionate.
   • Nessun Costo Inferenza: Il modello Teacher è utilizzato solo durante l'addestramento. Durante l'inferenza, viene utilizzato solo il modello studente, senza costi aggiuntivi.
   • Stabilità Temporale: La distillazione incoraggia assegnazioni di esperti semanticamente significative e coerenti fin dalle prime fasi.

3. Contributi Chiave

• Identificazione del Blocco del Gradiente: Il paper evidenzia come la mancanza di feedback informativo dalle route non selezionate sia la causa principale dell'instabilità nel routing MoE, un problema spesso sottovalutato rispetto al semplice sbilanciamento del carico.
• Framework TGR-MoE: Introduzione di un metodo semplice ma efficace che utilizza un router Teacher derivato da un modello denso per fornire supervisione pseudo-etichettata al router studente.
• Analisi della Consistenza: Dimostrazione che il metodo riduce drasticamente le fluttuazioni di routing e permette al modello di raggiungere una specializzazione stabile degli esperti molto prima rispetto ai metodi standard.
• Scalabilità: Il metodo mantiene i guadagni di prestazioni anche con un numero elevato di esperti, dove i metodi VMoE standard tendono a saturare.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-1K (pre-addestramento) e su dataset downstream come CIFAR-10/100 e Oxford-IIIT Pets.

• Prestazioni di Accuratezza:

  • TGR-MoE supera costantemente i modelli ViT densi, VMoE standard e varianti avanzate (come Expert Choice MoE, SoftMoE, e VMoE con z-loss).
  • Su ImageNet-1K, per il modello "Tiny" con K=1, l'accuratezza Top-1 sale dal 77.85% (VMoE) al 78.78% (TGR-MoE). Per il modello "Small", sale dal 82.63% all'83.34%.
  • I miglioramenti sono coerenti su tutte le scale (Tiny, Small, Base).

• Stabilità e Consistenza del Routing:

  • Analisi di Accordo: TGR-MoE raggiunge un accordo (>70%) con l'assegnazione finale degli esperti entro le prime 50 epoche, mentre VMoE mostra fluttuazioni significative (accordo ~40%) anche a metà addestramento.
  • Fine-tuning: Mantenere TGR-MoE durante il fine-tuning è cruciale. Se si usa TGR-MoE solo per l'inizializzazione e poi si passa a VMoE standard, le prestazioni crollano vicino alla baseline. L'uso continuo di TGR-MoE durante il fine-tuning su CIFAR-100 ha portato all'accuratezza più alta (91.07% per il modello Base).

• Scalabilità con il Numero di Esperti:

  • Aumentando il numero di esperti (da 4 a 256), TGR-MoE mostra miglioramenti continui, mentre VMoE standard smette di migliorare significativamente oltre un certo punto. Con 128 esperti, TGR-MoE mantiene un guadagno significativo (+0.98% rispetto a VMoE).

• Analisi dell'Upper Bound:

  • Un esperimento "upper bound" (usare direttamente il router Teacher per l'inferenza) ha mostrato un'accuratezza potenziale del 80.19%, indicando che c'è ancora spazio per migliorare l'allocazione degli esperti, ma TGR-MoE offre un compromesso pratico eccellente (77.81%) senza richiedere il modello Teacher in produzione.

5. Significato e Implicazioni

Il lavoro dimostra che l'ottimizzazione del routing negli MoE sparsi non può basarsi esclusivamente sui gradienti del task (che sono sparsi e discontinui). L'introduzione di segnali di supervisione esterni e stabili, derivati da modelli densi pre-addestrati, è essenziale per:

1. Risolvere il problema del blocco del gradiente.
2. Stabilizzare la dinamica di addestramento fin dall'inizio.
3. Permettere una scalabilità efficace all'aumentare del numero di esperti.

TGR-MoE offre una soluzione pratica che non richiede modifiche all'architettura di inferenza, rendendolo immediatamente applicabile per migliorare modelli di visione basati su MoE, con potenziali estensioni future per compiti di linguaggio naturale e multimodali.