Utviklingen av store språkmodeller (LLM) bygget fra decoder-kun transformermodeller har spilt en avgjørende rolle i å transformere det naturlige språkbehandlingsområdet (NLP), samt å fremme ulike dyp læringsapplikasjoner, inkludert forsterket læring, tidsserieanalyse, bildebehandling og mye mer. Likevel, til tross for deres skalerbarhet og sterke ytelse, møter LLM-bygget fra decoder-kun transformermodeller fremdeles betydelige begrensninger. Selv om de er uttrykksfulle, krever oppmerksomhetsmekanismen i transformer-avledede LLM betydelige beregningsressurser under både inferens og trening, og nødvendiggjør betydelig minne for sekvenslengden og kvadratisk FLOPs. Dette høye beregningskravet begrenser modellens kontekstlengde, gjør autoregressiv generering dyrt i forhold til skalaen, og hindrer modellens evne til å lære fra kontinuerlige datastrømmer og prosessere sekvenser av ubegrenset lengde effektivt.

I løpet av de siste årene har State Space-modeller (SSM) vist bemerkelsesverdige evner og ytelse, og konkurrerer med transformer-arkitekturmodeller i store modelleringstester, samt oppnår lineær minnekomplesitet som en funksjon av sekvenslengden. I tillegg har Mamba, en nylig lansert State Space-modell, vist fremragende ytelse i en rekke språkmodellering og langsekvensbehandlingsoppgaver. Samtidig har Mixture of Expert-modeller (MoE) også vist imponerende ytelse, og betydelig redusert latent og beregningskostnader for inferens, om enn på bekostning av en større minneavtrykk. Ved å bygge på Mamba og MoE-modeller, vil denne artikkelen diskutere BlackMamba, en ny arkitektur som kombinerer Mamba State Space-modellen med MoE-modeller for å utnytte fordelene til begge rammer. Eksperimenter med BlackMamba har demonstrert dens evne til å overgå den eksisterende Mamba-rammen og transformer-baselinjer i både trening FLOPs og inferens. Den eksepsjonelle ytelsen til BlackMamba-rammen viser at den kan kombinere evnene til Mamba- og MoE-rammene på en utmerket måte, og tilby rask og kostnadseffektiv inferens fra MoE med lineær-kompleksitet generering fra Mamba.

Denne artikkelen har som mål å dekke BlackMamba-rammen i dybden. Vi utforsker mekanismen, metodologien og arkitekturen til rammen, samt sammenligner den med state-of-the-art bilde- og videogenereringsrammer. La oss begynne.

BlackMamba: En introduksjon til MoE for State Space-modeller

Utviklingen av store språkmodeller (LLM), spesielt de som er basert på decoder-kun transformerarkitekturer, har hatt en betydelig innvirkning på det naturlige språkbehandlingsområdet (NLP) og har utvidet seg til ulike dyp læringsapplikasjoner, inkludert forsterket læring, tidsserieanalyse, bildebehandling og mer. Likevel, til tross for deres skalerbarhet og robuste ytelse, møter disse decoder-kun transformer-baserte LLM fremdeles betydelige utfordringer. Oppmerksomhetsmekanismen, en nøkkelkomponent i transformer-baserte LLM, krever betydelige beregningsressurser for både inferens og trening. Dette involverer et behov for minne som vokser med sekvenslengden og beregningsoperasjoner (FLOPs) som øker kvadratisk. Slike intense beregningskrav begrenser modellens kontekstlengde, øker kostnadene for autoregressiv generering som modellen skalerer, og hindrer modellens evne til å lære fra kontinuerlige datastrømmer eller prosessere sekvenser av ubegrenset lengde effektivt.

Betydelige anstrengelser har blitt gjort i de siste årene for å overvinne disse begrensningene, og oppmerksomheten har blitt rettet mot å utvikle alternative arkitekturer til de kanoniske tette oppmerksomhets transformermodellene med SSM og MoE-modeller som de mest lovende kandidatarkitekturer. Den viktigste fordelen ved å favorisere State Space-modeller over transformer-arkitekturmodeller er den lineære beregningskompleksiteten i forhold til inndata-sekvenslengden som tilbys av SSM, i motsetning til den kvadratiske kompleksiteten som tilbys av transformatorer. Teoretisk sett ermöglijer lineær beregningskompleksitet i forhold til inndata-sekvenslengden State Space-modeller å prosessere lengre sekvenser enn transformer-arkitekturmodeller for en gitt FLOPS eller flytende punkt-operasjoner per sekund-budsjett, og å gjøre autoregressiv generering konstant i beregning uten en KV-cache. Nylig utviklede State Space-modeller, inkludert RWKV og Mamba, har demonstrert effektiv langsekvens-inferens og trening, samt konkurrerende språkmodellering-ytelse med transformatorer med lignende skalerings egenskaper. På den andre siden har Mixture of Expert-modellarkitekturer blitt populære som en alternativ til tette transformatorer, ettersom de ermöglijer en betydelig reduksjon i inferens- og trening FLOPs, essensielt for å oppnå sammenlignbar kvalitet med en tett modell. MoE (Mixture of Experts) modeller opererer ved å aktivere bare en sparsom seleksjon av de totale parameterne under en enkelt fremover-passering. De bruker en ruting-funksjon til å bestemme hvilke “eksperter” som kalles inn i aksjon basert på den gitt konteksten. Dette tilnærmingen skaper en separasjon mellom beregningskostnaden for inferens og det totale antallet parametre, og ermöglijer forbedret ytelse innenfor et fast inferens-budsjett, om enn med en økt antall parametre og en større minneavtrykk.

