Utvecklingen av stora språkmodeller (LLM) byggda från transformatormodeller som endast avkodare har spelat en avgörande roll för att transformera NLP-domänen (Natural Language Processing), såväl som för att utveckla olika tillämpningar för djupinlärning, inklusive förstärkning lärande, tidsserieanalys, bildbehandling och mycket mer. Men trots sin skalbarhet och starka prestanda möter LLM:er byggda av transformatormodeller med endast dekoder fortfarande betydande brister. Även om den är uttrycksfull kräver uppmärksamhetsmekanismen i transformatorhärledda LLM: er höga beräkningsresurser under både slutledning och träning, vilket kräver betydande minne för sekvenslängden och kvadratiska FLOP:er. Detta höga beräkningskrav begränsar kontextlängden för transformatormodeller, vilket gör autoregressiva genereringsuppgifter proportionellt dyra med skalan, och hindrar lärande från kontinuerliga dataströmmar och möjligheten till verkligt obegränsad sekvensbearbetning.

Nyligen, Statliga rymdmodeller (SSM) har visat anmärkningsvärda egenskaper och prestanda, och konkurrerar med transformatorarkitekturmodeller i storskaliga modelleringsriktmärken samtidigt som de uppnår minneskomplexitet som en funktion av sekvenslängd och linjär tid. Dessutom har Mamba, en nyligen släppt State Space Model, visat enastående prestanda i en rad språkmodelleringar och långsekvensbearbetningsuppgifter. Samtidigt har Mixture of Expert (MoE)-modeller också visat imponerande prestanda samtidigt som de avsevärt minskat latensen och beräkningskostnaderna för slutledning, om än på bekostnad av ett större minnesfotavtryck. Med utgångspunkt i Mamba- och MoE-modeller kommer den här artikeln att diskutera BlackMamba, en ny arkitektur som kombinerar Mamba State Space Model med MoE-modeller för att dra nytta av fördelarna med båda ramverken. Experiment på BlackMamba har visat sin förmåga att överträffa det befintliga Mamba-ramverket och transformatorbaslinjerna i både tränings-FLOP och slutledning. BlackMamba-ramverkets exceptionella prestanda visar att det effektivt kan kombinera förmågorna hos Mamba- och MoE-ramverken, vilket ger snabb och kostnadseffektiv slutledning från MoE med generering av linjär komplexitet från Mamba.

Den här artikeln syftar till att täcka BlackMamba-ramverket på djupet. Vi utforskar ramverkets mekanism, metodik och arkitektur, tillsammans med dess jämförelse med toppmoderna ramverk för bild- och videogenerering. Låt oss börja.

BlackMamba: An Introduction to MoE for State Space Models

Utvecklingen av stora språkmodeller (LLM), särskilt de som är baserade på transformatorarkitekturer med endast dekoder, har särskilt påverkat Naturlig språkbehandling (NLP) och utökas till olika tillämpningar för djupinlärning, inklusive förstärkningsinlärning, tidsserieanalys, bildbehandling och mer. Trots sin skalbarhet och robusta prestanda möter dessa transformatorbaserade LLM:er endast av avkodare anmärkningsvärda utmaningar. Uppmärksamhetsmekanismen, en nyckelfunktion i transformatorbaserade LLMs, kräver omfattande beräkningsresurser för både slutledning och träning. 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 autoregressiva genereringsuppgifter när modellen skalas, och hindrar modellernas förmåga att effektivt lära av kontinuerliga dataströmmar eller processsekvenser av obegränsad längd. 

Betydande ansträngningar har gjorts under de senaste åren i ett försök att övervinna dessa begränsningar, och uppmärksamheten har flyttats mot att ta fram arkitektoniska alternativ till de kanoniska täta uppmärksamhetstransformatormodellerna med SSM- och MoE-modeller som de mest lovande kandidatarkitekturerna. Den viktigaste fördelen med att gynna tillståndsrymdmodeller framför transformatorarkitekturmodeller är den linjära beräkningskomplexiteten med avseende på ingångssekvenslängden som erbjuds av SSM:er i motsats till den kvadratiska komplexiteten som erbjuds av transformatorer. Teoretiskt sett möjliggör linjär beräkningskomplexitet med avseende på ingångssekvenslängd tillståndsrymdmodeller att bearbeta större sekvenser än transformatorarkitekturmodeller för en given FLOPS- eller Flytpunktsoperationer per sekund-budget, och att göra autoregressiv generering konstant vid beräkning utan en KV-cache. Nyligen utvecklade State Space Models inklusive Mamba, RetNet och några andra har visat effektiv långsekvens slutledning och träning, tillsammans med konkurrenskraftiga språkmodelleringsuppdrag för transformatorer med liknande skalningsegenskaper. Å andra sidan vinner Mixture of Expert-modellarkitekturer i popularitet som ett alternativ till täta transformatorer eftersom det underlättar en betydande minskning av slutledningar och tränings-FLOP:er som är nödvändiga för att uppnå jämförbar kvalitet med en tät modell. MoE-modeller (Mixture of Experts) fungerar genom att endast aktivera ett sparsamt urval av de totala parametrarna under en enda framåtpassning. De använder en routingfunktion för att avgöra vilka "experter" som kallas till handling baserat på det givna sammanhanget. Detta tillvägagångssätt skapar en separation mellan beräkningskostnaden för slutledning och det totala antalet parametrar, vilket möjliggör förbättrad prestanda inom en fast slutledningsbudget, om än med ett ökat antal parametrar och ett större minneskrav.

