Utvecklingen av stora språkmodeller (LLM) byggda från decoder-only transformermodeller har spelat en avgörande roll i att förändra det naturliga språkbehandlingsområdet (NLP) och främja olika djupinlärningsapplikationer, inklusive förstärkt inlärning, tidsserieanalys, bildbehandling och mycket mer. Men trots deras skalbarhet och starka prestanda möter LLM:er byggda från decoder-only transformermodeller fortfarande betydande brister. Även om de är uttrycksfulla kräver uppmärksamhetsmekanismen i transformer-derivade LLM:er höga beräkningsresurser under både inferens och utbildning, vilket kräver betydande minne för sekvenslängden och kvadratiska FLOPs. Detta höga beräkningskrav begränsar modellernas kontextlängd, gör autoregressiv generering proportionellt dyr med skalan och hindrar modellernas förmåga att lära sig från kontinuerliga dataströmmar och bearbeta sekvenser av obegränsad längd effektivt.

På senare tid har State Space-modeller (SSM) visat anmärkningsvärda förmågor och prestanda, i konkurrens med transformer-arkitekturmodeller i stora modelljämförelser medan de uppnår minneskomplexitet som en funktion av sekvenslängden och linjär tid. Dessutom har Mamba, en nyligen släppt State Space-modell, visat utmärkt prestanda i en rad språkmodellering och långa sekvensbearbetningsuppgifter. Samtidigt har Mixture of Expert-modeller (MoE) också visat imponerande prestanda medan de betydligt minskar svarstiden och beräkningskostnaderna för inferens, om än på bekostnad av en större minnesavtryck. Genom att bygga på Mamba och MoE-modeller kommer denna artikel att diskutera BlackMamba, en ny arkitektur som kombinerar Mamba State Space-modellen med MoE-modeller för att utnyttja fördelarna som båda ramverken erbjuder. Experiment på BlackMamba har visat dess förmåga att överträffa den befintliga Mamba-ramen och transformer-baslinjerna i både utbildnings-FLOPs och inferens. Den exceptionella prestandan hos BlackMamba-ramen visar att den kan kombinera förmågorna hos Mamba- och MoE-ramarna på ett utmärkt sätt, eftersom den kombinerar den billiga och snabba inferensen från MoE med linjärkomplexitetsgenerering från Mamba.

Denna artikel syftar till att täcka BlackMamba-ramen i detalj. Vi utforskar mekanismen, metodiken och arkitekturen i ramen, tillsammans med dess jämförelse med state-of-the-art-bild- och videogenereringsramar. Låt oss komma igång.

BlackMamba : En introduktion till MoE för State Space-modeller

Utvecklingen av stora språkmodeller (LLM), särskilt de som bygger på decoder-only transformerarkitekturer, har märkbart påverkat det naturliga språkbehandlingsområdet (NLP) och expanderat till olika djupinlärningsapplikationer, inklusive förstärkt inlärning, tidsserieanalys, bildbehandling och bortom. Trots deras skalbarhet och robusta prestanda möter dessa decoder-only transformer-baserade LLM:er betydande utmaningar. Uppmärksamhetsmekanismen, en nyckelfunktion i transformer-baserade LLM:er, kräver omfattande beräkningsresurser för både inferens och utbildning. Detta innebär ett behov av minne som växer med sekvenslängden och beräkningsoperationer (FLOPs) som ökar kvadratiskt. Sådana intensiva beräkningsbehov begränsar modellernas kontextlängd, höjer kostnaderna för autoregressiv generering som modellen skalar och hindrar modellernas förmåga att lära sig från kontinuerliga dataströmmar eller bearbeta sekvenser av obegränsad längd effektivt.

Betydande ansträngningar har gjorts under de senaste åren för att övervinna dessa begränsningar, och uppmärksamheten har skiftats mot att utveckla alternativa arkitekturer till de kanoniska täta uppmärksamhets-transformermodellerna med SSM och MoE-modeller som de mest lovande kandidatarkitekturerna. Den viktigaste fördelen med att föredra State Space-modeller framför transformer-arkitekturmodeller är den linjära beräkningskomplexiteten i förhållande till indatasekvenslängden som erbjuds av SSM, till skillnad från den kvadratiska komplexiteten som erbjuds av transformer. Teoretiskt sett möjliggör den linjära beräkningskomplexiteten i förhållande till sekvenslängden att State Space-modellerna kan bearbeta längre sekvenser än transformer-arkitekturmodeller för en given FLOPS- eller flyttalsoperationer per sekund-budget, och att göra autoregressiv generering konstant i beräkning utan en KV-cache. Nyligen utvecklade State Space-modeller, inklusive RWKV och Mamba, har visat effektiv långsekvensinferens och utbildning, tillsammans med konkurrenskraftig språkmodellering och prestanda jämförbar med transformer med liknande skalningsegenskaper. Å andra sidan har Mixture of Expert-modellarkitekturer blivit alltmer populära som ett alternativ till täta transformer, eftersom de möjliggör en betydande minskning av inferens- och utbildnings-FLOPs, vilket är nödvändigt för att uppnå jämförbar kvalitet med en tät modell. MoE (Mixture of Experts)-modeller fungerar genom att aktivera endast en sparsam selektion av totala parametrar under en enda framåtriktad passering. De använder en routerfunktion för att bestämma vilka “experter” som ska aktiveras baserat på den givna kontexten. Detta tillvägagångssätt skapar en separation mellan beräkningskostnaden för inferens och det totala antalet parametrar, vilket möjliggör förbättrad prestanda inom en fast inferensbudget, om än med en ökad parameterantal och ett större minneskrav.

