Mixture-of-Experts-teknikens uppgång för effektiva stora språkmodeller

I världen av naturlig språkbehandling (NLP) har jakten på att bygga större och mer kapabla språkmodeller varit en drivande kraft bakom många nyliga framsteg. Men när dessa modeller växer i storlek blir de beräkningsmässiga kraven för utbildning och inferens alltmer krävande, vilket trycker mot gränserna för tillgängliga hårdvaruresurser.

Här kommer Mixture-of-Experts (MoE) in, en teknik som lovar att lindra denna beräkningsbörda samtidigt som den möjliggör utbildning av större och mer kraftfulla språkmodeller. Nedan kommer vi att diskutera MoE, utforska dess ursprung, inre funktioner och dess tillämpningar i transformer-baserade språkmodeller.

Mixture-of-Experts ursprung

Begreppet Mixture-of-Experts (MoE) kan spåras tillbaka till början av 1990-talet när forskare undersökte idén om villkorlig beräkning, där delar av ett neuronnätverk selektivt aktiveras baserat på indata. Ett av de banbrytande arbetena inom detta område var “Adaptive Mixture of Local Experts“-artikeln av Jacobs et al. 1991, som föreslog en övervakad inlärningsram för en ensemble av neuronnätverk, var och en specialiserad på en annan region av indatautrymmet.

Den grundläggande idén bakom MoE är att ha flera “experter”-nätverk, var och en ansvarig för att bearbeta en delmängd av indata. En grindmekanism, vanligtvis ett neuronnätverk i sig, bestämmer vilken/vilka experter som ska bearbeta en given indata. Denna tillvägagångssätt möjliggör för modellen att allokera sina beräkningsresurser mer effektivt genom att endast aktivera relevanta experter för varje indata, snarare än att anställa full modellkapacitet för varje indata.

Under åren har olika forskare utforskat och utvidgat idén om villkorlig beräkning, vilket lett till utvecklingar som hierarkiska MoE, låg-rankapproximationer för villkorlig beräkning och tekniker för att uppskatta gradienter genom stokastiska neuroner och hårda tröskelaktiveringsfunktioner.

Mixture-of-Experts i Transformers

Mixture of Experts

Medan idén om MoE har funnits i decennier, dess tillämpning på transformer-baserade språkmodeller är relativt nylig. Transformers, som har blivit standarden för toppmoderna språkmodeller, består av flera lager, var och en innehållande en självuppmärksamhetsmekanism och ett feed-forward neuronnätverk (FFN).

Den nyckelinnovationen i att tillämpa MoE på transformers är att ersätta de täta FFN-lagren med glesa MoE-lager, var och en bestående av flera expert-FFN och en grindmekanism. Grindmekanismen bestämmer vilken/vilka experter som ska bearbeta varje indata-token, vilket möjliggör för modellen att selektivt aktivera endast en delmängd av experter för en given indatasekvens.

Ett av de tidiga arbetena som visade potentialen för MoE i transformers var “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”-artikeln av Shazeer et al. 2017. Detta arbete introducerade begreppet ett sparsamt grindat MoE-lager, som använde en grindmekanism som lade till sparsitet och brus till experturvalprocessen, vilket säkerställde att endast en delmängd av experter aktiverades för varje indata.

Sedan dess har flera andra arbeten ytterligare utvecklat tillämpningen av MoE på transformers, och adresserat utmaningar som utbildningsinstabilitet, lastbalansering och effektiv inferens. Noterbara exempel inkluderar Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) och GLaM (Du et al., 2022).

Fördelar med Mixture-of-Experts för språkmodeller

Den primära fördelen med att använda MoE i språkmodeller är förmågan att skala upp modellstorleken samtidigt som den relativa beräkningskostnaden under inferens hålls konstant. Genom att selektivt aktivera endast en delmängd av experter för varje indata-token kan MoE-modeller uppnå den uttrycksfulla kraften hos mycket större täta modeller samtidigt som de kräver avsevärt mindre beräkning.

Till exempel, överväg en språkmodell med ett tätt FFN-lager på 7 miljarder parametrar. Om vi ersätter detta lager med ett MoE-lager bestående av åtta experter, var och en med 7 miljarder parametrar, ökar det totala antalet parametrar till 56 miljarder. Men under inferens, om vi endast aktiverar två experter per token, är den beräkningsmässiga kostnaden ekvivalent med en 14 miljarders parameter tät modell, eftersom det beräknar två 7 miljarders parameter matris multiplikationer.

Denna beräkningsmässiga effektivitet under inferens är särskilt värdefull i distributions-scenarier där resurser är begränsade, såsom mobila enheter eller edge-beräkningsmiljöer. Dessutom kan de minskade beräkningskraven under utbildning leda till betydande energibesparingar och ett lägre koldioxidavtryck, i linje med den ökande betoningen på hållbara AI-praxis.

