Riepilogo Tecnico: Routing Sinkhorn Selettivo per un Miglioramento dei Mixture of Experts Sparsi

Problematica

I modelli Sparse Mixture-of-Experts (SMoE) offrono un approccio scalabile per aumentare la capacità del modello mantenendo l'efficienza computazionale, attivando solo un sottoinsieme di esperti per ogni token. Tuttavia, i modelli SMoE convenzionali che si affidano al routing basato su Softmax soffrono spesso di "collasso del routing" (routing collapse), in cui un piccolo sottoinsieme di esperti viene sovrautilizzato mentre altri rimangono sottoutilizzati.

Le strategie di mitigazione esistenti si affidano tipicamente a:

1. Obiettivi Ausiliari: Perdite di bilanciamento del carico (load-balancing losses) o z-loss per penalizzare l'uso non uniforme degli esperti. Queste possono introdurre disallineamento dell'obiettivo, instabilità dell'addestramento e maggiore complessità.
2. Rumore di Gating (Noisy Gating): Iniezione di rumore per incoraggiare l'esplorazione, il che aggiunge parametri addestrabili o iperparametri.
3. Routing basato su Sinkhorn: I metodi precedenti che utilizzano l'algoritmo di Sinkhorn per il trasporto ottimale (OT) spesso usano la mappa di trasporto esclusivamente per la selezione degli esperti (top-k) piuttosto che per l'assegnazione dei pesi. Inoltre, possono ridurre la flessibilità del routing perché la matrice di gating non viene ottimizzata direttamente tramite l'apprendimento basato sul gradiente nello stesso modo dei metodi Softmax, oppure comportano un overhead di addestramento sostanziale a causa della natura iterativa degli aggiornamenti Sinkhorn.

Metodologia

Il documento propone il Selective Sinkhorn Routing (SSR), un meccanismo leggero che riformula l'assegnazione token-esperto come un problema di trasporto ottimale regolarizzato dall'entropia.

Formulazione Centrale

Gli autori modellano l'assegnazione di $m$ token a $n$ esperti come un problema di trasporto ottimale a costo massimo con regolarizzazione entropica. L'obiettivo è trovare un piano di trasporto $\hat{\Pi}$ che massimizzi la compatibilità soddisfacendo specifici vincoli:

• Positività: $\Pi > 0$.
• Somma per Riga: Ogni token instrada la sua intera massa agli esperti ($\Pi \mathbf{1}_n = \mathbf{1}_m$).
• Somma per Colonna: Ogni esperto riceve un carico totale atteso uguale ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$).

La matrice dei costi $C$ è derivata dai punteggi di gating $S$. Gli autori propongono due varianti:

• Costo Lineare: $C = S$.
• Costo Softmax: $C_{i,:} = \text{softmax}(S_{i,:})$, utilizzato per prevenire l'overflow numerico durante le iterazioni di Sinkhorn.

A differenza dei precedenti metodi basati su OT che utilizzano il piano di trasporto solo per la selezione, l'SSR utilizza direttamente i valori in $\hat{\Pi}$ per calcolare i pesi di routing. Per ogni token, i primi $k$ esperti vengono selezionati in base alle voci più alte nella corrispondente riga di $\hat{\Pi}$, e i pesi sono normalizzati da queste voci.

Strategia Selettiva e Iniezione di Rumore

Una sfida chiave identificata è che il routing Sinkhorn non aggiorna la matrice dei pesi di gating $W_g$ poiché $W_g$ è scollata dall'obiettivo OT. Per affrontare questo problema, l'SSR impiega una strategia di addestramento ibrida:

• Applicazione Selettiva: Durante l'addestramento, ogni blocco MoE utilizza il routing Sinkhorn con una piccola probabilità $p$ (ad esempio, 0.001) e il normale gating Softmax altrimenti. Ciò consente al gating Softmax di aggiornare efficacemente $W_g$, mentre il routing Sinkhorn fornisce segnali di utilizzo bilanciato degli esperti.
• Iniezione di Rumore: Per prevenire il collasso degli esperti e incoraggiare l'esplorazione, viene aggiunto del rumore gaussiano alla matrice dei costi ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$). L'analisi teorica (Proposizione 4.3) dimostra che questo assicura che ogni esperto abbia una probabilità non nulla di selezione.

Comportamento in Inferenza

