🤖 machine learning

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Dit artikel introduceert Expert-Choice-routing als een superieur alternatief voor Token-Choice-routing in Diffusie-taalmodellen, waarbij deterministische load balancing, tijdstap-afhankelijke expertcapaciteit en de mogelijkheid om bestaande modellen te upgraden leiden tot hogere doorvoer, snellere convergentie en betere prestaties.

Oorspronkelijke auteurs: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

Gepubliceerd 2026-04-03

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

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

De Slimme Routeplanner voor Taalmodellen: Hoe een nieuwe methode AI sneller en slimmer maakt

Stel je voor dat je een enorm groot kantoor hebt met honderden specialisten (de "experts"). Hun taak is om zinnen te schrijven, alsof ze samen een boek schrijven. In de huidige wereld van AI (zoals ChatGPT) werken deze specialisten vaak volgens een oude regel: elke woordkeuze (token) kiest zelf zijn favoriete specialist.

Dit lijkt logisch, maar het heeft een groot probleem:

Soms kiezen 100 woorden allemaal dezelfde specialist. Die wordt dan overbelast en staat in de rij.
Andere specialisten staan de hele tijd te wachten met hun handen in de schoot.
Het hele kantoor moet wachten tot de langzaamste specialist klaar is. Dit kost veel tijd en energie.

Deze paper introduceert een slimme nieuwe manier om dit op te lossen, speciaal voor een nieuw type AI dat werkt met diffusie (een proces waarbij een tekst eerst als een wazige vlek begint en steeds scherper wordt, in plaats van woord voor woord te bouwen).

Hier is de uitleg in drie simpele stappen:

1. De Omgekeerde Keuze: "De Chef kiest de Werknemer"

In de oude methode (Token-Choice) kiezen de woorden hun eigen chef. In de nieuwe methode (Expert-Choice) draait het omgekeerd: elke specialist kiest zelf de woorden die hij wil helpen.

De Analogie: Stel je een restaurant voor.
- Oude manier: Gasten (woorden) rennen naar de kassa en kiezen hun favoriete kok. De ene kok krijgt 50 bestellingen, de andere 0. De keuken staat vast.
- Nieuwe manier (Expert-Choice): Elke kok (expert) heeft een vaste capaciteit, bijvoorbeeld "ik kan 5 bestellingen tegelijk aan". Ze kijken naar de rij en nemen de eerste 5 gasten die bij hen passen.
Het resultaat: Er is geen ruzie om werk, niemand staat in de rij, en niemand is lui. Het kantoor werkt als een goed geoliede machine. De onderzoekers laten zien dat dit 2 keer zo snel gaat als de oude methode.

2. De Slimme Werkrooster: "Meer energie op de moeilijke momenten"

Bij dit nieuwe type AI (diffusie) wordt een tekst stap voor stap scherper gemaakt.

Stap 1: De tekst is bijna volledig onleesbaar (veel "maskers").
Stap 100: De tekst is bijna klaar, alleen een paar woorden moeten nog worden ingevuld.

De onderzoekers ontdekten iets verrassends: De laatste stappen (waar de tekst al bijna klaar is) zijn eigenlijk het belangrijkst om op te focussen.

Waarom? Als er maar één woord ontbreekt in een zin, is het heel makkelijk om dat woord perfect te kiezen als je de hele context ziet. Als je daar extra "rekenkracht" (experts) op zet, leert het model daar het snelst van.
Als de tekst nog heel wazig is (veel maskers), helpt extra rekenkracht minder, omdat het model nog te veel moet raden.

De nieuwe methode past het aantal specialisten dynamisch aan:

Aan het begin (wazig): Weinig specialisten nodig.
Aan het einde (scherp): Veel specialisten inschakelen om de laatste details perfect te maken.

Dit noemen ze "Timestep-Adaptive Capacity". Het is alsof je een auto niet constant op 100 km/u rijdt, maar juist harder accelereert op de rechte stukken waar dat het meeste vooruitgang oplevert.

3. De "Retrofit": Oude auto's, nieuwe motor

De onderzoekers toonden ook aan dat je bestaande AI-modellen (die nog met de oude, trage methode werken) makkelijk kunt ombouwen. Je hoeft de hele motor niet te vervangen; je hoeft alleen de routeplanner (de router) te verwisselen.

Het resultaat? De oude modellen worden direct sneller in het leren en sneller in het beantwoorden van vragen, zonder dat je ze helemaal opnieuw hoeft te bouwen.

Samenvatting in één zin

Deze paper zegt: "Stop met het laten kiezen van de specialisten door de woorden zelf; laat de specialisten de woorden kiezen, en geef ze extra hulp op de momenten dat het echt uitmaakt."

De winst:

Snelheid: De AI leert 2 keer zo snel.
Efficiëntie: Geen tijdverlies door wachtende computers.
Kwaliteit: De AI wordt slimmer omdat hij zijn energie op de juiste momenten inzet.

Het is een bewijs dat we in de toekomst niet hoeven te bouwen aan grotere, zwaardere computers, maar slimmere manieren moeten vinden om de rekenkracht die we al hebben in te zetten.

Probleemstelling

