Sintesi Tecnica: DOT-MoE: Trasporto Ottimale Differenziabile per la MoEficazione

1. Definizione del Problema

La rapida scalabilità dei Large Language Models (LLM) ha creato un significativo collo di bottiglia nell'efficienza dell'inferenza. Sebbene le architetture Mixture-of-Experts (MoE) scollino la dimensione del modello dal costo di inferenza attivando solo un sottoinsieme di parametri per token, l'addestramento delle MoE da zero è notoriamente instabile, esigente in termini di dati e computazionalmente costoso.

Di conseguenza, la MoEficazione — la conversione di modelli densi pre-addestrati in MoE sparse — è emersa come un'alterna percorribile. Tuttavia, i metodi esistenti per convertire le reti Feed-Forward (FFN) dense in esperti sparsi si affidano pesantemente a strategie euristiche (ad esempio, suddivisione casuale, clustering basato sui pesi o clustering basato sulle attivazioni). Tali approcci soffrono di tre limitazioni primarie:

1. Separazione delle Responsabilità: Trattano tipicamente l'assegnazione dei neuroni (partizionamento della FFN) e l'addestramento del router come fasi separate e sequenziali, impedendo loro di co-adattarsi.
2. Ottimizzazione tramite Proxy: Ottimizzano basandosi su statistiche intermedie (pesi o attivazioni) piuttosto che sul reale contributo in uscita dei neuroni, portando spesso a una ricostruzione subottimale del comportamento del modello denso.
3. Mancanza di Differenziabilità: La natura discreta dell'assegnazione dei neuroni impedisce il flusso dei gradienti end-to-end, rendendo difficile l'ottimizzazione congiunta della decomposizione strutturale e della politica di routing.

2. Metodologia: DOT-MoE

Gli autori propongono DOT-MoE, un framework che formula la decomposizione dei livelli densi in esperti come un problema di Trasporto Ottimale Differenziabile (DOT). Questo approccio consente l'apprendimento end-to-end sia della decomposizione degli esperti che del meccanismo di routing.

2.1. Formulazione del Trasporto Ottimale

L'intuizione centrale è quella di vedere l'assegnazione dei neuroni come un problema di trasporto bilanciato in cui la "massa" (neuroni) deve essere trasportata da una distribuzione sorgente (i neuroni della FFN densa) a una distribuzione target (gli esperti) sotto rigorosi vincoli di capacità.

• Sorgente: $dffn$ neuroni, ciascuno con massa unitaria.
• Target: $E$ esperti, ciascuno dei quali richiede esattamente $s$ neuroni ($s \cdot E = dffn$).
• Funzione di Costo: Invece di una matrice di costo fissa, il metodo utilizza una matrice di affinità apprendibile $A$, dove $A_{i,e}$ rappresenta l'affinità di assegnare il neurone $i$ all'esperto $e$. L'obiettivo è massimizzare l'affinità totale soddisfacendo i vincoli marginali.

2.2. Relaxing Differenziabile tramite Sinkhorn-Knopp

Per rendere differenziabile il problema dell'assegnazione discreta, DOT-MoE impiega la Regolarizzazione Entropica.

• Il problema del trasporto ottimale viene rilassato utilizzando l'algoritmo di Sinkhorn-Knopp, che normalizza iterativamente le righe e le colonne della matrice del piano di trasporto.
• Ciò produce una matrice di assegnazione morbida $M_{soft} \in [0, 1]^{dffn \times E}$ che soddisfa i vincoli di capacità in termini di aspettativa.
• Man mano che il parametro di temperatura $\tau \to 0$, la soluzione si avvicina all'ottimo discreto.

2.3. Ottimizzazione Congiunta con Straight-Through Estimators (STE)

Una sfida critica è che l'implementazione richiede assegnazioni discrete, e il routing comporta la selezione top-k, entrambi non differenziabili. DOT-MoE affronta questo problema utilizzando gli Straight-Through Estimators (STE):

