De ontwikkeling van Large Language Models (LLM’s) op basis van decoder-only transformermodellen heeft een cruciale rol gespeeld bij het transformeren van het domein van Natural Language Processing (NLP) en het bevorderen van diverse deep learning-toepassingen, waaronder reinforcement learning, time-seriesanalyse, beeldverwerking en meer. Echter, ondanks hun schaalbaarheid en sterke prestaties, worden LLM’s op basis van decoder-only transformermodellen nog steeds geconfronteerd met aanzienlijke beperkingen. Hoewel ze uitdrukkingskrachtig zijn, vereist de aandachtmechanisme in transformer-gebaseerde LLM’s hoge computationele middelen tijdens zowel inferentie als training, waardoor een aanzienlijke hoeveelheid geheugen nodig is voor de sequentielengte en kwadratische FLOPs. Deze hoge computationele vereisten beperken de contextlengte van transformermodellen, maken autoregressieve generatietaken evenredig duur met de schaal en belemmeren het leren van continue gegevensstromen en de mogelijkheid om sequenties van onbeperkte lengte efficiënt te verwerken.

In recente tijden hebben State Space Modellen (SSM’s) opmerkelijke mogelijkheden en prestaties getoond, waarmee ze concurreren met transformer-architectuurmodellen in grote modelleringstests en een lineaire complexiteit bereiken met betrekking tot de sequentielengte. Bovendien heeft Mamba, een onlangs uitgebracht State Space Model, uitstekende prestaties getoond in een reeks taalmodellering- en lange-sequentieverwerkingstaken. Tegelijkertijd hebben Mixture of Expert (MoE) modellen eveneens indrukwekkende prestaties getoond, terwijl ze de latentie en computationele kosten van inferentie aanzienlijk verminderen, hoewel ten koste van een grotere geheugenvoetafdruk. Door te bouwen op Mamba en MoE-modellen, zal dit artikel BlackMamba bespreken, een novum architectuur dat het Mamba State Space Model combineert met MoE-modellen om de voordelen van beide kaders te benutten. Experimenten met BlackMamba hebben aangetoond dat het de bestaande Mamba-kader en transformer-baselines overtreft in zowel training FLOPs als inferentie. De uitzonderlijke prestaties van het BlackMamba-kader tonen aan dat het de capaciteiten van de Mamba- en MoE-kaders effectief kan combineren, waardoor snelle en kostenefficiënte inferentie van MoE met lineaire complexiteit generatie van Mamba mogelijk wordt.

Dit artikel heeft als doel het BlackMamba-kader diepgaand te bespreken. We onderzoeken de mechanismen, methodologie en architectuur van het kader, evenals de vergelijking met state-of-the-art beeld- en videogeneratiekaders. Laten we beginnen.

BlackMamba: Een Inleiding tot MoE voor State Space Modellen

De ontwikkeling van Large Language Models (LLM’s), met name die op basis van decoder-only transformerarchitectuur, heeft een aanzienlijke invloed gehad op het domein van Natural Language Processing (NLP) en heeft diverse deep learning-toepassingen uitgebreid, waaronder reinforcement learning, time-seriesanalyse, beeldverwerking en meer. Niettemin, ondanks hun schaalbaarheid en robuuste prestaties, worden deze decoder-only transformer-gebaseerde LLM’s geconfronteerd met aanzienlijke uitdagingen. De aandachtmechanisme, een cruciaal onderdeel van transformer-gebaseerde LLM’s, vereist uitgebreide computationele middelen tijdens zowel inferentie als training. Dit vereist een geheugen dat toeneemt met de sequentielengte en computationele operaties (FLOPs) die kwadratisch toenemen. Deze hoge computationele vereisten beperken de contextlengte van de modellen, verhogen de kosten van autoregressieve generatietaken naarmate de schaal toeneemt en belemmeren het leren van continue gegevensstromen en de mogelijkheid om sequenties van onbeperkte lengte efficiënt te verwerken.

