Post-transformerowe architektury sekwencyjne Mamba i SSM: Jak liniowa złożoność O(N) i selektywne bramkowanie eliminują kwadratowe wąskie gardło atencji Transformersów w kontekście miliona tokenów

Przez niemal dekadę od publikacji historycznego artykułu badaczy z zespołu Google pod tytułem Attention Is All You Need w 2017 roku, architektura Transformer stanowiła absolutny i niemal bezdyskusyjny monolit w świecie głębokiego uczenia. Od modeli językowych serii GPT i Claude, przez wizyjne modele ViT, aż po wielomodalne systemy generowania wideo, mechanizm samo-atencji Self-Attention oparty na iloczynie skalarnym Softmax(QK^T) stał się uniwersalnym silnikiem ekstrakcji korelacji w danych sekwencyjnych. Sukces ten opierał się na unikalnej cesze: w odróżnieniu od klasycznych rekurencyjnych sieci neuronowych RNN i pamięci LSTM, transformery umożliwiały pełną paralelizację treningu na tysiącach procesorów graficznych GPU, eliminując ograniczenia sekwencyjnego przepływu gradientów.
Jednak ten spektakularny triumf okupiony został fundamentalnym kosztem algorytmicznym, znanym jako przekleństwo kwadratowej złożoności obliczeniowej O(N^2). W klasycznym mechanizmie Multi-Head Attention każdy token w sekwencji wejściowej musi obliczyć iloczyn skalarny z każdym innym tokenem. O ile dla krótkich sekwencji rzędu dwóch czy czterech tysięcy tokenów narzut ten jest w pełni akceptowalny, o tyle dla kontekstów sięgających miliona tokenów (niezbędnych do analizy całych repozytoriów kodu, wielogodzinnych materiałów wideo w wysokiej rozdzielczości czy kompletnych genomów) liczba operacji zmiennoprzecinkowych FLOPs rośnie do astronomicznych wartości bilionów operacji, a pamięć podręczna kluczy i wartości KV Cache pożera setki gigabajtów pamięci HBM akceleratorów.
Z tego względu społeczność badawcza sztucznej inteligencji od lat poszukiwała 'świętego Graala' modelowania sekwencji: architektury łączącej pełną paralelizację treningu znaną z transformerów ze stałym kosztem inferencji O(1) na token i liniowym czasem przetwarzania O(N) charakterystycznym dla klasycznych automatów rekurencyjnych. Odpowiedzią na to wyzwanie stała się teoria modeli przestrzeni stanów State Space Models (SSM), wywodząca się z klasycznej teorii sterowania i fizyki matematycznej, której najbardziej dojrzałą i przełomową realizacją stała się rodzina architektur Mamba oraz Mamba-2, opracowana przez Alberta Gu i Tri Dao.
Przełom Mamby nie polega na prostym uproszczeniu obliczeń, lecz na fundamentalnym przeprojektowaniu interakcji pomiędzy matematyką równań różniczkowych a fizyczną hierarchią pamięci nowoczesnych akceleratorów graficznych GPU. Zastępując statyczne macierze transformacji dynamicznym mechanizmem selektywnym zależnym od wejścia Input-Dependent Selective Scan oraz wdrażając fuzję jąder obliczeniowych w szybkiej pamięci podręcznej SRAM, Mamba udowodniła, że liniowe modele sekwencyjne potrafią dorównać lub przewyższyć transformery w najbardziej wymagających zadaniach wnioskowania logicznego przy stukrotnie niższym zużyciu zasobów obliczeniowych.

Matematyczne fundamenty ciągłych przestrzeni stanów i dyskretyzacja

