Een introductie tot de Mamba LLM-architectuur: een nieuw paradigma in machine learning

Taalmodellen zijn een type machinelearningmodel dat getraind is om een waarschijnlijkheidsverdeling over natuurlijke taal te modelleren. Hun architectuur bestaat voornamelijk uit meerdere lagen neurale netwerken, zoals recurrente lagen, feedforwardlagen, embeddinglagen en attention-lagen. Deze lagen werken samen om een gegeven invoertekst te verwerken en uitvoer-voorspellingen te genereren.

Eind 2023 publiceerden onderzoekers van Carnegie Mellon en Princeton University een onderzoekspaper waarin een nieuwe architectuur voor Large Language Models (LLM’s) werd onthuld, genaamd Mamba. Mamba is een nieuwe state space model-architectuur gericht op sequentiemodellering. Het is ontwikkeld om enkele beperkingen van transformermodellen aan te pakken, vooral bij het verwerken van lange reeksen, en laat veelbelovende prestaties zien.

Laten we de Mamba LLM-architectuur en het belang ervan in machine learning verkennen.

Wat is Mamba?

Mamba is een nieuwe LLM-architectuur die het Structured State Space sequence (S4)-model integreert om lange datasequenties te beheren. Door de beste eigenschappen van recurrente, convolutionele en continu-tijdmodellen te combineren, kan S4 langdurige afhankelijkheden effectief en efficiënt simuleren. Dit stelt het in staat om onregelmatig bemonsterde data te verwerken, een onbegrensde context te hebben en de rekenkundige efficiëntie tijdens training en testen te behouden.

Voortbouwend op het S4-paradigma brengt Mamba verschillende opmerkelijke verbeteringen, vooral bij het afhandelen van tijdsvariantie in operaties. De architectuur draait om een speciaal selectiemechanisme dat de parameters van het structured state space model (SSM) aanpast op basis van de input.

Daardoor kan Mamba met succes minder belangrijke data eruit filteren door zich te richten op cruciale informatie binnen sequenties. Volgens Wikipedia: "Het model schakelt over van een tijdsinvariant naar een tijdsvariërend raamwerk, wat zowel de berekening als de efficiëntie van het systeem beïnvloedt."

Belangrijkste kenmerken en innovaties

Mamba wijkt af van conventionele attention- en MLP-blokken, wat het onderscheidt. Deze vereenvoudiging resulteert in een lichter, sneller model dat lineair schaalt met de lengte van de sequentie – iets wat geen van zijn voorgangers bereikte.

De kernonderdelen van Mamba zijn onder andere:

• Selective-State-Spaces (SSM): Recurrente modellen, die informatie selectief verwerken afhankelijk van de huidige input, vormen de basis van Mamba-SSM’s. Dit stelt ze in staat om irrelevante data te filteren en zich te concentreren op relevante informatie, wat kan leiden tot efficiëntere verwerking.
• Vereenvoudigde architectuur: Mamba vervangt de complexe attention- en MLP-blokken van Transformers door één samenhangend SSM-blok. Dit is bedoeld om inferentie te versnellen en de rekenkundige complexiteit te verlagen.
• Hardware-bewuste parallelisatie: Mamba kan nog beter presteren doordat het een recurrente modus gebruikt met een parallel algoritme dat speciaal is ontworpen voor hardware-efficiëntie.

Een ander essentieel onderdeel is Linear Time Invariance (LTI); LTI is een van de kernkenmerken van S4-modellen. Dit houdt in dat de modelparameters gedurende alle tijdstappen constant blijven, waardoor de dynamiek van het model consistent blijft. Het bouwen van sequentiemodellen wordt eenvoudiger en effectiever met LTI, de basis van recurrency en convoluties.

Details van de Mamba LLM-architectuur

De architectuur van Mamba benadrukt verder het belang van de vooruitgang die in machine learning is geboekt. Het wijzigt de manier waarop modellen sequenties verwerken door een geselecteerde state space model (SSM)-laag te introduceren. Hierdoor kan Mamba twee uiterst belangrijke dingen doen:

1. Focus op relevante informatie – Mamba kan meer voorspellende data voor de taak prioriteren door aan elke input een ander gewicht toe te kennen.

2. Dynamisch aanpassen aan input – Doordat het model zich kan aanpassen aan de input, kan Mamba gemakkelijk verschillende sequentiemodelleringstaken aan.

Daardoor kan Mamba sequenties met ongekende efficiëntie verwerken, wat het de ideale keuze maakt voor taken met lange datasequenties.