Diffusie-taalmodellen (DLM's) bieden een veelbelovend alternatief voor autoregressieve generatie door parallelle, niet-autoregressieve tekstgeneratie mogelijk te maken via iteratieve denoising. Om deze modellen te schalen naar honderden miljarden parameters, wordt vaak gebruikgemaakt van Mixture-of-Experts (MoE) architecturen. Echter, bestaande DLM-MoE-modellen hebben de Token-Choice (TC) routing overgenomen uit autoregressieve systemen.

Dit leidt tot twee fundamentele problemen:

Load Imbalance (Ongelijke belasting): Omdat tokens onafhankelijk hun favoriete experts kiezen, raken sommige experts overbelast terwijl anderen inactief blijven. Dit vereist een extra "load-balancing loss" die extra rekenkracht kost en kan interfereren met het primaire leerdoel, zonder echter een harde garantie op balans te bieden.
Rigide computatie-allokatie: In TC-routing is de belasting van experts een emergent gevolg van token-keuzes en kan niet direct worden gecontroleerd. Dit verhindert het toewijzen van meer rekenkracht aan specifieke fasen van het denoising-proces waar dit het meest effectief zou zijn.

Methodologie

De auteurs introduceren en evalueren Expert-Choice (EC) routing als een superieur paradigma voor DLM's, gebaseerd op drie kerncomponenten:

Expert-Choice Routing:
- In plaats van dat tokens experts kiezen, kiezen de experts zelf een vast aantal tokens (capaciteit $c$ ).
- Dit zorgt voor deterministische load balancing: elke expert verwerkt exact hetzelfde aantal tokens, wat load imbalance elimineert zonder extra loss-functies.
- DLM's zijn hierbij inherent compatibel omdat ze in elke denoising-stap een niet-causaal (bidirectioneel) overzicht van alle tokens hebben, wat nodig is voor expert-side selectie.
Timestep-afhankelijke Expert Capaciteit:
- Omdat de capaciteit in EC een externe hyperparameter is, kunnen auteurs deze dynamisch aanpassen per denoising-stap ( $t$ ).
- Ze testen verschillende scheduling-strategieën om te bepalen welke fasen (gebaseerd op het masker-ratio $r$ ) het meeste baat hebben bij extra experts.
- De strategieën variëren van lineair, cosinus en Gaussisch, zowel voorwaarts (meer capaciteit bij hoge maskers) als achteruit (meer capaciteit bij lage maskers).
Retrofitting van Bestaande Modellen:
- De auteurs tonen aan dat bestaande, voorgeïmplementeerde TC-DLM's kunnen worden omgezet naar EC door alleen de router (de gating-mechanisme) te vervangen, terwijl de expert-weights en embeddings ongewijzigd blijven.

Belangrijkste Bijdragen

Superioriteit van EC voor DLM's: EC-routing levert betere load balancing, hogere doorvoer (throughput) en snellere convergentie op dan TC-routing.
Adaptieve Computatie: De introductie van timestep-afhankelijke capaciteit, waarbij meer rekenkracht wordt toegewezen aan stappen met een laag masker-ratio.
Mechanistische Verklaring: Een analyse toont aan dat tokens in contexten met een laag masker-ratio (waar de meeste tokens al zichtbaar zijn) een leer-efficiëntie hebben die een orde van grootte hoger is dan bij hoge masker-ratio's. Het concentreren van compute op deze stappen levert de grootste marginale opbrengst op.
Praktische Toepasbaarheid: Bewijs dat bestaande modellen met één router-wissel kunnen worden verbeterd voor zowel training als inferentie.

Resultaten

De experimenten omvatten pre-training op schalen tot 8B parameters en fine-tuning op diverse taken (code, wiskunde, medisch):

Trainingssnelheid: EC bereikt een verlies van 3.75 in 10,6 uur, wat 2,0x sneller is dan de beste TC-configuratie (20,7 uur).
Doorvoer (Throughput): EC bereikt 52,1 TFLOP/s/GPU, wat 1,5x tot 2,1x hoger is dan alle TC-varianten. Dit komt door het elimineren van "straggler"-effecten (wachten op de langzaamste GPU).
Scheduling Strategieën: De Linear-Reverse scheduler (meer capaciteit bij lage masker-ratio's) presteert consistent het beste onder gelijke FLOPs, met een perplexiteit van 36,5 op OpenWebText, vergeleken met 37,1 voor een statische baseline.
Learnings Efficiency: Analyse toont aan dat stappen met een laag masker-ratio ([0, 0.25)) tot 7x sneller convergeren dan stappen met een hoog masker-ratio.
Inferentie en Fine-tuning: Bij het retrofitting van een voorgeïmplementeerd TC-model (LLaDA-MoE) resulteerde EC in snellere convergentie en een 1,3x tot 1,5x snellere inferentietijd, terwijl de nauwkeurigheid gelijk bleef of verbeterde (bijv. 58,6% Pass@1 op HumanEval vs. 53,9% voor TC).

Betekenis en Conclusie

Dit werk vestigt Expert-Choice routing als het superieure paradigma voor MoE-modellen binnen diffusie-taalmodellen. Het toont aan dat:

Load imbalance in DLM's structureel kan worden opgelost zonder extra loss-functies.
Computatie in DLM's niet als een vaste architecturale constante moet worden behandeld, maar als een adaptief beleid dat kan worden geoptimaliseerd per denoising-stap.
De grootste leerwinst in DLM's plaatsvindt in de latere fasen van denoising (lage masker-ratio's), wat in tegenspraak is met intuïties uit continue diffusie-modellen waar vaak wordt aangenomen dat hoge ruis (hoge masker-ratio) moeilijker is.

De bevindingen bieden een nieuwe route voor het efficiënter trainen en infereren van grote taalmodellen, met directe toepasbaarheid op bestaande systemen via minimale architecturale aanpassingen.