Aby zrozumieć rewolucję Mamby, należy cofnąć się do fundamentalnych równań liniowych przestrzeni stanów Linear State Space Models, opisujących ewolucję układu dynamicznego w czasie ciągłym. Ciągły model SSM mapuje jednowymiarowy sygnał wejściowy u(t) na sygnał wyjściowy y(t) za pośrednictwem N-wymiarowego ukrytego wektora stanu h(t). Ewolucja układu jest zdefiniowana przez układ liniowych równań różniczkowych: pochodna stanu h'(t) = A h(t) + B u(t), natomiast wyjście y(t) = C h(t) + D u(t), gdzie macierz A reprezentuje wewnętrzną dynamikę układu, macierz B odpowiada za sterowanie wejściowe, a macierz C dokonuje projekcji stanu na obserwowalne wyjście.
Ponieważ komputery cyfrowe operują na sekwencjach dyskretnych tokenów x_0, x_1, ..., x_L, równania te muszą zostać poddane dyskretyzacji numerycznej za pomocą kroku czasowego Delta. W architekturach SSM stosuje się najczęściej dyskretyzację metodą biliniową Tustina lub dokładną regułę zerowego rzędu Zero-Order Hold (ZOH). Po przekształceniach ciągłe macierze A i B zostają zastąpione macierzami dyskretnymi: A_bar = exp(Delta * A) oraz B_bar = (Delta * A)^-1 * (exp(Delta * A) - I) * Delta * B. Zdyskretyzowany układ przyjmuje postać rekurencji: h_t = A_bar * h_{t-1} + B_bar * x_t, a wyjście y_t = C * h_t.
Niezwykłą własnością liniowych modeli SSM jest ich podwójna natura: mogą być one traktowane zarówno jako liniowa rekurencja RNN podczas generowania kolejnych tokenów (gdzie stan h_t jest aktualizowany w stałym czasie O(1) bez konieczności buforowania całej historii), jak i równoważny globalny splot całkowy Global Convolution podczas treningu. Rozwijając rekurencję w czasie, wyjście y można zapisać jako splot wektora wejściowego x z jądrem splotowym K = (C*B_bar, C*A_bar*B_bar, C*A_bar^2*B_bar, ...). Ponieważ splot można obliczyć za pomocą szybkiej transformaty Fouriera FFT w czasie O(L log L), wcześniejsze modele takie jak S4 (Structured State Spaces) oferowały błyskawiczny trening na całych sekwencjach.

Jednak klasyczne modele S4 cierpiały na śmiertelną wadę: były układami liniowymi niezmienniczymi w czasie Linear Time-Invariant (LTI). Macierze A, B, C oraz krok dyskretyzacji Delta były stałe dla wszystkich tokenów w sekwencji, co uniemożliwiało modelowi selektywne decydowanie, które informacje należy zapamiętać, a które bezpowrotnie odrzucić. W przeciwieństwie do transformerów, modele LTI nie potrafiły zrealizować prostego zadania kopiowania czy asocjacyjnego wyszukiwania klucz-wartość, co skazywało je na rolę ciekawostki matematycznej.

Mechanizm selektywności: Jak dynamiczne bramkowanie przełamuje barierę LTI