Dette fremgangen i arkitektur tilbyr betydelige fordeler over tradisjonelle transformatorer og representerer en spennende retning for videre utvikling. Vi antar at integrering av disse forbedringene i en kombineret Mamba-MoE-modell kunne signifikant akselerere språkmodellering og effektivitet utover det som er mulig med standard transformer-modeller. De forventede fordelene med en Mamba-MoE-arkitektur i forhold til en tradisjonell tett transformer-modell inkluderer:

Mamba: Oppnår lineær beregningskompleksitet i forhold til inndata-sekvenslengden for både trening og inferens-faser. Den ermöglijer autoregressiv generering å skje i en konstant tidsramme og med konstant minnebruk.

MoE: Tilbyr inferens-hastighet og trening-beregnings-effektivitet sammenlignbar med en mindre, tett baseline-modell, mens den opprettholder en modellkvalitet som rivaliserer med den til en modell med en tilsvarende antall parametre som den tette versjonen.

Med det sagt, er det essensielt å påpeke at transformer-arkitekturmodeller fortsatt er state-of-the-art, og har demonstrert konsistent og bemerkelsesverdige sterke ytelse på språkmodellering og sekvensbehandlingsoppgaver. I kjernen anvender transformer-arkitekturen selv-oppmerksomhet som utfører en kvadratisk all-til-all-sammenligning av dot-produkt-lignende mellom innleggelsene av ulike token i en sekvens, og utfører en lineær kartlegging til en utgangsvektor. Transformer-modellen består av selv-oppmerksomhets-blokker stablet mellom MLP eller Multi-Layer Perceptron-blokker som videre består av en to-lags MLP med en gitt aktiveringsfunksjon.

BlackMamba: Arkitektur og Metodologi

State Space-modeller

State Space-modeller tilhører gruppen av sekvensmodeller med lineær kompleksitet i forhold til lengden av inndata-sekvensen. Arkitekturen til State Space-modeller ligner mer på Resurrent Neural Networks og Convolutional Neural Networks enn på oppmerksomhets-basert arkitektur, og er inspirert av et kontinuerlig dynamisk system som kartlegger en 1-dimensjonal funksjon gjennom en implisitt latent rom. Et lineært dynamisk system gjør parallell beregning effektiv ved å bruke enten en assosiativ eller en konvolusjonsskanning. I praktiske scenarier har den resurrente naturen til State Space-modeller vært årsaken til at de fortsatt ikke er adoptert på høyt parallell AI-hardware som GPU-er. Likevel har fremveksten av SSM-er som RWKV og Mamba brukt parallell skanning-kjerner til å kartlegge resurrente operasjoner effektivt til GPU-er, og ermöglijer trening av nye arkitekturer med effektivitet sammenlignbar med den som oppnås av transformer-modeller.

Den innebygde kvadratiske kompleksiteten i forhold til sekvenslengde innen transformatorer er en velkjent begrensning som hindrer forståelse og komprehensjon over svært lange kontekster. Nylige innovasjoner har introdusert ideen om å utvide kontekstlengden, og ermöglijer transformatorer å bli trent på en gjennomførbar skala før de blir anvendt på mye lengre kontekster under inferens. Likevel, krever inferens-prosessen fremdeles en betydelig mengde beregningsressurser og minne, spesielt for å opprettholde Key-Value (KV)-cachen, og gjør det til en ressurs-intensiv oppgave. Nylige forskningsinnsats har fokusert på å forbedre de uttrykksfulle evnene til State Space-modeller ved å inkorporere inndata-avhengige porterings-mekanismer, lignende de som finnes i oppmerksomhets-mekanismer.

Disse innsatsene søker å bevare den innebygde lineære fremdriften til State Space-rekursjon, og ermöglijer effektiv utførelse gjennom enten konvolusjon eller en selektiv skanning-prosess. Dette tilnærmingen reduserer betydelig ytelsesforskjellen med transformatorer i praktiske applikasjoner. Blant disse fremgangene, utmerker Mamba seg som en State Space-modell som speiler målene til tidligere forskning, og viser imponerende ytelse på nivå med transformatorer på skalaer opp til 2,8 milliarder parametre. Den oppnår dette ved å anvende inndata-avhengige porterings-mekanismer til inngangene til State Space-modell-rekursjonen, samtidig som den sikrer effektiv beregning gjennom å bruke spesialbygde selektive skanning-kjerner.

Mixture of Expert-modeller

