Technische Samenvatting: Expert Merging in Sparse Mixture of Experts met Nash-onderhandeling

Probleemstelling

Sparse Mixture of Experts (SMoE) architecturen zijn een hoeksteen geworden voor het efficiënt schalen van neurale netwerken door slechts een subset van expert-modules per input te activeren. Terwijl aanzienlijk onderzoek zich heeft gericht op het verbeteren van routing-mechanismen (bijv. load balancing, stabiliteit), blijft het complementaire probleem van expert merging — het combineren van alle expert-parameters tot een verenigd model voor deployment of inferentie — onderbelicht.

Bestaande merging-strategieën vertrouwen vaak op heuristische weegschema's, zoals eenvoudige input-onafhankelijke middeling of top-$k$ aggregatie. Deze methoden negeren vaak de complexe coöperatieve en competitieve dynamiek tussen experts. Bov furthermore hebben recente benaderingen zoals Curvature-aware Merging of Experts (CAMEx) en de dynamische variant EP-CAMEx, die natuurlijke gradiënten gebruiken om experts uit te lijnen, beperkingen getoond. Specifiek presteert EP-CAMEx vaak minder goed dan zijn statische tegenhanger, waarschijnlijk door onvoldoende coördinatie tussen de bijdragen van experts en een trage convergentie in propagatie. De kernuitdaging is om expert merging niet louter te interpreteren als een lineair gemiddelde, maar als een gestructureerde interactie die de unieke bijdragen van diverse experts balanceert.

Methodologie: NAMEx

De auteurs stellen NAMEx (Nash Merging of Experts) voor, een framework dat expert merging herinterpreteert door de lens van speltheorie, specifiek gebruikmakend van de Nash-onderhandelingsoplossing (NBS).

1. Expert Merging als een Onderhandelingsspel

De auteurs kaderen het merging-proces als een coöperatief-competitief spel tussen experts.

• Spelers: Elke expert $i$ in de SMoE-laag.
• Nutfunctie: Gedefinieerd als $u_i(\Delta E) = \tau_i^\top \Delta E$, waarbij $\tau_i = E_i - E_m$ de "domein-vector" is die de afwijking van de $i$-de expert van de basis-expert $E_m$ representeert. Deze vector legt de unieke bijdrage van de expert vast.
• Disagreement Point: Gesteld op 0, wat overeenkomt met geen update aan de basis-expert.
• Doelstelling: Vind een update-richting $\Delta E$ binnen een bal $B_\epsilon$ die het product van individuele winsten ten opzichte van het disagreement point maximaliseert:
  $$\Delta E^* = \arg\max_{\Delta E \in B_\epsilon} \prod_{i=1}^N (\Delta E^\top \tau_i)$$
  Onder de aanname van lineair onafhankelijke domein-vectoren, wordt de oplossing afgeleid als een gewogen som van domein-vectoren:
  $$\Delta E = \sum_{i=1}^N \alpha_i \tau_i$$
  waarbij de gewichten $\alpha$ voldoen aan de Nash-onderhandelingsvergelijking: $G^\top G \alpha = 1/\alpha$, waarbij $G$ de matrix van domein-vectoren is en $1/\alpha$ de elementgewijze reciproken zijn.

2. Integratie met Expert Propagatie

NAMEx wordt geïntegreerd in het EP-CAMEx-framework. In plaats van een vaste gradiëntafdaling-stap voor de basis-expert propagatie te gebruiken, berekent NAMEx de optimale update-richting $\Delta E^{(l)}$ op elke laag $l$ met behulp van de Nash-oplossing. De basis-expert $E_m$ wordt bijgewerkt als:
$$E_m^{(l+1)} = E_m^{(l)} + \gamma \sum_{i=1}^N \alpha_i^{(l)} \tau_i^{(l)}$$
waarbij $\gamma$ de stapgrootte is.

3. Complexe Momentum

Om de trage convergentie die in EP-CAMEx werd waargenomen aan te pakken, integreren de auteurs complexe momentum in het NAMEx-framework (NAMEx-Momentum).

• Een complexe momentum-buffer $\mu^{(j)} \in \mathbb{C}^d$ accumuleert directionele updates.
• De update-regel omvat een complexe momentum-coëfficiënt $\beta \in \mathbb{C}$:
  $$\mu^{(j+1)} = \beta \mu^{(j)} + \Delta E^{(j)}$$
  $$E_m^{(j+1)} = E_m^{(j)} + \Re(\gamma \mu^{(j+1)})$$
