Architektury Mamba-2, State Space Models i liniowa atencja: Jak modele SSM i hybrydowe sieci rekurencyjne eliminują kwadratową złożoność transformerów w przetwarzaniu kontekstów milionowych tokenów

Impas kwadratowej złożoności mechanizmu Softmax Attention

Od czasu publikacji przełomowej pracy Attention Is All You Need w 2017 roku, mechanizm samoatencji znormalizowanej funkcją Softmax stanowił niepodważalny dogmat głębokiego uczenia maszynowego. Transformery zdołały wyprzeć klasyczne rekurencyjne sieci neuronowe (RNN) oraz sieci ze sprzężeniem zwrotnym LSTM, oferując coś, czego wcześniejsze architektury nie potrafiły zapewnić: pełne zrównoleglenie procesu trenowania na procesorach graficznych. Zdolność do jednoczesnego przetwarzania wszystkich tokenów w sekwencji bez konieczności czekania na sekwencyjne przekazywanie stanu ukrytego pozwoliła na bezprecedensowe przeskalowanie modeli do setek miliardów parametrów i terabajtów danych tekstowych.

Jednakże w miarę jak apetyt inżynierów i badaczy rósł, a długość okna kontekstowego stała się kluczową metryką rynkową determinującą użyteczność modeli w analizie wielotomowych repozytoriów kodu, dokumentacji medycznej czy pełnometrażowych filmów wideo, fundamenty architektury Transformer zaczęły pękać. Podstawowym problemem jest mechanizm iloczynu skalarnego macierzy zapytań (Query) i kluczy (Key). Aby obliczyć wagi atencji dla sekwencji o długości N tokenów, konieczne jest wygenerowanie pełnej macierzy korelacji o wymiarach N na N. Oznacza to, że zarówno liczba operacji zmiennoprzecinkowych, jak i zapotrzebowanie na pamięć podręczną rosną z kwadratem długości sekwencji.

Kwadratowa złożoność staje się szczególnie dewastująca podczas etapu wnioskowania (Inference). W trybie autoregresyjnym, w którym model generuje odpowiedź token po tokenie, konieczne jest bezustanne utrzymywanie w pamięci HBM akceleratora wektorów kluczy i wartości wszystkich wcześniej przetworzonych słów w postaci tak zwanego bufora KV Cache. Dla kontekstu rzędu kilkuset tysięcy tokenów, rozmiar bufora KV Cache dla pojedynczej sesji użytkownika zaczyna wielokrotnie przekraczać rozmiar wag samego modelu neuronowego. Prowadzi to do drastycznego spadku przepustowości klastra obliczeniowego, ograniczenia liczby równolegle obsługiwanych zapytań oraz gigantycznego wzrostu kosztów infrastruktury serwerowej.

Warto podkreślić, że problem ten ma charakter fundamentalny i wynika bezpośrednio z natury algebry liniowej leżącej u podstaw klasycznych transformatorów. Żadne techniki inżynieryjne, takie jak kwantyzacja wektorów kluczy i wartości do formatów 4-bitowych czy zaawansowane metody stronicowania pamięci pokroju vLLM i PagedAttention, nie są w stanie zmienić wykładniczego nachylenia krzywej kosztu pamięciowego. Stanowią one jedynie doraźne zabiegi optymalizacyjne odsuwające w czasie zderzenie z fizycznym limitem przepustowości szyn pamięci HBM, co czyni poszukiwanie alternatywnych fundamentów matematycznych absolutną koniecznością dla dalszego rozwoju sztucznej inteligencji.

Matematyczne fundamenty modeli przestrzeni stanów SSM

W poszukiwaniu ucieczki przed kwadratowym wąskim gardłem badacze zwrócili się ku klasycznej teorii sterowania i równaniom różniczkowym, które od dziesięcioleci stanowiły fundament inżynierii lotniczej, przetwarzania sygnałów i automatyki przemysłowej. Modele przestrzeni stanów (State Space Models, SSM) opisują ewolucję układu dynamicznego w czasie ciągłym za pomocą układu liniowych równań różniczkowych. Układ ten mapuje jednowymiarowy sygnał wejściowy u(t) na N-wymiarowy stan ukryty h(t), a następnie rzutuje go na sygnał wyjściowy y(t) za pośrednictwem macierzy parametrów A, B, C i D.