Detta framsteg inom arkitektur erbjuder anmärkningsvärda fördelar jämfört med traditionella transformatorer och representerar en spännande riktning för vidare utveckling. Vi antar att integrering av dessa förbättringar i en kombinerad Mamba-MoE-modell kan avsevärt accelerera språkmodelleringsmöjligheter och effektivitet utöver standardtransformatormodeller. De förväntade fördelarna med en Mamba-MoE-arkitektur jämfört med en traditionell tät transformatormodell inkluderar:

Mamba: Uppnår linjär beräkningskomplexitet i förhållande till inmatningssekvenslängden för både tränings- och slutledningsfaser. Det gör att autoregressiv generering kan ske i en konstant tidsram och med konstant minnesanvändning.

MoE: Erbjuder slutledningshastigheten och beräkningseffektiviteten för träning som är jämförbar med en mindre, tät baslinjemodell samtidigt som den bibehåller en modellkvalitetsnivå som konkurrerar med en modell med ett motsvarande antal parametrar som den tätare versionen.

Med detta sagt är det viktigt att konstatera att transformatorarkitekturmodeller fortfarande är toppmoderna och har visat konsekvent och anmärkningsvärt stark prestanda på språkmodelleringsuppgifter och sekvensbearbetningsuppgifter. I sin kärna använder transformatorarkitekturen självuppmärksamhet som utför en kvadratisk allt-till-alla-jämförelse av punktproduktlikheterna mellan inbäddningar av olika tokens i en sekvens, och utför en linjär karta till en utdatavektor. Transformatormodellen består av självuppmärksamhetsblock staplade mellan MLP- eller Multi-Layer Perceptron-block som vidare består av en tvålagers MLP med en given aktiveringsfunktion. 

BlackMamba: Arkitektur och metodik

Statliga rymdmodeller

Tillståndsrymdmodeller tillhör gruppen av sekvensmodeller med linjär komplexitet med avseende på längden på ingångssekvensen. Arkitekturen för State Space Models ligger mer i linje med Recurrent Neural Networks och Convolutional Neural Networks snarare än uppmärksamhetsbaserad arkitektur, och är inspirerad av ett kontinuerligt dynamiskt system som kartlägger en 1-dimensionell funktion genom ett implicit latent utrymme. Ett linjärt dynamiskt system gör parallella beräkningar effektiva genom att använda antingen en associativ eller en faltningsskanning. I praktiska scenarier har den återkommande karaktären hos State Space Models varit anledningen till att de fortfarande ska användas på mycket parallell AI-hårdvara som GPU:er. Men uppkomsten av SSM som RWKV och Mamba har använt parallella skanningskärnor för att effektivt kartlägga återkommande operationer till GPU:er, vilket underlättar utbildningen av nya arkitekturer med effektivitet jämförbar med dem som uppnås med transformatormodeller. 

Den inneboende kvadratiska komplexiteten i förhållande till sekvenslängd inom transformatorer är en välkänd begränsning som försvårar resonemang och förståelse över mycket långa sammanhang. Nya innovationer har introducerat idén om att utöka kontextlängden, vilket gör att transformatorer kan tränas i en genomförbar skala innan de appliceras på mycket längre sammanhang under slutledning. Trots dessa framsteg kräver slutledningsprocessen fortfarande en avsevärd 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 strävan. Nyligen genomförda forskningsansträngningar har fokuserat på att förbättra de uttrycksfulla kapaciteterna hos stat-rymdmodeller genom att införliva input-beroende grindmekanismer, besläktade med Query, Key, Value (QKV)-matriser som finns i uppmärksamhetsmekanismer. 

Dessa ansträngningar syftar till att bevara den inneboende linjära progressionen av tillstånd-rymd-rekursion, vilket möjliggör effektiv exekvering genom antingen faltning eller en selektiv skanningsprocess. Detta tillvägagångssätt minskar avsevärt prestandaskillnaderna med transformatorer i praktiska tillämpningar. Bland dessa framsteg framstår Mamba som en stat-rymdmodell som speglar målen för tidigare forskning, och visar imponerande prestandanivåer jämförbara med transformatorer på skalor upp till 2.8 miljarder parametrar. Den uppnår detta genom att tillämpa ingångsberoende grindning på ingångarna för tillståndsrymdmodellens (SSM)-rekursion, samtidigt som den säkerställer effektiv beräkning genom användning av skräddarsydda selektiva skanningskärnor.

Blandning av expertmodeller

