Architektura Mamba-2 i hybrydowe modele State Space Models SSM: Jak State Space Duality SSD i liniowa złożoność eliminują kwadratowe wąskie gardło transformatorów przy przetwarzaniu milionów tokenów

Przez niemal dekadę, od momentu publikacji przełomowej pracy Attention Is All You Need przez zespół Google w 2017 roku, architektura transformatora dzieliła i rządziła światem głębokiego uczenia maszynowego. Wszystkie flagowe modele językowe, od serii GPT przez Claude aż po rodzinę Llama, zostały zbudowane wokół mechanizmu samoobsługi uwagi Softmax Attention. Jednak w miarę jak apetyt przemysłu sztucznej inteligencji przesunął się ku przetwarzaniu gigantycznych okien kontekstowych liczących setki tysięcy, a nawet miliony tokenów, transformatory zderzyły się z nieprzekraczalną ścianą czystej złożoności obliczeniowej. Matematyczna natura iloczynu skalarnego zapytań i kluczy narzuca kwadratową złożoność czasową i pamięciową O N kwadrat, sprawiając, że podwojenie długości przetwarzanego tekstu wymaga czterokrotnego wzrostu mocy obliczeniowej FLOPs oraz gigantycznej alokacji bufora pamięci podręcznej KV-Cache w pamięci HBM akceleratorów graficznych.

Ten fundamentalny paraliż skalowania sprawił, że badacze od lat poszukiwali alternatywnych modeli sekwencyjnych o liniowej złożoności O N, które potrafiłyby utrzymać zdolność długodystansowego rozumowania bez marnotrawstwa pamięciowego. Przełomem stała się koncepcja modeli przestrzeni stanów State Space Models SSM, wywodząca się z klasycznej teorii sterowania układami dynamicznymi. Po sukcesie architektury Mamba-1, która udowodniła, że selektywne modele SSM mogą dorównać transformatorom na standardowych benchmarkach, nadeszła prawdziwa rewolucja teoretyczna i sprzętowa w postaci Mamba-2 oraz sformułowania paradygmatu dwoistości przestrzeni stanów State Space Duality SSD przez Alberta Gu i Tri Dao. Odkrycie to położyło kres wieloletniej wojnie pomiędzy zwolennikami uwagi a entuzjastami rekurencji, dowodząc, że oba te światy są w istocie dwiema stronami tej samej monety matematycznej.

Kwadratowa klątwa mechanizmu Attention i kryzys bufora KV-Cache

Aby zrozumieć wagę rewolucji Mamba-2, należy precyzyjnie zdiagnozować przyczyny niewydolności klasycznego mechanizmu uwagi. W warstwie Multi-Head Attention każdy token w sekwencji wejściowej musi wejść w interakcję ze wszystkimi poprzedzającymi go tokenami za pośrednictwem iloczynu macierzy Query i Key transposed. Wynikowa macierz uwagi posiada wymiary N na N, gdzie N oznacza długość sekwencji. Dla kontekstu liczącego cztery tysiące tokenów narzut ten jest w pełni akceptowalny. Kiedy jednak sekwencja wzrasta do miliona tokenów, macierz uwagi zawiera bilion elementów zmiennoprzecinkowych, co wymaga terabajtów przepustowości pamięciowej dla pojedynczej warstwy sieci neuronowej.

Jeszcze bardziej destrukcyjny charakter ma ten problem podczas fazy autoregresywnego generowania tekstu token po tokenie. Aby uniknąć ponownego przeliczania reprezentacji dla całej historii konwersacji przy każdym nowym słowie, serwery produkcyjne utrzymują w pamięci VRAM tak zwany bufor kluczy i wartości KV-Cache. Rozmiar tego bufora rośnie ściśle liniowo z każdym wygenerowanym tokenem dla każdego użytkownika z osobna. Przy oknach rzędu stu tysięcy tokenów pojedyncza sesja potrafi zająć kilkadziesiąt gigabajtów pamięci HBM3e na akceleratorach Nvidia H100, drastycznie ograniczając współbieżność zapytań Batch Concurrency i windując koszt infrastruktury chmurowej do poziomów nieakceptowalnych biznesowo.

