Utviklingen av store språkmodeller (LLM) bygget fra transformatormodeller som kun er dekoder, har spilt en avgjørende rolle i å transformere NLP-domenet (Natural Language Processing), i tillegg til å fremme ulike dyplæringsapplikasjoner, inkludert forsterkning læring, tidsserieanalyse, bildebehandling og mye mer. Til tross for deres skalerbarhet og sterke ytelse, har LLM-er bygget fra transformatormodeller som kun er dekoder, fortsatt betydelige mangler. Selv om den er uttrykksfull, krever oppmerksomhetsmekanismen i transformatoravledede LLM-er høye beregningsressurser under både inferens og trening, noe som krever betydelig minne for sekvenslengden og kvadratiske FLOP-er. Dette høye beregningskravet begrenser kontekstlengden til transformatormodeller, noe som gjør autoregressive generasjonsoppgaver proporsjonalt dyre med skala, og hindrer læring fra kontinuerlige datastrømmer og muligheten for virkelig ubegrenset sekvensbehandling.

I det siste, Statlige rommodeller (SSM-er) har demonstrert bemerkelsesverdige evner og ytelse, og konkurrerer med transformatorarkitekturmodeller i storskala modelleringsstandarder, mens de oppnår minnekompleksitet som en funksjon av sekvenslengde og lineær tid. Dessuten har Mamba, en nylig utgitt State Space Model, vist enestående ytelse i en rekke språkmodellerings- og langsekvensbehandlingsoppgaver. Samtidig har Mixture of Expert (MoE)-modeller også vist imponerende ytelse samtidig som de har redusert ventetiden og beregningskostnadene ved slutninger betydelig, om enn på bekostning av et større minneavtrykk. Med utgangspunkt i Mamba- og MoE-modeller vil denne artikkelen diskutere BlackMamba, en ny arkitektur som kombinerer Mamba State Space Model med MoE-modeller for å utnytte fordelene som tilbys av begge rammeverkene. Eksperimenter på BlackMamba har vist dens evne til å utkonkurrere det eksisterende Mamba-rammeverket og transformatorbaselinjene både i trenings-FLOP-er og slutninger. Den eksepsjonelle ytelsen til BlackMamba-rammeverket viser at det effektivt kan kombinere evnene til Mamba- og MoE-rammeverket, og tilbyr rask og kostnadseffektiv slutning fra MoE med generering av lineær kompleksitet fra Mamba.

Denne artikkelen tar sikte på å dekke BlackMamba-rammeverket i dybden. Vi utforsker mekanismen, metodikken og arkitekturen til rammeverket, sammen med dets sammenligning med toppmoderne bilde- og videogenereringsrammeverk. La oss komme i gang.

BlackMamba: En introduksjon til MoE for State Space Models

Progresjonen av store språkmodeller (LLM), spesielt de som er basert på transformatorarkitekturer som kun er dekoder, har særlig påvirket Natural Language Processing (NLP) felt og utvidet til forskjellige dyplæringsapplikasjoner, inkludert forsterkende læring, tidsserieanalyse, bildebehandling og mer. Ikke desto mindre, til tross for deres skalerbarhet og robuste ytelse, møter disse transformatorbaserte LLM-ene som kun er dekoder, bemerkelsesverdige utfordringer. Oppmerksomhet mekanisme, en nøkkelfunksjon i transformator-basert LLM-ers, krever omfattende beregningsressurser for både inferens og trening. Dette innebærer et behov for minne som vokser med sekvenslengden og beregningsoperasjoner (FLOPs) som øker kvadratisk. Slike intensive beregningsbehov begrenser modellenes kontekstlengde, øker kostnadene ved autoregressive generasjonsoppgaver etter hvert som modellen skaleres, og hindrer modellenes evne til å lære av kontinuerlige datastrømmer eller prosesssekvenser med ubegrenset lengde effektivt. 