Blandning av expertmodeller (MoE) uppnår en separation mellan slutledningskostnaden och det totala parameterantalet genom att selektivt aktivera parametrar under framåtpassningen. Istället för att använda alla parametrar riktar dessa modeller tokens till specifika Multilayer Perceptron (MLP) experter. Helst är varje expert skräddarsydd för att bearbeta en viss typ av indata, med en routingmekanism, i huvudsak ett kompakt neuralt 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 uttryckskraften hos en modell med ett ekvivalent antal parametrar i en tätare konfiguration, men med avsevärt minskade beräkningskrav. Vanligtvis är routern en kartläggning av de linjära lagren från tokens till expertindex, där varje expert helt enkelt är en standardtransformator Multilayer Perceptron. Utvecklare har dock ännu inte räknat ut den optimala träningsmetoden för routern eftersom problemet med experttilldelningen inte är differentierbart, och Mixture of Expert-modeller kämpar ofta med lastbalansering och träningsstabilitet mellan olika experter för hårdvarueffektivitet. 

arkitektur

I sin kärna använder BlackMamba en standardtransformatormodell som består av interfolierade MLP-block och uppmärksamhetsblock som läggs till i sekvens längs en restström. Nu ersätter en majoritet av Mixture of Expert-modellerna helt enkelt flerskiktsperceptronblocken med ett dirigerat expertlager. Å andra sidan ersätter BlackMamba-ramverket inte bara flerskiktsperceptronblocket i transformatorn med ett dirigerat expertlager, utan ersätter också uppmärksamhetslagret med ett Mamba State Space Model-lager. Arkitekturen för BlackMamba-ramverket visas i följande figur. 

Utbildning och datauppsättning

BlackMamba-modellen är tränad på över 300 miljarder tokens på en anpassad datauppsättning och använder SwiGLU-aktiveringsfunktionen för expertens flerskiktsperceptroner. Ramverket tränar med 8 experter, ett antal som utvecklare fann vara den rätta balansen och avvägning mellan minnesfotavtrycket och slutsatskostnaden för modellen. Den anpassade datamängden som används för att träna BlackMamba-ramverket består av en blandning av redan befintliga datauppsättningar med öppen källkod inklusive Starcoder, SlimPajama, Pile och mer. Följande tabell visar vikterna för varje datauppsättning som används för att träna BlackMamba-ramverket. Sammantaget finns det 1.8 biljoner tokens i datamängden. 

BlackMamba : Resultat

För att säkerställa en rättvis jämförelse mellan Mamba och BlackMamba har utvecklare tränat båda modellerna med samma träningsparametrar på samma träningsdata. BlackMamba-ramverket kan överträffa både Mamba- och transformatormodeller för identisk frampassningsmodellstorlek vid slutledningstidpunkten samt träna flytande kommaoperationer per sekund. Följande figur visar hur lång tid det tar att generera en sekvens med en given längd autoregressivt från en initial en-token prompt som en funktion av sekvenslängden. 

Dessutom kombineras latensfördelarna med både Mixture of Expert- och Mamba-modellerna i BlackMamba-ramverket, vilket resulterar i betydligt snabbare slutledningstider jämfört med transformatormodeller, rena Mamba-modeller och MoE-modeller. Dessutom är inferensfördelen med BlackMamba-ramverket direkt proportionell mot sekvenslängderna, vilket gör BlackMamba extremt effektiv vid lång sekvensgenerering. Följande figur illustrerar antalet tokens som tilldelats BlackMamba-modellerna med 340 miljoner respektive 640 miljoner parametrar. Som det kan ses visar en majoritet av lagren en hög nivå av expertbalans som ett resultat av den förbättrade Sinkhorn-algoritmen implementerad av BlackMamba-modellerna. 

Följande tabell täcker utvärderingspoängen för BlackMamba-ramverket jämfört med en rad förtränade språkmodeller med öppen källkod. Som det kan observeras kan BlackMamba-ramverket konkurrera och överträffa med en majoritet av 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 förmågan hos BlackMamba-ramverket med färre parametrar. 

Avslutande tankar

I den här artikeln har vi pratat om BlackMamba, en ny arkitektur som kombinerar Mamba State Space Model med Mixture of Expert-modeller för att skörda fördelarna som erbjuds av båda dessa ramverk. Experiment på BlackMamba har visat att det överträffar det befintliga Mamba-ramverket och transformatorbaslinjerna i både tränings-FLOP och slutledning. BlackMamba-ramverkets exceptionella prestanda visar att det kan ärva och kombinera förmågorna hos Mamba- och MoE-ramverken exceptionellt bra eftersom det kombinerar den billiga och snabba slutsatsen från MoE med generering av linjär komplexitet från Mamba. Vi har pratat om hur arkitekturen i BlackMamba-ramverket kan överträffa starkt tränade stora språkmodeller, befintliga Mamba-ramverk och Mixture of Expert-modeller när det gäller tränings-FLOPs och slutledningskostnader. Dessutom ärver BlackMamba-ramverket också generationens FLOPs och minskad träning från både Mixture of Expert-modeller och Mamba-ramverket samtidigt.