Prawdziwym przełomem Mamby było zburzenie dogmatu niezmienniczości w czasie i wprowadzenie mechanizmu selektywnego Selective SSM. Autorzy zauważyli, że fundamentalna przewaga atencji nad rekurencją wynika z faktu, że wagi atencji są dynamicznie wyliczane na podstawie bieżącego tokenu wejściowego. W Mambie parametry B, C oraz krok dyskretyzacji Delta stają się jawnymi funkcjami bieżącego wektora wejściowego x_t, obliczanymi za pomocą prostych projekcji liniowych: B_t = Linear_B(x_t), C_t = Linear_C(x_t) oraz Delta_t = Softplus(Parameter + Linear_Delta(x_t)).
Fizyczna interpretacja kroku czasowego Delta_t jest kluczem do zrozumienia inteligencji Mamby. W teorii sterowania Delta_t reguluje stopień, w jakim układ dynamiczny skupia się na bieżącym sygnale wejściowym u(t), a w jakim stopniu podtrzymuje dotychczasowy stan h(t). Gdy model napotka słowo pozbawione znaczenia semantycznego (np. spójnik 'i' lub kropkę), sieć wyznacza wartość Delta_t bliską zera. Wówczas macierz A_bar zbliża się do macierzy jednostkowej, a B_bar dąży do zera – model całkowicie ignoruje bieżący token i bez strat przenosi dotychczasowy stan h_{t-1}.
I odwrotnie: gdy model przetwarza krytyczną informację (np. nazwisko bohatera czy klucz API w kodzie), wyznacza dużą wartość Delta_t. Wtedy dotychczasowy stan h zostaje zresetowany, a wektor wejściowy x_t zostaje głęboko zapisany w pamięci wewnętrznej. Dzięki temu mechanizmowi Mamba potrafi precyzyjnie selekcjonować informacje, filtrując szum i kompresując kontekst o dowolnej długości w skończonym wektorze stanu ukrytego.
Niestety, uczynienie macierzy B, C i Delta zależnymi od czasu t miało katastrofalny skutek matematyczny: model przestał być układem LTI, co całkowicie uniemożliwiło wykorzystanie globalnego splotu i transformaty Fouriera podczas treningu! Model powrócił do formy czysto sekwencyjnej, która na tradycyjnym oprogramowaniu głębokiego uczenia wykonywałaby się setki razy wolniej niż Transformer. Aby pokonać tę barierę, konieczne było zejście na poziom fizyki półprzewodników i mikroarchitektury GPU.

Sprzętowa świadomość: Algorytm równoległego skanowania w pamięci SRAM

Współczesne procesory graficzne, takie jak NVIDIA H100 czy B200, nie są ograniczone surową mocą obliczeniową rdzeni tensorowych (która sięga petaflopów), lecz przepustowością magistrali pamięci Memory Bandwidth pomiędzy główną pamięcią wysokoprzepustową HBM a ultraszybką pamięcią podręczną SRAM umieszczoną bezpośrednio na krzemie procesora. W klasycznych implementacjach rekurencyjnych stan ukryty h o rozmiarze (Batch, Długość, Wymiar, Stan) musiał być po każdym kroku czasowym zapisywany z SRAM do HBM i ponownie odczytywany w kolejnym kroku, co powodowało zjawisko Memory-Bound i drastyczny spadek wykorzystania procesora do poziomu poniżej dziesięciu procent.
Tri Dao, współtwórca rewolucyjnego algorytmu FlashAttention, zastosował w Mambie paradygmat sprzętowej optymalizacji Hardware-Aware Design. Zamiast sekwencyjnego wykonywania pętli for w pamięci HBM, stan ukryty h nigdy nie opuszcza szybkiej pamięci SRAM procesora podczas przetwarzania bloku sekwencji! Parametry B_t, C_t oraz Delta_t są ładowane do SRAM, gdzie dyskretyzacja oraz operacja rekurencyjna są wykonywane w jednym połączonym jądrze obliczeniowym Fused Kernel.
Kluczowym elementem algorytmicznym stało się wdrożenie równoległego skanowania asocjacyjnego Parallel Associative Scan (znanego jako algorytm Blellocha). Ponieważ operacja rekurencyjna h_t = A_t * h_{t-1} + B_t * x_t jest ściśle asocjacyjna: (x * y) * z = x * (y * z), proces jej obliczania można przedstawić w postaci zbalansowanego drzewa binarnego. Zamiast L kroków sekwencyjnych, GPU przetwarza sekwencję w zaledwie 2 * log_2(L) krokach równoległych w hierarchii wątków Warp i Threadblock.
Dzięki fuzji operacji w pamięci SRAM i algorytmowi Blellocha, Mamba osiąga ponad siedemdziesiąt procent teoretycznej wydajności sprzętowej procesorów graficznych, trenując się szybciej niż zoptymalizowany FlashAttention-2 dla dowolnych długości sekwencji powyżej dwóch tysięcy tokenów, przy jednoczesnym zerowym zużyciu pamięci na stan pośredni w fazie inferencji.

