Mamba: Omdefinerer sekvensmodellering og overgår Transformers-arkitektur

I denne artikel om Mamba undersøger vi, hvordan denne innovative tilstandsrumsmodel (SSM) ændrer sekvensmodellering. Mamba er udviklet af Albert Gu og Tri Dao og behandler komplekse sekvenser effektivt inden for sprog, genomik og lydanalyse. Selektive tilstandsrum giver tidsforbrug, der vokser lineært med sekvenslængden, og stærke resultater på tværs af disse modaliteter.

Vi ser på, hvordan Mamba håndterer de beregningsmæssige udfordringer, traditionelle Transformer-modeller møder ved lange sekvenser. Den selektive tilgang giver hurtigere inferens og lineær skalering med sekvenslængden, hvilket øger gennemstrømningen.

De særlige egenskaber er hurtig behandling, et selektivt SSM-lag og et hardwarevenligt design inspireret af FlashAttention. De gør Mamba i stand til at overgå mange eksisterende modeller, herunder Transformer-baserede modeller, og udgør et vigtigt fremskridt i maskinlæring.

Transformer-modeller og Mamba

Transformer-modeller som dem bag GPT-4 har sat standarder inden for naturlig sprogbehandling. Effektiviteten falder dog med længere sekvenser. Mamba behandler lange sekvenser mere effektivt med en arkitektur, der forenkler processen.

Transformer-modeller håndterer datasekvenser som tekst. I modsætning til tidligere modeller, der behandlede data strengt sekventielt, kan de behandle hele sekvenser parallelt og dermed opfange komplekse relationer i dataene.

En opmærksomhedsmekanisme lader modellen fokusere på forskellige dele af sekvensen under forudsigelser. Beregningen bruger queries, keys og values, som udledes af inputdata. Hvert element sammenlignes med de øvrige, og vægtene angiver deres relative betydning for forudsigelsen.

Den encoder-decoder-arkitektur, artiklen beskriver, har en encoder til input og en decoder til output. Encoderen består af flere lag med to dellag: selvopmærksomhed med flere hoveder og et fuldt forbundet feed-forward-netværk, som anvendes positionsvist. Normalisering og residualforbindelser understøtter træningen af dybe netværk.

Decoderen har tilsvarende dellag og desuden et tredje, som udfører opmærksomhed med flere hoveder over encoderens output. Ved autoregressiv generering kan en position kun bygge på tidligere positioner. Mamba bruger i stedet selektive tilstandsrum til at håndtere lange sekvenser.

Sådan fungerer en Transformer-model i hovedtræk:

1. Inputbehandling: Data kodes til repræsentationer, modellen kan forstå, ofte med embeddings, som også indeholder information om elementernes positioner.
2. Opmærksomhed: Der beregnes scorer for, hvor stor vægt andre dele af inputsekvensen skal have ved fortolkningen af det aktuelle element.
3. Encoder-decoder-arkitektur: Encoderen behandler input og decoderen genererer output. Flere lag forfiner repræsentationerne.
4. Opmærksomhed med flere hoveder: Modellen kan samtidigt fokusere på forskellige dele af sekvensen fra flere repræsentationsrum og lære fra forskellige kontekster.
5. Positionsvise feed-forward-netværk: Et neuralt netværk behandler hver position separat og på samme måde. Residualforbindelser og lagnormalisering kombinerer og bearbejder repræsentationerne.
6. Outputgenerering: Decoderen forudsiger outputsekvensen ud fra encoderens kontekst og det, den allerede har genereret.

Parallel sekvensbehandling og opmærksomhed gør Transformer-modeller effektive til oversættelse og tekstgenerering. Mambas selektive tilstandsrum adresserer deres beregningsmæssige ineffektivitet ved lange sekvenser. Hurtigere inferens og lineær skalering kan give en stadig større fordel, når sekvenserne vokser.

Mamba

Mamba adskiller sig fra traditionelle modeller ved at undvære opmærksomheds- og MLP-blokke. Forenklingen giver en lettere og hurtigere model, som skalerer lineært med sekvenslængden.

De vigtigste egenskaber er:

• Selektive SSM’er: De filtrerer irrelevante oplysninger fra og fokuserer på relevante data. Selektiviteten er afgørende for effektiv ræsonnering baseret på indhold.
• Hardwarebevidst algoritme: En parallel algoritme er optimeret til moderne hardware, især GPU’er. Den giver hurtigere beregning og mindre hukommelsesforbrug end traditionelle modeller.
• Forenklet arkitektur: Selektive SSM’er erstatter opmærksomheds- og MLP-blokke i en enklere og mere ensartet struktur, som forbedrer skalering og ydeevne.

Mamba har opnået stærke resultater inden for sprog, lyd og genomik, både ved fortræning og i domænespecifikke opgaver. I sprogmodellering matcher eller overgår den større Transformer-modeller.

Koden og de fortrænede modeller er offentligt tilgængelige på GitHub.

Almindelige kopieringsopgaver er enkle for lineære modeller. Selektiv kopiering og induction heads kræver dynamisk hukommelse, der tager højde for indholdet.

Strukturerede tilstandsrumsmodeller, S4, er en lovende klasse af sekvensmodeller, der kombinerer egenskaber fra RNN’er, CNN’er og klassiske tilstandsrumsmodeller. De er inspireret af kontinuerte systemer, som transformerer endimensionale funktioner eller sekvenser gennem en implicit latent tilstand. I deep learning giver de en metode til at udvikle effektive og tilpasningsdygtige sekvensmodeller.