Od ciągłych równań różniczkowych do selektywnej Mamby-1

Geneza modeli State Space Models wywodzi się z fizyki i inżynierii automatyki, gdzie stan układu ciągłego w czasie opisuje się parą równań różniczkowych: pochodna stanu ukrytego h od t równa się macierz A razy h od t plus macierz B razy wektor wejściowy x od t, podczas gdy sygnał wyjściowy y od t stanowi iloczyn macierzy C razy h od t plus macierz D razy x od t. W pionierskich pracach nad modelem S4 Structured State Spaces Albert Gu wykazał, że odpowiednia dyskretyzacja tych równań za pomocą metody biliniowej Tustina lub ekspotencjalnej transformacji macierzowej z użyciem specjalnych struktur macierzy HiPPO High-order Polynomial Projection Operators pozwala na uchwycenie zależności czasowych o długości setek tysięcy kroków.

Fundamentalnym ograniczeniem pierwszych modeli SSM była jednak ich niezmienniczość w czasie Linear Time-Invariance LTI. Macierze przejścia A, B i C były stałymi wagami sieci, niezależnymi od napływających danych, co uniemożliwiało modelowi selektywne filtrowanie informacji: sieć pamiętała wszystko z jednakową intensywnością, nie potrafiąc zapomnieć szumu ani skupić się na rzadkich faktach. Przełomem architektury Mamba-1 stało się wprowadzenie selektywności: parametry B, C oraz krok dyskretyzacji delta przestały być stałymi wagami, stając się funkcjami dynamicznie wyliczanymi z bieżącego tokenu x. Pozwoliło to modelowi na selektywne resetowanie stanu pamięci lub absorbowanie kluczowych informacji w zależności od kontekstu semantycznego.

Paradoks sprzętowy Mamby-1 i koncepcja State Space Duality SSD

Mimo bezdyskusyjnych zalet teoretycznych Mamba-1 napotkała na potężną barierę praktyczną związaną z architekturą nowoczesnych procesorów graficznych. Współczesne akceleratory, takie jak Nvidia A100 i H100, nie są zoptymalizowane pod kątem operacji wektorowych ani skanowania sekwencyjnego, lecz pod kątem gęstego mnożenia macierzy za pomocą dedykowanych rdzeni Tensor Cores. Architektura Mamba-1 opierała się na sekwencyjnym skanowaniu asocjacyjnym Selective Scan, które wymagało pisania skomplikowanych, niskopoziomowych kerneli w języku CUDA i Triton w celu przenoszenia stanów przez pamięć podręczną SRAM.

W rezultacie Mamba-1 osiągała zaledwie ułamek teoretycznej wydajności sprzętowej FLOP Utilization akceleratorów graficznych w porównaniu z mechanizmem Attention, który dzięki bibliotekom FlashAttention idealnie wysyca jednostki GEMM General Matrix Multiply. Autorzy Mamby stanęli przed dramatycznym pytaniem: czy liniowa złożoność musi oznaczać powolne wykonanie na krzemie? Odpowiedzią stała się koncepcja State Space Duality SSD, sformułowana w architekturze Mamba-2.

W ramach SSD udowodniono fundamentalne twierdzenie matematyczne: selektywny model przestrzeni stanów oraz mechanizm uwagi z maską przyczynową są dwoma wariantami tego samego przekształcenia liniowego na macierzach pół-separowalnych Semiseparable Matrices rzędu pierwszego. Jeśli macierz transformacji dynamicznej A w modelu SSM zostanie ograniczona do postaci skalarnej pomnożonej przez macierz jednostkową dla każdego kanału głowicy, całe równanie rekurencyjne można przepisać jako iloczyn macierzy wejściowej z dolnotrójkątną macierzą maskującą o strukturze pół-separowalnej.

Od naiwnej liniowej uwagi do RWKV i RetNet: Dlaczego wcześniejsze próby zawodziły

