Mamba a State Space Models: Architektúra, ktorá chce zvrhúť Transformer
Transformer dominuje AI od roku 2017 — no jeho Achilova päta, kvadratická pamäťová zložitosť, sa stáva čoraz väčším problémom pri spracovaní dlhých textov. State Space Models (SSM) a ich najznámejší predstaviteľ Mamba ponúkajú fundamentálne iný prístup: lineárnu zložitosť a pamäťovú efektívnosť, ktorú Transformer nedokáže dosiahnuť.
1. Čo sú State Space Models
State Space Models sú trieda matematických modelov s koreňmi v teórii riadenia z 60. rokov 20. storočia. Ich základná myšlienka: namiesto toho, aby model videl celú vstupnú sekvenciu naraz, udržiava kompaktný skrytý stav (hidden state), ktorý priebežne aktualizuje pri každom novom tokene.
Formálne SSM definujú dve rovnice:
- Stavová rovnica:
h'(t) = Ah(t) + Bx(t)— ako sa skrytý stav mení - Výstupná rovnica:
y(t) = Ch(t) + Dx(t)— čo model vypíše
kde h je skrytý stav, x je vstup, y je výstup a A, B, C, D sú naučiteľné matice.
Táto myšlienka nie je nová — rekurentné siete (RNN, LSTM) fungujú podobne. Problém bol, že klasické RNN sa ťažko trénujú (vanishing gradient) a neparalelizujú efektívne. Výskumníci ako Albert Gu a spolupracovníci zo Stanfordu ukázali, ako SSM vyriešiť tieto problémy pomocou efektívnych konvolúcií — modely S4 (2021), H3 (2022) a napokon Mamba (2023) posunuli SSM do centra pozornosti.
2. Prečo Transformer nestačí na všetko
Transformery sú výnimočné. Mechanizmus attention im umožňuje spájať ľubovoľné tokeny v sekvencii — čo je hlavný dôvod ich sily. No práve táto flexibilita má cenu:
- Kvadratická zložitosť: Výpočet attention rastie ako O(n²) vzhľadom na dĺžku sekvencie. Sekvenciu 2× dlhšiu spracuje 4× pomalšie.
- KV cache: Pri inferenci si Transformer musí pamätať kľúče a hodnoty všetkých predchádzajúcich tokenov. Kontext 128 000 tokenov môže zabrať desiatky gigabajtov pamäte.
- Obmedzenia pri dlhých kontextoch: Aj keď moderné Transformery dokážu spracovať milión tokenov a viac, praktická efektívnosť klesá a náklady rastú neúmerne.
Pre aplikácie ako spracovanie celých kníh, genomika, analýza dlhých časových radov alebo autonómne riadenie — kde sekvencie bývajú extrémne dlhé — je kvadratická zložitosť reálny bloker.
3. Mamba: Selektívne State Space Models
Mamba (2023, Gu & Dao) je kľúčový prelom v oblasti SSM. Oproti predchádzajúcim prístupom prináša tri inovácie:
Selektívny mechanizmus: Klasické SSM mali fixné matice A, B, C. Mamba ich robí závislými od vstupu — model sa sám naučí, čo si pamätať a čo zabudnúť. Toto je ekvivalent selektivity attention bez kvadratickej ceny.
Hardware-aware algoritmus: Mamba prepočítava stavy priamo v SRAM (rýchla pamäť GPU) namiesto DRAM, čo dramaticky znižuje pamäťové prenosy. Tento „parallel scan" umožňuje rýchle trénovanie napriek rekurentnej povahe modelu.
Rekurencia pri inferenci: Kým Transformer potrebuje celý KV cache, Mamba pri generovaní ďalšieho tokenu potrebuje len aktuálny skrytý stav — fixnej veľkosti, nezávisle od dĺžky kontextu. Inferencia je teda O(1) v pamäti.
Výsledok: Mamba je pri dlhých sekvenciách až 5× rýchlejšia ako porovnateľný Transformer s rovnakým počtom parametrov.
4. Porovnanie architektúr
| Vlastnosť | Transformer | RNN / LSTM | Mamba (SSM) |
|---|---|---|---|
| Trénovacia zložitosť | O(n²) | O(n) | O(n) |
| Pamäť pri inferenci | O(n) — rastie | O(1) — fixná | O(1) — fixná |
| Paralelizácia trénovania | Áno | Nie | Áno (scan) |
| In-context learning | Silné | Slabé | Stredné |
| Dlhé sekvencie (>100 K) | Nákladné | Degraduje | Efektívne |
| Najlepšie benchmark skóre | Áno | Nie | Čiastočne |
| Hybridné verzie | Transformer + MoE | — | Mamba + Transformer |
5. Modely a praktické nasadenie
Po vydaní Mamby vznikol ekosystém modelov a hybridných architektúr:
- Mamba-2 (2024): Vylepšená verzia so State Space Duality (SSD) — matematicky unifikuje SSM a attention, umožňuje ešte rýchlejší tréning.
- Jamba (AI21 Labs): Hybridný model kombinujúci Mamba bloky a Transformer bloky. Dosiahol výkon porovnateľný s Mistral-7B pri nižšej spotrebe pamäte.
- Falcon Mamba 7B (TII): Prvý plne SSM model vo veľkosti 7B parametrov vydaný ako open-source. Prekonal viaceré 7B Transformer modely na štandardných benchmarkoch.
- NVIDIA Hymba: Hybridná architektúra od NVIDIA kombinujúca SSM s „memory tokens" pre lepší in-context recall.
- Zamba: Open-source hybridný model zameraný na efektívnosť pre edge nasadenie.
Prakticky sa SSM modely uplatňujú v:
- Genomike a spracovaní DNA sekvencií (sekvencie s dĺžkou 30 000+ báz)
- Časových radoch — finančné dáta, priemyselné senzory
- Audio spracovaní bez nutnosti rozdeľovania na kratšie okná
- Edge AI zariadeniach s obmedzenou RAM, kde fixná pamäť pri inferenci je kritická výhoda
6. Limity a čo ďalej
Napriek sľubným výsledkom SSM nie sú bez slabín:
In-context learning: Transformer exceluje v schopnosti „naučiť sa" priamo z príkladov v kontexte. SSM s fixným skrytým stavom túto informáciu môžu stratiť pri kompresii dlhých sekvencií — je ťažšie spoľahlivo uchovať detaily zo začiatku veľmi dlhého vstupu.
Benchmarky: Na štandardných jazykových benchmarkoch (MMLU, GPQA) stále vedú Transformer modely. Výhoda SSM sa prejavuje najmä pri dlhých sekvenciách a pamäťovej efektivite, nie pri všeobecnom jazykovom porozumení.
Menší ekosystém: Knižnice ako HuggingFace Transformers sú optimalizované pre Transformer. SSM modely vyžadujú špeciálne CUDA kernely a nie vždy fungujú bez ďalšej konfigurácie.
Hybridný smer: Najperspektívnejší vývoj je kombinácia oboch prístupov — modely ako Jamba, Hymba alebo Zamba spájajú SSM efektívnosť s Transformer presnosťou. Táto hybridná architektúra môže byť optimálnou voľbou pre nasledujúcu generáciu modelov, kde ani čisto SSM, ani čisto Transformer nepostačuje sám o sebe.
Zhrnutie: State Space Models a architektúra Mamba prinášajú skutočnú alternatívu k Transformerom — lineárnu zložitosť, fixnú pamäť pri inferenci a efektívnosť pri extrémne dlhých sekvenciách. Hybridné prístupy naznačujú, že budúcnosť pravdepodobne nebude ani čisto SSM, ani čisto Transformer, ale to najlepšie z oboch svetov.