Mamba e State Space Models — L'Alternativa Lineare ai Transformer per Sequenze Lunghe

Cos'è: Mamba è un'architettura di sequence modeling pubblicata da Albert Gu (CMU) e Tri Dao (Princeton, autore di FlashAttention) a dicembre 2023. Si basa sui State Space Model (SSM) — una formulazione presa dalla teoria del controllo classica — con due innovazioni chiave: un meccanismo di selettività che permette ai parametri del modello di dipendere dall'input (cosa che gli SSM tradizionali non facevano), e un'implementazione hardware-aware che lo rende efficiente sulle GPU moderne. Risultato: complessità lineare O(n) per token rispetto alla lunghezza della sequenza, contro O(n²) dell'attenzione standard. A parità di parametri sotto i 3B raggiunge o supera i transformer; oltre quella soglia, il transformer puro mantiene un vantaggio.

State Space Models: la radice classica del paper

Gli State Space Model sono una rappresentazione matematica usata fin dagli anni '60 nella teoria del controllo per descrivere sistemi dinamici lineari. La formulazione discreta è: h(t) = A·h(t-1) + B·x(t) e y(t) = C·h(t), dove h è uno stato latente nascosto, x l'input, y l'output, e A, B, C sono matrici di parametri. Albert Gu, già nel 2021-2022, aveva iniziato a esplorare l'uso degli SSM come backbone per il sequence modeling nel deep learning, pubblicando S4 ("Efficiently Modeling Long Sequences with Structured State Spaces", 2021) che aveva ottenuto risultati eccezionali sul benchmark Long Range Arena, dove i transformer fallivano su sequenze da 16K+ token.

Il limite di S4 e di tutti gli SSM precedenti era che i parametri A, B, C erano fissi, indipendenti dall'input: il modello processava ogni token con la stessa dinamica, indipendentemente da cosa contenesse. Questo è il motivo per cui erano forti su task di dipendenze a lunga distanza con segnale uniforme (audio, time series) ma deboli sul linguaggio naturale, dove la rilevanza dei token cambia drasticamente in base al contesto. Un transformer, invece, usa l'attenzione per "guardare selettivamente" quali token del passato sono importanti per il token corrente.

L'idea Mamba: rendere gli SSM "selettivi"

L'innovazione centrale di Gu e Dao è rendere i parametri B, C e il passo temporale Δ dipendenti dall'input: B(x), C(x), Δ(x). In pratica, queste matrici diventano funzioni dell'input corrente attraverso layer lineari aggiuntivi. Questo permette al modello di "decidere" in funzione del token attuale quanto del passato far entrare nello stato (B), come proiettare lo stato in output (C), e quanto tempo "passare" su questo token (Δ). È l'equivalente concettuale dell'attenzione, ma realizzato come modulazione di una dinamica ricorrente invece che come prodotto scalare globale.

Il problema tecnico è che rendere i parametri dipendenti dall'input rompe la struttura matematica che permetteva agli SSM classici di essere convolution-equivalenti e quindi parallelizzabili in modo efficiente. La soluzione di Mamba è un parallel scan algorithm hardware-aware implementato in CUDA, che sfrutta la gerarchia di memoria SRAM/HBM nello stesso spirito di FlashAttention (non a caso Tri Dao è coautore di entrambi). L'algoritmo evita di materializzare in HBM lo stato latente espanso, mantenendo la complessità I/O sotto controllo. Risultato: training tempi competitivi con i transformer, e inferenza significativamente più veloce perché lo stato ricorrente è di dimensione costante invece di crescere con il contesto.

Performance: dove vince, dove perde

Sotto i 3B parametri, Mamba a parità di compute e dati eguaglia o supera i transformer standard sui benchmark di language modeling (Pile, LAMBADA, HellaSwag). Su sequenze molto lunghe — DNA, audio raw, time series — Mamba è chiaramente superiore: scala fino a 1 milione di token senza degrado, dove i transformer hanno problemi sia di memoria sia di qualità su context window estremi. L'inferenza è 5x più veloce dei transformer di dimensione equivalente perché il decoding è genuinamente ricorrente con stato costante, senza il KV cache che cresce linearmente.

Sopra i 3B parametri il quadro è meno chiaro. Mistral, Meta, Anthropic, OpenAI hanno tutti sperimentato con architetture SSM e ibride ma nessuno ha pubblicato un modello SSM puro frontier. La community open ha tentato Mamba 7B e Mamba 2.8B (Cartesia AI, la startup di Gu e Dao), ma in benchmark in-context learning e reasoning multi-step questi modelli sono ancora indietro rispetto a Llama 3 di dimensione equivalente. L'ipotesi corrente è che la dipendenza globale "ogni-token-a-ogni-token" dell'attenzione cattura pattern di ragionamento che lo stato ricorrente comprime troppo aggressivamente.

L'era degli ibridi: Jamba, Codestral Mamba, Zamba

Nel 2024 la community ha gradualmente convergato sull'idea che Mamba non sostituisce il transformer ma lo complementa. A marzo 2024 AI21 Labs rilascia Jamba, il primo modello ibrido a scala produzione: 52B parametri totali (12B attivi via Mixture of Experts), architettura alternata Mamba+Transformer in proporzione 7:1 (sette blocchi Mamba per ogni blocco Transformer). Jamba supporta context window di 256K token con un footprint di KV cache drammaticamente ridotto rispetto a un transformer puro equivalente. A luglio 2024 esce Jamba 1.5 con 398B parametri totali e context fino a 256K.

A luglio 2024 Mistral rilascia Codestral Mamba 7B, un modello specializzato in code generation basato su Mamba puro. La scelta è strategica: il code completion in editor richiede latenze bassissime su context lunghi (file di centinaia di righe), e l'inferenza con stato costante di Mamba è ideale per questo caso d'uso. Zamba di Zyphra e Falcon Mamba 7B di TII completano il panorama 2024 degli ibridi e dei Mamba puri commerciali. Il consenso tecnico emergente è che il futuro dei modelli a context window estremamente lunghi (1M+ token) sarà ibrido, con il transformer che fornisce reasoning globale e Mamba/SSM che fornisce ricorrenza efficiente sul long tail. I transformer puri continueranno a dominare le fasce frontier per qualche generazione ancora, ma la tassonomia delle architetture si sta diversificando per la prima volta dal 2017.

---