Hierarchical Copula-Gumbel-Top-\texorpdfstring{$K$}{K} Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws

Sintesi Tecnica: Routing Gerarchico Copula-Gumbel-Top-K

Definizione del Problema
Nei modelli sparse Mixture-of-Experts (MoE), il meccanismo di routing stocastico standard (Gumbel-Top-K) definisce una "legge di routing" per ogni token: una specifica distribuzione di probabilità su liste ordinate di esperti e pesi di miscelazione, condizionata dai logit del token. Mentre quasi tutto il lavoro esistente modifica questa legge per singolo token (ad esempio, tramite il fine-tuning dei logit o perdite ausiliarie), la distribuzione congiunta delle scelte di routing tra diversi token rimane ampiamente inesaminata. Il saggio si chiede: dato che la legge di routing di ogni singolo token è mantenuta esattamente fissa (come in un modello congelato), quali distribuzioni congiunte sulle scelte di routing di diversi token sono raggiungibili? In particolare, è possibile controllare la dipendenza tra i token per migliorare la coerenza degli esperti o il bilanciamento del carico senza alterare il comportamento marginale di alcun singolo token?

Metodologia: Hierarchical Copula-Gumbel-Top-K (H-CGA)
L'autore propone l'H-CGA, uno strato di dipendenza lato routing che opera sul vettore di rumore Gumbel prima della selezione Top-K. Il metodo si basa sul teorema di Sklar per separare le distribuzioni marginali dalle strutture di dipendenza. La costruzione procede in quattro fasi:

1. Raggruppamento Pre-Routing: I token sono partizionati in gruppi disgiunti (ad esempio, finestre fisse di token adiacenti). Un controllore legge le caratteristiche pre-routing congelate (ad esempio, stati nascosti, entropia del gate) per determinare le intensità di accoppiamento.
2. Accoppiamento Positivo Intra-Gruppo: Per i token all'interno di un gruppo $g$, le perturbazioni Gumbel per ogni coordinata di esperto sono generate tramite una copula Gaussiana scambiabile. I token condividono una variabile latente $\zeta_{ge}$, inducendo una correlazione positiva nelle loro perturbazioni casuali. Ciò aumenta la probabilità che token correlati selezionino gli stessi esperti (coerenza).
3. Accoppiamento Antitetico Cross-Gruppo: Coppie disgiunte di gruppi $(g, g')$ vengono accoppiate. Viene introdotta una forza di opposizione regolabile $\alpha_{g,g'} \in [0, 1]$. Il latente condiviso del secondo gruppo è costruito come una combinazione lineare del latente del primo gruppo e di un termine di rumore indipendente. Quando $\alpha_{g,g'} = 1$, i gruppi ricevono spinte condivise completamente opposte; quando $\alpha_{g,g'} = 0$, sono indipendenti. Ciò introduce una dipendenza negativa tra i gruppi per contrastare la burstiness del carico.
4. Preservazione delle Marginali: La trasformazione dai latenti condivisi alle varianti Gumbel assicura che, per qualsiasi singolo token $t$, il vettore delle perturbazioni $(\gamma_{t1}, \dots, \gamma_{tE})$ rimanga distribuito come Gumbel i.i.d., identico al routing indipendente.

Contributi Chiave

1. Architettura H-CGA: Uno strato di copula gerarchica che controlla la dipendenza cross-token in entrambe le direzioni (accoppiamento positivo intra-gruppo, dipendenza negativa regolabile cross-gruppo) senza modificare i logit del router, i pesi degli esperti o qualsiasi legge di routing per singolo token.
2. Invarianza Teorica: L'autore dimostra (Teorema 1) che l'intera gerarchia preserva la distribuzione condizionata della lista Top-K ordinata e dei pesi di miscelazione di ogni token. Di conseguenza, i conteggi attesi condizionati dell'inclusione degli esperti (Corollario 1) sono identici a quelli del routing indipendente.
3. Caratterizzazione del Trade-off Coerenza–Dispersione: Il saggio caratterizza il trade-off della varianza (Proposizione 1):
   • L'accoppiamento positivo intra-gruppo aumenta la varianza del carico realizzato rispetto al routing indipendente (un costo della coerenza locale).
   • L'opposizione cross-gruppo non negativa riduce questa varianza rispetto all'accoppiamento "piatto" (gruppi indipendenti) alla stessa intensità intra-gruppo, sebbene non riduca necessariamente la varianza al di sotto del baseline del routing indipendente.
