Oppsvinget til Mixture-of-Experts for effektive store språkmodeller

I verden av naturlig språkbehandling (NLP) har jakten på å bygge større og mer kapable språkmodeller vært en drivende kraft bak mange nylige fremgang. Imidlertid, når disse modellene vokser i størrelse, blir de komputasjonelle kravene for trening og inferens mer og mer krevende, og presser mot grensene for tilgjengelige hardware-resurser.

Her kommer Mixture-of-Experts (MoE) inn, en teknikk som lover å lettet denne komputasjonelle byrden samtidig som den muliggjør trening av større og mer powerful språkmodeller. Under skal vi diskutere MoE, utforske dens opphav, indre mekanismer og dens anvendelser i transformer-baserte språkmodeller.

Opphavet til Mixture-of-Experts

Begrepet Mixture-of-Experts (MoE) kan spores tilbake til tidlig på 1990-tallet, da forskere utforsket ideen om betinget beregning, der deler av et neuralt nettverk selektivt aktiveres basert på inndata. En av de pionerende arbeidene i dette feltet var “Adaptive Mixture of Local Experts“-artikkelen av Jacobs et al. i 1991, som foreslo et overvåket læring-rammeverk for en ensemble av neurale nettverk, hver spesialisert i en annen region av inndata-rommet.

Kjerneideen bak MoE er å ha flere “eksperter”-nettverk, hver ansvarlig for å prosessere en undergruppe av inndata. En porteringsmekanisme, vanligvis et neuralt nettverk selv, bestemmer hvilke eksperter som skal prosessere en gitt inndata. Dette tilnærmingen tillater modellen å allokere sine komputasjonelle ressurser mer effektivt ved å aktivere bare de relevante eksperter for hver inndata, i stedet for å anvende full modellkapasitet for hver inndata.

Over årene har forskere utforsket og utvidet ideen om betinget beregning, noe som har ført til utviklinger som hierarkiske MoE, lav-rang-approksimasjoner for betinget beregning og teknikker for å estimere gradienter gjennom stokastiske nerver og hard-terskel-aktiveringsfunksjoner.

Mixture-of-Experts i Transformers

Mixture of Experts

Selv om ideen om MoE har vært rundt i flere tiår, dens anvendelse i transformer-baserte språkmodeller er relativt ny. Transformers, som har blitt standarden for state-of-the-art språkmodeller, består av flere lag, hver med en selv-oppmerksomhetsmekanisme og en feed-forward neuralt nettverk (FFN).

Den nøkkelinnovasjonen i å anvende MoE til transformers er å erstatte de tette FFN-lagene med sparsomme MoE-lag, hver bestående av flere eksperter og en porteringsmekanisme. Porteringsmekanismen bestemmer hvilke eksperter som skal prosessere hver inndata-token, og muliggjør at modellen selektivt aktiverer bare en undergruppe av eksperter for en gitt inndata-sekvens.

En av de tidlige arbeidene som demonstrerte potensialet til MoE i transformers var “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”-artikkelen av Shazeer et al. i 2017. Dette arbeidet innførte konseptet om et sparsomt porterings-MoE-lag, som anvendte en porteringsmekanisme som la til sparsomhet og støy til eksperter-valgsprosessen, og sikret at bare en undergruppe av eksperter ble aktivert for hver inndata.

Siden da har flere andre arbeider videreutviklet anvendelsen av MoE til transformers, og har adressert utfordringer som trening-instabilitet, lastbalansering og effektiv inferens. Notable eksempler inkluderer Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) og GLaM (Du et al., 2022).

Fordelene med Mixture-of-Experts for språkmodeller

Den primære fordelen med å anvende MoE i språkmodeller er evnen til å skalerer opp modellstørrelsen samtidig som den holder en relativt konstant komputasjonell kostnad under inferens. Ved å selektivt aktivere bare en undergruppe av eksperter for hver inndata-token, kan MoE-modeller oppnå den uttrykksfulle kraften til mye større tette modeller, samtidig som de krever betydelig mindre komputasjon.

For eksempel, betrakt en språkmodell med et tett FFN-lag på 7 milliarder parametre. Hvis vi erstatter dette laget med et MoE-lag bestående av åtte eksperter, hver med 7 milliarder parametre, øker det totale antallet parametre til 56 milliarder. Imidlertid, under inferens, hvis vi bare aktiverer to eksperter per token, er den komputasjonelle kostnaden ekvivalent til en 14 milliarder parameter tett modell, ettersom det utfører to 7 milliarder parameter matrisemultiplikasjoner.

Denne komputasjonelle effisiensen under inferens er spesielt verdifull i distribusjonsscenarioer der ressurser er begrensede, som mobil enheter eller edge-computing-miljøer. I tillegg kan den reduserte komputasjonelle kostnaden under trening føre til betydelige energibesparelser og en lavere karbonavtrykk, i tråd med den økende fokuset på bærekraftig AI-praksis.