Kluczowym problemem w zastosowaniu równań ciągłych w architekturze cyfrowych procesorów jest konieczność ich dyskretyzacji. Do reprezentacji dyskretnej wykorzystuje się zaawansowane transformacje matematyczne, takie jak metoda ekstrapolacji zerowego rzędu (Zero-Order Hold, ZOH) lub transformacja biliniowa Tustina. Po zdyskretyzowaniu za pomocą kroku czasowego delta, ciągłe macierze przekształcają się w dyskretne operatory przejścia stanu A-bar oraz B-bar. Niezwykłą właściwością liniowych modeli przestrzeni stanów jest ich podwójna natura: mogą one być ewaluowane w sposób rekurencyjny krok po kroku (co zapewnia stały koszt pamięciowy O(1) podczas inferencji) lub sformułowane jako globalny splot dyskretny (co pozwala na pełne zrównoleglenie obliczeń na GPU za pomocą szybkiej transformaty Fouriera FFT podczas treningu).

Przełomem teoretycznym, który umożliwił modelom SSM zapamiętywanie długoterminowych zależności bez zjawiska zanikającego gradientu, była macierz HiPPO (High-order Polynomial Projection Operators), opracowana przez Alberta Gu. Macierz HiPPO kompresuje całą historię napływającego sygnału wejściowego do współczynników wielomianów Legendre'a, optymalizując rzutowanie przeszłości na skończoną przestrzeń stanu ukrytego. Dzięki temu model zyskuje matematycznie dowiedzioną zdolność do idealnego rekonstruowania informacji sprzed setek tysięcy kroków czasowych, co stanowiło piętę achillesową tradycyjnych sieci rekurencyjnych.

Od S4 do Mamba-1: Wprowadzenie mechanizmu selektywnego

Mimo matematycznej elegancji pierwotne modele oparte na strukturze HiPPO i architekturze S4 (Structured State Space for Sequence Modeling) posiadały fundamentalną słabość w porównaniu z transformerami: były systemami liniowymi niezmienniczymi w czasie (Linear Time-Invariant, LTI). Oznaczało to, że macierze A, B i C były statyczne i nie zależały od aktualnie przetwarzanej treści. Model traktował każdy napływający token z jednakową wagą wagową, nie potrafiąc selektywnie odfiltrować nieistotnego szumu ani skupić uwagi na kluczowych faktach semantycznych. W efekcie w testach syntetycznych wymagających zapamiętywania asocjacyjnego i selektywnego kopiowania informacji modele LTI ponosiły spektakularne porażki.

Rozwiązaniem tego problemu stała się zaprezentowana pod koniec 2023 roku architektura Mamba-1, stworzona przez Alberta Gu i Tri Dao. Kluczową innowacją Mamby było przekształcenie modelu LTI w system zależny od czasu (Time-Varying) poprzez mechanizm selektywnego skanowania (Selective State Space). W architekturze Mamba macierze parametrów B i C oraz krok czasowy delta stały się dynamicznymi funkcjami wektora wejściowego w danym kroku czasowym. Model zyskał zdolność do dynamicznego otwierania i zamykania bram pamięciowych: jeśli token niesie kluczową informację merytoryczną, krok delta ulega zwiększeniu, wprowadzając dane do stanu ukrytego; jeśli token jest pustym spójnikiem lub szumem, krok delta dąży do zera, całkowicie ignorując wejście.

Wprowadzenie selektywności zniszczyło możliwość stosowania prostego splotu globalnego, co groziło utratą zrównoleglenia na procesorach graficznych. Autorzy Mamby rozwiązali ten dylemat za pomocą wysoce zoptymalizowanego jądra sprzętowego (Hardware-Aware Parallel Scan). Zamiast materializować wielkie tensory w powolnej pamięci HBM, algorytm wykonuje prefiksowy skan równoległy bezpośrednio w ultraszybkiej pamięci podręcznej SRAM procesora graficznego, eliminując wąskie gardła transferu I/O i deklasując transformery pod względem szybkości trenowania.

