Mamba: de architectuur die AI-modellen een efficiënter geheugen geeft
Mamba is een architectuur voor het bouwen van AI-modellen die reeksen data verwerken, zoals tekst, spraak of DNA-sequenties. Het is een alternatief voor de zogeheten Transformer-architectuur, die al jaren de motor is achter chatbots zoals ChatGPT en Claude. Waar een Transformer bij elk woord dat hij genereert terugkijkt naar alle voorgaande woorden in een tekst, werkt Mamba anders: het model houdt een compact, voortdurend bijgewerkt 'geheugen' bij terwijl het de tekst doorleest, en hoeft dus niet steeds alles opnieuw te doorzoeken.
Een vergelijking maakt het duidelijk. Stel je een onderzoeker voor die een dik dossier moet doorspitten. De Transformer-aanpak is als iemand die bij elke nieuwe vraag het hele dossier opnieuw doorbladert om alle relevante pagina's te vinden: grondig, maar bij lange dossiers steeds trager en geheugenintensiever. De Mamba-aanpak lijkt meer op iemand die het dossier één keer doorleest en onderweg een samenvatting bijwerkt in een notitieboekje, waarbij hij zelf beslist wat belangrijk genoeg is om te onthouden en wat hij kan laten vallen. Dat maakt het proces sneller en zuiniger, vooral bij hele lange documenten.
Wat is het precies?
Om Mamba te begrijpen, helpt het om eerst te kijken naar het probleem dat het oplost. Transformer-modellen gebruiken een mechanisme dat 'attention' (aandacht) heet: voor elk woord in een zin berekent het model hoe relevant elk ander woord is. Dat werkt goed, maar de rekenkosten stijgen kwadratisch met de lengte van de tekst. Een tekst die twee keer zo lang is, kost dus ongeveer vier keer zoveel rekenkracht en geheugen. Bij boeken, lange gesprekken of hele genomen wordt dat al snel onhoudbaar.
Mamba is gebouwd op een oudere wiskundige familie van modellen, de zogeheten 'state space models' (toestandsruimtemodellen), een concept dat oorspronkelijk uit de regeltechniek komt, denk aan de wiskunde achter een thermostaat of een navigatiesysteem dat continu zijn positie bijhoudt. Zo'n model onderhoudt een 'toestand', een soort samenvatting van alles wat het tot dusver heeft gezien, en werkt die toestand bij telkens wanneer er nieuwe informatie binnenkomt. Dat kost, in tegenstelling tot attention, maar een vaste hoeveelheid rekenwerk per nieuw stukje tekst: de kosten groeien lineair (recht evenredig) met de lengte van de tekst in plaats van kwadratisch.
Een eerdere versie van dit idee, genaamd S4 (Structured State Space Sequence model, voorgesteld in 2021 door onderzoekers van Stanford), kon al goed omgaan met lange reeksen, maar had een belangrijke beperking: de manier waarop het model informatie bijwerkte, lag vooraf vast en hing niet af van de inhoud van de tekst zelf. Het model kon dus niet flexibel beslissen welke woorden belangrijk waren en welke genegeerd konden worden.
De kernvernieuwing van Mamba, gepresenteerd in december 2023 door onderzoekers Albert Gu en Tri Dao, is een 'selectiemechanisme'. De parameters die bepalen hoe de toestand wordt bijgewerkt, hangen nu af van de input zelf. Concreet betekent dit dat het model per woord of teken kan 'kiezen' of het die informatie in zijn geheugen opneemt of juist laat verdwijnen, vergelijkbaar met hoe een lezer die een tekst doorneemt automatisch meer aandacht besteedt aan een naam of getal dan aan een lidwoord. Deze selectiviteit maakte de rekenmethode in eerste instantie lastiger te versnellen op grafische processoren (GPU's), maar de onderzoekers ontwikkelden een hardware-bewust algoritme, gebaseerd op een techniek die 'parallel scan' heet, waarmee de berekeningen alsnog efficiënt parallel uitgevoerd kunnen worden. Een Mamba-blok vervangt bovendien in één keer zowel de attention-laag als de daaropvolgende neurale netwerklaag (MLP) die Transformers gebruiken, wat het bouwblok compacter maakt.
Wat wil men ermee bereiken?
Het belangrijkste doel is efficiëntie bij lange reeksen data. Transformers hebben last van een groeiend 'geheugen' tijdens het genereren van tekst, de zogeheten key-value cache, die met elk nieuw woord groter wordt en uiteindelijk veel werkgeheugen opslokt. Mamba-modellen houden daarentegen een vaste, compacte toestand aan, ongeacht hoe lang de tekst al is. Dat belooft lagere geheugenkosten en snellere verwerking, vooral bij taken met zeer lange context: denk aan het samenvatten van hele boeken, het analyseren van lange audio-opnames, of het doorzoeken van DNA-sequenties die miljoenen tekens lang kunnen zijn.
Een tweede drijfveer is toegankelijkheid. Als modellen met dezelfde prestaties minder rekenkracht en geheugen nodig hebben, worden ze goedkoper om te trainen en te draaien. Dat is aantrekkelijk voor bedrijven die AI willen inzetten op apparaten met beperkte capaciteit, zoals telefoons, of die grote hoeveelheden aanvragen tegelijk willen verwerken zonder torenhoge serverkosten.
Tot slot spelen wetenschappelijke motieven mee: onderzoekers willen weten of attention, het fundament van bijna alle moderne taalmodellen sinds 2017, wel de enige of beste manier is om taal te modelleren. Mamba is een van de meest serieuze pogingen om die aanname te toetsen.
Voorbeelden uit de praktijk
Het oorspronkelijke Mamba-onderzoek van Albert Gu en Tri Dao (december 2023) liet al zien dat de architectuur competitief presteerde op taalmodellering, audiogeneratie en het modelleren van DNA-sequenties, met een doorvoersnelheid die tot ongeveer vijf keer hoger lag dan vergelijkbare Transformer-modellen.
In 2024 volgde Mamba-2, een verbeterde versie die een theoretisch verband blootlegde tussen state space models en attention, de zogeheten 'structured state space duality'. Dit verband hielp onderzoekers de twee benaderingen beter met elkaar te verzoenen.
Mistral AI, het Franse AI-bedrijf, bracht in 2024 Codestral Mamba uit, een model van zeven miljard parameters (een maat voor de omvang van een AI-model) specifiek getraind voor het genereren van programmeercode, volledig gebaseerd op de Mamba-architectuur.
AI21 Labs uit Israël ontwikkelde Jamba (2024), een hybride model dat Mamba-lagen afwisselt met een beperkt aantal Transformer-lagen en met een 'mixture-of-experts'-structuur, een techniek waarbij slechts een deel van het model actief wordt voor elke taak. Jamba combineert zo de efficiëntie van Mamba met de sterke punten van attention.
Het Technology Innovation Institute in Abu Dhabi bracht in 2024 Falcon Mamba uit, een puur op Mamba gebaseerd taalmodel van zeven miljard parameters, als tegenhanger van hun eerdere Transformer-modellen uit de Falcon-reeks.
Hoe ver is de techniek?
Mamba is nog een jong onderzoeksgebied: het originele artikel verscheen eind 2023, en de meeste toepassingen dateren van 2024. Sindsdien is er een golf van vervolgonderzoek en varianten verschenen, wat wijst op grote wetenschappelijke belangstelling.
Toch is de techniek nog niet dominant. De meest gebruikte grote taalmodellen, zoals die van OpenAI, Anthropic en Google, zijn nog altijd overwegend op de klassieke Transformer gebaseerd. Puur op Mamba gebaseerde modellen presteren op de meeste taken vergelijkbaar met Transformers, maar blijken op sommige specifieke taken nog achter te lopen, met name taken die exact terughalen van een specifiek stukje informatie uit een lange context vereisen (in-context leren en precieze retrieval). Attention is daar van nature sterk in, omdat het letterlijk elk eerder woord apart kan 'opzoeken', terwijl Mamba alles samenperst in één vaste toestand.
Daarom wordt momenteel veel onderzoek gedaan naar hybride architecturen die Mamba-lagen combineren met een beperkt aantal attention-lagen, zoals Jamba en Zamba (van het bedrijf Zyphra) al laten zien. Deze hybrides proberen het beste van beide werelden te combineren: de snelheid en het lage geheugengebruik van Mamba, met de precisie van attention waar dat nodig is. Er is nog geen wetenschappelijke consensus of pure state space models de Transformer op termijn zullen vervangen, of dat hybride vormen de toekomst zijn.
Wie werken eraan?
De architectuur is bedacht door Albert Gu, verbonden aan Carnegie Mellon University, en Tri Dao, verbonden aan Princeton University. Beiden bouwden voort op eerder werk aan state space models, waaronder S4, dat mede werd ontwikkeld met onderzoekers van Stanford University. Gu en Dao richtten later ook het bedrijf Cartesia op, dat de technologie commercieel toepast, onder meer voor realtime spraaktoepassingen.
Daarnaast werken meerdere AI-bedrijven aan eigen Mamba-gebaseerde of hybride modellen: Mistral AI (Frankrijk), AI21 Labs (Israël), het Technology Innovation Institute (Verenigde Arabische Emiraten) en Zyphra (Verenigde Staten) zijn de bekendste voorbeelden. Verder publiceren onderzoeksgroepen wereldwijd op platforms als arXiv over verwante architecturen, zoals RWKV en Hyena, die net als Mamba proberen af te stappen van de kwadratische rekenkosten van klassieke attention.