Teacher-Guided Routing for Sparse Vision Mixture-of-Experts

1. Het Probleem: Instabiliteit bij Sparse MoE Training

Sparse Mixture-of-Experts (MoE) modellen zijn populair geworden omdat ze de rekenkracht kunnen schalen zonder de inferentiekosten evenredig te laten toenemen. Ze doen dit door voor elke input slechts een klein subset van "experts" (neuronale netwerken) te activeren via een router.

Hoewel effectief, ondervinden deze modellen aanzienlijke trainingsschalen:

• Gradient Blocking (Blokkering van Gradiënten): Omdat de router alleen gradiënten ontvangt van de experts die in de forward pass zijn geselecteerd, krijgt hij geen informatieve feedback van de niet-geselecteerde routes. Dit leidt tot een gebrek aan data-afhankelijke supervisie.
• Routing Fluctuatie: Door het gebrek aan globale supervisie, kan de toewijzing van experts aan specifieke inputs tijdens de training sterk fluctueren. Dezelfde sample kan in verschillende iteraties naar verschillende experts worden gestuurd.
• Gevolgen: Deze instabiliteit verhindert dat experts zich stabiel specialiseren, leidt tot suboptimale prestaties en maakt training moeilijk, vooral bij zeer sparse configuraties (veel experts, weinig actieve experts per token). Bestaande methoden zoals load-balancing losses zijn vaak niet voldoende om deze fluctuaties te onderdrukken.

2. Methodologie: TGR-MoE (Teacher-Guided Routing)

De auteurs stellen TGR-MoE voor, een framework dat de training stabiliseert door gebruik te maken van een voorgeladen, dichte (dense) teacher-model als externe gids.

Kerncomponenten:

• Teacher Router Constructie:
  • Er wordt een lichtgewicht "teacher router" toegevoegd aan een bevroren (frozen) dense teacher backbone (bijv. DeiT-III).
  • Deze teacher router wordt getraind met een combinatie van een load-balancing loss (voor gelijke verdeling van experts) en een entropy loss (voor vertrouwen in toewijzingen, zodat de router niet volledig uniform wordt).
  • Het doel is niet het minimaliseren van de taakfout direct, maar het genereren van een stabiele, gebalanceerde en semantisch zinvolle verdeling van expert-toewijzingen.
• Student Training met Distillatie:
  • De student (het MoE-model) leert zijn router te imiteren door de outputverdeling van de teacher router te distilleren.
  • De totale loss voor de student bestaat uit de standaard taakloss (bijv. classificatie) plus een Kullback-Leibler (KL) divergentie term die de student-router probeert dicht bij de teacher-router te houden.
  • De teacher backbone blijft bevroren; alleen de teacher-router en de student worden geoptimaliseerd.
• Voordeel: De teacher biedt een "globale prior" die informatieve feedback geeft aan de student-router, zelfs voor experts die in de huidige stap niet zijn geselecteerd. Dit doorbreekt het probleem van gradient blocking zonder extra kosten tijdens de inferentie (testtijd).

3. Belangrijkste Bijdragen

1. Nieuwe Training Paradigma: Introductie van TGR-MoE, een methode die router-stabiliteit bereikt door supervisie af te leiden van een dense teacher, in plaats van alleen te vertrouwen op de sparse gradiënten van het eigen model.
2. Oplossing voor Gradient Blocking: Het bieden van een externe prior die de router helpt om consistente expert-selectie te leren, zelfs in de vroege stadia van training wanneer specialisatie nog niet is opgetreden.
3. Scalabiliteit: Het aantonen dat deze methode de schaalbaarheid van MoE-modellen verbetert, zelfs bij een zeer groot aantal experts, waar traditionele methoden vaak verzadigen.
4. Analyse van Routing Consistentie: Uitgebreide analyse die aantoont dat TGR-MoE routing-fluctuaties aanzienlijk reduceert en sneller convergeert naar stabiele expert-toewijzingen dan standaard VMoE.

4. Resultaten

De auteurs hebben TGR-MoE geëvalueerd op ImageNet-1K (pre-training) en downstream taken zoals CIFAR-10/100 en Oxford-IIIT Pets.

• Prestatieverbetering: TGR-MoE overtreft consistent standaard VMoE en andere geavanceerde varianten (zoals Expert-Choice Routing, SoftMoE, en z-loss) op alle modelgroottes (Tiny, Small, Base).
  • Voorbeeld: Op ImageNet-1K verbetert TGR-MoE de Top-1 nauwkeurigheid van het Tiny-model van 77,85% (VMoE) naar 78,78%.
• Stabiliteit bij Fine-tuning: Het behouden van TGR-MoE tijdens fine-tuning leidt tot betere transferprestaties dan het alleen gebruiken van de voorgeladen weights. Dit suggereert dat de routing-kennis actief moet worden onderhouden.
• Effect van Aantal Experts: Terwijl de prestaties van standaard VMoE afvlakken bij een groot aantal experts (bijv. >64), blijft TGR-MoE lineair verbeteren. Bij 128 experts behaalde TGR-MoE een verbetering van +1,21% ten opzichte van VMoE.
• Routing Consistentie: Analyse toont aan dat TGR-MoE binnen de eerste 50 epochs al >70% overeenstemming bereikt met de finale toewijzingen, terwijl VMoE nog steeds grote fluctuaties vertoont.

5. Betekenis en Conclusie

TGR-MoE biedt een eenvoudige maar krachtige oplossing voor een fundamenteel probleem in het trainen van sparse MoE-modellen: de instabiliteit van de router.

• Efficiëntie: De methode vereist geen extra rekenkracht tijdens de inferentie, omdat de teacher alleen tijdens de training wordt gebruikt.
• Algemeen Inzicht: Het paper onderstreept dat effectieve MoE-routing supervisie nodig heeft die verder gaat dan alleen task-driven gradiënten. Het integreren van "teacher-informed priors" helpt experts zich sneller en consistenter te specialiseren.
• Toekomstperspectief: Hoewel het nu gericht is op beeldherkenning, biedt de methode een veelbelovende richting voor het uitbreiden van MoE-architecturen naar andere domeinen (zoals NLP) en grotere datasets, waar stabiliteit en schaalbaarheid cruciaal zijn.

Kortom, TGR-MoE maakt sparse Vision MoE-modellen stabieler, schaalbaarder en nauwkeuriger door slim gebruik te maken van kennis van een voorgeladen dichte teacher.