Rewolucja Mamba-2 i teoria State Space Duality SSD

Mimo sukcesu Mamby-1, wdrożenie modeli SSM na skalę przemysłową napotykało na barierę infrastrukturalną. Współczesne akceleratory obliczeniowe, takie jak NVIDIA H100, B200 czy układy TPU, są fizycznie zoptymalizowane pod kątem jednej, konkretnej operacji matematycznej: gęstego mnożenia macierzy (GEMM) na dedykowanych rdzeniach Tensor. Równoległy skan selektywny w Mambie-1, choć wydajny pamięciowo, wykorzystywał głównie tradycyjne jednostki zmiennoprzecinkowe CUDA, pozostawiając potężne rdzenie Tensor w stanie chronicznej niedoutilizacji.

Odpowiedzią na ten impas stała się architektura Mamba-2, oparta na monumentalnej teorii State Space Duality (SSD), sformułowanej przez Alberta Gu i Tri Dao w 2024 roku. Badacze udowodnili matematycznie, że ustrukturyzowane modele przestrzeni stanów oraz warianty liniowej atencji (Linear Attention) to dokładnie ten sam obiekt teoretyczny, rozpatrywany w dwóch różnych reprezentacjach numerycznych. Okazało się, że selektywny model przestrzeni stanów z diagonalną macierzą przejścia A może zostać zapisany w postaci operacji macierzowej z trójkątną dolną macierzą masującą o specyficznej strukturze półoddzielnej (Semi-Separable Matrix).

Odkrycie dualizmu SSD przyniosło gigantyczne konsekwencje praktyczne. Pozwoliło na sformułowanie algorytmu SSD jako blokowego mnożenia macierzy, które w osiemdziesięciu pięciu procentach wykonuje się bezpośrednio na rdzeniach Tensor procesorów graficznych. Wewnątrz pojedynczego bloku tokenów (np. 64 lub 128 słów) algorytm wykorzystuje reprezentację atencyjną zoptymalizowaną pod Tensor Cores, natomiast pomiędzy blokami przekazuje skompresowany stan ukryty w reprezentacji rekurencyjnej SSM. Mamba-2 osiągnęła dzięki temu od dwu- do ośmiokrotnego przyspieszenia treningu w stosunku do Mamby-1, oferując jednocześnie pełną zgodność z nowoczesnymi bibliotekami rozproszonego treningu Megatron-LM i parallelismem tensorowym.

Porównanie złożoności obliczeniowej i zachowania pamięci

Aby w pełni zrozumieć przełom, jaki niosą architektury SSM, konieczne jest zestawienie ich charakterystyk wydajnościowych z tradycyjnym mechanizmem Transformer. W klasycznym transformerze z pełną atencją złożoność obliczeniowa treningu wynosi O(N^2 * D), gdzie N to długość kontekstu, a D to wymiar ukryty modelu. W architekturach Mamba-2 i SSD złożoność ta ulega redukcji do ściśle liniowej O(N * D * S), gdzie S to stały wymiar stanu przestrzeni, niezależny od długości tekstu. Oznacza to, że przy wydłużeniu kontekstu z czterech tysięcy do miliona tokenów, zapotrzebowanie na operacje FLOPs w transformerze rośnie sześćdziesiąt cztery tysiące razy, podczas gdy w Mambie-2 rośnie dokładnie dwieście pięćdziesiąt razy.

Jeszcze bardziej spektakularna różnica dotyczy zapotrzebowania na pamięć operacyjną podczas generowania tekstu. W transformerze rozmiar bufora KV Cache rośnie w sposób ciągły z każdym wygenerowanym słowem. Dla modelu o rozmiarze siedemdziesięciu miliardów parametrów i kontekstu miliona tokenów, KV Cache pochłania setki gigabajtów pamięci HBM, wymagając alokacji całych szaf serwerowych wyłącznie na potrzeby przechowywania przeszłych wektorów atencji. W architekturze Mamba-2 stan ukryty ma stały rozmiar (zależny wyłącznie od wymiaru macierzy stanu), co oznacza, że zużycie pamięci podczas inferencji jest całkowicie niezmienne O(1) bez względu na to, czy generowany jest setny, czy milionowy token w sekwencji.

