Mixture-of-Experts-teknikkens opkomst 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 og mere krævende, og presser mod grænserne for de tilgængelige hardwareressourcer.

Indtræder 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 powerful 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å inputdataen. En af de pionerarbejder i dette felt var “Adaptive Mixture of Local Experts“-artiklen af Jacobs et al. i 1991, som foreslog et overvåget læringssystem for et ensemble af neurale netværk, hver specialiseret i en anden del af inputrummet.

Kernen i MoE er at have multiple “ekspert”-netværk, hver ansvarlig for at behandle en undermængde af inputdataen. En gateway-mekanisme, typisk et neuralt netværk i sig selv, bestemmer, hvilke ekspert(er) der skal behandle en given input. Dette tillader modellen at allocere sine computermæssige ressourcer mere effektivt ved at aktivere kun de relevante eksperter for hver input, i stedet for at anvende den fulde modelkapacitet for hver input.

Over årene har forskellige 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 at estime 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 dets anvendelse i transformer-baserede sprogmodeller relativt ny. Transformers, som er blevet standarden for state-of-the-art-sprogmodeller, består af multiple lag, hver indeholdende 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 sparse MoE-lag, hver bestående af multiple ekspert-FFN’er og en gateway-mekanisme. Gateway-mekanismen bestemmer, hvilke ekspert(er) der skal behandle hver input-token, hvilket muliggør, at modellen selektivt aktiverer kun en undermængde af eksperter for en given input-sekvens.

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 sparsely-gated MoE-lag, som anvendte en gateway-mekanisme, der tilføjede sparsomhed og støj til ekspert-selektionsprocessen, hvilket sikrede, at kun en undermængde af eksperter blev aktiveret for hver input.

Siden da har flere andre arbejder videreudviklet anvendelsen af MoE til transformers, hvor de 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

Det primære fordel ved at anvende MoE i sprogmodeller er evnen til at skala op modellens størrelse, mens man opretholder en relativt konstant computermæssig omkostning under inferens. Ved at selektivt aktivere kun en undermængde af eksperter for hver input-token, kan MoE-modeller opnå den udtrykskraft af langt større tætte modeller, mens de kræver betydeligt mindre beregning.

For eksempel kan man overveje en sprogmodel med et tæt FFN-lag på 7 milliarder parametre. Hvis vi erstatter dette lag med et MoE-lag bestående 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 pr. token, er den computermæssige omkostning lig med en 14 milliarder parametre tæt model, da den kun udfører to 7 milliarder parametre matrix-multiplicationer.

Denne computermæssige effektivitet under inferens er særligt værdifuld i deploymentscenarier, 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 kulstofaftryk, hvilket er i overensstemmelse med den voksende fokus på bæredygtig AI-praksis.

Udfordringer og overvejelser

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

1. Træningsinstabilitet: MoE-modeller er kendt for at være mere udsat for træningsinstabilitet sammenlignet med deres tætte modstykke. Dette problem opstår fra den sparse og betingede natur af ekspert-aktiveringerne, hvilket kan føre til udfordringer i gradient-propagation og konvergens. Teknikker som router z-loss (Zoph et al., 2022) er blevet foreslået for at mindske 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 den nedstrømsopgave har et relativt lille datasæt. Dette adfærd kan tilskrives den øgede kapacitet og sparsomhed af MoE-modeller, hvilket kan føre til overspecialisering på træningsdataen. Omhyggelig regularisering og finjusteringsstrategier er nødvendige for at mindske 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 ekspert-vægte skal indlæses i hukommelsen, selvom kun en undermængde af eksperter aktiveres for hver input. 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 eksperterne, således at ingen enkelt ekspert er overbelastet, mens andre forbliver underudnyttede. Denne belastningsbalance opnås typisk gennem auxiliary-loss under træning og omhyggelig finjustering af kapacitetsfaktoren, som bestemmer det maksimale antal tokens, der kan tildeles 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 aktiverings- 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 mindske denne overhead.

Trods disse udfordringer har potentialet for MoE-modeller i at muliggøre større og mere kapable sprogmodeller fremmet betydelige forskningsindsats for at adressere og mindske disse problemer.