Próby przełamania kwadratowej złożoności transformatora mają długą historię i obfitowały w liczne ślepe zaułki inżynieryjne. Pierwszą falę stanowiły modele z rodziny Linear Attention, zaproponowane między innymi przez Angelosa Katharopoulosa w 2020 roku. W modelach tych standardową nieliniową funkcję Softmax zastępowano iloczynem funkcji jądrowych phi od Q i phi od K, co pozwalało na zmianę kolejności mnożenia macierzy drogą łączności iloczynu i uzyskanie liniowej złożoności czasowej. Jednak naiwne modele liniowej uwagi cierpiały na dramatyczną degradację jakości: pozbawienie mechanizmu uwagi ostrej selektywności Softmaxu sprawiało, że model rozmywał reprezentacje semantyczne, nie potrafiąc skupić się na rzadkich, lecz kluczowych tokenach.

Kolejnym etapem ewolucji były architektury nawiązujące do sieci rekurencyjnych, takie jak RWKV Receptance Weighted Key Value oraz RetNet Retentive Network opracowany przez Microsoft Research. Modele te wprowadziły jawny mechanizm zaniku wykładniczego pamięci, symulując zachowanie klasycznego transformatora w trybie równoległym podczas treningu oraz zachowanie komórki RNN w trybie rekurencyjnym podczas wnioskowania. Choć RWKV i RetNet stanowiły potężny krok naprzód, brakowało im formalnego aparatu matematycznego gwarantującego stabilność numeryczną przy ekstremalnie długich kontekstach oraz możliwości bezpośredniej optymalizacji pod kątem wyspecjalizowanych koprocesorów macierzowych. Mamba-2 wypełniła tę lukę, dostarczając ścisły dowód równoważności z teorią układów dynamicznych i otwierając drogę do idealnej synchronizacji algorytmu z fizyczną strukturą krzemu.

Matematyka macierzy pół-separowalnych i dekompozycja blokowa SSD

Pół-separowalność oznacza, że każdy blok podprzekątniowy macierzy transformacji posiada rząd co najwyżej równy rzędowi przestrzeni stanów. W klasycznym mechanizmie Attention macierz relacji reprezentowana jest przez gęstą strukturę QK transposed, która może posiadać dowolnie wysoki rząd matematyczny. W architekturze Mamba-2 macierz ta zostaje zastąpiona strukturą o ściśle kontrolowanej strukturze pół-separowalnej, w której każdy element jest ważony przez iloczyn skumulowanych wag dyskretyzacji delta. Kluczowym uproszczeniem wprowadzonym w Mambie-2 jest zwężenie struktury macierzy dynamicznej A do postaci ściśle skalarnej dla każdego kanału głowicy. Podczas gdy w Mambie-1 macierz A była pełną macierzą diagonalną, co wymagało skomplikowanego śledzenia niezależnych stanów dla każdego wymiaru cech, postać skalarna redukuje liczbę parametrów ewolucji do pojedynczej wartości liczbowej. Posunięcie to radykalnie upraszcza jądro obliczeniowe z kwadratowego narzutu pamięciowego do operacji czysto liniowej, nie powodując przy tym jakiejkolwiek degradacji zdolności aproksymacyjnych sieci, co potwierdziły rygorystyczne dowody zbieżności oparte na teorii reprezentacji operatorów całkowych.

Ta elegancka tożsamość algebraiczna pozwala na zastosowanie genialnego algorytmu dekompozycji blokowej Block Decomposition. Cała sekwencja wejściowa o długości N zostaje podzielona na równe bloki o rozmiarze Q, na przykład sześćdziesięciu czterech lub stu dwudziestu ośmiu tokenów. Obliczenia wewnątrz każdego bloku są wykonywane w sposób całkowicie równoległy przy użyciu standardowego mnożenia macierzy GEMM na rdzeniach Tensor Cores, dokładnie tak jak w miniaturowym mechanizmie uwagi, co gwarantuje maksymalną utylizację sprzętową GPU rzędu osiemdziesięciu procent mocy szczytowej.

