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

Technische Samenvatting: Hiërarchische Copula-Gumbel-Top-K Routering

Probleemstelling
In sparse Mixture-of-Experts (MoE) modellen definieert het standaard stochastische routeringsmechanisme (Gumbel-Top-K) een "routeringswet" voor elk token: een specifieke waarschijnlijkheidsverdeling over geordende lijsten van experts en menggewichten, geconditioneerd op de logits van het token. Terwijl bijna al het bestaande werk deze per-token wet aanpast (bijvoorbeeld via het finetunen van logits of hulpfuncties/auxiliary losses), blijft de gezamenlijke verdeling van routeringskeuzes over verschillende tokens grotendeels ononderzocht. De paper vraagt: gegeven dat de routeringswet van elk individueel token exact vaststaat (zoals in een bevroren model), welke gezamenlijke verdelingen over de routeringskeuzes van verschillende tokens zijn bereikbaar? Kan men specif gezien de afhankelijkheid tussen tokens controleren om de expert-coherentie of de belastingbalans (load balancing) te verbeteren zonder het marginale gedrag van een enkel token te wijzigen?

Methodologie: Hiërarchische Copula-Gumbel-Top-K (H-CGA)
De auteur stelt H-CGA voor, een routeringszijde afhankelijkheidslaag die opereert op de Gumbel-ruisvector vóór de Top-K selectie. De methode steunt op de stelling van Sklar om marginale verdelingen te scheiden van afhankelijkheidsstructuren. De constructie verloopt in vier stadia:

1. Pre-Routing Groepering: Tokens worden verdeeld in disjuncte groepen (bijvoorbeeld vaste vensters van aangrenzende tokens). Een controller leest bevroren pre-routing kenmerken (bijvoorbeeld verborgen toestanden, gate-entropie) om koppelingssterktes te bepalen.
2. Binnen-groep Positieve Koppeling: Voor tokens binnen een groep $g$ worden de Gumbel-perturbaties voor elk expert-coördinaat gegenereerd via een uitwisselbare Gaussische copula. Tokens delen een latente variabele $\zeta_{ge}$, wat een positieve correlatie induceert in hun willekeurige perturbaties. Dit vergroot de waarschijnlijkheid dat gerelateerde tokens dezelfde experts selecteren (coherentie).
3. Cross-groep Antithese-koppeling: Disjuncte paren groepen $(g, g')$ worden gekoppeld. Een instelbare oppositiekracht $\alpha_{g,g'} \in [0, 1]$ wordt geïntroduceerd. De gedeelde latente variabele van de tweede groep wordt geconstrueerd als een lineaire combinatie van de latente variabele van de eerste groep en een onafhankelijke ruisterm. Bij $\alpha_{g,g'} = 1$ ontvangen de groepen volledig tegengestelde gedeelde impulsen; bij $\alpha_{g,g'} = 0$ zijn ze onafhankelijk. Dit introduceert negatieve afhankelijkheid tussen groepen om de last-pieken (load burstiness) tegen te gaan.
4. Preservering van Marginalen: De transformatie van gedeelde latente variabelen naar Gumbel-variaten zorgt ervoor dat voor elk enkel token $t$, de vector van perturbaties $(\gamma_{t1}, \dots, \gamma_{tE})$ identiek Gumbel-verdeeld blijft als bij onafhankelijke routering.

Belangrijkste Bijdragen

1. H-CGA Architectuur: Een hiërarchische copula-laag die de cross-token afhankelijkheid in beide richtingen controleert (positieve binnen-groep, instelbare negatieve cross-groep) zonder de router-logits, expert-gewichten of enige per-token routeringswet te wijzigen.
2. Theoretische Invariantie: De auteur bewijst (Theorem 1) dat de gehele hiërarchie de conditionele distributie van elke Top-K lijst van een token en de menggewichten behoudt. Bijgevolg zijn de conditionele verwachte expert-inclusie-aantallen (Corollary 1) identiek aan die van onafhankelijke routering.
3. Karakterisering van de Coherentie–Dispersie Trade-off: De paper karakteriseert de variantie-trade-off (Proposition 1):
   • Positieve binnen-groep koppeling verhoogt de variantie van de gerealiseerde expert-belastingen ten opzichte van onafhankelijke routering (een kostenpost van lokale coherentie).
   • Niet-negatieve cross-groep oppositie vermindert deze variantie ten opzichte van "platte" koppeling (onafhankelijke groepen) bij dezelfde binnen-groep sterkte, hoewel het de variantie niet noodzakelijkerwijs onder het niveau van de onafhankelijke routering-baseline brengt.