Mamba-2 i zunifikowana teoria SSD: Most łączący SSM i Attention

W 2024 roku ukazała się kolejna odsłona architektury – Mamba-2, która wniosła fundamentalny wkład teoretyczny w postaci strukturalnej dualności przestrzeni stanów Structured State Space Duality (SSD). Autorzy wykazali matematyczny dowód na to, że selektywne modele SSM oraz mechanizm atencji ze strukturą liniową Masked Attention są w istocie dwoma krańcami tego samego spektrum matematycznego, różniącymi się jedynie strukturą macierzy transformacji.
Mamba-2 przekształca obliczenia selektywnego skanowania w operacje mnożenia macierzy Matrix Multiplication (MatMul). Rdzenie tensorowe Tensor Cores w układach GPU są sprzętowo wyspecjalizowane w błyskawicznym mnożeniu gęstych macierzy, podczas gdy operacje wektorowe skanowania nie wykorzystywały pełnego potencjału jednostek systolitycznych. W Mambie-2 sekwencja jest dzielona na bloki o długości Q (np. 64 tokeny). Wewnątrz każdego bloku obliczenia są realizowane za pomocą czystego mnożenia macierzowego z wykorzystaniem Tensor Cores, a pomiędzy blokami stan jest propagowany za pomocą klasycznego skanowania rekurencyjnego.
Przejście na paradygmat SSD pozwoliło Mambie-2 na ośmiokrotne przyspieszenie treningu w porównaniu z pierwszą wersją Mamby, osiągając przepustowości rzędu kilkuset tysięcy tokenów na sekundę na pojedynczej karcie graficznej. Co więcej, Mamba-2 umożliwiła bezproblemową integrację technik wielogłowicowych Multi-Head SSM na wzór Multi-Head Attention, co radykalnie poprawiło zdolność modelu do jednoczesnego śledzenia wielu wątków i faktów w ultra-długim tekście.

Hybrydowe architektury Jamba i Samba: Fuzja najlepszych cech obu światów

Choć czyste modele SSM deklasują transformery w zadaniach przetwarzania długiego kontekstu i oszczędności pamięci, w klasycznych zadaniach wymagających precyzyjnego przywoływania faktów w krótkim oknie kontekstowym (tzw. in-context associative recall) pełny mechanizm atencji Softmax wciąż wykazuje subtelną przewagę. Wynika to z faktu, że atencja posiada nieograniczoną pojemność pamięci roboczej (kosztem rosnącej pamięci KV Cache), podczas gdy SSM kompresuje całą przeszłość do wektora o stałym wymiarze.
W odpowiedzi na to zjawisko czołowe laboratoria badawcze, w tym AI21 Labs z modelem Jamba oraz naukowcy z Microsoftu i Stanfordu z modelem Samba, stworzyły architektury hybrydowe. Jamba łączy warstwy Mamby z warstwami Transformer Attention oraz mechanizmem rzadkiej aktywacji Mixture-of-Experts (MoE) w stosunku jedna warstwa atencji na każde siedem warstw Mamby. Taka konfiguracja pozwala na redukcję rozmiaru pamięci KV Cache o ponad osiemdziesiąt pięć procent w porównaniu ze standardowym modelem Llama-3, przy jednoczesnym zachowaniu bezbłędnej precyzji w testach Needle In A Haystack na kontekstach sięgających dwustu pięćdziesięciu sześciu tysięcy tokenów.
Podobne sukcesy odnotowuje architektura RecurrentGemma opracowana przez Google DeepMind, wykorzystująca warstwy Griffin – hybrydę bramkowanej rekurencji liniowej Real-Gated Linear Recurrence z lokalną samo-atencją Sliding Window Attention. Modele te udowadniają, że przyszłość architektur wielkich modeli nie polega na ślepym trzymaniu się czystego transformera, lecz na inteligentnym łączeniu warstw o liniowej złożoności z rzadko rozmieszczonymi punktami globalnej koncentracji uwagi.