Ta niezmienniczość pamięciowa całkowicie eliminuje problem fragmentacji pamięci i pozwala na radykalne zwiększenie rozmiaru partii (Batch Size) w klastrach produkcyjnych. Pojedynczy serwer GPU jest w stanie obsługiwać dziesiątki równoległych strumieni zapytań o ekstremalnej długości kontekstu, osiągając przepustowość generowania tokenów na sekundę nieosiągalną dla jakiejkolwiek implementacji tradycyjnego transformera, nawet przy zastosowaniu agresywnej kwantyzacji KV Cache.

Hybrydyzacja architektoniczna: Modele Jamba, Zamba i połączenie światów

Mimo niepodważalnych przewag wydajnościowych, czyste modele przestrzeni stanów posiadają jedną specyficzną słabość wynikającą ze skończonej pojemności ich stanu ukrytego: wykazują gorsze rezultaty w zadaniach wymagających precyzyjnego przypominania faktów asocjacyjnych (In-Context Associative Recall) oraz zaawansowanego wnioskowania logicznego na odległych fragmentach tekstu w porównaniu z pełną atencją Softmax. Atencja Softmax działa jak bezstratna baza danych, w której każdy token może w dowolnym momencie bezpośrednio odpytać każdy inny token w sekwencji, podczas gdy model rekurencyjny kompresuje całą przeszłość do jednego wektora.

Najbardziej pragmatycznym i efektywnym kierunkiem rozwoju sztucznej inteligencji stała się hybrydyzacja architektoniczna. Zamiast budować modele oparte wyłącznie na Mambie lub wyłącznie na Transformerze, inżynierowie łączą obie technologie w spójne bloki warstwowe. Przykładem takiego podejścia są komercyjne modele z rodziny Jamba (opracowane przez AI21 Labs) oraz Zamba. W architekturze hybrydowej większość warstw sieci (np. siedemdziesiąt do osiemdziesięciu procent) stanowią bloki Mamba-2 o liniowej złożoności, które odpowiadają za płynną syntezę tekstu, filtrację szumu i lokalne modelowanie języka, natomiast co kilka warstw umieszcza się klasyczną warstwę atencji Transformer z mechanizmem Shared KV Cache.

Takie połączenie oferuje to, co najlepsze z obu światów. Rzadkie warstwy atencji gwarantują nieskazitelną precyzję wyszukiwania faktów w teście Needle In A Haystack na dystansie milionów tokenów, podczas gdy dominujące warstwy Mamby drastycznie redukują rozmiar całkowitego bufora pamięciowego i przyspieszają wnioskowanie. Co więcej, architektury hybrydowe doskonale integrują się ze strukturami mieszaniny ekspertów (Mixture of Experts, MoE), gdzie wybrane warstwy Mamba i atencji aktywują jedynie podzbiór dostępnych parametrów w zależności od semantyki zapytania.

Implementacja sprzętowa i implikacje dla centrów danych AI

Wdrożenie architektur liniowych wywołuje potężne konsekwencje dla projektowania sprzętu w centrach danych przyszłości. Przez lata dominacja transformerów wymuszała na producentach akceleratorów (NVIDIA, AMD, Google, cerebras) nieustanny wyścig o zwiększanie przepustowości magistrali pamięciowej HBM. W modelach o liniowej złożoności i stałym stanie ukrytym wąskie gardło pamięciowe ulega znacznemu złagodzeniu, przesuwając punkt ciężkości z powrotem ku czystej gęstości upakowania rdzeni obliczeniowych i optymalizacji transferu na poziomie chipletów.

Architektura Mamba-2 stwarza również bezprecedensowe możliwości dla systemów brzegowych (Edge AI) i urządzeń konsumenckich o rygorystycznych ograniczeniach poboru mocy. Smartfony, laptopy i autonomiczne roboty, dysponujące ograniczonymi zasobami pamięci LPDDR5, nie były w stanie uruchamiać transformerów z oknem kontekstowym przekraczającym kilka tysięcy tokenów. Dzięki modelom SSM urządzenia te mogą bez problemu przetwarzać ciągłe strumienie danych telemetrycznych, audio i wideo przez całe doby bez ryzyka przepełnienia pamięci RAM i przegrzania układu SoC.

