Stigningen af Mixture-of-Experts til effektive store sprogmodeller

I verden af naturlig sprogbehandling (NLP) har jagten på at bygge større og mere kapable sprogmodeller været en drivende kraft bag mange nyere fremskridt. Imidlertid, da disse modeller vokser i størrelse, bliver de computermæssige krav til træning og inferens mere krævende, og presser mod grænserne for de tilgængelige hardwareressourcer.

Indtræd Mixture-of-Experts (MoE), en teknik, der lover at lette denne computermæssige byrde, samtidig med at det muliggør træning af større og mere kraftfulde sprogmodeller. Herunder vil vi diskutere MoE, udforske dets oprindelse, indre mekanismer og dets anvendelser i transformer-baserede sprogmodeller.

Mixture-of-Experts’ Oprindelse

Begrebet Mixture-of-Experts (MoE) kan spores tilbage til begyndelsen af 1990’erne, da forskere udforskede ideen om betinget beregning, hvor dele af et neuralt netværk selektivt aktiveres baseret på inddata. En af de pionerarbejder i dette felt var “Adaptive Mixture of Local Experts“-artiklen af Jacobs et al. i 1991, der foreslog et overvåget læringssystem for et ensemble af neurale netværk, hvor hvert netværk specialiserer sig i en anden del af inddata-rummet.

Kernen i MoE er at have multiple “eksperter”-netværk, hvor hvert netværk er ansvarligt for at behandle en undermængde af inddata. En styringsmekanisme, typisk et neuralt netværk i sig selv, bestemmer, hvilke eksperter der skal behandle en given indgang. Dette tillader modellen at allokerer sine computermæssige ressourcer mere effektivt ved at aktiverer kun de relevante eksperter for hver indgang, i stedet for at anvende den fulde modelkapacitet for hver indgang.

Gennem årene har forskere udforsket og udvidet ideen om betinget beregning, hvilket har ført til udviklinger som hierarkiske MoE’er, lav-rang-approksimationer for betinget beregning og teknikker til estimering af gradienter gennem stokastiske neuroner og hard-threshold-aktiveringsfunktioner.

Mixture-of-Experts i Transformers

Mixture of Experts

Selvom ideen om MoE har været til stede i årtier, er dens anvendelse i transformer-baserede sprogmodeller relativt ny. Transformers, der er blevet standarden for state-of-the-art-sprogmodeller, består af multiple lag, hvor hvert lag indeholder en selv-attention-mekanisme og et feed-forward neuralt netværk (FFN).

Den nøgleinnovation i at anvende MoE til transformers er at erstatte de tætte FFN-lag med sparsomme MoE-lag, hvor hvert lag består af multiple eksperter og en styringsmekanisme. Styringsmekanismen bestemmer, hvilke eksperter der skal behandle hver indgangstoken, og tillader modellen at selektivt aktiverer kun en undermængde af eksperter for en given indgangssekvens.

En af de tidlige arbejder, der demonstrerede potentialet for MoE i transformers, var “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”-artiklen af Shazeer et al. i 2017. Dette arbejde introducerede begrebet om et sparsomt-styret MoE-lag, der anvendte en styringsmekanisme, der tilføjede sparsomhed og støj til eksperter-valgsprocessen, og sikrede, at kun en undermængde af eksperter blev aktiveret for hver indgang.

Siden da har flere andre arbejder videreudviklet anvendelsen af MoE til transformers, og har adresseret udfordringer som træningsinstabilitet, belastningsbalance og effektiv inferens. Bemærkelsesværdige eksempler inkluderer Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022), og GLaM (Du et al., 2022).

Fordele ved Mixture-of-Experts for Sprogmodeller

Den primære fordel ved at anvende MoE i sprogmodeller er evnen til at skalerer op modelstørrelsen, samtidig med at den relativt konstante computermæssige omkostning under inferens opretholdes. Ved at selektivt aktiverer kun en undermængde af eksperter for hver indgangstoken, kan MoE-modeller opnå den udtrykskraft af meget større tætte modeller, samtidig med at de kræver betydeligt mindre computation.

For eksempel kan vi overveje en sprogmodel med et tæt FFN-lag på 7 milliarder parametre. Hvis vi erstatter dette lag med et MoE-lag, der består af otte eksperter, hver med 7 milliarder parametre, øges det totale antal parametre til 56 milliarder. Imidlertid, under inferens, hvis vi kun aktiverer to eksperter per token, er den computermæssige omkostning lig med en 14 milliarder parameter tæt model, da den beregner to 7 milliarder parameter matrix-multiplicationer.

Denne computermæssige effektivitet under inferens er særligt værdifuld i installations-scenarier, hvor ressourcer er begrænsede, såsom mobile enheder eller edge-computing-miljøer. Derudover kan den reducerede computermæssige omkostning under træning føre til betydelige energibesparelser og en lavere carbonaftryk, i overensstemmelse med den voksende fokus på bæredygtige AI-praktikker.

