De opkomst van Mixture-of-Experts voor efficiënte grote taalmodellen

In de wereld van natuurlijke taalverwerking (NLP) is de ontwikkeling van grotere en krachtigere taalmodellen een belangrijke drijvende kracht achter veel recente vooruitgang. Echter, naarmate deze modellen in omvang toenemen, worden de computationele vereisten voor training en inferentie steeds veeleisender, waardoor de grenzen van de beschikbare hardwarebronnen worden bereikt.

Hier komt Mixture-of-Experts (MoE) om de hoek kijken, een techniek die belooft deze computationele last te verlichten en tegelijkertijd de training van grotere en krachtigere taalmodellen mogelijk te maken. Hieronder zullen we MoE bespreken, zijn oorsprong, interne werking en toepassingen in transformer-gebaseerde taalmodellen onderzoeken.

De oorsprong van Mixture-of-Experts

Het concept van Mixture-of-Experts (MoE) kan worden teruggevoerd tot de vroege jaren negentig, toen onderzoekers het idee van conditionele berekening onderzochten, waarbij delen van een neurale netwerk selectief worden geactiveerd op basis van de invoergegevens. Een van de pionierswerken op dit gebied was de paper “Adaptive Mixture of Local Experts” van Jacobs et al. in 1991, die een supervised learning-framework voor een ensemble van neurale netwerken voorstelde, waarbij elk netwerk zich specialiseerde in een ander deel van de invoerruimte.

Het kernidee achter MoE is om meerdere “expert”-netwerken te hebben, waarbij elk netwerk verantwoordelijk is voor het verwerken van een deel van de invoergegevens. Een poortmechanisme, meestal een neurale netwerk op zich, bepaalt welke expert(s) een bepaalde invoer moeten verwerken. Deze aanpak stelt het model in staat om zijn computationele middelen efficiënter te gebruiken door alleen de relevante experts te activeren voor elke invoer, in plaats van de volledige modelcapaciteit voor elke invoer te gebruiken.

In de loop der jaren hebben verschillende onderzoekers het idee van conditionele berekening verder onderzocht en uitgebreid, wat heeft geleid tot ontwikkelingen zoals hiërarchische MoE’s, lage-rangbenaderingen voor conditionele berekening en technieken voor het schatten van gradients via stochastische neuronen en hard-thresholdactivatiefuncties.

Mixture-of-Experts in Transformers

Mixture of Experts

Hoewel het idee van MoE al decennia bestaat, is de toepassing ervan in transformer-gebaseerde taalmodellen relatief recent. Transformers, die de facto de standaard zijn geworden voor state-of-the-art taalmodellen, bestaan uit meerdere lagen, waarbij elke laag een self-attention-mechanisme en een feed-forward neurale netwerk (FFN) bevat.

De sleutelinnovatie bij het toepassen van MoE in transformers is het vervangen van de dichte FFN-lagen door spaarzame MoE-lagen, elk bestaande uit meerdere expert-FFN’s en een poortmechanisme. Het poortmechanisme bepaalt welke expert(s) elk invoertoken moeten verwerken, waardoor het model selectief alleen een subset van experts kan activeren voor een bepaalde invoersequentie.

Een van de eerste werken die het potentieel van MoE in transformers demonstreerde, was de paper “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” van Shazeer et al. in 2017. Dit werk introduceerde het concept van een spaarzaam-poort-MoE-laag, die een poortmechanisme gebruikte dat sparsity en ruis toevoegde aan het expertselectieproces, waardoor alleen een subset van experts werd geactiveerd voor elke invoer.

Sindsdien hebben verschillende andere werken de toepassing van MoE in transformers verder uitgebreid, waarbij uitdagingen zoals trainingsinstabiliteit, load balancing en efficiënte inferentie werden aangepakt. Notabele voorbeelden zijn de Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) en GLaM (Du et al., 2022).

Voordelen van Mixture-of-Experts voor taalmodellen