Utmaningar och överväganden

Medan MoE-modeller erbjuder lockande fördelar, kommer deras antagande och distribution också med flera utmaningar och överväganden:

1. Utbildningsinstabilitet: MoE-modeller är kända för att vara mer benägna till utbildningsinstabilitet jämfört med deras täta motsvarigheter. Detta problem uppstår från den glesa och villkorliga naturen hos expertaktiveringarna, som kan leda till utmaningar i gradientpropagation och konvergens. Tekniker som router z-förlust (Zoph et al., 2022) har föreslagits för att mildra dessa instabiliteter, men ytterligare forskning behövs.
2. Finjustering och överanpassning: MoE-modeller tenderar att överanpassa sig mer lätt under finjustering, särskilt när nedströmsuppgiften har ett relativt litet datamängd. Detta beteende tillskrivs den ökade kapaciteten och glesheten hos MoE-modeller, som kan leda till överspecialisering på utbildningsdata. Försiktig regularisering och finjusteringsstrategier krävs för att mildra detta problem.
3. Minneskrav: Medan MoE-modeller kan minska beräkningskostnaderna under inferens, har de ofta högre minneskrav jämfört med täta modeller av liknande storlek. Detta beror på att alla expertvikter måste laddas in i minnet, även om endast en delmängd aktiveras för varje indata. Minnesbegränsningar kan begränsa skalbarheten hos MoE-modeller på resursbegränsade enheter.
4. Lastbalansering: För att uppnå optimal beräkningsmässig effektivitet är det avgörande att balansera lasten över experter, säkerställande att ingen enskild expert är överbelastad medan andra förblir underutnyttjade. Denna lastbalansering uppnås vanligtvis genom auxiliära förluster under utbildning och noggrann justering av kapacitetsfaktorn, som bestämmer det maximala antalet token som kan tilldelas varje expert.
5. Kommunikationsöverhuvud: I distribuerad utbildning och inferensscenarier kan MoE-modeller introducera ytterligare kommunikationsöverhuvud på grund av behovet av att utbyta aktiverings- och gradientinformation över experter som bor på olika enheter eller acceleratorer. Effektiva kommunikationsstrategier och hårdvaru-medveten modell-design är avgörande för att mildra detta överhuvud.

Trots dessa utmaningar har den potentiella fördelen med MoE-modeller i att möjliggöra större och mer kapabla språkmodeller väckt betydande forskningsinsatser för att adressera och mildra dessa problem.

Exempel: Mixtral 8x7B och GLaM

För att illustrera den praktiska tillämpningen av MoE i språkmodeller, låt oss överväga två noterbara exempel: Mixtral 8x7B och GLaM.

Mixtral 8x7B är en MoE-variant av Mistral-språkmodellen, utvecklad av Anthropic. Den består av åtta experter, var och en med 7 miljarder parametrar, vilket resulterar i totalt 56 miljarder parametrar. Men under inferens aktiveras endast två experter per token, vilket effektivt minskar den beräkningsmässiga kostnaden till den för en 14 miljarders parameter tät modell.

Mixtral 8x7B har visat imponerande prestanda, överträffande den 70 miljarders parameter Llama-modellen medan den erbjuder mycket snabbare inferenstider. En instruktionsjusterad version av Mixtral 8x7B, kallad Mixtral-8x7B-Instruct-v0.1, har också släppts, vilket ytterligare förbättrar dess förmågor i att följa naturliga språkinstruktioner.

Ett annat noterbart exempel är GLaM (Google Language Model), en stor MoE-modell utvecklad av Google. GLaM använder en decoder-endast transformer-arkitektur och tränades på en massiv 1,6 biljoner token datamängd. Modellen uppnår imponerande prestanda på få-skott och en-skott utvärderingar, matchande kvaliteten hos GPT-3 medan den använder endast en tredjedel av den energi som krävs för att träna GPT-3.

GLaMs framgång kan tillskrivas dess effektiva MoE-arkitektur, som möjliggjorde utbildning av en modell med ett stort antal parametrar samtidigt som den upprätthöll rimliga beräkningskrav. Modellen visade också potentialen för MoE-modeller att vara mer energi-effektiva och miljövänliga jämfört med deras täta motsvarigheter.

Grok-1-arkitekturen

GROK MIXTURE OF EXPERT

Grok-1 är en transformer-baserad MoE-modell med en unik arkitektur utformad för att maximera effektivitet och prestanda. Låt oss dyka in i de viktigaste specifikationerna:

1. Parametrar: Med en imponerande 314 miljarder parametrar är Grok-1 den största öppna LLM hittills. Men tack vare MoE-arkitekturen är endast 25% av vikterna (cirka 86 miljarder parametrar) aktiva vid varje given tidpunkt, vilket förbättrar bearbetningsförmågor.
2. Arkitektur: Grok-1 använder en Mixture-of-8-Experts-arkitektur, där varje token bearbetas av två experter under inferens.
3. Lager: Modellen består av 64 transformerlager, var och ett innehållande multi-huvuduppmärksamhet och täta block.
4. Tokenisering: Grok-1 använder en SentencePiece-tokenisator med en ordförrådsstorlek på 131 072 token.
5. Inbäddningar och positionell kodning: Modellen har 6 144-dimensionella inbäddningar och använder roterande positionella inbäddningar, vilket möjliggör en mer dynamisk tolkning av data jämfört med traditionella fasta positionella inbäddningar.
6. Uppmärksamhet: Grok-1 använder 48 uppmärksamhets-huvuden för frågor och 8 uppmärksamhets-huvuden för nycklar och värden, var och en med en storlek på 128.
7. Sammanhangslängd: Modellen kan bearbeta sekvenser upp till 8 192 token i längd, med bfloat16 precision för effektiv beräkning.

Prestanda och implementeringsdetaljer

Grok-1 har visat imponerande prestanda, överträffande LLaMa 2 70B och Mixtral 8x7B med en MMLU-poäng på 73%, vilket visar dess effektivitet och noggrannhet över olika tester.

Men det är viktigt att notera att Grok-1 kräver betydande GPU-resurser på grund av sin enorma storlek. Den nuvarande implementeringen i den öppna källkoden fokuserar på att validera modellens korrekthet och använder en ineffektiv MoE-lagerimplementering för att undvika behovet av anpassade kärnor.

Trots detta stöder modellen aktiverings-sharding och 8-bitars kvantisering, vilket kan optimera prestanda och minska minneskrav.

I en anmärkningsvärd drag, har xAI släppt Grok-1 under Apache 2.0-licensen, vilket gör dess vikter och arkitektur tillgängliga för den globala gemenskapen för användning och bidrag.

Den öppna källkodsutgåvan innehåller ett JAX-exempelkodsrepository som visar hur man laddar och kör Grok-1-modellen. Användare kan ladda ner checkpoint-vikterna med en torrent-klient eller direkt via HuggingFace Hub, vilket underlättar tillgång till denna banbrytande modell.

Mixture-of-Experts framtid i språkmodeller

Medan efterfrågan på större och mer kapabla språkmodeller fortsätter att växa, förväntas antagandet av MoE-tekniker att vinna ytterligare momentum. Pågående forskningsinsatser fokuserar på att adressera de återstående utmaningarna, såsom att förbättra utbildningsstabilitet, mildra överanpassning under finjustering och optimera minnes- och kommunikationskrav.

En lovande riktning är utforskningen av hierarkiska MoE-arkitekturer, där varje expert i sig består av flera underexperter. Denna tillvägagångssätt kan potentiellt möjliggöra ännu större skalbarhet och beräkningsmässig effektivitet samtidigt som den upprätthåller den uttrycksfulla kraften hos stora modeller.

Dessutom är utvecklingen av hårdvaru- och programvarusystem som är optimerade för MoE-modeller ett aktivt forskningsområde. Specialiserade acceleratorer och distribuerade utbildningsramverk som är utformade för att effektivt hantera de glesa och villkorliga beräkningsmönstren hos MoE-modeller kan ytterligare förbättra deras prestanda och skalbarhet.

Förutom detta kan integrationen av MoE-tekniker med andra framsteg inom språkmodellering, såsom gles uppmärksamhet, effektiv tokenisering och multimodala representationer, leda till ännu mer kraftfulla och mångsidiga språkmodeller som kan hantera en mängd olika uppgifter.

Slutsats

Mixture-of-Experts-tekniken har uppstått som ett kraftfullt verktyg i jakten på större och mer kapabla språkmodeller. Genom att selektivt aktivera experter baserat på indata, erbjuder MoE-modeller ett lovande svar på de beräkningsmässiga utmaningarna som är förknippade med att skala upp täta modeller. Medan det fortfarande finns utmaningar att övervinna, såsom utbildningsinstabilitet, överanpassning och minneskrav, gör de potentiella fördelarna med MoE-modeller i termer av beräkningsmässig effektivitet, skalbarhet och miljövänlighet dem till ett spännande forsknings- och utvecklingsområde.

Medan fältet för naturlig språkbehandling fortsätter att driva gränserna för vad som är möjligt, är det troligt att antagandet av MoE-tekniker kommer att spela en avgörande roll i att möjliggöra nästa generation av språkmodeller. Genom att kombinera MoE med andra framsteg inom modellarkitektur, utbildningstekniker och hårdvaruoptimering, kan vi se fram emot ännu mer kraftfulla och mångsidiga språkmodeller som kan förstå och kommunicera med människor på ett naturligt och sömlöst sätt.