Aanzienlijke inspanningen zijn geleverd in de afgelopen jaren om deze beperkingen te overwinnen, en de aandacht is verschoven naar het ontwikkelen van architecturale alternatieven voor de canonieke dense aandacht transformermodellen met SSM’s en MoE-modellen als de meest veelbelovende kandidaatarchitecturen. Het belangrijkste voordeel van het kiezen voor State Space Modellen boven transformer-architectuurmodellen is de lineaire computationele complexiteit met betrekking tot de lengte van de invoersequentie die door SSM’s wordt aangeboden, in tegenstelling tot de kwadratische complexiteit die door transformers wordt aangeboden. Theoretisch maakt lineaire computationele complexiteit met betrekking tot de lengte van de invoersequentie het mogelijk voor State Space Modellen om grotere sequenties te verwerken dan transformer-architectuurmodellen voor een gegeven FLOPS- of Floating-point operaties per seconde-budget, en maakt autoregressieve generatie constant in compute zonder een KV-cache. Onlangs ontwikkelde State Space Modellen, waaronder Mamba, RetNet en enkele anderen, hebben efficiënte lange-sequentie-inferentie en -training getoond, evenals concurrerende taalmodelleringstakenprestaties met transformers met vergelijkbare schaalbaarheidskenmerken. Aan de andere kant hebben Mixture of Expert modellenarchitecturen aan populariteit gewonnen als een alternatief voor dense transformers, aangezien ze een aanzienlijke reductie in inferentie- en trainings-FLOPs mogelijk maken, die essentieel zijn voor het bereiken van vergelijkbare kwaliteit met een dense model. MoE (Mixture of Experts) modellen werken door alleen een spaarzame selectie van de totale parameters te activeren tijdens een enkele forward pass. Ze gebruiken een routeringsfunctie om te bepalen welke ‘experts’ in actie worden geroepen op basis van de gegeven context. Deze benadering creëert een scheiding tussen de computationele kosten van inferentie en het totale aantal parameters, waardoor een betere prestatie mogelijk wordt binnen een vaste inferentiebegroting, hoewel met een grotere geheugenvoetafdruk.

Deze vooruitgang in architectuur biedt aanzienlijke voordelen ten opzichte van traditionele transformers en vertegenwoordigt een spannende richting voor verdere ontwikkeling. We stellen dat het integreren van deze verbeteringen in een gecombineerd Mamba-MoE-model de taalmodelleringcapaciteiten en -efficiëntie aanzienlijk kan versnellen en verder kan gaan dan die van standaard transformermodellen. De verwachte voordelen van een Mamba-MoE-architectuur in vergelijking met een traditioneel dense transformermodel zijn:

Mamba: bereikt lineaire computationele complexiteit ten opzichte van de lengte van de invoersequentie voor zowel de trainings- als inferentiefase. Het maakt autoregressieve generatie mogelijk in een constante tijdsframe en met constante geheugengebruik.

MoE: biedt inferentiesnelheid en trainingscomputationele efficiëntie die vergelijkbaar zijn met een kleinere, dense baseline-model, terwijl het een niveau van modelkwaliteit behoudt dat concurrerend is met dat van een model met een equivalent aantal parameters als de dichtere versie.

Met dat gezegd hebbende, is het essentieel om te vermelden dat transformer-architectuurmodellen nog steeds state-of-the-art zijn en consistent sterke prestaties hebben getoond op taalmodelleringtaken en sequentieverwerkingstaken. In het hart van de transformer-architectuur wordt zelfaandacht gebruikt, die een kwadratische all-to-all-vergelijking van de dotproductovereenkomsten tussen de embeddings van verschillende tokens in een sequentie uitvoert, en voert een lineaire kaart uit naar een uitvoervector. Het transformermodel bestaat uit zelfaandachtsblokken die tussen MLP- of Multilayer Perceptron-blokken zijn gestapeld, die verder bestaan uit een tweelaags MLP met een gegeven activatiefunctie.

BlackMamba: Architectuur en Methodologie

State Space Modellen