Het primaire voordeel van het gebruik van MoE in taalmodellen is de mogelijkheid om de modelgrootte op te schalen terwijl de computationele kosten tijdens inferentie relatief constant blijven. Door selectief alleen een subset van experts te activeren voor elk invoertoken, kunnen MoE-modellen de expressieve kracht van veel grotere dichte modellen bereiken terwijl ze significant minder berekening vereisen.

Bijvoorbeeld, overweeg een taalmodel met een dichte FFN-laag van 7 miljard parameters. Als we deze laag vervangen door een MoE-laag bestaande uit acht experts, elk met 7 miljard parameters, neemt het totale aantal parameters toe tot 56 miljard. Echter, tijdens inferentie, als we alleen twee experts per token activeren, is de computationele kost gelijk aan die van een 14 miljard parameter dicht model, aangezien het twee 7 miljard parameter matrixvermenigvuldigingen berekent.

Deze computationele efficiëntie tijdens inferentie is bijzonder waardevol in implementatiescenario’s waar resources beperkt zijn, zoals op mobiele apparaten of edge-computing-omgevingen. Bovendien kan de vermindering van computationele kosten tijdens training leiden tot aanzienlijke energibesparingen en een lager koolstofvoetafdruk, waardoor het zich aanpast aan de groeiende nadruk op duurzame AI-praktijken.

Uitdagingen en overwegingen

Hoewel MoE-modellen aantrekkelijke voordelen bieden, gaat hun adoptie en implementatie ook gepaard met verschillende uitdagingen en overwegingen:

1. Trainingsinstabiliteit: MoE-modellen zijn bekend om meer gevoelig te zijn voor trainingsinstabiliteit in vergelijking met hun dichte tegenhangers. Dit probleem ontstaat door de spaarzame en conditionele aard van de expertactivaties, die kan leiden tot uitdagingen in gradientpropagatie en convergentie. Technieken zoals de router z-verlies (Zoph et al., 2022) zijn voorgesteld om deze instabiliteit te mitigeren, maar verder onderzoek is nog steeds nodig.
2. Fine-tuning en overfitting: MoE-modellen hebben de neiging om sneller over te fitten tijdens fine-tuning, vooral wanneer de downstream-taak een relatief kleine dataset heeft. Dit gedrag wordt toegeschreven aan de toegenomen capaciteit en sparsity van MoE-modellen, die kan leiden tot overspecialisatie op de trainingsgegevens. Zorgvuldige regularisatie en fine-tuningstrategieën zijn vereist om dit probleem te mitigeren.
3. Geheugeneisen: Hoewel MoE-modellen de computationele kosten tijdens inferentie kunnen verlagen, hebben ze vaak hogere geheugeneisen in vergelijking met dichte modellen van vergelijkbare grootte. Dit komt doordat alle expertgewichten in het geheugen moeten worden geladen, zelfs als alleen een subset wordt geactiveerd voor elke invoer. Geheugensbeperkingen kunnen de schaalbaarheid van MoE-modellen op resource-beperkte apparaten beperken.
4. Load balancing: Om optimale computationele efficiëntie te bereiken, is het cruciaal om de belasting over experts te balanceren, waarbij wordt gegarandeerd dat geen enkele expert overbelast is terwijl anderen onderbenut worden. Deze load balancing wordt meestal bereikt door middel van auxiliare verliezen tijdens training en zorgvuldige afstemming van de capaciteitsfactor, die het maximum aantal tokens bepaalt dat aan elke expert kan worden toegewezen.
5. Communicatieoverhead: In gedistribueerde trainings- en inferentiescenario’s kunnen MoE-modellen extra communicatieoverhead introduceren vanwege de noodzaak om activatie- en gradientinformatie uit te wisselen tussen experts die op verschillende apparaten of accelerators resideren. Efficiënte communicatiestrategieën en hardwarebewuste modelontwerp zijn essentieel om deze overhead te mitigeren.

