Sintesi Tecnica: Fusione Esperta in Sparse Mixture of Experts con Nash Bargaining

Definizione del Problema

Le architetture Sparse Mixture of Experts (SMoE) sono diventate un pilastro per scalare le reti neurali in modo efficiente, attivando solo un sottoinsieme di moduli esperti per ogni input. Mentre la ricerca significativa si è concentrata sul miglioramento dei meccanismi di routing (ad esempio, bilanciamento del carico, stabilità), il problema complementare della fusione degli esperti (expert merging) — ovvero combinare tutti i parametri degli esperti in un modello unificato per il deployment o l'inferenza — rimane poco esplorato.

Le strategie di fusione esistenti si affidano tipicamente a schemi di ponderazione euristici, come la semplice media indipendente dall'input o l'aggregazione top-$k$. Questi metodi ignorano spesso le intricate dinamiche cooperative e competitive tra gli esperti. Inoltre, approcci recenti come Curvature-aware Merging of Experts (CAMEx) e la sua variante dinamica (EP-CAMEx), che utilizzano i gradienti naturali per allineare gli esperti, hanno mostrato dei limiti. Nello specifico, EP-CAMEx spesso performa peggioamente rispetto alla sua controparte statica, probabilmente a causa di una coordinazione insufficiente tra i contributi degli esperti e di una lenta convergenza nella propagazione. La sfida centrale è interpretare la fusione degli esperti non semplicemente come una media lineare, ma come un'interazione strutturata che bilancia i contributi unici di diversi esperti.

Metodologia: NAMEx

Gli autori propongono NAMEx (Nash Merging of Experts), un framework che reinterpreta la fusione degli esperti attraverso la lente della teoria dei giochi, utilizzando specificamente la Soluzione di Nash per la Trattativa (Nash Bargaining Solution - NBS).

1. Fusione degli Esperti come Gioco di Trattativa

Gli autori inquadrano il processo di fusione come un gioco cooperativo-competitivo tra esperti.

• Giocatori: Ogni esperto $i$ nel layer SMoE.
• Funzione di Utilità: Definita come $u_i(\Delta E) = \tau_i^\top \Delta E$, dove $\tau_i = E_i - E_m$ è il "vettore di dominio" che rappresenta la deviazione del $i$-esimo esperto dall'esperto base $E_m$. Questo vettore cattura il contributo unico dell'esperto.
• Punto di Disaccordo: Impostato a 0, corrispondente all'assenza di aggiornamento all'esperto base.
• Obiettivo: Trovare una direzione di aggiornamento $\Delta E$ all'interno di una palla $B_\epsilon$ che massimizzi il prodotto dei guadagni individuali rispetto al punto di disaccordo:
  $$\Delta E^* = \arg\max_{\Delta E \in B_\epsilon} \prod_{i=1}^N (\Delta E^\top \tau_i)$$
  Sotto l'assunzione di vettori di dominio linearmente indipendenti, la soluzione è derivata come una somma pesata dei vettori di dominio:
  $$\Delta E = \sum_{i=1}^N \alpha_i \tau_i$$
  dove i pesi $\alpha$ soddisfano l'equazione di Nash: $G^\top G \alpha = 1/\alpha$, con $G$ la matrice dei vettori di dominio e $1/\alpha$ che denota i reciproci elemento per elemento.

2. Integrazione con la Propagazione degli Esperti

NAMEx è integrato nel framework EP-CAMEx. Inveve di utilizzare un passo di discesa del gradiente fisso per la propagazione dell'esperto base, NAMEx calcola la direzione di aggiornamento ottimale $\Delta E^{(l)}$ ad ogni layer $l$ utilizzando la soluzione di Nash. L'esperto base $E_m$ viene aggiornato come:
$$E_m^{(l+1)} = E_m^{(l)} + \gamma \sum_{i=1}^N \alpha_i^{(l)} \tau_i^{(l)}$$
dove $\gamma$ è il passo (step size).

3. Momento Complesso

Per affrontare la lenta convergenza osservata in EP-CAMEx, gli autori integrano il momento complesso (complex momentum) nel framework NAMEx (NAMEx-Momentum).

• Un buffer di momento complesso $\mu^{(j)} \in \mathbb{C}^d$ accumula gli aggiornamenti direzionali.
• La regola di aggiornamento coinvolge un coefficiente di momento complesso $\beta \in \mathbb{C}$:
  $$\mu^{(j+1)} = \beta \mu^{(j)} + \Delta E^{(j)}$$
  $$E_m^{(j+1)} = E_m^{(j)} + \Re(\gamma \mu^{(j+1)})$$