Betydelig innsats har blitt gjort de siste årene i et forsøk på å overvinne disse begrensningene, og oppmerksomheten har blitt flyttet mot å utvikle arkitektoniske alternativer til de kanoniske tette oppmerksomhetstransformatormodellene med SSM-er og MoE-modeller som de mest lovende kandidatarkitekturene. Den viktigste fordelen høstet ved å favorisere State Space Models fremfor transformatorarkitekturmodeller er den lineære beregningskompleksiteten med hensyn til inngangssekvenslengden som tilbys av SSM-er i motsetning til den kvadratiske kompleksiteten som tilbys av transformatorer. Teoretisk sett gjør lineær beregningskompleksitet med hensyn til inngangssekvenslengde det mulig for State Space Models å behandle større sekvenser enn transformator-arkitekturmodeller for et gitt FLOPS- eller Floating-point-operasjoner per sekund-budsjett, og å gjengi autoregressiv generering konstant i beregning uten en KV-cache. Nylig utviklede State Space Models, inkludert Mamba, RetNet og noen få andre, har vist effektiv lang-sekvens inferens og trening, sammen med konkurrerende språkmodelleringsoppgaveytelse til transformatorer med lignende skaleringsegenskaper. På den annen side vinner Mixture of Expert-modellarkitekturer popularitet som et alternativ til tette transformatorer siden det muliggjør en betydelig reduksjon i inferens og trening av FLOP-er som er avgjørende for å oppnå sammenlignbar kvalitet med en tett modell. MoE (Mixture of Experts)-modeller fungerer ved å aktivere kun et sparsomt utvalg av de totale parametrene under en enkelt foroverpassering. De bruker en rutefunksjon for å bestemme hvilke "eksperter" som blir kalt til handling basert på den gitte konteksten. Denne tilnærmingen skaper et skille mellom beregningskostnadene ved slutning og det totale antallet parametere, noe som gir mulighet for forbedret ytelse innenfor et fast slutningsbudsjett, om enn med et økt antall parametere og et større minnebehov.

Dette fremskrittet innen arkitektur gir bemerkelsesverdige fordeler i forhold til tradisjonelle transformatorer og representerer en spennende retning for videre utvikling. Vi antar at integrering av disse forbedringene i en kombinert Mamba-MoE-modell kan betydelig akselerere språkmodelleringsevner og effektivitet utover standard transformatormodeller. De forventede fordelene med en Mamba-MoE-arkitektur sammenlignet med en tradisjonell tett transformatormodell inkluderer:

Mamba: Oppnår lineær beregningskompleksitet i forhold til inngangssekvenslengden for både trenings- og slutningsfaser. Den gjør det mulig for autoregressiv generering å skje i en konstant tidsramme og med konstant minnebruk.

MoE: Tilbyr inferenshastighet og treningsberegningseffektivitet som kan sammenlignes med en mindre, tett grunnlinjemodell, samtidig som den opprettholder et nivå av modellkvalitet som konkurrerer med en modell med et tilsvarende antall parametere som den tettere versjonen.

Når det er sagt, er det viktig å si at transformatorarkitekturmodeller fortsatt er toppmoderne, og har vist konsistent og bemerkelsesverdig sterk ytelse på språkmodelleringsoppgaver og sekvensbehandlingsoppgaver. I kjernen bruker transformatorarkitekturen selvoppmerksomhet som utfører en kvadratisk alt-til-alle-sammenligning av punktproduktlikhetene mellom innebyggingen av forskjellige tokens i en sekvens, og utfører et lineært kart til en utgangsvektor. Transformatormodellen består av selvoppmerksomhetsblokker stablet mellom MLP- eller Multi-Layer Perceptron-blokker som videre består av en to-lags MLP med en gitt aktiveringsfunksjon. 

BlackMamba: Arkitektur og metodikk

Statlige rommodeller

