Mixture-of-Experts as Soft Clustering: A Dual Jacobian-PCA Spectral Geometry Perspective

Probleemstelling

Mixture-of-Experts (MoE) architecturen worden voornamelijk ontwikkeld voor efficiëntie en conditionele berekening, waardoor grote modelcapaciteiten mogelijk zijn tegen vaste rekenkosten. Echter, de invloed van het routing-mechanisme op de geometrie van de geleerde functies en de interne representaties blijft slecht gekarakteriseerd.
In dichte (dense) netwerken moet één globaal gedeelde parameterisatie heterogeen gedrag ondersteunen, wat kan leiden tot sterke lokale gevoeligheid (high sensitivity) en verstrengelde representaties. Het is onduidelijk hoe MoE-routing deze lokale kromming, gevoeligheid en de dimensionaliteit van expert-lokale representaties herstructureert, onafhankelijk van downstream-taakprestaties.

Methodologie

De auteurs analyseren MoE's via een geometrische lens, waarbij routing wordt geïnterpreteerd als een vorm van soft partitioning in overlappende "expert-lokale kaarten" (charts). Om dit te onderzoeken, introduceren ze een Dual Jacobian-PCA Spectral Geometry Probe in een gecontroleerde setting:

1. Experimentele Opzet:

   • Gebruik van een gecontroleerde MLP-MoE architectuur (in plaats van volledige transformers) om exacte Jacobian-berekeningen en stabiele covariantie-schattingen mogelijk te maken.
   • Vergelijking van drie configuraties met gelijke capaciteit: Dense, Top-k routing (spars), en Fully-soft routing.
   • Training op synthetische, isotrope Gaussische invoer om verstorende factoren (zoals taalstructuur) uit te sluiten.

2. Analytische Probes:

   • Jacobian Spectral Probe (Functiegeometrie): Berekening van het singuliere-waarden spectrum van de expert-lokale Jacobian ($J_e(x)$). Dit meet lokale gevoeligheid (via de grootste singuliere waarde $\sigma_1$) en de verdeling van gevoeligheid over richtingen. Ook wordt de kruis-expert Jacobian-alignatie (cosine-similariteit) gemeten om functional overlap te kwantificeren.
   • Weighted PCA Probe (Representatiegeometrie): Gewogen Principal Component Analysis op de verborgen staten die de MoE-laag binnenkomen, waarbij de routing-weights ($g_e(x)$) dienen als sample-weights. Dit onthult de effectieve rang (effective rank) en de spreiding van variantie binnen expert-lokale representaties.

Belangrijkste Bijdragen

1. Geometrisch Kader: MoE's worden geïnterpreteerd als soft partitioning van de functieruimte, waarbij experts gespecialiseerde, overlappende kaarten vormen.
2. Spectrale Vergelijking: Systematische analyse van hoe de scherpte van routing (Top-k vs. soft) lokale gevoeligheid en variantie-structuur beïnvloedt.
3. Mechanistische Probes: Introductie van een combinatie van Jacobian-analyse en gewogen PCA om expert-lokale structuur te kwantificeren, wat leidt tot drie testbare hypothesen voor grote modellen.

Resultaten

1. Jacobian-analyse (Functiegeometrie)

• Verminderde lokale gevoeligheid: MoE-experts vertonen kleinere leidende singuliere waarden en een snellere spectrale verval dan dichte baselines. Dit duidt op vlatere lokale kromming (flatter local mappings).
• Lage kruis-expert alignatie: De cosine-similariteit tussen de gemiddelde Jacobian-matrices van verschillende experts ligt dicht bij nul. Experts leren dus fundamenteel verschillende transformaties met weinig overlap, in plaats van geschaalde varianten van dezelfde kaart.
• Concentratie van gevoeligheid: Hoewel de absolute waarden lager zijn, is de gevoeligheid van MoE-experts geconcentreerd in minder dominante richtingen vergeleken met dichte modellen.

2. Representatiegeometrie (PCA)

• Hogere effectieve rang: In de synthetische setting spreiden MoE-experts de variantie over meer hoofdcomponenten dan dichte modellen. Dichte modellen concentreren variantie snel (lage rang), terwijl MoE's een bredere, hogere-rang structuur behouden.
• Invloed van Routing-Scherpte:
  • Top-k routing: Leidt tot meer geconcentreerde variantie en een lagere effectieve rang binnen experts (strakkere partitionering).
  • Fully-soft routing: Resulteert in bredere variantiestructuren en een vlotter PCA-spectrum (hoge rang).

3. Validatie op Natuurlijke Taal (WikiText)

Bij toepassing op een 3-laags transformer met WikiText-data:

• Bevestiging: De vermindering van kromming (lagere Jacobian-normen) blijft bestaan. Top-k routing zorgt voor nog lagere kruis-expert alignatie dan soft routing.
• Inversie van PCA: In tegenstelling tot synthetische data, vereist het dichte model meer componenten om 90% variantie te bereiken dan de MoE-experts. Dit suggereert dat MoE's tokens partitioneren in eenvoudigere, lager-dimensionale sub-manifolds op het taalmanifold, terwijl dichte lagen heterogene contexten moeten mixen.

Betekenis en Implicaties

De studie biedt een mechanistische basis om te begrijpen hoe expert-routing de geometrie van modellen herschikt. De resultaten ondersteunen het idee dat MoE's fungeren als soft clustering van de functieruimte.

De auteurs formuleren drie testbare hypothesen voor grote schaalmodellen (bijv. 70B+):

1. Geometrische "Elbow": Er bestaat een optimaal aantal experts gebaseerd op variantie-concentratie (diminishing returns).
2. Onderdrukking van Hallucinaties: De vlatere lokale mapping (lagere Jacobian-normen) zou kunnen leiden tot minder gevoeligheid voor infinitesimale perturbaties, wat hallucinaties in lange contexten kan verminderen.
3. Ensemble Betrouwbaarheid: Top-k routing biedt superieure ensemble-diversiteit door de near-orthogonaliteit van expert-Jacobians, wat functionele diversiteit maximaliseert.

Conclusie:
Routing scherpte reguleert niet alleen welke experts actief zijn, maar ook hoe variantie binnen experts is georganiseerd. Deze inzichten bieden actievere ontwerpcriteria voor het balanceren van routing-scherpte en representatie-rang in toekomstige lange-context transformers.