Il documento sostiene che imporre vincoli di bilanciamento (come le somme per colonna uguali) durante l'inferenza su un singolo input sia controproducente, poiché forza un routing uniforme indipendentemente dai punteggi di compatibilità specifici dell'input. Pertate, sia il routing Sinkhorn che l'iniezione di rumore sono disabilitati durante l'inferenza, tornando a un routing Softmax deterministico per garantire predizioni coerenti e dipendenti dall'input.

Contributi Chiave

1. Nuovo Framework di Routing: Un metodo che integra il trasporto ottimale regolarizzato dall'entropia con l'iniezione di rumore stocastico per promuovere un utilizzo bilanciato degli esperti senza perdite di bilanciamento ausiliarie.
2. Approfondimenti Teorici: Dimostrazioni che il routing basato su Sinkhorn e l'iniezione di rumore aiutano l'addestramento incoraggiando l'esplorazione e il bilanciamento, ma dovrebbero essere disabilitati durante l'inferenza per evitare assegnazioni distorte.
3. Efficienza e Prestazioni: Valutazioni estese che mostrano come l'SSR migliori l'efficienza dell'addestramento, l'accuratezza e la robustezza rispetto alla corruzione dell'input rispetto ai principali modelli di riferimento.

Risultati Sperimentali

Gli autori hanno valutato l'SSR su compiti di modellazione del linguaggio (WikiText-103, Enwik-8) e classificazione di immagini (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R).

• Modellazione del Linguaggio:
  • Su WikiText-103, le varianti SSR superano costantemente Vanilla SMoE e i modelli di riferimento che utilizzano la perdita di bilanciamento del carico, la z-loss o il rumore.
  • SSR-L con rumore ha raggiunto una perplessità (PPL) di test di 34.367, una riduzione di 1.183 rispetto a Vanilla SMoE (35.550), superando il secondo miglior modello di riferimento con un margine significativo.
  • Fondamentalmente, l'SSR ha comportato un overhead di addestramento minimo (0.33% – 0.65%), mentre il pieno routing Sinkhorn-based SMoE presentava un overhead del 72.47%.
  • In configurazioni Momentum, l'SSR ha preservato l'evoluzione caratteristica della norma di Vanilla SMoE migliorandone al contempo le prestazioni, mentre altri metodi di bilanciamento hanno degradato le prestazioni.
• Modellazione a livello di Byte (Enwik-8):
  • SSR-S con rumore ha raggiunto il miglior valore di Bits-per-Character (BPC) di test di 1.128, migliorando Vanilla SMoE di 0.010 BPC.
  • Era 2.28x più veloce del pieno Sinkhorn-based SMoE.
• Compiti di Visione:
  • Su ImageNet-1K, SSR-L con rumore ha raggiunto un'accuratezza Top-1 di 77.420, un guadagno di +2.368 rispetto a Vanilla SMoE.
  • Ha inoltre mostrato una forte robustezza su ImageNet-O e ImageNet-R.
• Studi di Ablazione:
  • La probabilità $p$ di applicare il routing Sinkhorn è critica; valori molto piccoli (ad esempio, $10^{-3}$ o $10^{-4}$) sono stati sufficienti per fornire segnali di bilanciamento senza vincolare eccessivamente il router.
  • L'iniezione di rumore ha costantemente migliorato le prestazioni in tutte le impostazioni.
  • Disabilitare i meccanismi di bilanciamento in fase di inferenza ha prodotto i risultati migliori, confermando le affermazioni teoriche.

Significato e Rivendicazioni

Il documento afferma che il Selective Sinkhorn Routing offre una soluzione pratica ed efficace per la progettazione di SMoE. Sostituendo i complessi obiettivi ausiliari con un'ottimizzazione intermittente ed efficiente basata su Sinkhorn, l'SSR ottiene:

• Migliore Stabilità dell'Addestramento: Attraverso un bilanciamento degli esperti integrato senza disallineamento dell'obiettivo.
• Maggiore Accuratezza: Dimostrata attraverso benchmark di testo e visione.
• Robustezza: Prestazioni migliorate contro la corruzione dell'input e gli esempi avversari.
• Efficienza: Overhead computazionale minimo rispetto sia ai metodi con perdita ausiliaria che al pieno routing Sinkhorn.

Gli autori sottolineano che il metodo è leggero, non richiede parametri addestrabili aggiuntivi (nella variante con rumore) e mantiene la flessibilità della selezione degli esperti garantendo al contempo un utilizzo bilanciato durante la fase di addestramento.