State Space Models tilhører gruppen av sekvensmodeller med lineær kompleksitet med hensyn til lengden på inngangssekvensen. Arkitekturen til State Space Models er mer på linje med tilbakevendende nevrale nettverk og konvolusjonelle nevrale nettverk i stedet for oppmerksomhetsbasert arkitektur, og er inspirert fra et kontinuerlig dynamisk system som kartlegger en 1-dimensjonal funksjon gjennom et implisitt latent rom. Et lineært dynamisk system gjør parallelle beregninger effektive ved å bruke enten en assosiativ eller en konvolusjonsskanning. I praktiske scenarier har den tilbakevendende naturen til State Space Models vært grunnen til at den fortsatt skal tas i bruk på svært parallell AI-maskinvare som GPUer. Imidlertid fremveksten av SSM-er som RWKV og Mamba har brukt parallelle skannekjerner for å kartlegge tilbakevendende operasjoner effektivt til GPUer, og dermed lette opplæringen av nye arkitekturer med effektivitet som kan sammenlignes med de som oppnås med transformatormodeller. 

Den iboende kvadratiske kompleksiteten i forhold til sekvenslengde i transformatorer er en velkjent begrensning som hindrer resonnement og forståelse over svært lange sammenhenger. Nylige innovasjoner har introdusert ideen om å utvide kontekstlengden, slik at transformatorer kan trenes i en gjennomførbar skala før de brukes i mye lengre sammenhenger under inferens. Til tross for disse fremskrittene, krever slutningsprosessen fortsatt en betydelig mengde beregningsressurser og minne, spesielt for å opprettholde Key-Value (KV)-cachen, noe som gjør det til en ressurskrevende innsats. Nyere forskningsinnsats har fokusert på å forbedre uttrykksevnen til stat-rom-modeller ved å inkorporere inngangsavhengige portmekanismer, beslektet med Query, Key, Value (QKV)-matrisene som finnes i oppmerksomhetsmekanismer. 

Disse anstrengelsene tar sikte på å bevare den iboende lineære progresjonen av tilstand-rom-rekursjon, og muliggjør effektiv utførelse gjennom enten konvolusjon eller en selektiv skanneprosess. Denne tilnærmingen reduserer ytelsesforskjellen betydelig med transformatorer i praktiske applikasjoner. Blant disse fremskrittene skiller Mamba seg ut som en stat-rom-modell som gjenspeiler målene for tidligere forskning, og viser imponerende ytelsesnivåer som kan sammenlignes med transformatorer på skalaer opp til 2.8 milliarder parametere. Den oppnår dette ved å bruke inngangsavhengig gating på inngangene til state-space model (SSM) rekursjon, samtidig som den sikrer effektiv beregning gjennom bruk av skreddersydde selektive skannekjerner.

Blanding av ekspertmodeller

Blanding av ekspertmodeller (MoE) oppnår et skille mellom slutningskostnaden og det totale parameterantallet ved selektivt å aktivere parametere under fremføringen. I stedet for å bruke alle parametere, sender disse modellene tokens til spesifikke Multilayer Perceptron (MLP)-eksperter. Ideelt sett er hver ekspert skreddersydd for å behandle en bestemt type input, med en rutingmekanisme, i hovedsak et kompakt nevralt nettverk, som bestemmer den mest passende eksperten for hvert token. Denne tilnærmingen tar sikte på å bevare den omfattende uttrykkskraften til en modell med et tilsvarende antall parametere i en tettere konfigurasjon, men med betydelig reduserte beregningskrav. Vanligvis er ruteren en kartlegging av de lineære lagene fra tokens til ekspertindekser, der hver ekspert ganske enkelt er en standard transformator Multilayer Perceptron. Utviklere har imidlertid ennå ikke funnet ut den optimale treningsmetoden for ruteren siden ekspertoppdragsproblemet er ikke-differensierbart, og Mixture of Expert-modeller sliter ofte med belastningsbalansering og treningsstabilitet mellom forskjellige eksperter for maskinvareeffektivitet. 

arkitektur

I kjernen bruker BlackMamba en standard transformatormodell som består av sammenflettede MLP-blokker og oppmerksomhetsblokker lagt til i rekkefølge langs en gjenværende strøm. Nå erstatter et flertall av Mixture of Expert-modellene ganske enkelt flerlags perceptronblokkene med et rutet ekspertlag. På den annen side erstatter BlackMamba-rammeverket ikke bare flerlags perceptronblokken i transformatoren med et rutet ekspertlag, men erstatter også oppmerksomhetslaget med et Mamba State Space Model-lag. Arkitekturen til BlackMamba-rammeverket er demonstrert i følgende figur. 