Denna arkitektoniska utveckling erbjuder betydande fördelar jämfört med traditionella transformer. Vi hävdar att integrationen av dessa förbättringar i en kombinerad Mamba-MoE-modell kan avsevärt accelerera språkmodellering och effektivitet bortom standardtransformermodeller. De förväntade fördelarna med en Mamba-MoE-arkitektur jämfört med en traditionell tät transformermodell inkluderar:

Mamba: Uppnår linjär beräkningskomplexitet i förhållande till indatasekvenslängden för både utbildnings- och inferensfasen. Det möjliggör autoregressiv generering som sker i en konstant tidsram och med konstant minnesanvändning.

MoE: Erbjuder inferenshastighet och utbildningsberäknings-effektivitet jämförbar med en mindre, tät baslinjemodell medan den upprätthåller en modellkvalitet som kan mäta sig med en modell med ett lika stort antal parametrar som den tätare versionen.

Med det sagt är det viktigt att påpeka att transformer-arkitekturmodeller fortfarande är state-of-the-art och har visat konsekvent och anmärkningsvärd stark prestanda på språkmodellering och sekvensbearbetningsuppgifter. I sin kärna använder transformer-arkitekturen självuppmärksamhet som utför en kvadratisk all-till-all-jämförelse av prickproduktlikheterna mellan inbäddningarna av olika token i en sekvens och utför en linjär karta till en utgångsvektor. Transformer-modellen består av självuppmärksamhetsblock staplade mellan MLP- eller multilayerperceptronblock som ytterligare består av ett tvålagers MLP med en given aktiveringsfunktion.

BlackMamba : Arkitektur och metodik

State Space-modeller

State Space-modeller tillhör gruppen av sekvensmodeller med linjär komplexitet i förhållande till indatasekvenslängden. Arkitekturen för State Space-modeller överensstämmer mer med återkopplande neurala nätverk och konvolutionsneurala nätverk än uppmärksamhetsbaserad arkitektur och är inspirerad av ett kontinuerligt dynamiskt system som kartar en 1-dimensionell funktion genom ett implicit latentspace. Ett lineärt dynamiskt system gör parallella beräkningar effektiva med hjälp av antingen en associativ eller en konvolutionsgenomsökning. I praktiska scenarier har den återkopplande naturen av State Space-modeller varit anledningen till att de fortfarande inte har antagits på högt parallell AI-hårdvara som GPU:er. Men uppkomsten av SSM:er som RWKV och Mamba har använt parallella skanningskärnor för att kartlägga återkommande operationer effektivt till GPU:er, vilket möjliggör utbildning av nya arkitekturer med effektivitet jämförbar med den som uppnås av transformer-modeller.

Den inneboende kvadratiska komplexiteten i förhållande till sekvenslängd inom transformer är en välkänd begränsning som hindrar resonemang och förståelse över mycket långa kontexter. Nya innovationer har introducerat idén att förlänga kontextlängden, vilket möjliggör för transformer att utbildas på en rimlig skala innan de tillämpas på mycket längre kontexter under inferens. Trots dessa framsteg kräver inferensprocessen fortfarande en betydande mängd beräkningsresurser och minne, särskilt för att upprätthålla Key-Value (KV)-cachen, vilket gör det till en resurskrävande uppgift. Nya forskningsinsatser har fokuserat på att förbättra de uttrycksfulla förmågorna hos state-space-modeller genom att införa indataberoende styrningsmekanismer, liknande de fråga-nyckel-värde (QKV)-matriser som finns i uppmärksamhetsmekanismer.

Dessa insatser syftar till att bevara den inneboende linjära progressionen av state-space-rekursion, vilket möjliggör effektiv körning genom antingen konvolution eller en selektiv sökprocess. Detta tillvägagångssätt minskar betydligt prestandaskillnaden med transformer i praktiska tillämpningar. Bland dessa framsteg utmärker sig Mamba som en state-space-modell som speglar målen för tidigare forskning, visar imponerande prestandanivåer jämförbara med transformer på skalor upp till 2,8 miljarder parametrar. Det uppnår detta genom att tillämpa indataberoende styrning på indata till state-space-modellens (SSM) rekursion, samtidigt som det säkerställer effektiv beräkning genom användning av specialutvecklade selektiva skanningskärnor.

Mixture of Expert-modeller