State Space Modellen behoren tot de groep van sequentiemodellen met lineaire complexiteit ten opzichte van de lengte van de invoersequentie. De architectuur van State Space Modellen komt meer overeen met Recurrent Neural Networks en Convolutional Neural Networks dan met aandacht-gebaseerde architectuur en is geïnspireerd op een continu dynamisch systeem dat een 1-dimensionale functie door een impliciete latent ruimte kaart. Een lineair dynamisch systeem maakt parallelle berekeningen efficiënt met behulp van een associatief of convolutioneel scanproces. In praktische scenario’s is de recurrente aard van State Space Modellen de reden waarom het nog steeds moet worden geadopteerd op hoogwaardige AI-hardware zoals GPU’s. Echter, de opkomst van SSM’s zoals RWKV en Mamba hebben parallelle scan-kernels gebruikt om recurrente operaties efficiënt te kaarten naar GPU’s, waardoor de training van nieuwe architectuur met efficiëntie vergelijkbaar met die van transformermodellen mogelijk wordt.

De inherente kwadratische complexiteit in relatie tot sequentielengte binnen transformers is een bekende beperking die het redeneren en begrijpen van zeer lange contexten belemmert. Recentelijk zijn innovaties geïntroduceerd om de contextlengte uit te breiden, waardoor transformers kunnen worden getraind op een haalbare schaal voordat ze worden toegepast op veel langere contexten tijdens inferentie. Ondanks deze vooruitgang, vereist het inferentieproces nog steeds een aanzienlijke hoeveelheid computationele middelen en geheugen, vooral voor het onderhoud van de Key-Value (KV) cache, waardoor het een middelenintensieve onderneming wordt. Recent onderzoek heeft zich gericht op het verbeteren van de expressieve capaciteiten van state-space modellen door input-afhankelijke gating-mechanismen te incorporeren, vergelijkbaar met de Query, Key, Value (QKV) matrices die in aandachtmechanismen worden gevonden.

Deze inspanningen zijn gericht op het behoud van de inherente lineaire progressie van state-space recursie, waardoor efficiënte uitvoering mogelijk wordt door middel van convolution of een selectief scanproces. Deze benadering verkleint de prestatieverschillen met transformers in praktische toepassingen aanzienlijk. Onder deze vooruitgang staat Mamba uit als een state-space model dat de doelstellingen van eerder onderzoek weerspiegelt, met indrukwekkende prestatieniveaus die vergelijkbaar zijn met die van transformers tot schaalniveaus van 2,8 miljard parameters. Het bereikt dit door input-afhankelijke gating toe te passen op de ingangen van de state-space model (SSM) recursie, terwijl het efficiënte berekening waarborgt door het gebruik van speciale selectieve scan-kernels.

Mixture of Expert Modellen

Mixture of Expert (MoE) modellen bereiken een scheiding tussen de inferentiekosten en het totale aantal parameters door parameters selectief te activeren tijdens de forward pass. In plaats van alle parameters te gebruiken, leiden deze modellen tokens naar specifieke Multilayer Perceptron (MLP) experts. Ideaal gesproken is elke expert aangepast om een specifiek type invoer te verwerken, met een routeringsmechanisme, dat in wezen een compact neuronaal netwerk is, dat bepaalt welke expert het meest geschikt is voor elke token. Deze benadering is gericht op het behoud van de algehele expressieve kracht van een model met een equivalent aantal parameters in een dichtere configuratie, maar met aanzienlijk verminderde computationele eisen. Typisch is de router een afbeelding van de lineaire lagen van tokens naar expert-indices, waarbij elke expert simpelweg een standaard transformer Multilayer Perceptron is. Echter, ontwikkelaars moeten nog steeds de optimale trainingsmethode voor de router vinden, aangezien het expert-toewijzingsprobleem niet-differentieerbaar is, en Mixture of Expert modellen vaak worstelen met belastingsbalans en trainingsstabiliteit voor hardware-efficiëntie.

Architectuur

In het hart van BlackMamba wordt een standaard transformermodel gebruikt, bestaande uit tussenliggende MLP-blokken en aandachtsblokken die in sequentie langs een residustroom zijn toegevoegd. Nu vervangen de meeste Mixture of Expert modellen simpelweg de multilayer perceptron-blokken door een geroute expertlaag. Aan de andere kant vervangt het BlackMamba-kader niet alleen het multilayer perceptron-blok in de transformer door een geroute expertlaag, maar vervangt het ook de aandachtslaag door een Mamba State Space Model-laag. De architectuur van het BlackMamba-kader wordt in de volgende figuur getoond.