Ondanks deze uitdagingen heeft het potentieel van MoE-modellen in het mogelijk maken van grotere en krachtigere taalmodellen aanzienlijke onderzoeksinspanningen gegenereerd om deze problemen aan te pakken en te mitigeren.

Voorbeeld: Mixtral 8x7B en GLaM

Om de praktische toepassing van MoE in taalmodellen te illustreren, laten we twee opvallende voorbeelden bekijken: Mixtral 8x7B en GLaM.

Mixtral 8x7B is een MoE-variant van het Mistral-taalmodel, ontwikkeld door Anthropic. Het bestaat uit acht experts, elk met 7 miljard parameters, wat resulteert in een totaal van 56 miljard parameters. Echter, tijdens inferentie worden alleen twee experts per token geactiveerd, waardoor de computationele kost effectief wordt verlaagd tot die van een 14 miljard parameter dicht model.

Mixtral 8x7B heeft indrukwekkende prestaties getoond, waarbij het de 70 miljard parameter Llama-model overtrof terwijl het veel snellere inferentietijden bood. Een instructie-gefinetune versie van Mixtral 8x7B, genaamd Mixtral-8x7B-Instruct-v0.1, is ook uitgebracht, waardoor zijn mogelijkheden om natuurlijke taalinstructies te volgen verder zijn verbeterd.

Een ander opvallend voorbeeld is GLaM (Google Language Model), een grote MoE-model ontwikkeld door Google. GLaM gebruikt een decoder-only transformer-architectuur en is getraind op een massive 1,6 biljoen token dataset. Het model bereikt indrukwekkende prestaties op few-shot en one-shot evaluaties, waarbij het de kwaliteit van GPT-3 evenaart terwijl het slechts een derde van de energie gebruikt die nodig is om GPT-3 te trainen.

Het succes van GLaM kan worden toegeschreven aan zijn efficiënte MoE-architectuur, die het mogelijk maakte om een model met een groot aantal parameters te trainen terwijl de computationele vereisten redelijk bleven. Het model toonde ook het potentieel van MoE-modellen om energie-efficiënter en milieuvriendelijker te zijn in vergelijking met hun dichte tegenhangers.

De Grok-1-architectuur

GROK MIXTURE OF EXPERT

Grok-1 is een transformer-gebaseerd MoE-model met een unieke architectuur ontworpen om efficiëntie en prestaties te maximaliseren. Laten we de belangrijkste specificaties onderzoeken:

1. Parameters: Met een verbijsterende 314 miljard parameters is Grok-1 het grootste open LLM tot nu toe. Echter, dankzij de MoE-architectuur zijn slechts 25% van de gewichten (ongeveer 86 miljard parameters) actief op elk moment, waardoor de verwerking wordt verbeterd.
2. Architectuur: Grok-1 gebruikt een Mixture-of-8-Experts-architectuur, waarbij elk token door twee experts wordt verwerkt tijdens inferentie.
3. Lagen: Het model bestaat uit 64 transformer-lagen, elk met multihead-attention en dichte blokken.
4. Tokenisatie: Grok-1 gebruikt een SentencePiece-tokenizer met een vocabulairegrootte van 131.072 tokens.
5. Embeddings en positionele codering: Het model heeft 6.144-dimensionale embeddings en gebruikt rotary positionele codering, waardoor een meer dynamische interpretatie van gegevens mogelijk is in vergelijking met traditionele vaste positionele codering.
6. Aandacht: Grok-1 gebruikt 48 aandachtshoofden voor queries en 8 aandachtshoofden voor sleutels en waarden, elk met een grootte van 128.
7. Contextlengte: Het model kan sequenties verwerken tot 8.192 tokens in lengte, met bfloat16-nauwkeurigheid voor efficiënte berekening.

Prestaties en implementatiedetails

Grok-1 heeft indrukwekkende prestaties getoond, waarbij het LLaMa 2 70B en Mixtral 8x7B overtrof met een MMLU-score van 73%, waardoor zijn efficiëntie en nauwkeurigheid over verschillende tests werden aangetoond.