• Theoretische analyse (Propositie 3.5 en Stelling C.3) biedt convergentiegaranties, waarbij wordt aangetoond dat onder milde voorwaarden de spectrale radius van de systeemmatrix kleiner dan 1 kan worden gehouden, wat lineaire convergentie verzekert.

Belangrijkste Bijdragen

1. NAMEx Framework: Een nieuwe expert merging-methode die het Nash-onderhandelingsoptimalisatie-framework integreert in expert-propagatie, waarbij expert-domein-vectoren als nutfuncties worden gebruikt om eerlijke en efficiënte merging-coëfficiënten af te leiden.
2. Integratie van Complexe Momentum: De integratie van complexe momentum in NAMEx om convergentie te versnellen en stabiliteit te verbeteren, met name in scenario's met adversariële of conflicterende expert-dynamieken. De auteurs leveren theoretische bewijzen voor de convergentiesnelheid van NAMEx-Momentum.
3. Schaalbaarheid en Empirische Validatie: Demonstratie van de effectiviteit van NAMEx over diverse taken (taalmodellering, tekstclassificatie, beeldclassificatie) en de schaalbaarheid naar grootschalige systemen (Qwen1.5-MoE 14B, DeepSeek-MoE 16B).

Experimentele Resultaten

De auteurs evalueerden NAMEx tegenover baselines (SMoE, CAMEx, EP-CAMEx) over meerdere benchmarks:

• Taalmodellering (WikiText-103): NAMEx-Full-Mom behaalde de laagste validatie- en test-perplexiteit voor zowel kleine als middelgrote schaal pre-training, waarmee het alle baselines overtrof.
• Tekstclassificatie (GLUE): Op de GLUE-benchmark met T5-Base behaalde NAMEx-Full-Mom de beste resultaten over alle taken (SST-2, MRPC, CoLA, etc.), wat consistente verbeteringen laat zien ten opzichte van SMoE en curvature-aware methoden.
• Beeldclassificatie (ImageNet-1k & Robuustheid): Op Swin-MoE presteerde NAMEx-Mom beter dan alle baselines in Top-1 accuraatheid. Cruciaal is dat het superieure robuustheid toonde op gecorrumpeerde datasets (ImageNet-A, ImageNet-O, ImageNet-R), waarbij momentum-varianten de sterkste winsten lieten zien.
• Grootschalige Modellen: Wanneer toegepast op DeepSeek-MoE (16B) en Qwen1.5-MoE (14B), presteerde NAMEx-Full consequent beter dan baselines en EP-CAMEx in zowel zero-shot als fine-tuned (SmolTalk) settings op MMLU, GSM8K en ARC benchmarks.
• Ablatie-studies:
  • Momentum: Complexe momentum verbeterde de convergentiesnelheid en prestaties aanzienlijk vergeleken met real-valued momentum.
  • Iteraties: Het oplossen van het Nash-systeem met 2 iteraties per laag bleek efficiënt en effectief, met afnemende meeropbrengsten bij hogere iteratie-aantallen.
  • Disagreement Point: Het gebruik van 0 als het disagreement point bleek robuust en stabiel vergeleken met het gebruik van het gemiddelde.

Betekenis en Claims

Het artikel beweert dat NAMEx een geprincipiëerd alternatief biedt voor heuristische expert merging. Door expert-interacties te modelleren als een onderhandelingsspel, zorgt NAMEx voor een Pareto-efficiënt merging-proces dat rekening houdt met zowel coöperatie als competitie tussen experts. De integratie van complexe momentum adresseert convergentie-bottlenecks in dynamische merging-scenario's.

De auteurs benadrukken dat NAMEx:

• Prestaties Verbetert: Consistent beter presteert dan bestaande merging-strategieën over visuele en talige taken.
• Robuustheid Versterkt: Superieure generalisatie vertoont onder data-corruptie en distributieverschuivingen.
• Effectief Schalt: Bewijst levensvatbaar te zijn voor grootschalige MoE-architecturen (tot 16B parameters) zonder de computationele efficiëntie aanzienlijk in gevaar te brengen, aangezien de inferentiekosten ongewijzigd blijven en de trainingsoverhead beheersbaar is.
• Theoretische Onderbouwing Biedt: Convergentiegaranties levert voor de momentum-verbeterde variant, waarmee het verder gaat dan louter empirische heuristieken.

Het werk suggereert dat het bekijken van expert merging door de lens van speltheorie nieuwe wegen opent voor het optimaliseren van modulaire deep learning-systemen, met potentiële toepassingen in federated learning en transfer learning waar eerlijke expert-combinatie cruciaal is.