Utviklingen av store språkmodeller (LLM) bygget fra decoder-only transformer-modeller har spilt en avgjørende rolle i å transformere det naturlige språkbehandlingsområdet (NLP), samt å fremme diverse dype læringsapplikasjoner inkludert forsterkingslæring, tidsserieanalyse, bildebehandling og mye mer. Likevel, til tross for deres skalerbarhet og sterke ytelse, møter LLM-bygget fra decoder-only transformer-modeller fortsatt betydelige mangler. Selv om de er uttrykksfulle, krever oppmerksomhetsmekanismen i transformer-avledede LLM høye beregningsressurser under både inferens og trening, og nødvendiggjør betydelig minne for sekvenslengden og kvadratisk FLOPs. Dette høye beregningskravet begrenser kontekstlengden til transformer-modellene, gjør autoregressiv generering dyrt i forhold til skalaen, og hindrer læring fra kontinuerlige datastrømmer og evnen til å behandle 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 skalamodell-benchmark, samt oppnår minnekomplessitet som en funksjon av sekvenslengden og lineær tid. I tillegg har Mamba, en nylig lansert State Space-modell, vist fremragende ytelse i en rekke språkmodellering- og lange-sekvensbehandlingsoppgaver. Samtidig har Mixture of Expert-modeller (MoE) også vist imponerende ytelse, samtidig som de betydelig reduserer latency og beregningskostnader for inferens, om enn på bekostning av en større minneavtrykk. Bygget 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 på 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 effektivt, 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, metoden og arkitekturen til rammen, samt sammenligningen med bilde- og video-genereringsrammer. 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-only transformer-arkitekturer, har hatt en betydelig innvirkning på det naturlige språkbehandlingsområdet (NLP) og utvidet til diverse dype læringsapplikasjoner, inkludert forsterkingslæring, tidsserieanalyse, bildebehandling og mye mer. Likevel, til tross for deres skalerbarhet og robuste ytelse, møter disse decoder-only transformer-baserte LLM betydelige utfordringer. Oppmerksomhetsmekanismen, en nøkkelkomponent i transformer-baserte LLM, krever omfattende beregningsressurser for både inferens og trening. Dette involverer et behov for minne som vokser med sekvenslengden og beregningsoperasjoner (FLOPs) som øker kvadratisk. Slike intensive beregningsbehov begrenser modellens kontekstlengde, øker kostnadene for autoregressiv generering oppgaver i forhold til skalaen, og hindrer modellens evne til å lære fra kontinuerlige datastrømmer eller behandle sekvenser av ubegrenset lengde effektivt.

Betydelige anstrengelser har blitt gjort i de siste årene for å overvinne disse begrensningene, og fokus har blitt flyttet mot å utvikle alternative arkitekturer til de kanoniske tette oppmerksomhets-transformer-modellene med SSM og MoE-modeller som de mest lovende kandidat-arkitekturer. Den viktigste fordelen med å 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 transformer. Teoretisk sett tillater den lineære beregningskompleksiteten State Space-modeller å behandle 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 Mamba, RetNet og noen andre, har demonstrert effektiv lang-sekvens inferens og trening, samt konkurrerende språkmodellering-oppgave-ytelse til transformer med lignende skalerings-egenskaper. På den andre siden har Mixture of Expert-modell-arkitekturer også vist imponerende ytelse, samtidig som de betydelig reduserer latency og beregningskostnader for inferens, om enn på bekostning av en større minneavtrykk.

Dette fremgangen i arkitektur tilbyr betydelige fordeler over tradisjonelle transformer og representerer en spennende retning for videre utvikling. Vi påstår at integrering av disse forbedringene i en kombineret Mamba-MoE-modell kunne betydelig akselerere språkmodellering-evner og effektivitet utover det som er mulig med standard transformer-modeller. De forventede fordelene med en Mamba-MoE-arkitektur sammenlignet med 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 gjør autoregressiv generering mulig i en konstant tidsramme og med konstant minnebruk.

MoE: Tilbyr inferens-hastighet og trening-beregningseffektivitet sammenlignet med en mindre, tett baseline-modell, samtidig som den opprettholder en modellkvalitet som rivaliserer med 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 bemerkelsesverdig sterk ytelse på språkmodellering-oppgaver og sekvensbehandlingsoppgaver. I kjernen anvender transformer-arkitekturen selv-oppmerksomhet som utfører en kvadratisk all-til-all sammenligning av dot-produkt-ligninger mellom innleggelsene av forskjellige token i en sekvens, og utfører en lineær kartlegging til en utgangsvektor.

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 hjelp av enten en assosiativ eller en konvolusjons-scan. I praktiske scenarier har den resurrente naturen til State Space-modeller vært grunnen til at de fortsatt ikke er adoptert på høyt parallell AI-hardware som GPU-er. Likevel har fremveksten av SSM som RWKV og Mamba brukt parallell scan-kjerner til å kartlegge resurrente operasjoner effektivt til GPU-er, og dermed muliggjort trening av nye arkitekturer med effektivitet sammenlignet med transformer-modeller.

