Mamba: co potrafią modele przestrzeni stanów

Jest taki sposób czytania, który brzmi jak kara: przy każdym kolejnym słowie wracasz na początek i przebiegasz wzrokiem wszystko, co już przeczytałeś. Zdanie dalej — znowu od zera. Brzmi absurdalnie, a mniej więcej tak właśnie zwykły transformer obchodzi się z długim tekstem. Im dłuższy kontekst wrzucasz, tym więcej pamięci schodzi, i to nie po równo, tylko coraz szybciej. Wśród ludzi, co odpalają to w domu, wraca to jak natrętny wątek: „mam 24 GB, czemu przy 32k kontekstu i tak leci OOM?". Odpowiedź brzmi zawsze tak samo — bo attention, ta sama uwaga, która zrobiła z transformerów króla, ma pewien brzydki nawyk. Za każdy token płaci kwadratem.

I przez lata żyliśmy z tym jak z hałaśliwym sąsiadem: da się, tylko trzeba akceptować. Aż w grudniu 2023 Albert Gu i Tri Dao wrzucili na arXiv pracę o wdzięcznej nazwie Mamba, w której powiedzieli coś w stylu „a gdyby tak w ogóle wyrzucić attention?".¹ Nie ograniczyć, nie zoptymalizować — wyrzucić. Cały mechanizm uwagi, a przy okazji i bloki MLP, i zastąpić je czymś, co skaluje się liniowo, a przy generowaniu trzyma stan o stałym rozmiarze. Brzmi jak sprzedażowa bajka. A jednak papier ma liczby.

Skąd w ogóle ten kwadrat w attention

Zanim zachwycimy się Mambą, warto wiedzieć, od czego ona ucieka. Uwaga w transformerze działa tak, że każdy token patrzy na każdy inny token w sekwencji. Masz 1000 tokenów? To milion par do policzenia. Masz 10 000? Sto milionów. Rośnie z kwadratem długości — stąd to słynne O(n²), które przy długim kontekście robi się głównym żłobem na pamięć i czas.

Do tego dochodzi druga rzecz, boleśnie znajoma każdemu, kto odpala model lokalnie: KV-cache. Żeby przy generowaniu nie liczyć uwagi od zera dla każdego nowego tokenu, transformer trzyma w pamięci klucze i wartości wszystkich dotychczasowych tokenów. I ten bufor rośnie liniowo z długością rozmowy — im dłużej gadasz, tym więcej VRAM-u zjada sam cache, niezależnie od wag modelu. Rozgryzaliśmy ten mechanizm bliżej tutaj, ale sedno jest proste: transformer ma pamięć, która puchnie. Zawsze puchnie.

Mamba wychodzi z zupełnie innej rodziny — z modeli przestrzeni stanów (SSM, state space models).¹ To pomysł stary jak teoria sterowania: zamiast patrzeć na wszystko naraz, przesuwasz się przez sekwencję krok po kroku i utrzymujesz jeden stan — wektor o stałym rozmiarze, który streszcza wszystko, co do tej pory widziałeś. Nowy token wpada, aktualizuje stan, idziesz dalej. Trochę jak RNN, tylko policzony sprytniej.

Dlaczego wcześniejsze SSM nie zabiły transformera

Bo miały jedną wadę, przez którą były głupsze niż attention. Klasyczny SSM — jak S4, poprzednik Mamby od tych samych autorów — ma stałe parametry. Te same macierze przetwarzają każdy token identycznie, bez względu na to, co ten token właściwie mówi. To pozwala policzyć całość jako jeden wielki splot (konwolucję), błyskawicznie i równolegle — ale odbiera modelowi zdolność, którą attention ma za darmo: rozumowanie oparte na treści.

Wyobraź sobie zdanie, w którym ważne jest jedno słowo gdzieś na początku, a reszta to wypełniacz. Attention po prostu skupi się na tym jednym słowie. Stary SSM nie umie — traktuje wypełniacz tak samo poważnie jak sedno, bo nie potrafi wybrać. Nie umie powiedzieć „to zapamiętaj, a tamto olej".

Na czym polega „selekcja", cały myk Mamby

I tu wchodzi jedna zmiana, która robi całą różnicę — selekcja. Gu i Dao uzależnili parametry SSM od wejścia. Krok dyskretyzacji Δ oraz macierze B i C przestają być stałe i stają się funkcjami aktualnego tokenu.¹ Innymi słowy: model przy każdym kroku sam decyduje, ile z nowej informacji wpuścić do stanu, a ile ze starego stanu zapomnieć — w zależności od tego, co właśnie czyta.

To jest ten moment „aha". Selektywny SSM potrafi zrobić to, czego stary nie umiał: przepuścić ważny token dalej wzdłuż sekwencji, a nieważny wygasić. Dostaje content-based reasoning, którym attention się chwalił — tylko bez płacenia kwadratem.

Cena selekcji: znika droga na skróty