Training en Dataset

Het BlackMamba-model is getraind op meer dan 300 miljard tokens op een aangepast dataset en gebruikt de SwiGLU-activatiefunctie voor de expert multilayer perceptrons. Het kader traint met 8 experts, een aantal dat ontwikkelaars hebben gevonden als de juiste balans en compromis tussen de geheugenvoetafdruk en de inferentiekosten van het model. De aangepaste dataset die wordt gebruikt om het BlackMamba-kader te trainen, bestaat uit een mengsel van reeds bestaande open-source datasets, waaronder Starcoder, SlimPajama, Pile en meer. De volgende tabel toont de gewichten van elke dataset die wordt gebruikt voor het trainen van het BlackMamba-kader. Over het algemeen zijn er 1,8 biljoen tokens in de dataset.

BlackMamba: Resultaten

Om een eerlijke vergelijking te garanderen tussen Mamba en BlackMamba, hebben ontwikkelaars beide modellen getraind met dezelfde trainingsparameters op dezelfde trainingsgegevens. Het BlackMamba-kader is in staat om zowel Mamba als transformermodellen te overtreffen voor identieke forward pass modelgrootte op zowel inferentie- als trainings Floating-point operaties per seconde. De volgende figuur toont de tijd die nodig is om een sequentie van een bepaalde lengte autoregressief te genereren vanuit een initiële een-token prompt als een functie van de sequentielengte.

Bovendien combineren de latentievoordelen van zowel de Mixture of Expert- als de Mamba-modellen in het BlackMamba-kader, waardoor significant snellere inferentietijden mogelijk zijn in vergelijking met transformermodellen, pure Mamba-modellen en MoE-modellen. Bovendien is het inferentievoordeel van het BlackMamba-kader recht evenredig met de sequentielengte, waardoor BlackMamba uiterst effectief is bij lange sequentiegeneratie. Verder toont de volgende figuur het aantal tokens dat is toegewezen aan de BlackMamba-modellen met 340 miljoen en 640 miljoen parameters. Zoals te zien is, vertonen de meeste lagen een hoge mate van expertbalans als gevolg van de verbeterde Sinkhorn-algoritme die door de BlackMamba-modellen wordt geïmplementeerd.

De volgende tabel toont de evaluatiescores van het BlackMamba-kader in vergelijking met een reeks open-source voorgetrainde taalmodellen. Zoals te zien is, is het BlackMamba-kader in staat om te concurreren en te overtreffen met de meeste kaders op alle baselines. Bovendien is het de moeite waard om op te merken dat de modellen die BlackMamba overtreffen, aanzienlijk meer parameters hebben, en de prestatieverschillen minimaal zijn, wat aangeeft dat het BlackMamba-kader met minder parameters in staat is om concurrerende prestaties te leveren.

Slotgedachten

In dit artikel hebben we het over BlackMamba gehad, een novum architectuur dat het Mamba State Space Model combineert met MoE-modellen om de voordelen van beide kaders te benutten. Experimenten met BlackMamba hebben aangetoond dat het de bestaande Mamba-kader en transformer-baselines overtreft in zowel training FLOPs als inferentie. De uitzonderlijke prestaties van het BlackMamba-kader tonen aan dat het de capaciteiten van de Mamba- en MoE-kaders effectief kan combineren, waardoor snelle en kostenefficiënte inferentie van MoE met lineaire complexiteit generatie van Mamba mogelijk wordt. We hebben besproken hoe de architectuur van het BlackMamba-kader in staat is om sterke getrainde Large Language Models, bestaande Mamba-kaders en Mixture of Expert modellen te overtreffen in termen van training FLOPs en inferentiekosten. Bovendien erfde het BlackMamba-kader de generatie FLOPs en de vermindering van training van zowel Mixture of Expert modellen als Mamba-kaders tegelijkertijd.