Mixture of experts (MoE)

Wat is een mixture of experts?

Een mixture of experts, meestal afgekort tot MoE, is een modelontwerp dat een groot neuraal netwerk opsplitst in veel kleinere, gespecialiseerde onderdelen die experts heten. Voor elk stukje input beslist een klein routernetwerk welke expert of twee experts het afhandelen, terwijl de rest niks doet.

Het gevolg is dat een model een enorm aantal parameters kan bevatten en toch maar een klein deel ervan gebruikt voor elke input. Je krijgt de capaciteit van een heel groot model tegen bijna de rekenkost van een klein model.

Een druk ziekenhuis is een handige vergelijking. Een patiënt aan de balie wordt niet door elke arts in het gebouw gezien. Een verpleegkundige aan de triage stuurt hem naar de een of twee specialisten die zijn geval nodig heeft. MoE stuurt elke token op dezelfde manier, naar de experts die er het best bij passen.

Het probleem dat een mixture of experts oplost

In een gewoon "dense" model gebruikt elke parameter mee voor elke input. Een model slimmer maken betekent meestal het groter maken, en een groter dense model kost evenredig meer voor elk woord dat het verwerkt. Die kost loopt snel op, zowel tijdens de training als in dagelijks gebruik.

MoE breekt de band tussen hoeveel een model weet en hoeveel rekenwerk het per token doet. Het totale aantal parameters kan oplopen tot honderden miljarden, terwijl het werk per token ongeveer gelijk blijft. Daarom zijn veel van de grootste recente taalmodellen op deze manier gebouwd.

Hoe een mixture of experts werkt

De experts. Elke expert is een klein feed-forward netwerk, hetzelfde soort blok dat je in een standaard transformer vindt. Een MoE-laag houdt er meerdere naast elkaar. Experts krijgen geen onderwerpen toegewezen met de hand. Tijdens de training schuift elke expert vanzelf op naar de patronen die hij het best aankan.

De router. Ook wel het gating-netwerk genoemd. Dit is een klein netwerk dat naar elke binnenkomende token kijkt en elke expert een score geeft. Daarna stuurt het de token naar de hoogst scorende paar, meestal de top twee.

Sparse activatie. Omdat enkel de gekozen experts draaien, blijft het grootste deel van het model uitgeschakeld voor een bepaalde token. Dat houdt de rekenkost laag bij inferentie, zelfs als het volledige model gigantisch is. De uitkomsten van de actieve experts worden samengevoegd, gewogen volgens de scores van de router, en doorgegeven aan de volgende laag.

Een uitgewerkt voorbeeld

Mixtral 8x7B van Mistral AI is een duidelijk geval. Elke MoE-laag bevat 8 experts, en de router kiest er 2 voor elke token. Het volledige model telt 46,7 miljard parameters, maar slechts ongeveer 12,9 miljard zijn actief voor een enkele token. Zo draait het aan de snelheid en kost van een model van ongeveer 13 miljard parameters, terwijl het de kennis van een veel groter model draagt.

Mixture of experts versus een dense model

Een dense model gebruikt al zijn parameters voor elke token. Het is eenvoudiger om te trainen en te draaien, en het geheugengebruik is makkelijk te voorspellen, maar het betaalt de volle prijs voor zijn volle omvang bij elk woord.

Een MoE-model draait goedkoper per token voor dezelfde hoeveelheid opgeslagen kennis. De keerzijde is dat al die experts wel in het geheugen moeten staan, ook al draaien er maar een paar tegelijk. MoE-modellen hebben dus veel geheugen nodig om te draaien. Ze zijn ook lastiger om goed te trainen, want de router moet leren om het werk gelijk te verdelen over de experts.

Waar moet je op letten bij het gebruik van een mixture of experts

Geheugen wordt de limiet, niet rekenkracht. Elke expert moet geladen en klaar staan, ook al vuurt hij zelden. Een MoE-model dat goedkoop is per token kan toch veel meer geheugen nodig hebben dan het aantal actieve parameters doet vermoeden.

De router kan voorkeuren ontwikkelen. Als je de routing niet in de gaten houdt, krijgt een handvol experts het meeste verkeer terwijl de rest amper traint. Bouwers voegen tijdens de training technieken toe om de belasting gespreid te houden.

Actieve parameters kunnen misleiden. Een model dat met zijn aantal actieve parameters wordt aangeprezen, klinkt klein, maar je hebt nog altijd de hardware nodig om elke expert te bewaren. Lees zowel het totale als het actieve cijfer voor je oordeelt wat er nodig is om het te draaien.