4. Adattamento Solo-Routing: L'autore formula un metodo di adattamento in cui un piccolo controllore imposta i dial di dipendenza ($\rho_g, \alpha_{g,g'}$) basandosi su caratteristiche congelate. Questo controllore è addestrato utilizzando uno stimatore della funzione di punteggio (REINFORCE), richiedendo solo pass in avanti attraverso il modello base congelato e gradienti confinati al controllore.

Risultati e Studio Pilota
Il saggio presenta un piccolo studio pilota su un decoder-only MoE da 15.8M di parametri addestrato su TinyStories, con il modello base congelato durante la valutazione.

• Verifiche di Invarianza: Test sintetici hanno confermato che la frequenza empirica delle liste Top-K ordinate e le probabilità di inclusione degli esperti sotto H-CGA corrispondevano al routing indipendente entro i margini di errore Monte Carlo, validando l'invarianza teorica.
• Effetti Statistici:
  • L'aumento dell'accoppiamento intra-gruppo ($\rho$) ha aumentato significativamente la similitudine di Jaccard degli insiemi di esperti all'interno delle finestre e ha ridotto il numero di esperti distinti utilizzati per finestra.
  • L'aumento dell'opposizione cross-gruppo ($\alpha$) ha ridotto significativamente la similitudine di Jaccard ai confini dei gruppi, confermando la capacità di disaccoppiare le richieste dei gruppi.
• Apprendimento: Un controllore scalare appreso (3 parametri) è stato addestrato stabilmente ma non ha mostrato miglioramenti nella cross-entropia di validazione rispetto al baseline indipendente. Le intensità di accoppiamento apprese erano basse, in linea con l'osservazione del saggio che la preservazione delle marginali limita il segnale disponibile per l'ottimizzazione quando viene utilizzata solo la perdita per singolo token.
• Varianza del Carico: Il pilota ha osservato cambiamenti minimi nel coefficiente di variazione (CV) del carico aggregato, notando che questo metrica riassuntiva non testa direttamente i limiti della varianza condizionata proposti nella Proposizione 1.

Significatività e Rivendicazioni
Il saggio sostiene che l'H-CGA espone un nuovo grado di libertà nella routing delle MoE: la struttura di dipendenza congiunta del rumore di routing. La sua principale significatività è teorica e meccanicistica:

• Sicurezza: Fornisce una "garanzia di sicurezza" che la modifica della dipendenza cross-token non altera il comportamento di routing marginale di alcun token, assicurando che le preferenze apprese dal modello base rimangano intatte.
• Controllo: Offre due "dial" complementari per scambiare la coerenza locale dell'insieme di esperti con la dispersione globale del carico.
• Efficienza: Permette l'adattamento solo-routing dove il modello base viene valutato solo in direzione forward, evitando il costo computazionale di aggiornare i pesi degli esperti.

Limitazioni e Modestia
L'autore è esplicito riguardo all'ambito e ai limiti:

• Garanzia Locale dello Strato: I risultati di invarianza sono condizionati agli stati nascosti e ai logit che entrano in uno specifico strato di routing. Non garantiscono l'invarianza end-to-end in una rete multi-strato, poiché l'accoppiamento può alterare gli stati nascosti a valle.
• Nessun Guadagno a Livello di Task: Il pilota non stabilisce che l'adattamento solo-routing migliori le prestazioni sui task a valle. L'autore nota che la preservazione delle marginali implica che gli obiettivi che dipendono esclusivamente dalle distribuzioni del singolo token non possono identificare il controllore; segnali utili devono derivare da effetti congiunti (ad esempio, interazioni di attenzione), che possono essere deboli.
• Varianza vs. Overflow: Sebbene il metodo caratterizzi teoricamente la varianza del carico, il pilota non ha misurato direttamente la varianza del carico condizionato o gli overflow di capacità. Il saggio non afferma di risolvere l'overflow di capacità, ma suggerisce che il dial di opposizione sia un meccanismo per affrontarlo.
• Overhead: Il metodo richiede la generazione e la coordinazione delle perturbazioni Gumbel per tutti gli esperti, il che può introdurre un overhead computazionale rispetto al routing standard.

In sintesi, il saggio stabilisce un quadro matematicamente rigoroso per controllare la dipendenza di routing cross-token in MoE congelate, dimostrando che la coerenza e la dispersione del carico possono essere manipolate senza alterare le leggi di routing per singolo token, pur rimanendo cauti riguardo ai benefici pratici immediati per le prestazioni dei task.