Mixture of Expert (MoE)-modeller uppnår en separation mellan inferenskostnaden och det totala parameterantalet genom att selektivt aktivera parametrar under den framåtriktade passeringen. Istället för att använda alla parametrar dirigerar dessa modeller token till specifika multilayerperceptron (MLP)-experter. Idealiskt är varje expert anpassad för att bearbeta en specifik typ av indata, med en routermekanism, i princip en kompaktn neural nätverk, som bestämmer den mest lämpliga experten för varje token. Detta tillvägagångssätt syftar till att bevara den omfattande uttrycksfulla kraften hos en modell med ett lika stort antal parametrar i en tätare konfiguration, men med betydligt minskade beräkningskrav. Typiskt sett är routern en mappning av linjära lager från token till expertindex, med varje expert som enkelt är en standardtransformer Multilayer Perceptron. Men utvecklare har ännu inte funnit den optimala utbildningsmetoden för routern, eftersom experttilldelningsproblemet är icke-differentierbart, och Mixture of Expert-modeller ofta kämpar med belastningsbalans och utbildningsstabilitet mellan olika experter för hårdvarueffektivitet.

Arkitektur

I sin kärna använder BlackMamba en standardtransformermodell som består av växlande MLP-block och uppmärksamhetsblock som läggs till i sekvens längs en residualström. Nu, en majoritet av Mixture of Expert-modeller ersätter enkelt multilayerperceptronblocken med ett dirigerat expertlager. Å andra sidan ersätter BlackMamba-ramen inte bara multilayerperceptronblocket i transformer med ett dirigerat expertlager, utan ersätter också uppmärksamhetslagret med ett Mamba State Space-modelllager. Arkitekturen för BlackMamba-ramen visas i följande figur.

Utbildning och dataset

BlackMamba-modellen utbildas på över 300 miljarder token på ett anpassat dataset och använder SwiGLU-aktiveringsfunktionen för expertmultilayerperceptronerna. Ramen utbildas med 8 experter, ett antal som utvecklare fann vara rätt balans och avvägning mellan minnesavtryck och inferenskostnad för modellen. Det anpassade dataset som används för att utbilda BlackMamba-ramen består av en blandning av redan existerande öppen källkodsdataset, inklusive Starcoder, SlimPajama, Pile och fler. Följande tabell visar viktningen av varje dataset som används för att utbilda BlackMamba-ramen. Sammanlagt finns det 1,8 biljoner token i datasetet.

BlackMamba : Resultat

För att säkerställa en rättvis jämförelse mellan Mamba och BlackMamba har utvecklare utbildat båda modellerna med samma utbildningsparametrar på samma utbildningsdata. BlackMamba-ramen kan överträffa både Mamba och transformer-modeller för identisk framåtriktad passningsmodellstorlek vid inferenstiden samt utbildningsflyttalsoperationer per sekund. Följande figur visar den tid som krävs för att generera en sekvens av en given längd autoregressivt från en initial en-token-prompt som en funktion av sekvenslängden.

Dessutom kombineras svarstidsfördelarna hos både Mixture of Expert- och Mamba-modellerna i BlackMamba-ramen, vilket resulterar i betydligt snabbare inferenstider jämfört med transformer-modeller, rena Mamba-modeller och MoE-modeller. Dessutom är inferensfördelen hos BlackMamba-ramen direkt proportionell mot sekvenslängderna, vilket gör BlackMamba extremt effektiv för lång sekvensgenerering. Fortsättning, följande figur visar antalet token som tilldelats BlackMamba-modellerna med 340 miljoner respektive 640 miljoner parametrar. Som det kan ses visar de flesta lagren en hög nivå av expertbalans som ett resultat av den förbättrade Sinkhorn-algoritmen som implementerats av BlackMamba-modellerna.

Följande tabell täcker utvärderingspoängen för BlackMamba-ramen jämfört med en rad öppen källkods-förutbildade språkmodeller. Som det kan observeras kan BlackMamba-ramen tävla och överträffa de flesta ramverken över alla baslinjer. Dessutom är det värt att notera att modellerna som överträffar BlackMamba har betydligt fler parametrar, och gapet i prestanda är minimalt, vilket indikerar BlackMamba-ramens förmåga att prestera med färre parametrar.

Slutliga tankar

I denna artikel har vi talat om BlackMamba, en ny arkitektur som kombinerar Mamba State Space-modellen med Mixture of Expert-modeller för att utnyttja fördelarna som båda ramverken erbjuder. Experiment på BlackMamba har visat dess förmåga att överträffa den befintliga Mamba-ramen och transformer-baslinjerna i både utbildnings-FLOPs och inferens. Den exceptionella prestandan hos BlackMamba-ramen visar att den kan kombinera förmågorna hos Mamba- och MoE-ramarna på ett utmärkt sätt, eftersom den kombinerar den billiga och snabba inferensen från MoE med linjärkomplexitetsgenerering från Mamba. Vi har talat om hur arkitekturen för BlackMamba-ramen kan överträffa starkt utbildade stora språkmodeller, befintliga Mamba-ram och Mixture of Expert-modeller i termer av utbildnings-FLOPs och inferenskostnad. Dessutom ärver BlackMamba-ramen också genererings-FLOPs och reducerad utbildning från både Mixture of Expert-modeller och Mamba-ram samtidigt.