4. Routerings-alleen Adaptatie: De auteur formuleert een adaptatiemethode waarbij een kleine controller de afhankelijkheids-knoppen ($\rho_g, \alpha_{g,g'}$) instelt op basis van bevroren kenmerken. Deze controller wordt getraind met een score-functie estimator (REINFORCE), wat alleen forward passes door het bevroren basismodel vereist en waarbij gradiënten beperkt blijven tot de controller.

Resultaten en Pilotstudie
De paper presenteert een kleine schaal pilot op een 15.8M parameter decoder-only MoE getraind op TinyStories, waarbij het basismodel bevroren is tijdens evaluatie.

• Invariantie Checks: Synthetische tests bevestigden dat de empirische frequentie van geordende Top-K lijsten en expert-inclusie-waarschijnlijkheden onder H-CGA overeenkwamen met onafhankelijke routering binnen de Monte Carlo-foutmarges, wat de theoretische invariantie valideert.
• Statistische Effecten:
  • Het verhogen van de binnen-groep koppeling ($\rho$) verhoogde de Jaccard-gelijkenis van expert-sets binnen vensters significant en verminderde het aantal gebruikte distincte experts per venster.
  • Het verhogen van de cross-groep oppositie ($\alpha$) verminderde de Jaccard-gelijkenis bij de grenzen van groepen significant, wat de bekwaamheid bevestigt om de eisen van groepen te ontkoppelen.
• Leren: Een geleerde scalaire controller (3 parameters) trainde stabiel maar toonde geen verbetering in validatie cross-entropy vergeleken met de onafhankelijke baseline. De geleerde koppelingssterktes waren laag, wat consistent is met de observatie van de paper dat marginale preservering de signalen voor optimalisatie beperkt wanneer alleen per-token loss wordt gebruikt.
• Load Variantie: De pilot observeerde minimale veranderingen in de geaggregeerde load coefficient of variation (CV), waarbij werd opgemerkt dat deze samenvattende metriek de conditionele variantie-bounds uit Proposition 1 niet direct test.

Betekenis en Claims
De paper claimt dat H-CGA een nieuwe vrijheidsgraad blootlegt in MoE-routering: de gezamenlijke afhankelijkheidsstructuur van de routeringsruis. De primaire betekenis is theoretisch en mechanistisch:

• Veiligheid: Het biedt een "veiligheidsgarantie" dat het wijzigen van cross-token afhankelijkheid de marginale routeringsgedrag van een token niet verandert, waardoor de geleerde voorkeuren van het basismodel intact blijven.
• Controle: Het biedt twee complementaire "knoppen" om de afruil tussen lokale expert-set coherentie en globale last-dispersie te beheren.
• Efficiëntie: Het maakt routerings-alleen adaptatie mogelijk waarbij het basismodel alleen in de forward direction wordt geëvalueerd, wat de computationele kosten van het updaten van expert-gewichten vermijdt.

Beperkingen en Bescheidenheid
De auteur is expliciet over de reikwijdte en beperkingen:

• Laag-Lokale Garantie: De invariantie-resultaten zijn conditioneel op de verborgen toestanden en logits die een specifieke routeringslaag binnenkomen. Ze garanderen geen end-to-end invariantie in een multi-layer netwerk, aangezien koppeling de downstream verborgen toestanden kan wijzigen.
• Geen Taak-Niveau Winsten: De pilot stelt niet vast dat routerings-alleen adaptatie de prestaties op downstream taken verbetert. De auteur merkt op dat marginale preservering impliceert dat doelstellingen die uitsluitend afhangen van single-token distributies de controller niet kunnen identificeren; nuttige signalen moeten voortkomen uit gezamenlijke effecten (bijv. aandacht-interacties), die echter zwak kunnen zijn.
• Variantie vs. Overloop: Hoewel de methode de load-variantie theoretisch karakteriseert, heeft de pilot de conditionele load-variantie of capaciteits-overstromingen (overflows) niet direct gemeten. De paper claimt niet het probleem van capaciteitsoverloop op te lossen, maar suggereert dat de oppositie-knop een mechanisme is om dit aan te pakken.
• Overhead: De methode vereist het genereren en coördineren van Gumbel-perturbaties voor alle experts, wat een computationele overhead kan introduceren vergeleken met standaard routering.

Samenvattend stelt de paper een wiskundig rigoureus kader vast voor het controleren van de cross-token afhankelijkheid in bevroren MoEs, waarbij bewezen wordt dat coherentie en last-dispersie gemanipuleerd kunnen worden zonder de per-token routeringswetten te wijzigen, terwijl er voorzichtig wordt gedaan over de directe praktische voordelen voor taakprestaties.