State Space Models (SSM)

Un modello a spazio di stato (SSM) mappa una sequenza di input in una sequenza di output tramite un vettore di stato nascosto che aggiorna a ogni passaggio, perciò la sua memoria mantiene le stesse dimensioni a prescindere dalla lunghezza della sequenza. L'idea deriva dalla teoria del controllo, in cui un sistema lineare è espresso come aggiornamento dello stato h' = Ah + Bx e output y = Ch + Dx, e in cui il filtro di Kalman stima stati di questo tipo da misurazioni rumorose fin dal 1960. Gli SSM di deep learning apprendono A, B e C dai dati e impilano molti di questi livelli; dal 2021, con S4, e dal 2023, con Mamba, competono con i transformer su sequenze lunghe di testo, audio e dati di serie temporali, mentre il calcolo cresce linearmente con la lunghezza della sequenza.

Come funzionano gli SSM#

Un livello SSM neurale converte le equazioni continue in una ricorrenza discreta con un passo Δ appreso, usando la regola del mantenimento di ordine zero Ā = exp(ΔA). Ogni passaggio calcola quindi h_t = Ā h_(t-1) + B̄ x_t e y_t = C h_t. Quando A, B e C restano costanti nel tempo, lo svolgimento della ricorrenza mostra che l'output è una convoluzione dell'input con il kernel (CB̄, CĀB̄, CĀ²B̄, ...). Lo stesso livello può quindi funzionare come convoluzione durante l'addestramento, elaborando l'intera sequenza in parallelo su una GPU, e come ricorrenza durante la generazione, con un aggiornamento dello stato di dimensione costante per token. In PyTorch, le due rappresentazioni producono gli stessi output:

import torch

torch.manual_seed(0)
L, N = 64, 16  # lunghezza della sequenza, dimensione dello stato
x = torch.randn(L)

# SSM diagonale a tempo continuo, discretizzato con passo delta (mantenimento di ordine zero)
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# Rappresentazione ricorrente: un aggiornamento dello stato per passaggio, come nella generazione
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# Rappresentazione convoluzionale: precalcola il kernel K_k = C . (A_bar^k * B_bar), come nell'addestramento
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Stampa Same outputs: True. La ricorrenza mantiene solo 16 numeri di stato a ogni passaggio, mentre il kernel precalcolato permette all'addestramento di calcolare tutti i 64 output in parallelo.

Da S4 a Mamba#

I primi SSM profondi erano troppo lenti e richiedevano troppa memoria per essere addestrati su larga scala. S4, sviluppato da Albert Gu, Karan Goel e Christopher Ré, ha dato alla matrice di stato una struttura che rende economico il calcolo del kernel di convoluzione e un'inizializzazione progettata per ricordare storie lunghe. Ha stabilito lo stato dell'arte in tutte le attività del benchmark Long Range Arena, è diventato il primo modello a risolvere l'attività Path-X con sequenze da 16K passaggi e generava 60 volte più velocemente di un transformer.

Un kernel fisso tratta ogni input allo stesso modo, quindi non può decidere di ignorare il riempitivo e conservare un fatto chiave. Mamba, sviluppato da Albert Gu e Tri Dao, rende Δ, B e C funzioni del token corrente, permettendo al modello di scegliere cosa scrivere nel proprio stato e cosa dimenticare. Questi parametri selettivi interrompono la rappresentazione convoluzionale, perciò Mamba calcola la ricorrenza con una scansione parallela consapevole dell'hardware. Mamba ha raggiunto un throughput di inferenza cinque volte superiore a quello dei transformer e Mamba-3B ha superato i transformer delle stesse dimensioni, eguagliando quelli grandi il doppio. Ogni blocco Mamba avvolge l'SSM selettivo con una convoluzione breve e un gate SiLU:

Mamba-2 ha mostrato che gli SSM selettivi e una forma di attenzione lineare sono due rappresentazioni dello stesso calcolo, un framework che gli autori chiamano dualità spazio di stato, e lo ha usato per rendere il livello principale da 2 a 8 volte più veloce.

SSM e architetture correlate#

• Transformer: la self-attention confronta ogni token con tutti gli altri, quindi può copiare esattamente qualsiasi token precedente, ma il suo costo cresce in modo quadratico e la cache KV cresce con la finestra di contesto. Un SSM comprime la cronologia in uno stato fisso, più economico ma con perdita d'informazione; il paper del 2024 Repeat After Me ha mostrato che i transformer copiano e recuperano meglio contenuti da input lunghi rispetto agli SSM con stato fisso.
• Reti neurali ricorrenti (RNNs): una RNN o una LSTM applica un aggiornamento non lineare, perciò l'addestramento deve elaborare la sequenza un token alla volta. Gli aggiornamenti degli SSM sono lineari rispetto allo stato, il che consente di usare una convoluzione o una scansione parallela.
• Modelli di Markov nascosti (HMMs): un HMM ha un insieme finito di stati discreti con probabilità di transizione. Il modello lineare-gaussiano a spazio di stato alla base del filtro di Kalman ne è la controparte continua e gli SSM profondi mantengono questo vettore di stato continuo.

Dove si usano gli SSM#

La maggior parte degli impieghi in produzione riguarda modelli linguistici di grandi dimensioni (LLMs) ibridi, che mantengono alcuni livelli di attenzione per recuperare informazioni esatte. Jamba di AI21 alterna livelli transformer e Mamba a livelli di miscela di esperti (MoE), entra in una singola GPU da 80 GB e riporta risultati solidi con contesti fino a 256K token; i modelli ibridi Granite 4.0-H di IBM combinano livelli Mamba-2 e transformer in rapporto 9:1. Vision Mamba applica scansioni selettive alle patch delle immagini e gli SSM modellano anche sequenze grezze di audio e DNA, troppo lunghe per un'attenzione completa. La forma classica è usata nel tracciamento di Ultralytics YOLO26: il tracker predefinito TrackTrack, come ByteTrack e BoT-SORT, prevede il riquadro successivo di ciascun oggetto con un filtro di Kalman il cui stato a 8 dimensioni contiene il centro e le dimensioni del riquadro, oltre alle relative velocità; l'aggiornamento usa una matrice di transizione a velocità costante e una matrice di osservazione lineare che legge le quattro coordinate del riquadro.