Eksempel: Mixtral 8x7B og GLaM

For at illustrere den praktiske anvendelse af MoE i sprogmodeller, lad os 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 pr. token, hvilket reducerer den computermæssige omkostning til det af en 14 milliarder parametre tæt model.

Mixtral 8x7B har demonstreret imponerende præstationer, overgående den 70 milliarder parametre Llama-model, mens den tilbyder langt hurtigere inferenstider. En instruktions-justeret version af Mixtral 8x7B, kaldet Mixtral-8x7B-Instruct-v0.1, er også blevet udgivet, hvilket yderligere forbedrer dets evner i at følge naturligt sprog-instruktioner.

Et andet bemærkelsesværdigt eksempel er GLaM (Google (GOOGL ) Language Model), en stor MoE-model udviklet af Google. GLaM anvender en decoder-only transformer-arkitektur og blev trænet på en massiv 1,6 billion token-datasæt. Modellen opnår imponerende præstationer på few-shot og one-shot-evalueringer, svarende til kvaliteten af GPT-3, mens den kun kræver en tredjedel af den energi, der er nødvendig til at træne GPT-3.

GLaM’s succes kan tilskrives dets effektive MoE-arkitektur, som muliggjorde træning af en model med et enormt antal parametre, mens 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 (ca. 86 milliarder parametre) aktive på ethvert givent tidspunkt, hvilket forbedrer behandlingskapaciteten.
2. Arkitektur: Grok-1 anvender en Mixture-of-8-Experts-arkitektur, hvor hver token behandles af to eksperter under inferens.
3. Lag: Modellen består af 64 transformer-lag, hver indeholdende multihead-attention og dense-blokke.
4. Tokenisering: Grok-1 anvender en SentencePiece-tokenizer med en ordforrådsstørrelse på 131.072 tokens.
5. Indlejring og positionskodning: Modellen har 6.144-dimensionale indlejring og anvender rotary positionskodning, hvilket muliggør en mere dynamisk fortolkning af data sammenlignet med traditionelle faste positionskodninger.
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, anvendende bfloat16-præcision for effektiv beregning.

Præstations- og implementeringsdetaljer

Grok-1 har demonstreret imponerende præstationer, overgående LLaMa 2 70B og Mixtral 8x7B med en MMLU-score på 73%, hvilket viser dets 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 kildeudgave fokuserer på at validere modellens korrekthed og anvender en ineffektiv MoE-lag-implementering for at undgå behovet for brugerdefinerede kernel.

Nonetheless, modellen understøtter aktiverings-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 dens vægte og arkitektur tilgængelige for den globale fællesskab til brug og bidrag.

Den åbne kildeudgave 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 muliggør let adgang til denne banebrydende model.

Mixture-of-Experts’ fremtid i sprogmodeller

Da efterspørgslen efter større og mere kapable sprogmodeller fortsætter med at vokse, forventes adoptionen af MoE-teknikker at vinde yderligere momentum. Gående forskningsindsatser fokuserer på at adressere de resterende udfordringer, såsom forbedring af træningsstabilitet, mindske overfitting under finjustering og optimere hukommelses- og kommunikationskrav.

En lovende retning er udforskningen af hierarkiske MoE-arkitekturer, hvor hver ekspert selv består af multiple under-eksperter. Dette kan muliggøre endnu større skalerbarhed og computermæssig effektivitet, mens det 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 sparse og betingede beregningsmønstre af MoE-modeller, kan yderligere forbedre deres præstation og skalerbarhed.

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

Konklusion

Mixture-of-Experts-teknikken er blevet en kraftfuld værktøj i jagten på større og mere kapable sprogmodeller. Ved at selektivt aktivere eksperter baseret på inputdataen tilbyder MoE-modeller en lovende løsning på de computermæssige udfordringer, der er forbundet med at skala 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ømæssig bæredygtighed dem til et spændende forsknings- og udviklingsområde.

Da feltet af naturlig sprogbehandling fortsætter med at presse 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 powerful og fleksible sprogmodeller, der kan forstå og kommunikere med mennesker på en naturlig og ubrudt måde.