Udfordringer og overveielser

Selv om MoE-modeller tilbyr overbevisende fordeler, kommer deres tilpasning og distribusjon også med flere udfordringer og overveielser:

1. Trening-instabilitet: MoE-modeller er kjent for å være mer utsatt for trening-instabilitet sammenlignet med deres tette motparter. Dette problem oppstår fra den sparsomme og betingede naturen til eksperter-aktivering, som kan føre til utfordringer i gradient-propagasjon og konvergens. Teknikker som router z-loss (Zoph et al., 2022) har blitt foreslått for å mildne disse instabilitetene, men videre forskning er fortsatt nødvendig.
2. Finjustering og overfitting: MoE-modeller har en tendens til å overfitte mer lett under finjustering, spesielt når nedstrøms-oppgaven har et relativt lite datasett. Dette atferd skyldes den økte kapasiteten og sparsomheten til MoE-modeller, som kan føre til overspesialisering på treningdata. Omsorgsfulle regularisering og finjustering-strategier er nødvendige for å mildne dette problemet.
3. Minnekrav: Selv om MoE-modeller kan redusere komputasjonelle kostnader under inferens, har de ofte høyere minnekrav sammenlignet med tette modeller av samme størrelse. Dette skyldes at alle eksperter-vektorer må lastes inn i minnet, selv om bare en undergruppe aktiveres for hver inndata. Minnebegrensninger kan begrense skalerbarheten til MoE-modeller på ressurssvake enheter.
4. Lastbalansering: For å oppnå optimal komputasjonell effisiens, er det avgjørende å balansere lasten over eksperter, og sikre at ingen enkelt ekspert er overbelastet mens andre forblir underutnyttet. Dette lastbalanseringen oppnås vanligvis gjennom hjelpe-tap under trening og omsorgsfulle justering av kapasitetsfaktoren, som bestemmer det maksimale antallet token som kan tildeles hver ekspert.
5. Kommunikasjons-overhead: I distribuert trening- og inferens-scenarier kan MoE-modeller introdusere ekstra kommunikasjons-overhead på grunn av behovet for å utveksle aktivering- og gradient-informasjon over eksperter som bor på forskjellige enheter eller akseleratorer. Effektive kommunikasjons-strategier og hardware-orientert modell-design er essensielle for å mildne denne overheaden.

Til tross for disse udfordringene, har det potensielle fordelene til MoE-modeller i å muliggjøre større og mer kapable språkmodeller, spurt til betydelige forskningsinnsats for å adresse og mildne disse problemene.

Eksempel: Mixtral 8x7B og GLaM

For å illustrere den praktiske anvendelsen av MoE i språkmodeller, la oss betrakte to notabile eksempler: Mixtral 8x7B og GLaM.

Mixtral 8x7B er en MoE-variant av Mistral-språkmodellen, utviklet av Anthropic. Den består av åtte eksperter, hver med 7 milliarder parametre, noe som resulterer i en total på 56 milliarder parametre. Imidlertid, under inferens, aktiveres bare to eksperter per token, og den komputasjonelle kostnaden reduseres til den til en 14 milliarder parameter tett modell.

Mixtral 8x7B har demonstrert imponerende ytelse, og overgår 70 milliarder parameter Llama-modellen samtidig som den tilbyr mye raskere inferenstider. En instruksjons-justert versjon av Mixtral 8x7B, kalt Mixtral-8x7B-Instruct-v0.1, har også blitt utgitt, og ytterligere forbedret dens evner til å følge naturlige språkinstruksjoner.

Et annet notabelt eksempel er GLaM (Google (GOOGL ) Language Model), en stor-skala MoE-modell utviklet av Google. GLaM anvender en decoder-only transformer-arkitektur og ble trent på et massivt 1,6 billion token datasett. Modellen oppnår imponerende ytelse på few-shot og one-shot-evalueringer, og matcher kvaliteten til GPT-3 samtidig som den bruker bare en tredjedel av energien som trengs for å trene GPT-3.

GLaMs suksess kan tilskrives dens effektive MoE-arkitektur, som muliggjorde trening av en modell med et stort antall parametre samtidig som den holdt en rimelig komputasjonell kostnad. Modellen demonstrerte også potensialet til MoE-modeller til å være mer energieffektive og miljøvennlige sammenlignet med deres tette motparter.

Grok-1-arkitekturen

GROK MIXTURE OF EXPERT

Grok-1 er en transformer-basert MoE-modell med en unik arkitektur designet for å maksimere effisiens og ytelse. La oss dykke ned i de viktigste spesifikasjonene:

1. Parametre: Med en imponerende 314 milliarder parametre, er Grok-1 den største åpne LLM til dags dato. Imidlertid, takket være MoE-arkitekturen, er bare 25% av vektene (cirka 86 milliarder parametre) aktive på ethvert gitt tidspunkt, og forbedrer prosesseringskapasiteten.
2. Arkitektur: Grok-1 anvender en Mixture-of-8-Experts-arkitektur, hvor hver token prosesseres av to eksperter under inferens.
3. Lag: Modellen består av 64 transformer-lag, hver med multihead-oppmerksomhet og tette blokker.
4. Tokenisering: Grok-1 anvender en SentencePiece-tokeniserer med en ordforråd på 131 072 token.
5. Embetinger og posisjonskoding: Modellen har 6 144-dimensjonale embetinger og anvender rotary-posisjonskoding, som muliggjør en mer dynamisk tolkning av data sammenlignet med tradisjonelle faste posisjonskodinger.
6. Oppmerksomhet: Grok-1 anvender 48 oppmerksomhets-hoder for forespørsler og 8 oppmerksomhets-hoder for nøkler og verdier, hver med en størrelse på 128.
7. Kontekstlengde: Modellen kan prosessere sekvenser opp til 8 192 token i lengde, og anvender bfloat16-presisjon for effektiv beregning.

Ytelse og implementeringsdetaljer

Grok-1 har demonstrert imponerende ytelse, og overgår LLaMa 2 70B og Mixtral 8x7B med en MMLU-poengsum på 73%, og viser sin effisiens og nøyaktighet over ulike tester.

Imidlertid er det viktig å merke seg at Grok-1 krever betydelige GPU-resurser på grunn av sin enorme størrelse. Den nåværende implementeringen i den åpne kildekoden fokuserer på å validere modellens riktighet og anvender en ineffektiv MoE-lag-implementering for å unngå behovet for tilpassede kjerner.

Likevel støtter modellen aktiverings-sharding og 8-bit-kvantifisering, som kan optimalisere ytelse og redusere minnekrav.

I en bemerkelsesverdig handling har xAI utgitt Grok-1 under Apache 2.0-lisensen, og gjort vektene og arkitekturen tilgjengelig for den globale samfunnet for bruk og bidrag.

Den åpne kildekoden inkluderer et JAX-eksempel-kode-repo som demonstrerer hvordan man laster og kjører Grok-1-modellen. Brukere kan laste ned checkpoint-vektene ved hjelp av en torrent-klient eller direkte gjennom HuggingFace Hub, og det er lett å få tilgang til denne banebrytende modellen.

Fremtiden for Mixture-of-Experts i språkmodeller

Ettersom etterspørselen etter større og mer kapable språkmodeller fortsetter å vokse, forventes det at tilpasningen av MoE-teknikker vil øke ytterligere. Gående forskningsinnsats fokuserer på å adresse de gjenværende utfordringene, som å forbedre trening-stabilitet, mildne overfitting under finjustering og optimalisere minne- og kommunikasjonskrav.

En lovende retning er utforskningen av hierarkiske MoE-arkitekturer, hvor hver ekspert selv består av flere under-eksperter. Dette tilnærmingen kan potensielt muliggjøre enda større skalerbarhet og komputasjonell effisiens samtidig som den holder den uttrykksfulle kraften til store modeller.

I tillegg er utviklingen av hardware- og programvare-systemer som er optimalisert for MoE-modeller en aktivt forskningsområde. Spesialiserte akseleratorer og distribuert trening-rammeverk som er designet for å håndtere den sparsomme og betingede beregningsmønsteret til MoE-modeller kan ytterligere forbedre deres ytelse og skalerbarhet.

Fremover kan integreringen av MoE-teknikker med andre fremgang i språkmodellering, som sparsom oppmerksomhet-mekanismer, effektive tokenisering-strategier og multi-modale representasjoner, føre til enda mer powerful og fleksible språkmodeller som kan håndtere en rekke oppgaver.

Konklusjon

Mixture-of-Experts-teknikken har oppstått som et kraftfullt verktøy i jakten på større og mer kapable språkmodeller. Ved å selektivt aktivere eksperter basert på inndata, tilbyr MoE-modeller et løftende løsning på de komputasjonelle utfordringene forbundet med å skalerer opp tette modeller. Selv om det fortsatt er utfordringer å overvinne, som trening-instabilitet, overfitting og minnekrav, er de potensielle fordelene til MoE-modeller i form av komputasjonell effisiens, skalerbarhet og miljøvennligheit, og gjør dem til et spennende forskningsområde.

Ettersom feltet naturlig språkbehandling fortsetter å presse grensene for hva som er mulig, vil tilpasningen av MoE-teknikker sannsynligvis spille en avgjørende rolle i å muliggjøre den neste generasjonen av språkmodeller. Ved å kombinere MoE med andre fremgang i modell-arkitektur, treningsteknikker og hardware-optimalisering, kan vi se frem til enda mer powerful og fleksible språkmodeller som kan forstå og kommunisere med mennesker på en naturlig og sømløs måte.