Mamba: Het herdefiniëren van sequentiële modellering en het overtreffen van Transformers-architectuur

In dit artikel over Mamba, zullen we onderzoeken hoe dit innovatieve state-space model (SSM) de sequentiële modellering revolutioneert. Ontwikkeld door Albert Gu en Tri Dao, wordt Mamba onderscheiden door zijn efficiëntie in het verwerken van complexe sequenties in domeinen zoals taalverwerking, genomica en audio-analyse. De sequentiemodellering in lineaire tijd met selectieve state spaces zorgt voor uitzonderlijke prestaties in deze diverse modaliteiten.

We zullen ingaan op Mamba’s vermogen om de computationele uitdagingen die traditionele Transformers tegenkomen, te overwinnen, vooral bij lange sequenties. Zijn selectieve benadering in state space modellen maakt snellere inferentie en lineaire schaling met sequentielengte mogelijk, waardoor de doorvoer aanzienlijk verbetert.

Mamba’s uniekheid ligt in zijn snelle verwerking, selectieve SSM-laag en hardware-vriendelijke ontwerp geïnspireerd door FlashAttention. Deze functies stellen Mamba in staat om veel bestaanhet modellen te overtreffen, waaronder die op basis van de transformer-benadering, waardoor het een opmerkelijke vooruitgang in machine learning is.

Transformers vs Mamba

Transformers, zoals GPT-4, hebben benchmarks gezet in natuurlijke taalverwerking. Echter, hun efficiëntie daalt bij langere sequenties. Hier springt Mamba vooruit, met zijn vermogen om lange sequenties efficiënter te verwerken en zijn unieke architectuur die het hele proces vereenvoudigt.

Transformers zijn bedreven in het verwerken van sequenties van gegevens, zoals tekst voor taalmodellen. In tegenstelling tot eerdere modellen die gegevens sequentieel verwerkten, verwerken Transformers hele sequenties tegelijk, waardoor ze complexe relaties binnen de gegevens kunnen vastleggen.

Zij gebruiken een aandachtsmechanisme, dat het model in staat stelt zich op verschillende delen van de sequentie te richten bij het maken van voorspellingen.

Deze aandacht wordt berekend met behulp van drie sets van gewichten: queries, keys en values, die afgeleid zijn van de invoergegevens. Elk element in een sequentie wordt vergeleken met elk ander element, waardoor een gewicht wordt gegenereerd dat het belang aangeeft, of ‘aandacht’, die elk element zou moeten ontvangen bij het voorspellen van het volgende element in de sequentie.

Transformers behouden twee hoofdblokken: de encoder, die de invoergegevens verwerkt, en de decoder, die de uitvoer genereert. De encoder bestaat uit meerdere lagen, elk met twee sublagen: een multi-head self-attention mechanisme en een eenvoudige, positiegewijze volledig verbonden feed-forward netwerk. Normalisatie en residuele verbindingen worden gebruikt bij elke sublaag om te helpen bij het trainen van diepe netwerken.

De decoder heeft ook lagen met twee sublagen die vergelijkbaar zijn met de encoder, maar voegt een derde sublaag toe die multi-head aandacht uitvoert over de uitvoer van de encoder. De sequentiële aard van de decoder zorgt ervoor dat voorspellingen voor een positie alleen eerdere posities kunnen overwegen, waardoor de autoregressieve eigenschap wordt behouden.

In tegenstelling tot Transformers, neemt het Mamba-model een andere benadering. Terwijl Transformers het probleem van lange sequenties proberen op te lossen door complexere aandachtsmechanismen te gebruiken, gebruikt Mamba selectieve state spaces, waardoor een meer efficiënte verwerking mogelijk wordt.

Hier is een overzicht van hoe een transformer werkt:

1. Invoer verwerking: Transformers verwerken eerst de invoergegevens in een formaat dat het model kan begrijpen, vaak met behulp van embeddings die ook de positie van elk element in de sequentie incorporeren.
2. Aandachtsmechanisme: In zijn kern berekent het aandachtsmechanisme een score die aangeeft hoeveel aandacht moet worden besteed aan andere delen van de invoersequentie bij het begrijpen van een huidig element.
3. Encoder-decoder architectuur: Het Transformer-model bestaat uit een encoder om de invoer te verwerken en een decoder om de uitvoer te genereren. Elk bestaat uit meerdere lagen die het model’s begrip van de invoer verfijnen.
4. Multi-head aandacht: Binnen zowel de encoder als de decoder, stelt multi-head aandacht het model in staat om tegelijkertijd aandacht te besteden aan verschillende delen van de sequentie vanuit verschillende representatieruimten, waardoor het zijn vermogen om te leren van diverse contexten verbetert.
5. Positie-afhankelijke feed-forward netwerken: Na aandacht, verwerkt een eenvoudig neuronaal netwerk de uitvoer van elke positie afzonderlijk en identiek. Dit wordt gecombineerd met de invoer via een residuele verbinding en gevolgd door laagnormalisatie.
6. Uitvoer generatie: De decoder voorspelt vervolgens een uitvoersequentie, beïnvloed door de context van de encoder en wat het tot nu toe heeft gegenereerd.

De transformer’s vermogen om sequenties parallel te verwerken en zijn robuuste aandachtsmechanisme maken het krachtig voor taken zoals vertaling en tekstgeneratie.

Daarentegen werkt het Mamba-model anders door selectieve state spaces te gebruiken om sequenties te verwerken. Deze benadering adresseert de computationele inefficiëntie in Transformers bij het omgaan met lange sequenties. Mamba’s ontwerp maakt snellere inferentie en lineaire schaling met sequentielengte mogelijk, waardoor een nieuw paradigma voor sequentiële modellering ontstaat dat meer efficiënt kan zijn, vooral bij het omgaan met lange sequenties.

Mamba

Mamba laat traditionele aandachts- en MLP-blokken weg. Die vereenvoudiging levert een lichter en sneller model op dat lineair schaalt met de sequentielengte.

De belangrijkste eigenschappen zijn:

• Selectieve SSM’s: filteren irrelevante informatie en richten zich op relevante gegevens, zodat efficiënt redeneren op basis van inhoud mogelijk wordt.
• Hardwarebewust algoritme: een parallel algoritme, geoptimaliseerd voor moderne hardware en vooral GPU’s, versnelt berekeningen en vermindert het geheugengebruik.
• Eenvoudige architectuur: selectieve SSM’s vervangen aandacht en MLP-blokken in een homogener geheel met betere schaalbaarheid en prestaties.

Mamba heeft sterke resultaten behaald bij taal, audio en genomica, zowel bij voortraining als bij domeinspecifieke taken. Bij taalmodellering evenaart of overtreft het ook grotere Transformer-modellen.

Code en voorgetrainde modellen zijn openbaar beschikbaar op GitHub.

Standaard kopieertaken zijn eenvoudig voor lineaire modellen. Selectief kopiëren en induction heads vereisen dynamisch, inhoudsafhankelijk geheugen.

Gestructureerde toestandsruimtemodellen, S4, combineren eigenschappen van RNN’s, CNN’s en klassieke toestandsruimtemodellen. Ze zijn geïnspireerd op continue systemen die eendimensionale functies of sequenties via een impliciete latente toestand transformeren. Binnen deep learning bieden ze een methode voor efficiënte en aanpasbare sequentiemodellen.

De dynamiek van S4-modellen

SSM (S4): het basismodel zet sequentie x om in uitvoer y met geleerde parameters A, B en C en tijdparameter Δ. De transformatie discretiseert de continue parameters en past de SSM-bewerking toe. Bij tijdinvariantie blijven de parameters gelijk over de tijdstappen.

Het belang van discretisatie

Discretisatie zet continue parameters met vaste formules om in discrete parameters en behoudt de verbinding met continue-tijdsystemen. Dit geeft resolutie-invariantie en correcte normalisatie, ten gunste van stabiliteit en prestaties. Er zijn ook overeenkomsten met de gating-mechanismen in RNN’s die de informatiestroom regelen.

Lineariteit en tijdinvariantie (LTI)

Traditionele S4-modellen zijn lineair en tijdinvariant: hun dynamiek blijft gelijk en de parameters staan vast voor alle tijdstappen. LTI ondersteunt recurrente berekeningen en convoluties en biedt een eenvoudig maar krachtig kader voor sequentiemodellering.

Fundamentele beperkingen overwinnen