Wnioskowanie w czasie rzeczywistym i ekonomika centrów danych AI

Wdrożenie architektur Mamba i SSM wywołuje sejsmiczne zmiany w ekonomice chmurowych centrów danych AI. W tradycyjnym modelu opartym na transformerach koszt inferencji rośnie lawinowo wraz z długością sesji dialogowej użytkownika. Bufor KV Cache dla każdego aktywnego użytkownika musi być stale przetrzymywany w ultradrogiej pamięci HBM, co drastycznie ogranicza maksymalny rozmiar partii Batch Size i zmusza dostawców do instalowania klastrów o wielomilionowej wartości wyłącznie na potrzeby obsługi ruchu sieciowego.
W architekturze Mamba stan każdego użytkownika jest reprezentowany przez zwarty wektor o stałym rozmiarze kilkunastu kilobajtów, niezależnie od tego, czy rozmowa trwa sto tokenów, czy sto tysięcy tokenów. W efekcie pojedynczy serwer GPU potrafi obsłużyć dziesięciokrotnie większą liczbę równoległych strumieni zapytań Concurrency przy stałym opóźnieniu generowania Time-to-First-Token (TTFT) oraz stałej przepustowości tokenów na sekundę. Dla przedsiębiorstw wdrażających autonomicznych agentów AI pracujących w pętli 24/7 oznacza to redukcję kosztów chmurowych o ponad siedemdziesiąt procent.
Podsumowując, rewolucja architektur Mamba i modeli przestrzeni stanów SSM to najważniejszy przełom w algorytmice głębokiego uczenia od czasu narodzin samego Transformera. Przełamując barierę kwadratowej złożoności bez utraty ekspresywności reprezentacji, Mamba otwiera przed sztuczną inteligencją horyzonty przetwarzania nieskończonych sekwencji – od analizy całego dorobku literackiego cywilizacji w jednym oknie inferencyjnym, po ciągłe, bezbłędne sterowanie robotami humanoidalnymi w czasie rzeczywistym.
Warto szczegółowo przeanalizować matematyczną strukturę macierzy A w sformułowaniu HiPPO (High-order Polynomial Projection Operators), które stanowi teoretyczny rdzeń inicjalizacji wag w modelach SSM. Albert Gu udowodnił, że aby ciągły stan ukryty h(t) mógł optymalnie rekonstruować historię wejścia u(t) ze skończoną pamięcią, rzutowanie musi odbywać się na bazę ortogonalnych wielomianów Legendre'a. Macierz HiPPO-LegS posiada ściśle zdefiniowane elementy analityczne, w których wartości diagonalne i poddiagonalne wymuszają geometryczny rozkład czasów zaniku pamięci. Dzięki temu model nie ulega zjawisku zanikającego gradientu Vanishing Gradient, zachowując idealną pamięć o tokenach przetworzonych dziesiątki tysięcy kroków wcześniej bez konieczności stosowania bramek zapominania znanych z architektur LSTM.
Na poziomie implementacji kompilatorowej kluczowym narzędziem dla modeli Mamba stał się język OpenAI Triton oraz autorskie jądra w CUDA C++. W odróżnieniu od klasycznych operacji PyTorch, które generują narzut wywołań kerneli PyTorch Dispatcher, jądra Triton w Mambie dokonują fuzji operacji nieliniowych aktywacji SiLU (Swish), dyskretyzacji ZOH, skanowania prefiksowego oraz projekcji liniowej w pojedynczym bloku pamięci współdzielonej Shared Memory na procesorze strumieniowym SM. Pozwala to na uniknięcie jakichkolwiek pośrednich transferów do pamięci globalnej DRAM, co przekłada się na osiągnięcie niemal stuprocentowej wydajności arytmetycznej jednostek Compute Units.
Wielkim polem badawczym pozostaje także kwantyzacja wag i aktywacji modeli SSM do formatów niskoprecyzyjnych FP8, FP4 oraz INT4. O ile transformery wykazują silną obecność tak zwanych wartości odstających Outlier Features w kanałach aktywacji mechanizmu Softmax (co drastycznie utrudnia ich kwantyzację poniżej 8 bitów bez zaawansowanych technik SmoothQuant), o tyle liniowe przestrzenie stanów Mamby charakteryzują się znacznie gładszym, jednorodnym rozkładem statystycznym stanu ukrytego h. Umożliwia to agresywną kwantyzację wag macierzy A, B i C do formatu 4-bitowego bez zauważalnej degradacji perplexity, co toruje drogę do uruchamiania modeli liczących dziesiątki miliardów parametrów bezpośrednio na urządzeniach krawędziowych Edge AI, smartfonach i pokładowych komputerach dronów.
Nie można pominąć także zastosowania modeli Mamba w bioinformatyce i genomice obliczeniowej. Badania nad modelem Caduceus – dwukierunkowym modelem Mamba wytrenowanym na kompletnym genomie człowieka – wykazały spektakularną przewagę nad transformerami w przewidywaniu wpływu mutacji niekodujących w sekwencjach DNA liczących setki tysięcy par zasad. Zdolność Mamby do jednoczesnego modelowania obu nici helisy DNA z zachowaniem symetrii rewersyjno-komplementarnej Reverse-Complement Equivariance stanowi dowód na uniwersalność paradygmatu przestrzeni stanów w naukach przyrodniczych, daleko wykraczającą poza przetwarzanie ludzkiego języka naturalnego.
Kolejnym triumfem paradygmatu SSM stała się domena wizji komputerowej i przetwarzania wideo, gdzie modele Vision Mamba (Vim) oraz VMamba zastępują klasyczne transformery wizyjne ViT. W modelach ViT podział obrazu 4K lub 8K na mikroskopijne patche generuje dziesiątki tysięcy tokenów wizyjnych, co paraliżuje mechanizm samo-atencji. Architektura VMamba wprowadza dwukierunkowy mechanizm skanowania Visual State Space (VSS), który skanuje piksele w czterech kierunkach (horyzontalnie, wertykalnie i po przekątnych), zachowując dwuwymiarową spójność przestrzenną. Dzięki liniowej złożoności O(N) modele SSM potrafią przetwarzać niekompresowane strumienie wideo w czasie rzeczywistym z częstotliwością sześćdziesięciu klatek na sekundę na pojedynczym układzie wbudowanym, co stanowi rewolucję w autonomii dronów i robotów mobilnych.
W dalszej perspektywie badawczej wyłania się fuzja przestrzeni stanów z architekturami neuromorficznymi Spiking Neural Networks (SNN). Ciągłe równania różniczkowe Mamby wykazują bezpośrednią izomorficzność z dynamiką potencjału błonowego biologicznych neuronów typu Leaky Integrate-and-Fire (LIF). Implementacja modeli SSM na asynchronicznych procesorach neuromorficznych (takich jak Intel Loihi czy SynSense) pozwala na realizację wnioskowania sekwencyjnego przy poborze mocy rzędu pojedynczych miliwatów, otwierając drogę do stworzenia w pełni autonomicznych, miniaturowych systemów sztucznej inteligencji zdolnych do nieprzerwanej pracy operacyjnej przez lata na zasilaniu bateryjnym.
Podsumowując, mariaż rygorystycznej teorii układów dynamicznych z bezwzględną optymalizacją sprzętową wyznaczył nowy standard dla inżynierii sztucznej inteligencji. Przeniesienie ciężaru obliczeniowego z kwadratowej eksplozji pamięciowej na zoptymalizowaną rekurencję w pamięci SRAM to dowód na to, że prawdziwy postęp technologiczny rodzi się na styku algorytmiki i fizyki krzemu. Architektury Mamba, Mamba-2 oraz ich hybrydowe warianty bezdyskusyjnie otwierają nową epokę skalowania modeli sekwencyjnych, w której horyzont kontekstu przestaje być barierą inżynierską.