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

In de wereld van natuurlijke taalverwerking (NLP) is de zoektocht naar het bouwen van grotere en krachtigere taalmodellen een belangrijke drijfveer geweest achter veel recente vooruitgang. Echter, naarmate deze modellen in omvang toenemen, worden de computatievereisten voor training en inferentie steeds veeleisender, waardoor ze tegen de limieten van de beschikbare hardwarebronnen aanlopen.

Komt Mixture-of-Experts (MoE) in beeld, een techniek die belooft deze computatiebezwaren te verlichten en het trainen 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 op de vroege jaren 90, toen onderzoekers het idee van conditionele berekening verkenden, 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 begeleid leerframework voor een ensemble van neurale netwerken voorstelde, waarbij elk netwerk zich specialiseerde in een verschillend gebied van de invoerruimte.

De kernidee achter MoE is om meerdere “expert”-netwerken te hebben, waarbij elk 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 computatiebronnen efficiënter toe te wijzen 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 verkend en uitgebreid, wat heeft geleid tot ontwikkelingen zoals hiërarchische MoE’s, lage-rangbenaderingen voor conditionele berekening en technieken voor het schatten van gradiënten via stochastische neuronen en hard-thresholdactiveringsfuncties.

Mixture-of-Experts in Transformers

Mixture of Experts

Hoewel het idee van MoE al decennialang bestaat, is de toepassing ervan op 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 op transformers is om de dichte FFN-lagen te vervangen door spaarzame MoE-lagen, elk bestaande uit meerdere expert-FFN’s en een poortmechanisme. Het poortmechanisme bepaalt welke expert(s) elke 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 aantoonde, 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 op transformers verder ontwikkeld, waarbij uitdagingen zoals trainingsinstabiliteit, belastingsbalans en efficiënte inferentie werden aangepakt. Opvallende 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 computatiekosten tijdens inferentie relatief constant blijven. Door selectief alleen een subset van experts te activeren voor elke invoertoken, kunnen MoE-modellen de expressieve kracht van veel grotere dichte modellen bereiken terwijl ze aanzienlijk minder computatie 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, zijn de computatiekosten gelijk aan die van een 14 miljard parameter dicht model, aangezien het twee 7 miljard parameter matrixvermenigvuldigingen berekent.

Deze computatie-efficiëntie tijdens inferentie is bijzonder waardevol in implementatiescenario’s waar resources beperkt zijn, zoals op mobiele apparaten of edge-computing-omgevingen. Bovendien kunnen de verlaagde computatievereisten tijdens training aanzienlijke energiewinsten en een lager koolstofvoetafdruk opleveren, waarmee ze aansluiten bij de groeiende nadruk op duurzame AI-praktijken.

Uitdagingen en overwegingen

Hoewel MoE-modellen aantrekkelijke voordelen bieden, komen hun adoptie en implementatie ook 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 gradiëntpropagatie 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 computatiekosten tijdens inferentie kunnen verlagen, hebben ze vaak hogere geheugeneisen in vergelijking met dichte modellen van vergelijkbare grootte. Dit komt omdat 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. Belastingsbalans: Om optimale computatie-efficiëntie te bereiken, is het cruciaal om de belasting over experts te balanceren, ervoor te zorgen dat geen enkele expert overbelast wordt terwijl anderen onderbenut worden. Deze belastingsbalans wordt typisch bereikt door middel van hulpverliezen tijdens training en zorgvuldige afstemming van de capaciteitsfactor, die het maximum aantal tokens bepaalt dat kan worden toegewezen aan elke expert.
5. Communicatieoverhead: In gedistribueerde trainings- en inferentiescenario’s kunnen MoE-modellen extra communicatieoverhead introduceren vanwege de noodzaak om activatie- en gradiëntinformatie uit te wisselen tussen experts die op verschillende apparaten of accelerators resideren. Efficiënte communicatiestrategieën en hardware-georiënteerd modelontwerp zijn essentieel om deze overhead te mitigeren.

Ondanks deze uitdagingen hebben de potentiële voordelen van MoE-modellen in het mogelijk maken van grotere en krachtigere taalmodellen aanzienlijke onderzoeksinspanningen gestimuleerd 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 beschouwen: 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 computatiekosten effectief worden 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-gefinetuneerde versie van Mixtral 8x7B, genaamd Mixtral-8x7B-Instruct-v0.1, is ook uitgebracht, waardoor zijn capaciteiten om natuurlijke taalinstructies te volgen verder zijn verbeterd.

Een ander opvallend voorbeeld is GLaM (Google Language Model), een groot MoE-model ontwikkeld door Google. GLaM gebruikt een decoder-only transformerarchitectuur en werd getraind op een enorme dataset van 1,6 biljoen tokens. Het model behaalt 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 enorm aantal parameters te trainen terwijl de computatievereisten 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 verbazingwekkende 314 miljard parameters is Grok-1 het grootste open LLM tot nu toe. Echter, dankzij de MoE-architectuur zijn alleen 25% van de gewichten (ongeveer 86 miljard parameters) actief op elk moment, waardoor de verwerkingscapaciteiten worden verbeterd.
2. Architectuur: Grok-1 gebruikt een Mixture-of-8-Experts-architectuur, waarbij elke token door twee experts wordt verwerkt tijdens inferentie.
3. Lagen: Het model bestaat uit 64 transformerlagen, elk met multihead-attention en dense blokken.
4. Tokenisatie: Grok-1 gebruikt een SentencePiece-tokenizer met een woordenschatgrootte 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%, waarmee het zijn efficiëntie en nauwkeurigheid op verschillende tests aantoonde.

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.

Desalniettemin ondersteunt het model activatie-sharding en 8-bits kwantificatie, waardoor de prestaties kunnen worden geoptimaliseerd en de 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-voorbeeldcode repository die aantoont hoe u Grok-1 kunt laden en uitvoeren. 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 de trainingsstabiliteit, het mitigeren van overfitting tijdens fine-tuning en het optimaliseren van geheugen- en communicatievereisten.

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

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 in staat zijn om een breed scala aan taken aan te pakken.

Conclusie

De Mixture-of-Experts-techniek is opgekomen als een krachtig instrument 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 computatie-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 computatie-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.