💻 computer science

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

Dit artikel introduceert Selective Sinkhorn Routing (SSR), een lichtgewicht mechanisme dat de toewijzing van tokens aan experts frameert als een geconstrueerd optimaal transportprobleem om een gebalanceerde expertbenutting en verbeterde modelprestaties te bereiken zonder afhankelijk te zijn van hulpmatige balanseringsverliezen of complexe trainbare componenten.

Oorspronkelijke auteurs: Duc Anh Nguyen, Huu Binh Ta, Nhuan Le Duc, Tan Minh Nguyen, Toan Tran

Gepubliceerd 2026-06-05

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Duc Anh Nguyen, Huu Binh Ta, Nhuan Le Duc, Tan Minh Nguyen, Toan Tran

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorm, hoogtechnologisch callcenter runt. Je hebt duizenden binnenkomende gesprekken (tokens) en een team van 100 gespecialiseerde agenten (experts). Je doel is om elk gesprek naar de beste agent te leiden om het probleem snel op te lossen.

In een standaardopstelling gebruik je een eenvoudige regel: "Stuur het gesprek naar de agent die op dit moment het meest gekwalificeerd lijkt." Dit is een soort Softmax-router. Het probleem? Dezelfde paar "super-agenten" krijgen alle gesprekken, terwijl de andere 90 agenten niets doen. Dit wordt routing collapse genoemd. Het callcenter wordt hierdoor inefficiënt omdat je niet je volledige team gebruikt.

Om dit op te lossen, probeerden eerdere methoden de balans te forceren. Ze voegden een "manager" toe die constant het systeem bestookte met: "Hé, Agent 5 heeft al een uur geen gesprek gehad, stuur er een naar hem toe!" of "Agent 1 is te druk, stop met het sturen van gesprekken naar hem!" Dit zijn de auxiliary losses die in het paper worden genoemd. Hoewel ze helpen, zijn ze lomp, vergen ze extra rekenkracht van de computer en verwarren ze het systeem soms over wat het eigenlijk probeert te leren.

Het Nieuwe Idee: De "Perfect Gebalanceerde" Toewijzing

De auteurs van dit paper stellen een slimmere manier voor om gesprekken toe te wijzen, gebruikmakend van een wiskundig concept genaamd Optimal Transport (specifiek het Sinkhorn-algoritme).

Zie dit niet als een manager die agenten lastigvalt, maar als een perfect gechoreografeerde dans.

Het Doel: Elke agent moet over een bepaalde tijd precies hetzelfde aantal gesprekken krijgen, en elke beller moet worden gekoppeld aan een agent die goed is in zijn werk.
De Methode: In plaats van simpelweg de "beste" agent voor elk gesprek te kiezen, berekent het systeem een globale kaart. Het kijkt naar alle gesprekken en alle agenten tegelijk en bepaalt de meest efficiënte manier om de werkdruk te verdelen, zodat niemand overbelast is en niemand zich verveelt.

Het Probleem met de "Perfecte Dans"

Er is een addertje onder het gras. Als je deze perfecte balans afdwingt bij elk individueel gesprek dat binnenkomt, raakt het systeem in de war. Het zou een gesprek over "programmeren" naar een agent kunnen sturen die geweldig is in "koken", puur om de cijfers gelijk te houden. Dit schaadt de prestaties.

De doorbraak van het paper is Selective Sinkhorn Routing (SSR).

Hoe SSR Werkt: De "Hybride" Strategie

In plaats van de complexe "perfecte dans" voor elk gesprek te gebruiken, gebruikt SSR een slimme mix:

De meeste tijd (99%+): Het gebruikt de standaard, snelle methode (Softmax) om gesprekken te routeren. Dit laat het systeem leren waar agenten echt goed in zijn.
Zelden (0,1% tot 1% van de tijd): Het pauzeert en voert de "perfecte dans" uit (het Sinkhorn-algoritme).
- Waarom? Deze kleine hoeveelheid "perfecte balans" werkt als een zachte duw. Het herinnert het systeem eraan: "Vergeet de andere agenten niet!" zonder dat het bij elk gesprek een slechte match afdwingt.
- Het Resultaat: Het systeem leert zichzelf op natuurlijke wijze te balanceren, zonder dat er een plagerige manager (auxiliary loss) nodig is of een enorme hoeveelheid extra rekenkracht.