De LTI-eigenschap beperkt het modelleren van gegevens die aanpasbare dynamiek vereisen. Het Mamba-onderzoek introduceert tijdsafhankelijke parameters en heft die beperking op, zodat meer soorten sequenties en taken kunnen worden verwerkt.

De term toestandsruimtemodel omvat in brede zin recurrente processen met een latente toestand en wordt in verschillende vakgebieden gebruikt. In deep learning verwijzen S4 en gestructureerde SSM’s naar een specifieke klasse die efficiënt rekent en complexe sequenties kan modelleren.

S4-modellen kunnen als zelfstandige sequentietransformaties worden opgenomen in neurale netwerken die van begin tot eind worden getraind. Ze zijn vergelijkbaar met convolutielagen in CNN’s en kunnen de basis van verschillende sequentiearchitecturen vormen.

SSM tegenover SSM met selectie

Waarom selectiviteit belangrijk is

Gestructureerde SSM’s

Volgens het onderzoek is het comprimeren van context tot een beheersbare toestand een fundamenteel onderdeel van sequentiemodellering. Modellen die invoer selectief benadrukken of filteren, houden die gecomprimeerde toestand beter bij en worden efficiënter en krachtiger. Selectiviteit regelt de informatiestroom adaptief langs de sequentie, wat essentieel is voor complexe taaltaken en andere toepassingen.

Selectieve SSM’s maken hun parameters afhankelijk van de invoer en krijgen daarmee aanpassingsvermogen dat tijdinvariante modellen missen. Ze kunnen niet meer op dezelfde manier convoluties gebruiken voor efficiënt rekenen en steunen in plaats daarvan op een lineair recurrent mechanisme.

SSM met selectie (S6): deze variant maakt B, C en Δ afhankelijk van de invoer en kan zich richten op bepaalde delen van sequentie x. De parameters worden met inachtneming van de selectie gediscretiseerd. Een scan-bewerking past de tijdsafhankelijke SSM toe langs de sequentie en stelt de focus dynamisch bij.

Belangrijkste prestatieresultaten van Mamba

Evaluatieresultaten van Mamba zoals weergegeven in het artikel

Mamba combineert hoge inferentiesnelheid met nauwkeurigheid. Het ontwerp benut langere contexten beter, zoals blijkt uit DNA- en audiomodellering. Het overtreft eerdere modellen bij complexe taken met afhankelijkheden over grote afstanden. De veelzijdigheid blijkt ook uit evaluaties zonder taakspecifieke voorbeelden bij meerdere taken, met voordelen voor efficiëntie en schaalbaarheid.

Aan de slag met Mamba

De technische vereisten ten tijde van het artikel zijn Linux, een NVIDIA-GPU, PyTorch 1.12 of nieuwer en CUDA 11.6 of nieuwer. Installatie gebruikt pip-opdrachten voor pakketten uit de Mamba-repository. Bij compatibiliteitsproblemen met PyTorch kan –no-build-isolation helpen. Modellen getraind op onder meer Pile en SlimPajama voorzien in verschillende rekenbehoeften en prestatievereisten.

De interfaces lopen van de selectieve SSM-laag via het Mamba-blok tot complete taalmodellen. Het Mamba-blok is de hoofdmodule, gebruikt een causale Conv1d-laag en kan in andere neurale netwerken worden geïntegreerd. Het artikel beschrijft een Python-gebruiksvoorbeeld waarin een model wordt aangemaakt en gegevens verwerkt, als illustratie van eenvoud en flexibiliteit.

Voorgetrainde modellen zijn beschikbaar op Hugging Face, met 130 miljoen tot 2,8 miljard parameters en training op Pile en SlimPajama. Ze beantwoorden aan uiteenlopende reken- en prestatiebehoeften en volgen de dimensionele standaarden van GPT-3. Hoge doorvoer en nauwkeurigheid maken ze concurrerend voor taalmodellering en andere toepassingen.

De betekenis van Mamba

Mamba biedt een krachtig alternatief voor Transformer-architecturen bij informatierijke gegevens. Het ontwerp sluit aan op moderne hardware en optimaliseert geheugen en parallelle verwerking. Open toegang tot code en voorgetrainde modellen maakt het systeem toegankelijk voor onderzoekers en ontwikkelaars in AI en deep learning.