• Forward Pass: Utilizza decisioni hard (arrotondamento greedy di $M_{soft}$ per l'assegnazione; selezione top-k per il routing).
• Backward Pass: Instrada i gradienti attraverso le controparti soft ($M_{soft}$ e i logit del router $P$).
• Ciò permette ai pesi del router ($W_r$) e alla matrice di affinità di assegnazione ($A$) di co-adattarsi durante l'addestramento.

2.4. Obiettivo di Addestramento

Il framework viene addestrato tramite una Fase di Allineamento in cui i pesi del modello denso sono congelati. Il modello minimizza una combinazione di perdite:

1. Ricostruzione dell'Output: Misurata tramite la divergenza KL tra l'output del docente denso e dello studente MoE sparso.
2. Modellazione del Linguaggio: Perdita di cross-entropy sull'obiettivo di modellazione del linguaggio.
3. Ausiliari di Stabilità: Router $z$-loss (per prevenire l'esplosione dei logit) e perdita di bilanciamento del carico (per prevenire il collasso degli esperti).

3. Contributi Chiave

1. Framework di Trasporto Ottimale: Il paper introduce una nuova formulazione della conversione da denso a MoE come un problema di trasporto bilanciato, risolto tramite iterazioni Sinkhorn differenziabili per imporre rigorosi vincoli di capacità degli esperti.
2. Co-Adattamento a Doppio Livello: A differenza dei metodi precedenti che fissano le assegnazioni prima di addestrare i router, DOT-MoE ottimizza congiuntamente la decomposizione neurone-esperto e il routing token-esperto utilizzando STE complementari.
3. Costruzione Consapevole dell'Output: Modellando direttamente il contributo del neurone al risultato finale (anziché ottimizzare proxy come pesi o attivazioni intermedie), il metodo raggiunge una maggiore fedeltà di ricostruzione.

4. Risultati Sperimentali

Gli autori hanno valutato DOT-MoE su tre famiglie di modelli (LLaMA-2, LLaMA-3, Qwen2.5) e sei benchmark (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Performance vs. Pruning: Su LLaMA-2 7B con un budget di parametri del 50%, DOT-MoE ha raggiunto una perplexity di 7.99 su WikiText-2, superando significativamente il metodo di pruning strutturato allo stato dell'arte DISP-LLM (9.84).
• Performance vs. Esistente MoEficazione: In scenari zero-shot, DOT-MoE ha superato sostanzialmente metodi come CMoE e LLaMA-MoE-v2. Ad esempio, su LLaMA-2 7B, DOT-MoE ha ottenuto una accuratezza media del 61.5% rispetto al 44.5% di CMoE.
• Efficienza: DOT-MoE mantiene circa il 90% delle prestazioni del modello denso originale riducendo i parametri attivi del 50%.
• Fine-Tuning: Con soli 1.2B di token di fine-tuning continuo, i modelli DOT-MoE hanno recuperato prestazioni vicine a quelle dei controparti densi (ad esempio, 66.6% vs 76.6% su LLaMA-2 7B), superando i baseline che richiedevano dati simili o maggiori.
• Studi di Ablazione:
  • Granularità: Le prestazioni migliorano con l'aumento della granularità degli esperti fino a un punto di saturazione, dopo il quale i guadagni diminuiscono.
  • Velocità di Inferenza: Nonostante l'elevato numero di esperti, il throughput di inferenza rimane stabile quando si utilizzano kernel MoE fusi (es. vLLM), poiché la dimensione intermedia totale fusa rimane costante.
  • Sparsità: I modelli addestrati a livelli di sparsità più elevati (75%) generalizzano robustamente anche a regimi di sparsità ancora più estremi durante l'inferenza.

5. Significato e Rivendicazioni

Il paper sostiene che DOT-MoE rappresenti un avanzamento significativo nell'efficienza dell'inferenza degli LLM, fornendo un metodo differenziabile e principista per convertire modelli densi in MoE sparse.

• Cambio di Paradigma: Si passa da una partizione euristica e statica a una decomposizione dinamica e apprendibile che rispetta la coerenza funzionale degli esperti.
• Praticità: Il metodo consente la sparsificazione post-hoc efficiente di modelli pre-addestrati senza la necessità di un intero ri-addestramento da zero, mantenendo un'alta capacità pur riducendo drasticamente i parametri attivi.
• Robustezza: Il framework dimostra che preservare la fedeltà della FFN richiede la modellazione esplicita del contributo del neurone all'output, un principio che i metodi basati su proxy esistenti non riescono a catturare.

Gli autori concludono che DOT-MoE offre un'inizializzazione superiore per i modelli sparsi, colmando il divario di prestazioni tra architetture dense e sparse con un minimo di dati di addestramento aggiuntivi.