Z kolei transfer informacji pomiędzy poszczególnymi blokami odbywa się za pośrednictwem ultrakompaktowego stanu rekurencyjnego o wymiarach odpowiadających przestrzeni stanów. Zamiast przekazywać całą historię tokenów, model syntetyzuje skumulowany stan bloku w jeden wektor stanu ukrytego i przekazuje go do kolejnego bloku za pomocą szybkiego skanowania sekwencyjnego o mikroskopijnym narzucie. Połączenie lokalnego mnożenia macierzy z globalnym skanowaniem liniowym daje model o idealnie liniowej złożoności obliczeniowej O N, który na poziomie sprzętowym wykonuje się od trzech do ośmiu razy szybciej niż Mamba-1 i FlashAttention-2 przy długich sekwencjach.

Architektura Mamba-2: Wielogłowicowe stany SSM i unifikacja z MLP

Wprowadzenie paradygmatu SSD pozwoliło na radykalne przeprojektowanie całej warstwy sieciowej w modelu Mamba-2. Podczas gdy Mamba-1 stosowała pojedynczą, wielką projekcję stanów, Mamba-2 wprowadza strukturę wielogłowicową Multi-Head SSM Pattern, analogiczną do Multi-Head Attention w transformatorach. Pojedynczy blok Mamba-2 dzieli przestrzeń cech na dziesiątki niezależnych głowic, z których każda posiada własną skalarną macierz ewolucji stanu A oraz dedykowane parametry B i C.

Taka organizacja pozwala na równoległe modelowanie różnych skal czasowych i relacji semantycznych w tekście: część głowic specjalizuje się w wychwytywaniu lokalnych zależności gramatycznych z gwałtownie wygasającym stanem pamięci, podczas gdy inne głowice utrzymują niemal bezstratną pamięć długodystansową, śledząc wątki narracyjne przez setki stron dokumentu. Dodatkowo Mamba-2 integruje blok nieliniowej transformacji cech MLP bezpośrednio do wnętrza komórki SSM, eliminując zbędne warstwy pośrednie i redukując liczbę operacji odczytu i zapisu do pamięci HBM.

Najbardziej spektakularną korzyścią inżynieryjną jest jednak eliminacja bufora KV-Cache podczas generowania tekstu. W fazie autoregresywnej Mamba-2 nie musi utrzymywać rosnącej tablicy wektorów kluczy i wartości. Cała pamięć przeszłości modelu jest skondensowana w stałym, niezmiennym w czasie wektorze stanu ukrytego o rozmiarze zaledwie kilkunastu kilobajtów na głowicę. Złożoność pamięciowa generowania pojedynczego tokenu wynosi absolutne O 1: koszt pamięciowy i czasowy wygenerowania milionowego tokenu jest dokładnie taki sam, jak wygenerowania tokenu drugiego, co pozwala na uruchamianie gigantycznych procesów wnioskowania na pojedynczych kartach graficznych klasy konsumenckiej.

Architektury hybrydowe: Jamba, Samba i Zamba łączą światy SSM i Attention

Pomimo geniuszu matematycznego czystych modeli SSM, transformatory zachowały jedną przewagę kognitywną: niezrównaną zdolność do bezwzględnego, punktowego przypominania informacji z odległego kontekstu, zwaną testem igły w stogu siana Needle-in-a-Haystack. Ponieważ mechanizm Softmax Attention posiada bezpośredni, nieograniczony dostęp do każdego wcześniejszego tokenu, potrafi bezbłędnie zlokalizować pojedynczą liczbę ukrytą w milionie słów. Modele czysto rekurencyjne, kompresujące całą historię do wektora o skończonym rozmiarze, mogą w skrajnych przypadkach doświadczać utraty mikroskopijnych detali faktograficznych.