W wielkich klastrach chmurowych przejście na modele hybrydowe SSM-Transformer oznacza wielokrotne obniżenie całkowitego kosztu posiadania (TCO) infrastruktury. Operatorzy centrów danych mogą obsłużyć tę samą liczbę użytkowników przy użyciu znacznie mniejszej liczby szaf serwerowych, redukując zużycie energii elektrycznej i wody chłodzącej. Jest to kluczowy czynnik w obliczu narastającego kryzysu energetycznego napędzanego przez wykładniczy wzrost zapotrzebowania na moc obliczeniową AI.

Alternatywne architektury liniowe: RWKV v6 i mechanizmy RetNet

Choć Mamba-2 skupia na sobie największą uwagę środowiska akademickiego i przemysłowego, nie jest jedyną architekturą rzucającą wyzwanie transformerom. Równolegle rozwija się rodzina modeli RWKV (Receptance Weighted Key Value), której najnowsze iteracje v6 Eagle i Finch zyskały status pełnoprawnych modeli bazowych o architekturze rekurencyjnej. W przeciwieństwie do modeli SSM wywodzących się z teorii układów dynamicznych, RWKV zostało zaprojektowane jako bezpośrednia modyfikacja transformera, w której funkcja Softmax została zastąpiona mechanizmem liniowej interpolacji czasowej (Time-Mixing).

Kluczową innowacją w RWKV v6 jest wprowadzenie dynamicznego mechanizmu mieszania kanałów (Channel-Mixing) oraz wielogłowicowego mechanizmu akumulacji wykładniczej ze zmiennym współczynnikiem zaniku. Model zachowuje strukturę rekurencyjną w fazie inferencji, umożliwiając generowanie tekstu ze stałym zużyciem pamięci RAM, a jednocześnie trenuje się z prędkością tradycyjnego transformera. Podobną ścieżką podąża architektura RetNet (Retentive Network), zaproponowana przez inżynierów Microsoft Research, która łączy jawną rekurencję z mechanizmem obrotowych kodowań pozycyjnych (RoPE), oferując trzy reprezentacje obliczeniowe: równoległą do treningu, rekurencyjną do inferencji oraz blokową do przetwarzania długich sekwencji.

Wszystkie te podejścia łączy wspólny fundament matematyczny: rezygnacja z globalnej macierzy korelacji Softmax na rzecz kumulatywnych stanów pośrednich. O ile tradycyjna atencja zmusza model do bezustannego porównywania każdego nowego słowa z każdym słowem zapisanym w historii, o tyle architektury nowej generacji akumulują informację semantyczną w spójnym, gęstym wektorze pamięci dynamicznej, aktualizowanym za pomocą operacji addytywnych i multiplikatywnych w locie.

Wyzwania uczenia rozproszonego i równoległości sekwencyjnej

Skalowanie modeli o liniowej złożoności do rozmiarów setek miliardów parametrów wymagało przezwyciężenia unikalnych wyzwań w obszarze inżynierii uczenia rozproszonego. Klasyczne biblioteki takie jak Megatron-LM czy DeepSpeed przez lata optymalizowały techniki równoległości tensorowej (Tensor Parallelism) i potokowej (Pipeline Parallelism) pod kątem specyfiki transformerów. W przypadku transformerów podział macierzy atencji pomiędzy węzły klastra był zadaniem dobrze zrozumianym i ustandaryzowanym.

W architekturze SSM wprowadzenie sekwencyjnego stanu ukrytego początkowo utrudniało podział pojedynczej sekwencji pomiędzy wiele oddzielnych kart graficznych (Sequence Parallelism). Gdy kontekst liczy miliony tokenów, pojedynczy procesor graficzny nie jest w stanie pomieścić wejściowych wektorów aktywacji w pamięci podręcznej, co wymusza pocięcie sekwencji na fragmenty przydzielane do różnych GPU. W klasycznym transformerze stosowano algorytmy takie jak Ring Attention, w których karty graficzne przekazywały sobie bloki pamięci podręcznej w kołowym potoku transmisyjnym.