• L'analisi teorica (Proposizione 3.5 e Teorema C.3) fornisce garanzie di convergenza, mostrando che, in condizioni moderate, il raggio spettrale della matrice di sistema può essere mantenuto inferiore a 1, garantendo una convergenza lineare.

Contributi Chiave

1. Framework NAMEx: Un nuovo metodo di fusione degli esperti che integra il framework di ottimizzazione di Nash Bargaining nella propagazione degli esperti, trattando i vettori di dominio degli esperti come funzioni di utilità per derivare coefficienti di fusione equi ed efficienti.
2. Integrazione del Momento Complesso: L'incorporazione del momento complesso in NAMEx per accelerare la convergenza e migliorare la stabilità, particolarmente in scenari con dinamiche di esperti avversarie o conflittuali. Gli autori forniscono prove teoriche per il tasso di convergenza di NAMEx-Momentum.
3. Scalabilità e Validazione Empirica: Dimostrazione dell'efficacia di NAMEx attraverso diversi compiti (modellazione del linguaggio, classificazione del testo, classificazione delle immagini) e la sua scalabilità a sistemi su larga scala (Qwen1.5-MoE 14B, DeepSeek-MoE 16B).

Risultati Sperimentali

Gli autori hanno valutato NAMEx rispetto ai baseline (SMoE, CAMEx, EP-CAMEx) attraverso molteplici benchmark:

• Modellazione del Linguaggio (WikiText-103): NAMEx-Full-Mom ha ottenuto le perplexity di validazione e di test più basse sia per il pre-training su piccola che media scala, superando tutti i baseline.
• Classificazione del Testo (GLUE): Sul benchmark GLUE utilizzando T5-Base, NAMEx-Full-Mom ha ottenuto i migliori risultati in tutti i task (SST-2, MRPC, CoLA, ecc.), mostrando miglioramenti costanti rispetto a SMoE e ai metodi basati sulla curvatura.
• Classificazione delle Immagini (ImageNet-1k & Robustezza): Su Swin-MoE, NAMEx-Mom ha superato tutti i baseline in accuratezza Top-1. Fondamentalmente, ha dimostrato una robustezza superiore su dataset corrotti (ImageNet-A, ImageNet-O, ImageNet-R), con le varianti di momentum che mostravano i guadagni più forti.
• Modelli su Larga Scala: Quando applicato a DeepSeek-MoE (16B) e Qwen1.5-MoE (14B), NAMEx-Full ha costantemente superato i baseline e EP-CAMEx sia in scenari zero-shot che fine-tuned (SmolTalk) su benchmark MMLU, GSM8K e ARC.
• Studi di Ablazione:
  • Momento: Il momento complesso ha migliorato significativamente la velocità di convergenza e le prestazioni rispetto al momento a valori reali.
  • Iterazioni: Risolvere il sistema di Nash con 2 iterazioni per layer è stato trovato efficiente ed efficace, con rendimenti decrescenti a conteggi di iterazioni più elevati.
  • Punto di Disaccordo: L'uso di 0 come punto di disaccordo si è rivelato robusto e stabile rispetto all'uso della media.

Significato e Rivendicazioni

Il paper sostiene che NAMEx offra un'alternativa fondata alla fusione euristica degli esperti. Modellando le interazioni tra esperti come un gioco di trattativa, NAMEx garantisce un processo di fusione Pareto-efficiente che tiene conto sia della cooperazione che della competizione tra gli esperti. L'integrazione del momento complesso affronta i colli di bottiglia della convergenza negli scenari di fusione dinamica.

Gli autori sottolineano che NAMEx:

• Migliora le Prestazioni: Supera costantemente le strategie di fusione esistenti in compiti di visione e linguaggio.
• Aumenta la Robustezza: Dimostra una superiore generalizzazione sotto corruzione dei dati e shift della distribuzione.
• Scala Efficacemente: Dimostra di essere praticabile per architetture MoE su larga scala (fino a 16B di parametri) senza compromettere significativamente l'efficienza computazionale, poiché i costi di inferenza rimangono invariati e l'overhead di addestramento è gestibile.
• Fondamento Teorico: Fornisce garanzie di convergenza per la variante potenziata dal momento, andando oltre le pure euristiche empiriche.

Il lavoro suggerisce che vedere la fusione degli esperti attraverso la teoria dei giochi apra nuove strade per l'ottimizzazione di sistemi di deep learning modulari, con potenziali applicazioni nel federated learning e nel transfer learning dove la combinazione equa degli esperti è critica.