Praktyczną odpowiedzią globalnego przemysłu AI stały się architektury hybrydowe, które łączą warstwy Mamba-2 z rzadko rozmieszczonymi warstwami Attention. Pionierem tego podejścia stał się model Jamba opracowany przez firmę AI21 Labs, a także otwarte architektury Samba i Zamba. W modelach tych standardowy blok powtarzalny składa się z sekwencji sześciu lub ośmiu warstw Mamba-2, po których następuje jedna warstwa klasycznego transformatora z mechanizmem uwagi oraz modułami Mixture of Experts MoE.

Taka symbioza pozwala na uzyskanie tego, co najlepsze z obu światów. Warstwy Mamba-2 przejmują na siebie dziewięćdziesiąt procent obciążenia obliczeniowego, przetwarzając długie sekwencje z liniową prędkością i filtrując szum informacyjny, podczas gdy sporadyczne warstwy Attention zapewniają stuprocentową precyzję wyszukiwania asocjacyjnego i precyzyjnego wnioskowania relacyjnego. Rozmiar bufora KV-Cache w modelu hybrydowym spada o ponad osiemdziesiąt pięć procent w porównaniu ze standardowym transformatorem, umożliwiając jednoczesną obsługę setek użytkowników przy oknach kontekstowych przekraczających dwieście pięćdziesiąt tysięcy tokenów bez konieczności rozbudowy infrastruktury serwerowej.

Inżynieria niskopoziomowa: Kernelle Triton i jednostki Tensor Memory Accelerator

Sukces wdrożeniowy architektury Mamba-2 na nowoczesnych klastrach obliczeniowych wynika bezpośrednio z bezprecedensowego dopasowania matematyki SSD do mikroarchitektury najnowszych procesorów graficznych Nvidia Hopper H100 oraz Blackwell B200. W klasycznym podejściu sekwencyjnym wąskim gardłem była nieustanna wymiana danych pomiędzy wolną pamięcią globalną HBM a ultraszybką, lecz ciasną pamięcią podręczną SRAM na poziomie pojedynczych klastrów przetwarzania tekstur TPC. Mamba-1 wymagała zaawansowanych technik ponownego przeliczania aktywacji Recomputation, aby zmieścić stany w pamięci podręcznej rejestrów.

Mamba-2 całkowicie rekonfiguruje ten przepływ danych za pomocą dedykowanych bibliotek w języku Triton. Dekompozycja blokowa pozwala na bezpośrednie wykorzystanie jednostek Tensor Memory Accelerator TMA obecnych w układach architektury Hopper. Jednostka TMA realizuje asynchroniczny transfer wielowymiarowych tensorów pomiędzy pamięcią HBM a pamięcią współdzieloną Shared Memory w sposób całkowicie sprzętowy, odciążając jednostki cieniujące ALU. Rdzenie Tensor Cores mogą dzięki temu nieprzerwanie wykonywać operacje mnożenia macierzy blokowych w precyzji FP8 lub FP16, podczas gdy w tle realizowany jest kolejny transfer danych. Efektywność wykorzystania mocy obliczeniowej GPU osiąga w Mambie-2 poziom przekraczający osiemdziesiąt procent wydajności szczytowej, co stawia ją na absolutnym szczycie inżynierii wydajnościowej obok najgęstszych implementacji GEMM.

Ewaluacja empiryczna: Przetwarzanie genomu, analiza kodu i prawa skalowania

Weryfikacja empiryczna modeli Mamba-2 przyniosła dowody na to, że architektury SSM podlegają dokładnie takim samym prawom skalowania Scaling Laws, jak klasyczne transformatory. Wykresy zależności błędu predykcji od liczby parametrów i wielkości korpusu treningowego wykazują idealnie prostoliniowy spadek straty w skali logarytmicznej, co dowodzi, że Mamba-2 nie posiada żadnych ukrytych barier pojemności kognitywnej.

Na standardowych zestawach testowych języka naturalnego, takich jak MMLU, GSM8k, ARC oraz HumanEval, modele oparte na Mambie-2 osiągają wyniki w pełni porównywalne z najnowocześniejszymi transformatorami o tej samej liczbie aktywnych parametrów, deklasując je jednak bezapelacyjnie w zadaniach wymagających przetwarzania ekstremalnie długich sekwencji. W benchmarku RULER, badającym zachowanie modeli na kontekstach do stu dwudziestu ośmiu tysięcy tokenów, Mamba-2 utrzymała niemal stuprocentową retencję informacji, podczas gdy tradycyjne transformatory wykazywały gwałtowny spadek precyzji wywołany degradacją pozycji obrotowych RoPE.