Tyle że za wszystko się płaci. Uzależnienie parametrów od wejścia ma paskudny efekt uboczny: zabija formę splotową. Skoro macierze zmieniają się z każdym tokenem, nie policzysz już całości jednym szybkim splotem, jak w S4. Musisz wrócić do trybu rekurencyjnego — krok po kroku — a to na GPU brzmi jak przepis na wolny, sekwencyjny dramat.

Autorzy rozwiązali to algorytmem, który nazwali selective scan, i zaprojektowali go „świadomie sprzętowo" (hardware-aware), czerpiąc wprost z pomysłów, które znasz z FlashAttention.¹ Cały trik siedzi w tym, gdzie na GPU lądują dane. Karta graficzna ma pamięć szybką i małą (SRAM) oraz wolną i dużą (HBM) — a wąskim gardłem prawie zawsze jest przepychanie danych między nimi, nie samo liczenie.

Selective scan robi trzy rzeczy, żeby to obejść:¹

• Nie materializuje rozdmuchanego stanu w wolnej pamięci HBM — dyskretyzację i rekurencję liczy w szybkim SRAM.
• Zlepia operacje w jeden kernel (kernel fusion), zamiast latać z wynikami tam i z powrotem.
• Przy propagacji wstecznej przelicza stany od nowa, zamiast trzymać je w pamięci (recomputation) — bo na nowoczesnym GPU policzyć jest taniej niż odczytać.

Efekt? Model, który w teorii jest rekurencyjny i „powinien" być wolny, w praktyce śmiga. To ta sama filozofia, co przy flash attention: nie zmieniaj matematyki, zmień to, jak dane wędrują przez poziomy pamięci karty.

Co z tego wychodzi w liczbach

Papier nie owija w bawełnę. Kilka twardych wyników:¹

To „5× higher throughput" to nie chochlik — Mamba generuje tokeny pięć razy szybciej niż transformer podobnej klasy, właśnie dzięki temu, że nie musi za każdym krokiem przemielać rosnącego cache’u.¹ A jakościowo Mamba-3B dorównuje transformerom dwukrotnie większym i bije te swojej wielkości — zarówno w pretreningu, jak i w ewaluacji downstream.

I jedna rzecz, która dla długiego kontekstu jest najsmaczniejsza: wydajność na realnych danych poprawia się aż do sekwencji rzędu milionów tokenów.¹ Nie „nie psuje się" — poprawia. Tam, gdzie transformer się dusi, Mamba dopiero się rozkręca. Do tego papier pokazuje wyniki state-of-the-art nie tylko dla języka, ale i dla audio oraz genomiki — jako ogólny backbone sekwencyjny, nie wąska sztuczka do jednej modalności.

Zrób to sam (a przynajmniej dotknij)

To nie jest zamknięty papier z ładnym wykresem. Kod i checkpointy leżą w oficjalnym repo state-spaces/mamba — z gotowymi modelami od 130M do 2.8B parametrów, trenowanymi między innymi na 300 miliardach tokenów (wariant SlimPajama na 600 mld).² Model z abstraktu opisywany jako „Mamba-3B" to na Hugging Face ten sam checkpoint co mamba-2.8b — nazewnictwo się rozjeżdża, liczba parametrów ta sama.³

Jak chcesz zajrzeć pod maskę bloku Mamba, warto znać trzy pokrętła z implementacji:²

# kluczowe hiperparametry modułu Mamba (Mamba-1)
d_state = 16   # rozmiar stanu SSM — ile "pamięci" trzyma ukryty wektor
d_conv  = 4    # szerokość lokalnej konwolucji przed SSM
expand  = 2    # o ile blok rozszerza wymiar wewnętrzny

d_state to wielkość tego stałego stanu, o którym cała gadka — domyślnie 16 dla Mamby-1. d_conv to szerokość krótkiej lokalnej konwolucji (Mamba wciąż lubi lokalny kontekst tuż-tuż), a expand mówi, o ile blok pompuje wymiar wewnętrzny. Trzy liczby, a definiują charakter całej architektury.

Jedno zastrzeżenie na koniec, żeby nie wpaść w hype: Mamba nie wysyła transformerów na emeryturę. Papier pokazuje przewagę na modelowaniu języka, audio i genomice — nie ogłasza, że uwaga jest bezużyteczna wszędzie. Późniejsze prace tych samych autorów (Mamba-2 i tak zwane structured state space duality) pokazują wręcz, że SSM i attention są matematycznie spokrewnione bliżej, niż się na pierwszy rzut oka wydaje — to nie dwa wrogie plemiona, raczej dwie gałęzie tego samego drzewa.

I może właśnie o to chodzi. Przez kilka lat myśleliśmy, że długi kontekst to podatek, który płaci się kwadratem i pęczniejącym cache’em — po prostu prawo natury lokalnego LLM-owania. Mamba przyszła i pokazała, że to była tylko jedna droga przez las, nie jedyna. Ten sam worek, do którego dorzucałeś kamień za każdym tokenem, nagle okazał się plecakiem o stałej wadze — i można iść dalej, aż po horyzont milionów tokenów, bez zaciskania zębów przy liczniku VRAM-u.

Przypisy · źródła