SSM
Jak działa
1) Punktem wyjścia jest ciągły liniowy układ stanu z parametrami (A, B, C, D), gdzie ukryty stan x(t) ewoluuje liniowo w funkcji wejścia u(t). 2) Układ dyskretyzowany jest krokiem czasu Delta (np. metoda Zero-Order Hold lub bilinear), co daje rekurencję x_k = A_bar x_{k-1} + B_bar u_k oraz y_k = C x_k. 3) Przy odpowiedniej strukturze A (np. diagonal plus low-rank, HiPPO-LegS) rekurencję można wyrazić jako globalną konwolucję 1D z kernelem K, wyliczanym efektywnie w dziedzinie częstotliwości (FFT). 4) Trening odbywa się w trybie konwolucyjnym (równoległym po sekwencji), a autoregresyjna inferencja w trybie rekurencyjnym (O(1) pamięci na token). 5) W Mambie macierze B, C oraz krok Delta stają się funkcją wejścia (selective SSM), co usuwa ograniczenie LTI i wymaga niestandardowego, hardware-aware selective scan w pamięci SRAM GPU.
Rozwiązany problem
Kwadratowa złożoność self-attention w Transformerach (O(L^2) czasu i pamięci) ogranicza skalowanie kontekstu do dziesiątek tysięcy tokenów oraz uniemożliwia tanie modelowanie sekwencji o długości milionów kroków (audio, genomika, długie dokumenty, sterowanie robotem). SSM oferują liniową złożoność wzgledem długości oraz stały koszt pamięci na token podczas inferencji autoregresyjnej.
Komponenty
Macierz definiująca dynamikę ukrytego stanu. W S4 używana jest struktura HiPPO-LegS (diagonal plus low-rank), w S4D/S5 wariant diagonalny, co znacząco upraszcza obliczenia.
Parametr (uczony lub zależny od wejścia) zamieniający ciągły układ na rekurencję dyskretną metodą zero-order hold lub bilinear.
B rzutuje wejście u_k na przestrzeń stanu, C odczytuje wyjście z ukrytego stanu. W Mambie obie zależą od wejścia (selective SSM).
Hardware-aware algorytm rekurencji równoległej (parallel scan) trzymający stan w SRAM GPU, niezbędny gdy parametry zależą od wejścia i nie da się użyć trybu konwolucyjnego.
Implementacja
Krok Delta zainicjalizowany poza zalecanym log-uniform zakresem powoduje wybuchanie aktywacji lub martwy stan, szczególnie w głębokich stosach SSM.
Gdy B, C lub Delta zależą od wejścia (Mamba), nie istnieje globalny kernel konwolucyjny — naiwny trening rekurencyjny jest wolny.
Niewielki stan ukryty (np. N=16) ogranicza zdolnosc dokladnego kopiowania dlugich fragmentow z kontekstu, w czym Transformery z KV-cache pozostaja mocniejsze.
Ewolucja
Gu et al. wprowadzają HiPPO, framework dla operatorów stanu rekonstruujących sygnał wejściowy w bazie ortogonalnej. Podstawa teoretyczna późniejszych SSM.
Pierwsza warstwa SSM dla deep learning, ujawniająca dualność rekurencyjno-konwolucyjną; była jednak kosztowna obliczeniowo.
Wprowadzenie struktury DPLR i parametryzacji w dziedzinie częstotliwości; SOTA na Long Range Arena, sekwencje do 16k tokenów.
Pokazano, że diagonalna macierz A osiąga jakość zbliżoną do pełnego S4 przy znacznie prostszej implementacji.
Hungry Hungry Hippos pokazuje, że SSM zaczynają być konkurencyjne dla Transformerów w modelowaniu języka.
Gu i Dao wprowadzają zależne od wejścia B, C, Delta oraz hardware-aware selective scan; pierwszy SSM dorównujący Transformerom na benchmarkach językowych przy liniowej złożoności.
Mamba-2 unifikuje SSM i attention przez State Space Duality; Jamba (AI21) i Zamba mieszają warstwy Mamba i Transformera w produkcyjnych LLM.
Hiperparametry (konfigurowalne osie)
Wymiar ukrytego stanu na kanał. Typowo 16-64 w S4/Mambie. Kontroluje pojemność pamięci długoterminowej.
Zakres inicjalizacji kroku dyskretyzacji Delta (zwykle log-uniform). Krytyczny dla stabilności i jakości modelowania zależności o różnej skali.
Wybór struktury A: HiPPO-LegS (S4), diagonal (S4D, S5, Mamba), DPLR. Wpływa na wydajność i własności pamięci.
Czy parametry B, C, Delta zależą od wejścia (Mamba) czy są stałe (S4/S4D/S5).
Złożoność obliczeniowa
Złożoność czasowa: O(L) inferencja rekurencyjna; O(L log L) trening konwolucyjny (FFT); O(L) trening Mamby z selective scan. Złożoność przestrzenna: O(1) na token podczas autoregresyjnej inferencji; O(L) podczas treningu.
Paradygmat wykonania
Klasyczne SSM (S4, S4D) są LTI — parametry nie zależą od wejścia. Mamba wprowadza zależność od wejścia (selektywność), pozostając gęsta obliczeniowo.
Równoległość
W trybie konwolucyjnym SSM jest w pełni równoległy po sekwencji podczas treningu. W trybie rekurencyjnym (inferencja autoregresyjna) jest z natury sekwencyjny. Mamba odzyskuje równoległość treningu przez parallel scan.
Wymagania sprzętowe
Mamba i S4 mają oficjalne kernele CUDA (selective scan, FFT) wykorzystujące SRAM GPU i Tensor Cores.
Stała pamięć stanu na token czyni SSM atrakcyjnymi do inferencji na CPU/edge; brak jednak oficjalnych kerneli AVX równie zoptymalizowanych jak GPU.