Prawdziwym triumfem architektury Mamba-2 stały się jednak zastosowania poza domeną tekstu naturalnego, w szczególności w bioinformatyce i genomice. Modele sekwencjonowania DNA, takie jak Caduceus, przetwarzają ciągłe łańcuchy nukleotydów o długości milionów par zasad. Dla klasycznego transformatora zadanie to było całkowicie niewykonalne ze względów pamięciowych, podczas gdy liniowa Mamba-2 przetwarza całe chromosomy w jednym przebiegu, modelując długozasięgowe oddziaływania epigenetyczne i struktury chromatyny z precyzją niedostępną dla jakiejkolwiek wcześniejszej metody bioinformatycznej.

Ciągła percepcja wizyjna: Modele Mamba-Vision i analiza strumieni wideo

Poza domeną tekstu naturalnego architektura Mamba-2 rewolucjonizuje przetwarzanie obrazów i wideo w modelach widzenia komputerowego Vision-Mamba. Tradycyjne transformatory wizyjne ViT dzielą obraz na siatkę łatek patches, po czym aplikują do nich pełny mechanizm uwagi. W przypadku zdjęć wysokiej rozdzielczości lub sekwencji wideo liczba łatek rośnie lawinowo, powodując natychmiastowe załamanie wydajnościowe klasycznych transformatorów, które nie są w stanie przetwarzać wideo w pełnej rozdzielczości w czasie rzeczywistym.

Modele oparte na Mambie-2, takie jak Mamba-Vision czy Video-Mamba, traktują sekwencję klatek wideo jako ciągły strumień czasoprzestrzenny. Dzięki liniowej złożoności i stałemu stanowi pamięci model może przetwarzać strumień wideo o częstotliwości sześćdziesięciu klatek na sekundę w rozdzielczości 4K bez konieczności buforowania całej sekwencji w pamięci VRAM. Ma to fundamentalne znaczenie dla autonomicznych robotów humanoidalnych, pojazdów autonomicznych oraz systemów monitoringu miejskiego, które wymagają natychmiastowej, bezopóźnieniowej interpretacji otoczenia i ciągłego śledzenia dynamicznych obiektów bez ryzyka przepełnienia pamięci podręcznej akceleratora pokładowego.

Początek ery post-transformatorowej: Inteligencja w czasie rzeczywistym

Architektura Mamba-2 oraz teoria State Space Duality wyznaczają cezurę w rozwoju sztucznej inteligencji, którą historycy technologii będą wspominać jako początek ery post-transformatorowej. Przełamanie dogmatu kwadratowej złożoności i udowodnienie, że rekurencja i uwaga są dwoma obliczami tego samego aparatu algebraicznego, ostatecznie zdejmuje z modeli językowych ograniczenia pamięciowe i energetyczne.

Możliwość przetwarzania milionów tokenów w czasie rzeczywistym przy stałym zużyciu pamięci otwiera drogę do systemów sztucznej inteligencji, które będą mogły żyć i uczyć się w sposób ciągły. Agenci autonomiczni wyposażeni w pamięć Mamba-2 mogą nieustannie przetwarzać strumienie danych wideo z kamer robotycznych, analizować pełne repozytoria kodu oprogramowania w ułamku sekundy oraz prowadzić wielomiesięczne dialogi bez utraty spójności i bez konieczności resetowania sesji. Liniowa złożoność staje się nowym standardem nowoczesnej inteligencji maszynowej, a Mamba-2 kładzie nienaruszalny fundament pod nową generację systemów kognitywnych o skali, elegancji i efektywności energetycznej dorównującej biologicznemu mózgowi człowieka, redefiniując granice tego, co osiągalne w dziedzinie sztucznej inteligencji.