Den innebygde kvadratiske kompleksiteten i forhold til sekvenslengde innenfor transformer er en velkjent begrensning som hemmer forståelse og resonnering over svært lange kontekster. Nylige innovasjoner har introdusert ideen om å utvide kontekstlengden, og muliggjort at transformer kan trenes på en gjennomførbart skala før de anvendes på mye lengre kontekster under inferens. Likevel krever inferens-prosessen fortsatt 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 gating-mekanismer, lignende de som finnes i oppmerksomhets-mekanismer.

Disse innsatsene har som mål å bevare den lineære fremgangen til State Space-rekursjon, og å muliggjøre effektiv utførelse gjennom enten konvolusjon eller en selektiv scan-prosess. Dette tilnærminger nærmere ytelsesforskjellen med transformer 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å niveau med transformer på skalaer opptil 2,8 milliarder parametre. Den oppnår dette ved å anvende inndata-avhengig gating til inngangene til State Space-modellens (SSM) rekursjon, samtidig som den sikrer effektiv beregning gjennom bruk av spesialtilpassede selektive scan-kjerner.

Mixture of Expert-modeller

Mixture of Expert-modeller (MoE) oppnår en separasjon mellom inferens-kostnaden og det totale antall parametre ved å selektivt aktivere parametre under en enkelt fremad-gående passering. I stedet for å bruke alle parametre, dirigerer disse modellene token til bestemte Multilayer Perceptron (MLP)-eksperter. Ideelt sett er hver ekspert tilpasset å prosessere en bestemt type inndata, med en router-mekanisme, som i realiteten er et kompaktnettverk, som bestemmer den mest passende eksperten for hver token. Dette tilnærminger har som mål å bevare den omfattende uttrykksfulle kraften til en modell med en tilsvarende antall parametre i en tettere konfigurasjon, men med betydelig reduserte beregningskrav. Vanligvis er routeren en mapping av lineære lag fra token til ekspert-indeks med hver ekspert som en standard transformer Multilayer Perceptron. Likevel er utviklerne fortsatt i prosessen med å finne den optimale treningsmetoden for routeren, siden ekspert-tildelingsproblemet er ikke-differensierbart, og Mixture of Expert-modeller ofte sliter med last-balansering og treningsstabilitet mellom forskjellige eksperter for hardware-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 routed ekspert-lag. På den andre siden erstatter BlackMamba-rammen ikke bare Multilayer Perceptron-blokken i transformer med en routed 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 trenet på over 300 milliarder token på et eget dataset, og anvender 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 eget dataset som brukes til å trene BlackMamba-rammen består av en blanding av allerede eksisterende åpne kilde-dataset, inkludert Starcoder, SlimPajama, Pile og mer. Følgende tabell demonstrerer vektene til hvert av dataset som brukes til å trene BlackMamba-rammen. Totalt er det 1,8 billioner token i dataset.

BlackMamba : Resultater

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

I tillegg kombinerer latency-fordelene til både Mixture of Expert og Mamba-modeller i BlackMamba-rammen, og resulterer i betydelig raskere inferens-tider sammenlignet med transformer-modeller, rene Mamba-modeller og MoE-modeller. I tillegg er inferens-fordelen til BlackMamba-rammen direkte proporsjonal med sekvenslengdene, og gjør BlackMamba ekstremt effektiv for lang sekvens-generering. Videre illustrerer følgende figur antallet token som er tildelt 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.

Følgende tabell dekker evalueringsscorene til BlackMamba-rammen sammenlignet med en rekke åpne kilde-forhånds-trente språkmodeller. Som det kan observeres, er BlackMamba-rammen i stand til å konkurrere og overgå de fleste rammer over alle baselinjer. I tillegg 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 Mixture of Expert-modeller for å utnytte fordelene til begge rammer. Eksperimenter på 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 eksepsjonelt godt, 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å sterkt trenede store språkmodeller, eksisterende Mamba-ramme og Mixture of Expert-modeller i både trening FLOPs og inferens-kostnad. I tillegg arver BlackMamba-rammen også genererings-FLOPs og redusert trening fra både Mixture of Expert-modeller og Mamba-ramme samtidig.