Udfordringer og Overvejelser

Selvom MoE-modeller tilbyder overbevisende fordele, følger deres adoption og installation også flere udfordringer og overvejelser:

1. Træningsinstabilitet: MoE-modeller er kendt for at være mere udsatte for træningsinstabilitet sammenlignet med deres tætte modstykke. Dette problem opstår fra den sparsomme og betingede natur af eksperter-aktivering, hvilket kan føre til udfordringer i gradient-propagation og konvergens. Teknikker som router z-tab (Zoph et al., 2022) er blevet foreslået for at mildne disse instabiliteter, men yderligere forskning er stadig nødvendig.
2. Finjustering og Overfitting: MoE-modeller har tendens til at overfitte mere let under finjustering, især når nedstrøms-opgaven har et relativt lille dataset. Dette adfærd kan tilskrives den øgede kapacitet og sparsomhed af MoE-modeller, hvilket kan føre til overspecialisering på træningsdata. Omhyggelig regularisering og finjusteringsstrategier er nødvendige for at mildne dette problem.
3. Hukommelseskrav: Selvom MoE-modeller kan reducere computermæssige omkostninger under inferens, har de ofte højere hukommelseskrav sammenlignet med tætte modeller af samme størrelse. Dette skyldes, at alle eksperter-vægte skal indlæses i hukommelsen, selvom kun en undermængde af eksperter aktiveres for hver indgang. Hukommelsesbegrænsninger kan begrænse skalerbarheden af MoE-modeller på ressource-begrænsede enheder.
4. Belastningsbalance: For at opnå optimal computermæssig effektivitet er det afgørende at balancere belastningen over eksperter, således at ingen enkelt ekspert er overbelastet, mens andre forbliver underudnyttede. Denne belastningsbalance opnås typisk gennem auxiliary-tab under træning og omhyggelig finjustering af kapacitetsfaktoren, der bestemmer det maksimale antal tokens, der kan tildeles til hver ekspert.
5. Kommunikations-overhead: I distribueret træning og inferens-scenarier kan MoE-modeller introducere ekstra kommunikations-overhead på grund af behovet for at udveksle aktivering- og gradient-information over eksperter, der befinder sig på forskellige enheder eller acceleratorer. Effektive kommunikationsstrategier og hardware-orienteret model-design er afgørende for at mildne denne overhead.

Trods disse udfordringer har de potentielle fordele ved MoE-modeller i at muliggøre større og mere kapable sprogmodeller ført til betydelige forskningsindsats for at adressere og mildne disse problemer.

Eksempel: Mixtral 8x7B og GLaM

For at illustrere den praktiske anvendelse af MoE i sprogmodeller, kan vi overveje to bemærkelsesværdige eksempler: Mixtral 8x7B og GLaM.

Mixtral 8x7B er en MoE-variant af Mistral-sprogmodellen, udviklet af Anthropic. Den består af otte eksperter, hver med 7 milliarder parametre, hvilket resulterer i en total på 56 milliarder parametre. Imidlertid, under inferens, aktiveres kun to eksperter per token, hvilket reducerer den computermæssige omkostning til det niveau, der svarer til en 14 milliarder parameter tæt model.

Mixtral 8x7B har demonstreret imponerende præstationer, og overgår den 70 milliarder parameter Llama-model, samtidig med at den tilbyder meget hurtigere inferenstider. En instruktions-justeret version af Mixtral 8x7B, kaldet Mixtral-8x7B-Instruct-v0.1, er også blevet udgivet, hvilket yderligere forbedrer dets kapaciteter i at følge naturlige sprog-instruktioner.

Et andet bemærkelsesværdigt eksempel er GLaM (Google Language Model), en stor MoE-model udviklet af Google. GLaM anvender en decoder-kun transformer-arkitektur og blev trænet på et massivt 1,6 billion token-dataset. Modellen opnår imponerende præstationer på few-shot og one-shot-evalueringer, og matcher kvaliteten af GPT-3, samtidig med at den kun anvender en tredjedel af den energi, der kræves til at træne GPT-3.

GLaM’s succes kan tilskrives dens effektive MoE-arkitektur, der muliggjorde træning af en model med et enormt antal parametre, samtidig med at den opretholdt rimelige computermæssige krav. Modellen demonstrerede også potentialet for MoE-modeller til at være mere energi-effektive og miljøvenlige sammenlignet med deres tætte modstykke.

Grok-1-Arkitekturen

GROK MIXTURE OF EXPERT

Grok-1 er en transformer-baseret MoE-model med en unik arkitektur designet til at maksimere effektivitet og præstation. Lad os dykke ned i de nøgle-specifikationer:

1. Parametre: Med en imponerende 314 milliarder parametre er Grok-1 den største åbne LLM til dato. Imidlertid, takket være MoE-arkitekturen, er kun 25% af vægtene (cirka 86 milliarder parametre) aktive på et given tidspunkt, hvilket forbedrer processorkapaciteten.
2. Arkitektur: Grok-1 anvender en Mixture-of-8-Experts-arkitektur, hvor hvert token behandles af to eksperter under inferens.
3. Lag: Modellen består af 64 transformer-lag, hvor hvert lag inkorporerer multihead-attention og dense-blokke.
4. Tokenisering: Grok-1 anvender en SentencePiece-tokenisator med en ord-af-vokabular-størrelse på 131.072 tokens.
5. Indlejring og Positional Encoding: Modellen har 6.144-dimensionale indlejring og anvender rotary positional encoding, hvilket muliggør en mere dynamisk fortolkning af data sammenlignet med traditionelle faste positional encoding.
6. Attention: Grok-1 anvender 48 attention-hoveder til forespørgsler og 8 attention-hoveder til nøgler og værdier, hver med en størrelse på 128.
7. Kontekstlængde: Modellen kan behandle sekvenser op til 8.192 tokens i længde, og anvender bfloat16-præcision for effektiv beregning.

Præstation og Implementeringsdetaljer

Grok-1 har demonstreret imponerende præstationer, og overgår LLaMa 2 70B og Mixtral 8x7B med en MMLU-score på 73%, hvilket viser dens effektivitet og nøjagtighed på tværs af forskellige tests.

Imidlertid er det vigtigt at bemærke, at Grok-1 kræver betydelige GPU-ressourcer på grund af sin enorme størrelse. Den nuværende implementering i den åbne kildekode fokuserer på at validere modellens korrekthed og anvender en ineffektiv MoE-lag-implementering for at undgå behovet for brugerdefinerede kernel.

Alligevel understøtter modellen aktivering-sharding og 8-bit-kvantificering, hvilket kan optimere præstation og reducere hukommelseskrav.

I en bemærkelsesværdig bevægelse har xAI udgivet Grok-1 under Apache 2.0-licensen, hvilket gør vægt og arkitektur tilgængelige for den globale fællesskab til brug og bidrag.

Den åbne kildekode-udgivelse inkluderer en JAX-eksempelkode-repository, der demonstrerer, hvordan man kan indlæse og køre Grok-1-modellen. Brugere kan downloade checkpoint-vægtene ved hjælp af en torrent-klient eller direkte gennem HuggingFace Hub, hvilket faciliterer let adgang til denne banebrydende model.

Fremtiden for Mixture-of-Experts i Sprogmodeller

Da kravet til større og mere kapable sprogmodeller fortsætter med at vokse, forventes adoptionen af MoE-teknikker at vinde yderligere momentum. Fremtidige forskningsindsats er fokuseret på at adressere de resterende udfordringer, såsom forbedring af træningsstabilitet, mildning af overfitting under finjustering og optimering af hukommelses- og kommunikationskrav.

En lovende retning er udforskningen af hierarkiske MoE-arkitekturer, hvor hver ekspert i sig selv består af multiple under-eksperter. Dette tillader muligvis endnu større skalerbarhed og computermæssig effektivitet, samtidig med at den opretholder den udtrykskraft af store modeller.

Derudover er udviklingen af hardware- og software-systemer, der er optimeret til MoE-modeller, et aktivt forskningsområde. Specialiserede acceleratorer og distribuerede trænings-rammer, der er designet til at håndtere de sparsomme og betingede beregningsmønstre af MoE-modeller, kan yderligere forbedre deres præstation og skalerbarhed.

Endelig kan integrationen af MoE-teknikker med andre fremskridt i sprogmodellering, såsom sparsom attention-mekanismer, effektive tokenisering-strategier og multi-modale repræsentationer, føre til endnu mere kraftfulde og fleksible sprogmodeller, der kan tackle en bred vifte af opgaver.

Konklusion

Mixture-of-Experts-teknikken er opstået som et kraftfuldt værktøj i jagten på større og mere kapable sprogmodeller. Ved at selektivt aktiverer eksperter baseret på inddata, tilbyder MoE-modeller en lovende løsning på de computermæssige udfordringer, der er forbundet med at skalerer op tætte modeller. Selvom der stadig er udfordringer at overvinde, såsom træningsinstabilitet, overfitting og hukommelseskrav, gør de potentielle fordele ved MoE-modeller i forhold til computermæssig effektivitet, skalerbarhed og miljøvenlighed dem til et spændende forsknings- og udviklingsområde.

Da feltet af naturlig sprogbehandling fortsætter med at skubbe grænserne for, hvad der er muligt, er det sandsynligt, at adoptionen af MoE-teknikker vil spille en afgørende rolle i at muliggøre den næste generation af sprogmodeller. Ved at kombinere MoE med andre fremskridt i model-arkitektur, træningsteknikker og hardware-optimering kan vi se frem til endnu mere kraftfulde og fleksible sprogmodeller, der kan forstå og kommunikere med mennesker på en naturlig og ubesværet måde.