Mamba’s ontwerpfilosofie is gebaseerd op een goed begrip van moderne hardwarecapaciteiten. Het is ontworpen om optimaal gebruik te maken van GPU-rekenkracht en garandeert:

• Geoptimaliseerd geheugengebruik: Kortere datatransmissietijden en snellere verwerking worden bereikt door Mamba’s state-expansie zo te ontwerpen dat deze in het high-bandwidth memory (HBM) van GPU’s past.

• Maximale parallelle verwerking: Mamba behaalt een prestatieniveau dat een nieuwe norm stelt voor sequentiemodellen door zijn berekeningen af te stemmen op het parallelle karakter van GPU-computing.

Mamba vs Transformers

De introductie van Transformers, zoals GPT-4, bracht het vakgebied van natural language processing (NLP) vooruit en stelde maatstaven voor verschillende taalverwerkingstaken. Lange sequenties waren echter al lange tijd een zwakke plek voor transformers, omdat ze hun efficiëntie aanzienlijk belemmeren.

Precies op dit punt blinkt Mamba uit. Mamba kan namelijk lange sequenties sneller verwerken dan transformers en doet dit op een eenvoudigere manier dankzij zijn unieke architectuur.

Transformer-architectuur

Transformers zijn zeer bedreven in het werken met datasequenties, zoals tekst voor taalmodellen. In tegenstelling tot eerdere modellen die data sequentieel verwerkten, verwerken ze volledige sequenties gelijktijdig. Deze eigenschap stelt ze in staat complexe relaties in de data vast te leggen.

Ze gebruiken een attention-mechanisme waarmee het model zich tijdens het genereren van voorspellingen op verschillende segmenten van de sequentie kan concentreren. Drie sets gewichten worden gebruikt om deze attention te berekenen: values, keys en queries, die uit de inputdata worden verkregen.

Elk element in een sequentie wordt afgezet tegen elk ander element om aan te geven hoeveel gewicht — of "attention" — het moet hebben om het volgende element in de reeks te voorspellen.

Transformers bestaan uit twee primaire blokken: de decoder, die de uitvoer creëert, en de encoder, die de inputdata verwerkt.

De encoder bestaat uit meerdere lagen – elk met twee sublagen: een eenvoudige, positiegewijze, volledig verbonden feedforwardnetwerk en een multi-head self-attention-mechanisme. Om de training van diepe netwerken te ondersteunen, worden bij elke sublaag residual connections en normalisatie gebruikt.

Net als de encoder bestaat de decoder uit twee lagen met twee sublagen, maar voegt ook een derde sublaag toe die multi-head attention over de uitvoer van de encoder afhandelt. De autoregressieve eigenschap van de decoder blijft behouden door het sequentiële karakter van de decoder, waardoor voorspellingen voor een positie alleen rekening houden met eerdere posities.

Zo proberen Transformers het probleem van lange sequenties aan te pakken met meer complexe attention-mechanismen, maar Mamba kiest een andere benadering.

Mamba-architectuur

Mamba maakt gebruik van selectieve state spaces. Deze methode pakt de rekeninefficiënties van Transformers bij lange sequenties aan.

De architectuur van Mamba maakt snellere inferentie en lineaire schaling met de sequentielengte mogelijk, en creëert zo een nieuw paradigma voor sequentiemodellering dat effectiever kan blijken naarmate sequenties langer worden.

Aangezien we hierboven al dieper op Mamba’s architectuur zijn ingegaan, doen we dat hier niet opnieuw.

Hier is een grafiek van Wikipedia om beter te visualiseren hoe Mamba en Transformers zich tot elkaar verhouden:

Het is het vermelden waard dat, ondanks de vele voordelen die SSM’s hebben ten opzichte van Transformers, die laatsten veel langere sequenties aankunnen dan SSM’s in het geheugen kunnen opslaan, veel minder data nodig hebben om vergelijkbare taken te leren, en SSM’s overtreffen bij taken die het ophalen of kopiëren van de inputcontext vereisen, zelfs met minder parameters. 

Aan de slag met Mamba

Als je met Mamba wilt experimenteren of het in een project wilt gebruiken, heb je het volgende nodig:

• Linux
• NVIDIA GPU
• PyTorch 1.12+
• CUDA 11.6+

Om de vereiste pakketten uit de Mamba-repository te installeren, gebruik je een paar eenvoudige pip-opdrachten:

• [Optioneel] pip install causal-conv1d>=1.2.0: een efficiënte implementatie van een eenvoudige causale Conv1d-laag die binnen het Mamba-blok wordt gebruikt.
• pip install mamba-ssm: het kernpakket van Mamba.

Je kunt het ook vanaf de bron bouwen met pip install . vanuit deze repository.

Als PyTorch-versies compatibiliteitsproblemen veroorzaken, kan pip worden gebruikt met de switch --no-build-isolation om te helpen. Deze modellen zijn getraind op grote datasets zoals de Pile- en SlimPajama-datasets en zijn gebouwd om aan verschillende rekeneisen en prestatiebenchmarks te voldoen.

Het Mamba-model heeft meerdere interface­niveaus, maar de hoofdmodule is het Mamba-architectuurblok dat de selectieve SSM omhult.

Zodra alles is geïnstalleerd, kun je het als volgt gebruiken:

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Toepassingen van Mamba

De introductie van Mamba LLM is een belangrijke potentiële verschuiving in het LLM-architectuurlandschap. Sneller, efficiënter en schaalbaar: Mamba verwerkt moeiteloos lange sequenties met hoge prestatienormen, wat verklaart waarom het een cruciale rol zal spelen in de vormgeving van de toekomst van geavanceerde AI-systemen.

De volgende golf van AI-innovaties kan namelijk worden aangejaagd door de efficiëntie en prestaties van Mamba, die de weg vrijmaken voor de creatie van steeds complexere modellen en applicaties. De potentiële impact is enorm en omvat audio- en spraakverwerkingstoepassingen, analyse van langlopende teksten, contentcreatie, realtime vertaling en meer.

Sectoren die hierdoor kunnen worden getransformeerd, zijn onder andere:

• Gezondheidszorg: Mamba kan het ontwikkelen van gepersonaliseerde geneesmiddelen versnellen door genetische data snel te analyseren.

• Financiën: Mamba kan worden ingezet om langetermijnmarkttrends te analyseren, wat resulteert in nauwkeurigere aandelenvoorspellingen.

• Klantenservice: Mamba kan chatbots aandrijven die langdurige gesprekken monitoren en zo klantcommunicatie verbeteren.

De weg vooruit voor Mamba

Mamba springt eruit als een innovatieve baken dat de weg wijst naar nieuwe mogelijkheden voor het aanpakken van complexe sequentiemodelleringsproblemen. De introductie is een stap richting intelligentere, schaalbare en efficiënte systemen die sequenties met ongekende diepgang kunnen begrijpen en verwerken.

Hoewel Mamba een belangrijke technische mijlpaal is gebleken, hangt het succes niet alleen af van de technologische mogelijkheden – samenwerking in onderzoek, gedeelde resources en open-source-bijdragen spelen allemaal een essentiële rol:

• Open-source-bijdragen: Door onderzoekers en ontwikkelaars aan te moedigen bij te dragen aan de Mamba-codebase, kunnen veerkrachtigere en flexibelere modellen ontstaan.

• Gedeelde resources: De community kan de vooruitgang versnellen door kennis en middelen te bundelen en voorgetrainde modellen te delen.

• Gezamenlijk onderzoek: Samenwerkingen tussen academische instellingen en bedrijven kunnen de mogelijkheden van Mamba uitbreiden.

Conclusie

Mamba biedt niet slechts een incrementele verbetering van huidige sequentiemodellen; het herdefinieert wat mogelijk is. Met de introductie ervan gaat er een nieuw hoofdstuk open in de geschiedenis van kunstmatige intelligentie, waarin rekeninefficiënties en beperkingen in sequentielengte eindelijk tot het verleden gaan behoren.

De afgelopen jaren zagen we de evolutie van AI van RNN’s naar Transformers en nu Mamba, telkens een stap dichter bij AI die diep kan redeneren en informatie kan verwerken op een manier die met mensen vergelijkbaar is. Mamba belichaamt de revolutionaire geest die het vakgebied van AI vooruitstuwt met zijn geselecteerde state space-benadering en lineaire tijdsschaalbaarheid.

Mamba markeert het begin van een veelbelovende ontwikkeling in kunstmatige intelligentie. Het is een paradigma dat voor de toekomst is ontworpen en dat met zijn grenzeloze potentieel een aanzienlijke impact op AI zal hebben.

Wil je verder leren? Bekijk dan:

• Hoe je LLM-applicaties bouwt met LangChain
• Hoe je een LLM traint met PyTorch