Echter, het is belangrijk op te merken dat Grok-1 aanzienlijke GPU-bronnen vereist vanwege zijn enorme grootte. De huidige implementatie in de open-source release is gericht op het valideren van de correctheid van het model en gebruikt een inefficiënte MoE-laagimplementatie om de noodzaak van aangepaste kernels te vermijden.

Nonetheless, het model ondersteunt activatie-sharding en 8-bit kwantificatie, waardoor prestaties kunnen worden geoptimaliseerd en geheugeneisen kunnen worden verlaagd.

In een opmerkelijke stap heeft xAI Grok-1 uitgebracht onder de Apache 2.0-licentie, waardoor zijn gewichten en architectuur toegankelijk zijn voor de wereldwijde gemeenschap voor gebruik en bijdragen.

De open-source release bevat een JAX-voorbeeldcodebibliotheek die laat zien hoe het Grok-1-model kan worden geladen en uitgevoerd. Gebruikers kunnen de checkpoint-gewichten downloaden met een torrent-client of rechtstreeks via de HuggingFace Hub, waardoor toegang tot dit baanbrekende model wordt vergemakkelijkt.

De toekomst van Mixture-of-Experts in taalmodellen

Aangezien de vraag naar grotere en krachtigere taalmodellen blijft groeien, zal de adoptie van MoE-technieken naar verwachting verder toenemen. Lopende onderzoeksinspanningen zijn gericht op het aanpakken van de resterende uitdagingen, zoals het verbeteren van trainingsstabiliteit, het mitigeren van overfitting tijdens fine-tuning en het optimaliseren van geheugen- en communicatievereisten.

Een veelbelovende richting is het onderzoeken van hiërarchische MoE-architecturen, waarbij elke expert zelf bestaat uit meerdere sub-experts. Deze aanpak kan mogelijk nog grotere schaalbaarheid en computationele efficiëntie mogelijk maken terwijl de expressieve kracht van grote modellen wordt behouden.

Bovendien is de ontwikkeling van hardware- en software-systemen die zijn geoptimaliseerd voor MoE-modellen een actief onderzoeksgebied. Gespecialiseerde accelerators en gedistribueerde trainingsframeworks die zijn ontworpen om de spaarzame en conditionele berekeningspatronen van MoE-modellen efficiënt te verwerken, kunnen hun prestaties en schaalbaarheid verder verbeteren.

Verder kan de integratie van MoE-technieken met andere vooruitgang in taalmodellering, zoals spaarzame aandachtsmechanismen, efficiënte tokenisatiestrategieën en multimodale representaties, leiden tot nog krachtigere en veelzijdigere taalmodellen die een breed scala aan taken kunnen aanpakken.

Conclusie

De Mixture-of-Experts-techniek is een krachtig instrument gebleken in de zoektocht naar grotere en krachtigere taalmodellen. Door selectief experts te activeren op basis van de invoergegevens, bieden MoE-modellen een veelbelovende oplossing voor de computationele uitdagingen die samenhangen met het opschalen van dichte modellen. Hoewel er nog uitdagingen moeten worden overwonnen, zoals trainingsinstabiliteit, overfitting en geheugeneisen, maken de potentiële voordelen van MoE-modellen in termen van computationele efficiëntie, schaalbaarheid en milieuvriendelijkheid ze een spannend onderzoeks- en ontwikkelingsgebied.

Aangezien het veld van natuurlijke taalverwerking de grenzen van wat mogelijk is blijft verleggen, zal de adoptie van MoE-technieken naar verwachting een cruciale rol spelen bij het mogelijk maken van de volgende generatie taalmodellen. Door MoE te combineren met andere vooruitgang in modelarchitectuur, trainingsmethoden en hardware-optimalisatie, kunnen we uitkijken naar nog krachtigere en veelzijdigere taalmodellen die echt kunnen begrijpen en communiceren met mensen op een natuurlijke en naadloze manier.