Het Geheime Ingrediënt: Een Beetje Ruis Toevoegen

Het paper suggereert ook om een klein beetje willekeurige ruis (zoals statische elektriciteit op een radio) toe te voegen aan het besluitvormingsproces tijdens de training.

Analogie: Stel je voor dat de agenten een beetje dronken zijn of dat de telefoonlijnen een beetje krakend zijn. Het systeem kan niet 100% zeker zijn wie de "beste" agent is, dus probeert het een paar verschillende mensen.
Voordeel: Dit voorkomt dat het systeem in een sleur terechtkomt waarbij het altijd dezelfde top 3 agenten kiest. Het dwingt het systeem om te verkennen en te ontdekken dat andere agenten ook heel goed zijn.
Belangrijke Opmerking: Het paper benadrukt dat je deze ruis uitzet wanneer het systeem daadwerkelijk werkt (inference). Je wilt niet dat je callcenter willekeurig is wanneer een klant wacht; je wilt dat het snel en voorspelbaar is.

Wat Ze Hebben Gevonden

De auteurs hebben dit getest op twee hoofdtaken:

Taalmodellering (Schrijven): Ze testten het op datasets zoals WikiText-103.
- Resultaat: Hun methode (SSR) schreef betere teksten (lagere "perplexity", wat een score is voor hoe verward de AI is) dan eerdere methoden.
- Snelheid: Het was veel sneller te trainen omdat het geen zware "plagerige" losses nodig had. Het gebruikte de complexe wiskunde slechts voor een fractie van de tijd.
Beeldclassificatie (Visie): Ze testten het op ImageNet (het herkennen van afbeeldingen).
- Resultaat: Hun methode herkende afbeeldingen nauwkeuriger en kon beter omgaan met vreemde, beschadigde of "adversariële" afbeeldingen (afbeeldingen die ontworpen zijn om AI te misleiden).

De Kernboodschap

Het paper beweert dat Selective Sinkhorn Routing een lichtgewicht, efficiënte manier is om het "routing collapse"-probleem in Sparse Mixture of Experts (SMoE) modellen op te lossen.

De Oude Manier: Gebruik zware, complexe wiskunde of bestraffende strafpunten om balans af te dwingen. (Traag, soms instabiel).
De Nieuwe Manier (SSR): Gebruik de complexe wiskunde slechts zelden om het systeem te sturen, en voeg een beetje willekeur toe om het interessant te houden tijdens de training.
Resultaat: Je krijgt een slimmere, gebalanceerdere AI die sneller traint en beter werkt, zonder de extra ballast.

Cruciaal is dat het paper benadrukt dat de "perfecte balans" en de "ruis" alleen bedoeld zijn voor de training. Wanneer het model daadwerkelijk in de echte wereld wordt gebruikt, schakelt het terug naar een standaard, snelle en deterministische modus. Dit zorgt ervoor dat het eindproduct zowel van hoge kwaliteit als efficiënt is.

Technische Samenvatting: Selective Sinkhorn Routing voor Verbeterde Sparse Mixture of Experts

Probleemstelling

Sparse Mixture-of-Experts (SMoE) modellen bieden een schaalbare aanpak om de modelcapaciteit te vergroten terwijl de computationele efficiëntie behouden blijft door slechts een subset van experts per token te activeren. Echter, conventionele SMoE-modellen die vertrouwen op Softmax-gebaseerde routing lijden vaak aan "routing collapse", waarbij een kleine subset van experts overmatig wordt gebruikt terwijl andere onderbenut blijven.

Bestaande mitigatiestrategieën vertrouwen doorgaans op:

Auxiliary Objectives: Load-balancing losses of z-losses om ongelijke expert-utilisatie te bestraffen. Deze kunnen leiden tot objectief-misalignment, trainingsinstabiliteit en verhoogde complexiteit.
Noisy Gating: Het injecteren van ruis om exploratie te stimuleren, wat extra trainbare parameters of hyperparameters toevoegt.
Sinkhorn-gebaseerde Routing: Eerdere methoden die het Sinkhorn-algoritme voor optimal transport (OT) gebruiken, gebruiken de transportkaart vaak uitsluitend voor expertselectie (top-k) in plaats van gewichtsassignatie. Bovendien kunnen ze de routingflexibiliteit verminderen omdat de gating-matrix niet direct via gradiëntgebaseerd leren wordt geoptimaliseerd op dezelfde manier als Softmax-methoden, of ze veroorzaken aanzienlijke trainingsoverhead door de iteratieve aard van Sinkhorn-updates.

Methodologie

Het artikel stelt Selective Sinkhorn Routing (SSR) voor, een lichtgewicht mechanisme dat de toewijzing van token-naar-expert herformuleert als een entropie-geregulariseerd optimal transport probleem.

Kernformulering

De auteurs modelleren de toewijzing van $m$ tokens aan $n$ experts als een maximum-cost optimal transport probleem met entropische regularisatie. Het doel is om een transportplan $\hat{\Pi}$ te vinden dat compatibiliteit maximaliseert terwijl aan specifieke restricties wordt voldaan:

Positiviteit: $\Pi > 0$ .
Rijsom: Elk token routeert zijn volledige massa naar experts ( $\Pi \mathbf{1}_n = \mathbf{1}_m$ ).
Kolomsom: Elke expert ontvangt een gelijke verwachte totale belasting ( $\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$ ).

De kostenmatrix $C$ wordt afgeleid van de gating scores $S$ . De auteurs stellen twee varianten voor:

Lineaire Kosten: $C = S$ .
Softmax Kosten: $C_{i,:} = \text{softmax}(S_{i,:})$ , gebruikt om numerieke overflow tijdens Sinkhorn-iteraties te voorkomen.

In tegen tegenstelling tot eerdere OT-gebaseerde methoden die de transportkaart alleen voor selectie gebruiken, gebruikt SSR de waarden in $\hat{\Pi}$ direct om routinggewichten te berekenen. Voor elk token worden de top- $k$ experts geselecteerd op basis van de hoogste entries in de overeenkomstige rij van $\hat{\Pi}$ , en de gewichten worden genormaliseerd vanuit deze entries.

Selectieve Strategie en Ruisinjectie

Een belangrijke uitdaging die is geïdentificeerd, is dat Sinkhorn-routing de gating-gewichtmatrix $W_g$ niet bijwerkt omdat $W_g$ ontkoppeld is van het OT-doel. Om dit aan te pakken, gebruikt SSR een hybride trainingsstrategie:

Selectieve Toepassing: Tijdens de training gebruikt elke MoE-block met een kleine waarschijnlijkheid $p$ (bijv. 0,001) Sinkhorn-routing en standaard Softmax-gating in de overige gevallen. Dit stelt Softmax-gating in staat om $W_g$ effectief bij te werken, terwijl Sinkhorn-routing gebalanceerde expert-utilisatie-signalen levert.
Ruisinjectie: Om expert-collapse te voorkomen en exploratie te stimuleren, wordt Gaussische ruis toegevoegd aan de kostenmatrix ( $\tilde{C} = C + \alpha_{noise} \cdot \epsilon$ ). Theoretische analyse (Propositie 4.3) toont aan dat dit ervoor zorgt dat elke expert een niet-nul kans heeft op selectie.

Gedrag tijdens Inferentie

Het artikel betoogt dat het afdwingen van gebalanceerde restricties (zoals gelijke kolomsommen) tijdens inferentie op een enkele input contraproductief is, omdat het uniforme routing afdwingt ongeacht de specifieke compatibiliteitsscores van de input. Daarom worden zowel Sinkhorn-routing als ruisinjectie uitgeschakeld tijdens inferentie, waarbij wordt teruggevallen op deterministische Softmax-routing om consistente, input-afhankelijke voorspellingen te garanderen.

