ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization

Titel

ERMoE: Eigen-Reparameterized Mixture-of-Experts voor Stabiele Routing en Interpreteerbare Specialisatie

1. Het Probleem

Mixture-of-Experts (MoE) modellen zijn ontwikkeld om de capaciteit van neurale netwerken te vergroten zonder de rekentijd per token lineair te laten toenemen, door slechts een subset van "experts" (subnetwerken) per invoer te activeren. Echter, bestaande MoE-architecturen kampen met fundamentele structurele tekortkomingen:

• Misalignement tussen Router en Experts: De router (die bepaalt welke expert wordt gekozen) leert vaak los van de interne structuur van de experts. Dit leidt tot onstabiele routing, waarbij tokens naar suboptimale experts worden gestuurd, wat de specialisatie van experts ondermijnt en de voorspellingsnauwkeurigheid verlaagt.
• Onbalans in Belasting (Load Imbalance): Zonder correctie neigen routers ertoe om een "rijk worden rijker"-dynamiek te creëren, waarbij een paar experts overbelast raken (stragglers) terwijl andere onderbenut blijven.
• Nadelige Effecten van Hulpverliezen: Om deze onbalans op te lossen, gebruiken bestaande methoden vaak een auxiliary load-balancing loss (LBL). Hoewel dit de verdeling verbetert, introduceert deze loss interferentiegradiënten die de taakprestaties verzwakken en kan het leiden tot een te uniforme routing, waardoor de gewenste specialisatie van experts verdwijnt.

2. Methodologie: ERMoE

De auteurs stellen ERMoE (Eigen-Reparameterized Mixture-of-Experts) voor, een nieuwe architectuur die de routingmechaniek fundamenteel verandert door deze te koppelen aan de representatieruimte van de experts.

Kerncomponenten:

• Eigenbasis Reparameterisatie: In plaats van standaard gewichten, wordt de weight-matrix van elke expert $W^{(e)}$ herschreven in een orthonormale eigenbasis:
  $$W^{(e)} = U^{(e)} \text{diag}(s^{(e)}) V^{(e)\top}$$
  Hierbij zijn $U^{(e)}$ en $V^{(e)}$ orthonormale bases en $s^{(e)}$ trainbare coëfficiënten. Een lichte orthogonaliteitsstraf zorgt ervoor dat deze bases decorreleerde richtingen behouden.
• Eigenbasis Score (Routing): De router gebruikt geen geleerde logits meer. In plaats daarvan wordt de score berekend op basis van de cosinus-ähnlichkeit tussen de token-features (en hun context) en de eigenbasis van de expert.
  • De token en de context worden geprojecteerd op de basis van de expert.
  • De score is de cosinus tussen deze twee projecties.
  • Een hoge score betekent dat de token goed past binnen het representatieruimte van die specifieke expert.
• Gedrempelde Top-k Routing:
  • Alleen experts met een score boven een drempelwaarde $T$ komen in aanmerking.
  • Uit deze set worden de top-$k$ experts geselecteerd.
  • Als er minder dan $k$ experts boven de drempel zitten, wordt er een fallback-mechanisme gebruikt.
  • Dit elimineert de noodzaak voor een externe load-balancing loss, omdat de geometrische alignatie en de drempel zorgen voor een stabiele en evenwichtige verdeling.

Toepassing op 3D Neurobeeldvorming (ERMoE-ba):
Voor het voorspellen van hersenleeftijd uit 3D MRI-scans (T1-gewogen) wordt een variant ERMoE-ba ontwikkeld. Deze gebruikt:

• Een 3D ViT-achtergrond.
• Specifieke "region experts" (voorgedefinieerd op Witte Stof, Grijze Stof, en Cerebrospinaal Vloeistof) en "free experts".
• Dit zorgt voor anatomisch interpreteerbare specialisatie.

3. Belangrijkste Bijdragen

1. Nieuwe Architectuur: Introductie van ERMoE, waarbij experts worden geparameteriseerd in een orthonormale eigenbasis en routing plaatsvindt via content-bewuste alignatie-scores in plaats van vrije logits.
2. Oplossing voor Structurele Defecten: Het ontwerp lost het probleem van router-expert misalignement en onbalans op zonder de noodzaak van auxiliary load-balancing losses, waardoor interferentiegradiënten worden vermeden en specialisatie wordt behouden.
3. 3D Neuroimaging Variant: Ontwikkeling van ERMoE-ba voor 3D hersenbeeldvorming, wat leidt tot verbeterde hersenleeftijdvoorspelling en anatomisch interpreteerbare expert-specialisaties.

4. Resultaten

De methode is getest op zowel natuurlijke afbeeldingen als medische 3D-beelden:

• Beeldclassificatie (ImageNet):
  • ERMoE bereikt 88.03% Top-1 nauwkeurigheid op ImageNet-1K, wat een verbetering is ten opzichte van sterke baselines zoals V-MoE (87.41%).
  • Het presteert beter zonder gebruik van load-balancing losses.
• Few-Shot Learning (CIFAR & Tiny-ImageNet):
  • ERMoE toont een overtuigende superioriteit in few-shot scenario's (bijv. +4% op CIFAR-100 5-shot t.o.v. V-MoE), wat wijst op betere, minder redundante features.
• Multimodaal Retrieval (COCO & Flickr30K):
  • Als visuele encoder in een CLIP-achtig framework (CLIP-ERMoE) behaalt het state-of-the-art resultaten voor Image-to-Text retrieval (65.4% R@1 op COCO).
• 3D Hersenleeftijdvoorspelling (ADNI):
  • ERMoE-ba verbetert de voorspelling van hersenleeftijd met meer dan 7% (verlaging van MAE van ~3.13 naar 2.31 jaar) ten opzichte van bestaande 3D CNN's en Transformers.
  • De expert-specialisatie blijkt anatomisch betekenisvol: specifieke experts specialiseren zich in Witte Stof, Grijze Stof of CSF, wat zichtbaar is in de routingpatronen tijdens training.

5. Betekenis en Impact

Dit paper biedt een paradigmaverschuiving in het ontwerp van MoE-modellen:

• Stabiliteit zonder Trade-offs: Het bewijst dat stabiele routing en evenwichtige expertgebruik mogelijk zijn zonder de schadelijke bijwerkingen van load-balancing losses.
• Interpreteerbaarheid: Door routing te koppelen aan de intrinsieke representatieruimte van de experts, worden de beslissingen van het model transparanter. In de medische context betekent dit dat men kan zien welke anatomische structuren bijdragen aan de voorspelling.
• Efficiëntie: ERMoE bereikt betere prestaties met minder trainbare parameters en lagere FLOPs vergeleken met vergelijkbare MoE-architecturen, wat het zeer geschikt maakt voor schaalbare en klinische toepassingen.

Kortom, ERMoE lost het fundamentele probleem van instabiele routing in MoE-modellen op door de router te "verankeren" in de geometrie van de experts zelf, wat leidt tot robuustere, nauwkeurigere en interpreteerbaarder modellen.