Mixture of Expert (MoE) modeller oppnår en separasjon mellom inferens-kostnaden og det totale antallet parametre ved å selektivt aktivere parametre under fremover-passeringen. I stedet for å bruke alle parametre, retter disse modellene token til bestemte Multilayer Perceptron (MLP) eksperter. Ideelt sett er hver ekspert tilpasset å prosessere en bestemt type inndata, med en ruting-mekanisme, essensielt en kompakt neural nettverk, som bestemmer den mest passende eksperten for hver token. Dette tilnærmingen søker å bevare den omfattende uttrykksfulle kraften til en modell med en tilsvarende antall parametre i en tetter konfigurasjon, men med betydelig redusert beregningskostnader. Vanligvis er ruterne en mapping av lineære lag fra token til ekspert-indekser med hver ekspert som en standard transformer Multilayer Perceptron. Likevel, er utviklerne fremdeles i ferd med å finne den optimale treningsmetoden for ruterne, ettersom ekspert-tildelings-problemet er ikke-differensierbart, og Mixture of Expert-modeller ofte sliter med last-balansering og trenings-stabilitet for hårdvaru-effektivitet.

Arkitektur

I kjernen anvender BlackMamba en standard transformer-modell bestående av vekselvis MLP-blokker og oppmerksomhets-blokker lagt til i sekvens langs en residual-strøm. Nå, erstatter de fleste Mixture of Expert-modeller enkelt og rett Multilayer Perceptron-blokkene med en ruter-ekspert-lag. På den andre siden, erstatter BlackMamba-rammen ikke bare Multilayer Perceptron-blokken i transformator med en ruter-ekspert-lag, men også oppmerksomhets-laget med en Mamba State Space-modell-lag. Arkitekturen til BlackMamba-rammen demonstreres i figuren nedenfor.

Trening og Dataset

BlackMamba-modellen er trent på over 300 milliarder token på et eget dataset, og bruker SwiGLU-aktiveringsfunksjonen for ekspert-Multilayer Perceptron. Rammen trener med 8 eksperter, et antall som utviklerne fant å være riktig balanse og kompromiss mellom minneavtrykk og inferens-kostnad for modellen. Det egendefinerte datasett som brukes til å trene BlackMamba-rammen består av en blanding av allerede eksisterende åpne kilde-datasett, inkludert Starcoder, SlimPajama, Pile og mer. Tabellen nedenfor demonstrerer vektene til hvert av datasettene som brukes til å trene BlackMamba-rammen. Totalt er det 1,8 billioner token i datasettet.

BlackMamba: Resultater

For å sikre en rettferdig sammenligning mellom Mamba og BlackMamba, har utviklerne trent begge modellene med samme treningsparametre på samme treningsdata. BlackMamba-rammen er i stand til å overgå både Mamba og transformer-modeller for identisk fremover-passering modell-størrelse på inferens-tid samt trening Floating-point operasjoner per sekund. Figuren nedenfor demonstrerer tiden det tar å generere en sekvens av en gitt lengde autoregressivt fra en initial en-token-prompt som en funksjon av sekvenslengden.

Videre, kombinerer latency-fordelene til både Mixture of Expert og Mamba-modeller i BlackMamba-rammen, og resulterer i signifikant raskere inferens-tider sammenlignet med transformer-modeller, rene Mamba-modeller og MoE-modeller. Videre, er inferens-fordelene til BlackMamba-rammen direkte proporsjonale med sekvenslengdene, og gjør BlackMamba ekstremt effektiv for lang sekvens-generering. I tillegg, illustrerer figuren nedenfor antallet token som er tilordnet BlackMamba-modellene med 340 millioner og 640 millioner parametre henholdsvis. Som det kan ses, viser de fleste lagene en høy grad av ekspert-balansering som et resultat av den forbedrede Sinkhorn-algoritmen implementert av BlackMamba-modellene.

Tabellen nedenfor dekker evalueringsscorene til BlackMamba-rammen sammenlignet med en rekke åpne forhånds-trente språkmodeller. Som det kan observeres, er BlackMamba-rammen i stand til å konkurrere og overgå de fleste rammer over alle baselinjer. Videre, er det verdt å merke seg at modellene som overgår BlackMamba har betydelig flere parametre, og gapet i ytelse er minimalt, og indikerer evnen til BlackMamba-rammen med færre parametre.

Slutt-tanker

I denne artikkelen, har vi diskutert BlackMamba, en ny arkitektur som kombinerer Mamba State Space-modellen med MoE-modeller for å utnytte fordelene til begge rammer. Eksperimenter med BlackMamba har demonstrert dens evne til å overgå den eksisterende Mamba-rammen og transformer-baselinjer i både trening FLOPs og inferens. Den eksepsjonelle ytelsen til BlackMamba-rammen viser at den kan kombinere evnene til Mamba- og MoE-rammene på en utmerket måte, og tilby rask og kostnadseffektiv inferens fra MoE med lineær-kompleksitet generering fra Mamba. Vi har diskutert hvordan arkitekturen til BlackMamba-rammen kan overgå sterke trent Large Language Models, eksisterende Mamba-ramme og Mixture of Expert-modeller i forhold til trening FLOPs og inferens-kostnad. Videre, arver BlackMamba-rammen også genererings-FLOPs og redusert trening fra både Mixture of Expert-modeller og Mamba-ramme samtidig.