Dynamikken i S4-modeller

SSM (S4): Grundmodellen omformer en sekvens x til et output y med de lærte parametre A, B og C samt tidsparameteren Δ. Parametrene diskretiseres, så kontinuerte størrelser bliver diskrete, og SSM-operationen anvendes. I den tidsinvariante model ændres parametrene ikke fra tidsskridt til tidsskridt.

Betydningen af diskretisering

Diskretisering omdanner kontinuerte parametre til diskrete gennem faste formler og bevarer forbindelsen til systemer med kontinuert tid. Det giver egenskaber som invarians over for opløsning og korrekt normalisering, hvilket forbedrer stabilitet og ydeevne. Processen har også ligheder med gating-mekanismer i RNN’er, som styrer informationsstrømmen.

Linearitet og tidsinvarians (LTI)

De traditionelle S4-modeller er lineære og tidsinvariante. Dynamikken forbliver den samme, og parametrene er faste ved alle tidsskridt. Denne egenskab understøtter rekurrens og foldning og giver en enkel, stærk ramme for sekvensmodellering.

At overvinde grundlæggende begrænsninger

LTI-egenskaben begrænser modelleringen af data, der kræver tilpasningsdygtig dynamik. Mamba-arbejdet introducerer tidsvarierende parametre og fjerner dermed begrænsningen. Det udvider de sekvenser og opgaver, modellerne kan håndtere.

Begrebet tilstandsrumsmodel dækker generelt rekurrente processer med en latent tilstand og bruges i mange fagområder. Inden for deep learning betegner S4 eller strukturerede SSM’er en bestemt klasse, der er optimeret til effektiv beregning og modellering af komplekse sekvenser.

S4-modeller kan indgå som selvstændige sekvenstransformationer i neurale netværk, der trænes samlet fra input til output. De kan sammenlignes med foldningslag i CNN’er og fungere som grundstruktur i forskellige netværksarkitekturer.

SSM sammenlignet med SSM med selektion

Hvorfor selektivitet er vigtig

Strukturerede SSM’er

Artiklen argumenterer for, at sekvensmodellering grundlæggende handler om at komprimere konteksten til en håndterbar tilstand. Modeller, som kan fokusere selektivt på eller filtrere input, bevarer denne tilstand mere effektivt. De bliver dermed mere effektive og stærkere. Selektiviteten gør det muligt at styre informationsstrømmen adaptivt langs sekvensen, hvilket er vigtigt for komplekse sprogopgaver og andre anvendelser.

Selektive SSM’er gør parametrene afhængige af input og får dermed en tilpasningsevne, som tidsinvariante modeller mangler. De tidsvarierende modeller kan ikke længere bruge foldning på samme måde til effektiv beregning, men anvender i stedet en lineær rekurrent mekanisme.

SSM med selektion (S6): Varianten tilføjer inputafhængighed til B, C og Δ. Modellen kan fokusere på udvalgte dele af inputsekvensen x. Parametrene diskretiseres under hensyntagen til selektionen, og en scan-operation anvender den tidsvarierende SSM-beregning langs sekvensen og tilpasser fokus dynamisk.

Mambas vigtigste resultater

Mambas evalueringsresultater som vist i artiklen

Mamba kombinerer høj inferenshastighed med nøjagtighed. Modellen udnytter længere kontekster bedre, hvilket ses i DNA- og lydmodellering, hvor den overgår tidligere modeller på komplekse opgaver med langtrækkende afhængigheder. Alsidigheden fremgår også af evalueringer uden opgavespecifikke eksempler på tværs af flere opgaver, med forbedret effektivitet og skalerbarhed.

Kom i gang med Mamba

De tekniske krav angivet ved artiklens udgivelse er Linux, en NVIDIA-GPU, PyTorch 1.12 eller nyere og CUDA 11.6 eller nyere. Installationen bruger pip-kommandoer til pakker fra Mambas kodearkiv. Ved kompatibilitetsproblemer med PyTorch kan pip-flaget –no-build-isolation hjælpe. Modeller trænet på blandt andet Pile og SlimPajama er udformet til forskellige beregningsbehov og præstationskrav.

Grænsefladerne spænder fra det selektive SSM-lag til Mamba-blokken og hele sprogmodeller. Mamba-blokken er hovedmodulet, bruger et kausalt Conv1d-lag og kan integreres i andre neurale netværk. Artiklen omtaler et Python-eksempel, hvor en Mamba-model oprettes og behandler data, som illustration af systemets enkelhed og fleksibilitet.

Fortrænede modeller findes på Hugging Face i størrelser fra 130 millioner til 2,8 milliarder parametre, trænet på Pile og SlimPajama. De dækker forskellige beregnings- og præstationsbehov og følger GPT-3’s dimensionsstandarder. Høj gennemstrømning og nøjagtighed gør dem konkurrencedygtige til blandt andet sprogmodellering.

Mambas betydning

Mamba tilbyder et stærkt alternativ til Transformer-arkitekturer til informationsrige data. Designet passer til moderne hardware og optimerer både hukommelse og parallel behandling. Åben adgang til kode og fortrænede modeller gør systemet tilgængeligt for forskere og udviklere inden for AI og deep learning.