Belangrijkste Bijdragen

Nieuw Routing-framework: Een methode die entropie-geregulariseerd optimal transport integreert met stochastische ruisinjectie om gebalanceerde expert-utilisatie te bevorderen zonder hulp-balanceringsverliezen.
Theoretische Inzichten: Bewijzen die aantonen dat Sinkhorn-gebaseerde routing en ruisinjectie helpen bij de training door exploratie en balans te stimuleren, maar uitgeschakeld moeten worden tijdens inferentie om verstoorde toewijzingen te voorkomen.
Efficiëntie en Prestaties: Uitgebreide evaluaties tonen aan dat SSR de trainings-efficiëntie, nauwkeurigheid en robuustheid vergeleken met state-of-the-art baselines verbetert.

Experimentele Resultaten

De auteurs evalueerden SSR op taalmodellering (WikiText-103, Enwik-8) en beeldclassificatie (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R).

Taalmodellering:
- Op WikiText-103 presteerden SSR-varianten consequent beter dan Vanilla SMoE en baselines die gebruikmaken van load-balancing loss, z-loss of ruis.
- SSR-L w/ noise bereikte een test perplexiteit (PPL) van 34,367, een reductie van 1,183 ten opzichte van Vanilla SMoE (35,550), waarmee het de eerstvolgende beste baseline met een aanzienlijke marge verslaat.
- Cruciaal is dat SSR minimale trainingsoverhead veroorzaakte (0,33% – 0,65%), terwijl een volledige Sinkhorn-gebaseerde SMoE 72,47% overhead veroorzaakte.
- Onder Momentum-instellingen behield SSR de karakteristieke norm-evolutie van Vanilla SMoE terwijl de prestaties werden verbeterd, terwijl andere balanceringsmethoden de prestaties verslechterden.
Byte-level Modellering (Enwik-8):
- SSR-S w/ noise bereikte de beste test Bits-per-Character (BPC) van 1,128, een verbetering van 0,010 BPC ten opzichte van Vanilla SMoE.
- Het was 2,28x sneller dan volledige Sinkhorn-gebaseerde SMoE.
Visuele Taken:
- Op ImageNet-1K bereikte SSR-L w/ noise een Top-1 nauwkeurigheid van 77,420, een winst van +2,368 ten opzichte van Vanilla SMoE.
- Het vertoonde ook sterke robuustheid op ImageNet-O en ImageNet-R.
Ablatie-studies:
- De waarschijnlijkheid $p$ van het toepassen van Sinkhorn-routing is cruciaal; zeer kleine waarden (bijv. $10^{-3}$ of $10^{-4}$ ) waren voldoende om balancerende signalen te leveren zonder de router te veel te beperken.
- Ruisinjectie verbeterde de prestaties consistent over verschillende instellingen.
- Het uitschakelen van balanceringsmechanismen tijdens inferentie leverde de beste resultaten op, wat de theoretische claims bevestigt.

Betekenis en Claims

Het artikel claimt dat Selective Sinkhorn Routing een praktische en effectieve oplossing biedt voor SMoE-ontwerp. Door complexe hulp-verliezen te vervangen door efficiënte, intermitterende Sinkhorn-gebaseerde optimalisatie, bereikt SSR:

Verbeterde Trainingsstabiliteit: Door ingebouwde expert-balans zonder objectief-misalignment.
Hogere Nauwkeurigheid: Gedemonstreerd over tekst- en visuele benchmarks.
Robuustheid: Verbeterde prestaties tegen inputcorruptie en adversariële voorbeelden.
Efficiëntie: Minimale computationele overhead vergeleken met zowel hulp-verlies methoden als volledige Sinkhorn-routing.

De auteurs benadrukken dat de methode lichtgewicht is, geen extra trainbare parameters vereist (in de ruis-variant), en de flexibiliteit van expert-selectie behoudt terwijl het een gebalanceerde utilisatie tijdens de trainingsfase waarborgt.