Opplæring og datasett

BlackMamba-modellen er trent på over 300 milliarder tokens på et tilpasset datasett, og bruker SwiGLU-aktiveringsfunksjonen for ekspertens flerlags-perseptroner. Rammeverket trener med 8 eksperter, et antall som utviklere fant å være den rette balansen og bytte mellom minnefotavtrykket og slutningskostnadene til modellen. Det tilpassede datasettet som brukes til å trene BlackMamba-rammeverket består av en blanding av allerede eksisterende åpen kildekode-datasett inkludert Starcoder, SlimPajama, Pile og mer. Tabellen nedenfor viser vektene til hvert av datasettet som brukes til å trene BlackMamba-rammeverket. Totalt sett er det 1.8 billioner tokens i datasettet. 

BlackMamba : Resultater

For å sikre en rettferdig sammenligning mellom Mamba og BlackMamba, har utviklere trent begge modellene med samme treningsparametere på samme treningsdata. BlackMamba-rammeverket er i stand til å utkonkurrere både Mamba- og transformatormodeller for identisk foroverpass-modellstørrelse på slutningstidspunktet, i tillegg til å trene flytepunktoperasjoner per sekund. Følgende figur viser tiden det tar å generere en sekvens med en gitt lengde autoregressivt fra en innledende ett-token-ledetekst som en funksjon av sekvenslengden. 

Videre er latenstidsfordelene til både Mixture of Expert- og Mamba-modellene kombinert i BlackMamba-rammeverket, noe som resulterer i betydelig raskere slutningstider sammenlignet med transformatormodeller, rene Mamba-modeller og MoE-modeller. Videre er slutningsfordelen med BlackMamba-rammeverket direkte proporsjonal med sekvenslengdene, noe som gjør BlackMamba ekstremt effektiv ved lang sekvensgenerering. Følgende figur illustrerer antall tokens som er tildelt BlackMamba-modellene med henholdsvis 340 millioner og 640 millioner parametere. Som det kan sees, viser et flertall av lagene et høyt nivå av ekspertbalanse som et resultat av den forbedrede Sinkhorn-algoritmen implementert av BlackMamba-modellene. 

Tabellen nedenfor dekker evalueringsresultatene til BlackMamba-rammeverket sammenlignet med en rekke forhåndstrente språkmodeller med åpen kildekode. Som det kan observeres, er BlackMamba-rammeverket i stand til å konkurrere og utkonkurrere med et flertall av rammene på tvers av alle grunnlinjer. Videre er det verdt å merke seg at modellene som overgår BlackMamba har betydelig høyere antall parametere, og gapet i ytelse er minimalt, noe som indikerer evnen til BlackMamba-rammeverket med færre parametere. 

Final Thoughts

I denne artikkelen har vi snakket om BlackMamba, en ny arkitektur som kombinerer Mamba State Space Model med Mixture of Expert-modeller for å høste fordelene som tilbys av begge disse rammeverkene. Eksperimenter på BlackMamba har vist at den overgår det eksisterende Mamba-rammeverket og transformatorbaselinjene i både trenings-FLOP-er og slutninger. Den eksepsjonelle ytelsen til BlackMamba-rammeverket viser at det er i stand til å arve og kombinere evnene til Mamba- og MoE-rammeverket eksepsjonelt godt siden det kombinerer den billige og raske slutningen fra MoE med generering av lineær kompleksitet fra Mamba. Vi har snakket om hvordan arkitekturen til BlackMamba-rammeverket er i stand til å utkonkurrere sterke trente store språkmodeller, eksisterende Mamba-rammeverk og Mixture of Expert-modeller når det gjelder trenings-FLOPs og slutningskostnader. Videre arver BlackMamba-rammeverket også generasjons-FLOP-ene og redusert trening fra både Mixture of Expert-modeller og Mamba-rammeverket samtidig.