ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Sintesi Tecnica: ProbMoE: Routing Probabilistico Differenziabile per Mixture-of-Experts

Problematica
I modelli Mixture-of-Experts (MoE) scalano in modo efficiente attivando solo un piccolo sottoinsieme di esperti per ogni token. Tuttavia, l'addestramento di questi modelli rimane complesso perché il routing standard si basa su un operatore top-$k$ discreto e non differenziabile. Le procedure di addestramento convenzionali aggirano questo problema ignorando la dipendenza della selezione top-$k$ dai logit del router e propagando i gradienti solo attraverso le probabilità softmax degli esperti selezionati. Questa approssimazione limita il segnale di apprendimento per sottoinsiemi di esperti alternativi, portando spesso a distribuzioni di routing concentrate, scarsa utilizzazione degli esperti e dinamiche di addestramento instabili. I tentativi esistenti per mitigare ciò, come la backpropagation densa o gli stimatori di gradiente euristici, spesso non riescono a modellare esplicitamente la distribuzione sui sottoinsiemi di esperti fattibili o forniscono meccanismi limitati per esplorare combinazioni alternative durante l'addestramento.

Metodologia
Gli autori propongono ProbMoE, un framework che riformula il routing degli esperti come un problema di inferenza probabilistica su sottoinsiemi di esperti con vincolo di cardinalità. Invece di trattare il routing come un operatore deterministico, ProbMoE modella la selezione degli esperti come una distribuzione su sottoinsiemi discreti.

1. Formulazione Probabilistica:

   • Routing Exact-$k$: Il framework definisce una distribuzione su sottoinsiemi $S$ dove $|S|=k$. Condiziona le probabilità di selezione Bernoulli indipendenti (derivate dai logit del router tramite sigmoid) al vincolo che siano selezionati esattamente $k$ esperti. La costante di normalizzazione per questa distribuzione, $Z_k$, è calcolata esattamente tramite programmazione dinamica, rendendo la formulazione trattabile.
   • Routing Dynamic-$k$: Il framework si generalizza a un contesto con vincolo di intervallo dove la dimensione del sottoinsieme $|S|$ ricade nell'intervallo $[k_{min}, k_{max}]$. La distribuzione è condizionata su questo intervallo, permettendo al modello di inferire congiuntamente l'identità degli esperti e la cardinalità del routing per ogni token.

2. Forward Pass:

   • Durante l'addestramento, ProbMoE campiona un sottoinsieme di esperti $S$ dalla distribuzione definita (exact-$k$ o a vincolo di intervallo).
   • Lo strato MoE calcola l'output utilizzando solo gli esperti campionati, preservando il pattern di esecuzione sparsa dei modelli MoE standard.
   • In fase di inferenza, viene selezionato l'insieme Maximum-A-Posteriori (MAP), garantendo che i costi di inferenza corrispondano al routing MoE convenzionale.

3. Backward Pass:

   • Per consentire l'ottimizzazione basata sul gradiente, ProbMoE impiega lo stimatore di gradiente SIMPLE (Ahmed et al., 2023).
   • Invece di differenziare attraverso il campione discreto, il backward pass utilizza le probabilità marginali di ogni esperto di essere selezionato sotto la distribuzione vincolata.
   • I pesi di routing sono costruiti utilizzando un estimatore straight-through: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, dove $z$ è la maschera campionata, $m$ è il vettore delle probabilità marginali esatte e $\pi$ sono le probabilità di routing softmax.
   • Ciò consente ai gradienti di fluire attraverso le probabilità marginali, fornendo aggiornamenti informativi al router che riflettono come i cambiamenti nelle probabilità di selezione influenzino la perdita attesa, mantenendo al contempo l'esecuzione sparsa degli esperti.

Contributi Chiave

1. Formulazione di Inferenza Probabilistica: Il paper inquadra il routing MoE come inferenza probabilistica su sottoinsiemi di esperti con vincolo di cardinalità, andando oltre le approssimazioni euristiche del gradiente.
2. Ottimizzazione Trattabile Basata sulle Marginali: Sfruttando lo stimatore SIMPLE, gli autori accoppiano il campionamento stocastico di sottoinsiemi con gradienti del router esatti basati sulle marginali, consentendo un addestramento differenziabile senza computazione densa lato esperti.
3. Framework Unificato Dynamic-$k$: Il metodo si estende naturalmente al routing dinamico, dove il modello inferisce sia l'identità degli esperti che il numero di esperti attivi per token entro un intervallo predefinito, senza fare affidamento su regole di gating euristiche.
4. Miglioramenti Empirici:
   • Exact-$k$: ProbMoE raggiunge prestazioni elevate su benchmark (GSM8K, Law, Translation, MBPP, MMLU) su backbone OLMoE e Qwen, spesso eguagliando o superando baseline competitive come DenseMixer e il Top-$k$ Convenzionale.
   • Dynamic-$k$: ProbMoE Dynamic-$k$ raggiunge prestazioni comparabili a Exact-$k$ attivando mediamente meno esperti, dimostrando un'allocazione adattiva efficace.
5. Diversità di Routing: L'analisi mostra che ProbMoE induce distribuzioni di routing più ampie e un utilizzo più equilibrato degli esperti rispetto ai baseline, mitigando il collasso degli esperti.

Risultati e Significato
Il paper dimostra che ProbMoE migliora l'utilizzo degli esperti e la diversità del routing. Negli studi di ablazione, la combinazione di campionamento stocastico e ottimizzazione basata sulle marginali si rivela il principale motore dei guadagni di prestazione, superando le configurazioni che utilizzano solo uno di questi componenti.

Gli autori affermano che ProbMoE fornisce un metodo rigoroso per la selezione degli esperti che migliora la stabilità dell'addestramento e l'efficienza computazionale. Modellando il routing probabilisticamente, il framework consente l'esplorazione di combinazioni di esperti alternative durante l'addestramento, portando a distribuzioni di routing meglio calibrate. La variante Dynamic-$k$ mostra ulteriormente come la cardinalità adattiva possa essere appresa per adattarsi alla complessità del livello di token, riducendo i costi effettivi di routing senza sacrificare le prestazioni del compito. Il lavoro suggerisce che le direzioni future includano l'estensione di ProbMoE al pre-training e alle ottimizzazioni a livello di sistema per sfruttare appieno la sparsità in contesti dinamici.