Uppgången av Mixture-of-Experts 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 senaste 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 kraftfullare 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 ett övervakat lärande-ramverk för en ensemble av neuronnätverk, var och en specialiserad på en annan region av indatautrymmet.

Kärnidé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 expert(er) som ska bearbeta en given indata. Denna approach möjliggör för modellen att allokera sina beräkningsresurser mer effektivt genom att aktivera endast de relevanta experterna 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 har 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 hard-tröskelaktiveringsfunktioner.

Mixture-of-Experts i Transformer

Mixture of Experts

Medan idén om MoE har funnits i decennier, dess tillämpning på transformer-baserade språkmodeller är relativt ny. Transformer, 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 viktigaste innovationen i att tillämpa MoE på transformer ä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 expert(er) 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 transformer var “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer“-artikeln av Shazeer et al. 2017. Detta arbete introducerade begreppet ett sparsely-gated MoE-lager, som använde en grindmekanism som lade till sparsitet och brus i experturvalsprocessen, 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å transformer, och har hanterat 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 betydligt 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 miljarder parameter tät modell, eftersom det beräknar två 7 miljarder 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 en lägre koldioxidavtryck, i linje med den växande 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 att utbildningsinstabilitet jämfört med deras täta motsvarigheter. Denna fråga uppstår från den glesa och villkorliga naturen av expertaktiveringar, 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 den 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äkningskostnader under inferens, har de ofta högre minneskrav jämfört med täta modeller av liknande storlek. Detta beror på att alla expertviktningar 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, för att säkerställa att ingen enda expert är överbelastad medan andra förblir underutnyttjade. Denna lastbalansering uppnås vanligtvis genom hjälp av auxiliära förluster under utbildning och noggrann finjustering av kapacitetsfaktorn, som bestämmer det maximala antalet token som kan tilldelas varje expert.
5. Kommunikationsöverföring: I distribuerad utbildning och inferensscenarier kan MoE-modeller införa ytterligare kommunikationsöverföring 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 denna överföring.

Trots dessa utmaningar har de potentiella fördelarna med MoE-modeller i att möjliggöra större och mer kapabla språkmodeller sporrat betydande forskningsinsatser för att hantera och mildra dessa frågor.

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 miljarder parameter tät modell.

Mixtral 8x7B har visat imponerande prestanda, överträffande den 70 miljarder parameter Llama-modellen samtidigt som den erbjuder mycket snabbare inferenstider. En instruktionsfinjusterad 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å few-shot och one-shot utvärderingar, matchande kvaliteten på GPT-3 samtidigt som den använder endast en tredjedel av den energi som krävs för att träna GPT-3.

GLaM:s framgång kan tillskrivas dess effektiva MoE-arkitektur, som möjliggjorde utbildning av en modell med ett enormt antal parametrar samtidigt som den upprätthöll rimliga beräkningskrav. Modellen visade också potentialen för MoE-modeller att vara mer energisnåla 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 designad för att maximera effektivitet och prestanda. Låt oss dyka in i de viktigaste specifikationerna:

1. Parametrar: Med en förbluffande 314 miljarder parametrar är Grok-1 den största öppna LLM hittills. Men tack vare MoE-arkitekturen är endast 25% av viktorna (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 transformer-lager, var och en som inkorporerar multihead-uppmä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 implementationen i den öppna källkods-versionen fokuserar på att validera modellens korrekthet och använder en ineffektiv MoE-lagerimplementation för att undvika behovet av anpassade kärnor.

Ändå stöder modellen aktiverings-sharding och 8-bitars kvantifiering, vilket kan optimera prestanda och minska minneskrav.

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

Den öppna källkods-versionen innehåller ett JAX-exempelkod-repository som visar hur man kan ladda och köra Grok-1-modellen. Användare kan ladda ner checkpoint-viktningar med en torrent-klient eller direkt via HuggingFace Hub, vilket underlättar tillgång till denna banbrytande modell.

Framtiden för Mixture-of-Experts 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 öka ytterligare. Pågående forskningsinsatser fokuserar på att hantera 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 approach 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 mjukvarusystem som är optimerade för MoE-modeller ett aktivt forskningsområde. Specialiserade acceleratorer och distribuerade utbildningsramverk som är designade för att hantera de glesa och villkorliga beräkningsmönstren hos MoE-modeller kan ytterligare förbättra deras prestanda och skalbarhet.

Ytterligare är integrationen av MoE-tekniker med andra framsteg inom språkmodellering, såsom glesa uppmärksamhetsmekanismer, effektiva tokeniseringsstrategier och multimodala representationer, kan leda till ännu mer kraftfulla och mångsidiga språkmodeller som kan hantera en bred range av 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 lösning 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 område för forskning och utveckling.

Medan fältet för naturlig språkbehandling fortsätter att trycka gränserna för vad som är möjligt, kommer antagandet av MoE-tekniker sannolikt 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.