W Mambie-2 problem ten rozwiązano dzięki wykorzystaniu właściwości State Space Duality. Ponieważ operacja SSM może być zapisana w postaci blokowej macierzy trójkątnej dolnej, możliwe stało się zastosowanie hierarchicznego skanowania równoległego. Każdy procesor graficzny oblicza lokalną ewolucję stanu dla swojego fragmentu sekwencji, po czym następuje błyskawiczna operacja redukcji All-Reduce o znikomym narzucie komunikacyjnym, która propaguje stany brzegowe pomiędzy węzłami. Pozwala to na skalowanie treningu na klastry liczące dziesiątki tysięcy układów scalonych bez degradacji współczynnika skalowalności liniowej (Scaling Efficiency).

Zastosowania w biologii obliczeniowej i modelowaniu genomu

Jednym z najbardziej spektakularnych obszarów, w których modele przestrzeni stanów udowodniły swoją bezdyskusyjną wyższość nad transformerami, jest biologia obliczeniowa i bioinformatyka. Kod genetyczny zapisany w kwasach nukleinowych DNA i RNA nie jest podzielony na zwarte akapity ani zdania. Informacja biologiczna – regulacja ekspresji genów, interakcje chromatyny czy struktury przestrzenne białek – rozciąga się na miliony par zasad w sposób ciągły i nieliniowy.

Próby zastosowania klasycznych transformerów do modelowania całych genomów kończyły się fiaskiem ze względu na konieczność drastycznego obcinania sekwencji lub stosowania prymitywnego próbkowania. Modele oparte na architekturze SSM, takie jak HyenaDNA czy Caduceus, zdołały przetworzyć całe chromosomy ludzkie liczące setki milionów par zasad w pojedynczym oknie kontekstowym z rozdzielczością na poziomie pojedynczego nukleotydu. Zdolność modeli SSM do wykrywania długodystansowych korelacji genetycznych bez utraty szczegółowości na poziomie pojedynczych mutacji punktowych otwiera bezprecedensowe możliwości w projektowaniu spersonalizowanych terapii genowych, odkrywaniu leków onkologicznych oraz inżynierii enzymów syntetycznych.

Podobny przełom obserwujemy w przetwarzaniu surowych sygnałów fizycznych: sejsmicznych fal sprężystych, wielokanałowych zapisów elektroencefalografii (EEG) oraz danych telemetrycznych z silników odrzutowych i turbin wiatrowych. Wszędzie tam, gdzie dane napływają w postaci ciągłych, wielomilowych szeregów czasowych o wysokiej częstotliwości próbkowania, modele przestrzeni stanów deklasują dotychczasowe algorytmy, łącząc precyzję modeli matematycznych fizyki newtonowskiej z plastycznością głębokich sieci neuronowych.

Podsumowanie: Nowy horyzont architektury głębokiego uczenia

Ewolucja od klasycznych transformerów, przez modele S4 i Mamba-1, aż po sformułowanie teorii State Space Duality w Mamba-2 dowodzi, że rozwój sztucznej inteligencji nie polega wyłącznie na bezrefleksyjnym zwiększaniu skali obliczeniowej w myśl praw skalowania (Scaling Laws). Fundamentalne innowacje algorytmiczne i matematyczne potrafią przełamać bariery fizyczne, które jeszcze niedawno wydawały się nie do pokonania.

Złamanie kwadratowej złożoności mechanizmu samoatencji otwiera zupełnie nowy rozdział w historii uczenia maszynowego. Modele nowej ery, operujące z liniową złożonością obliczeniową i stałym kosztem pamięciowym, pozwolą na płynne przetwarzanie wielomiesięcznych strumieni interakcji, bezstratną analizę całych bibliotek wiedzy ludzkiej oraz konstrukcję prawdziwie autonomicznych agentów kognitywnych zdolnych do ciągłego uczenia się w czasie rzeczywistym bez ograniczeń sztucznie narzuconych okien kontekstowych.