State Space Models (SSM)

Um modelo de espaço de estados (SSM) mapeia uma sequência de entrada numa sequência de saída através de um vetor de estado oculto que atualiza a cada etapa, pelo que a sua memória mantém o mesmo tamanho, independentemente do crescimento da sequência. A ideia vem da teoria do controlo, onde um sistema linear é escrito como uma atualização de estado h' = Ah + Bx e uma saída y = Ch + Dx, e onde o filtro de Kalman estima estes estados a partir de medições ruidosas desde 1960. Os SSMs de aprendizagem profunda aprendem A, B e C a partir dos dados e empilham muitas destas camadas; desde o S4, em 2021, e o Mamba, em 2023, competem com os Transformers em sequências longas de texto, áudio e dados de séries temporais, enquanto o seu custo computacional cresce linearmente com o comprimento da sequência.

Como funcionam os SSMs#

Uma camada SSM neural converte as equações contínuas numa recorrência discreta com um passo aprendido Δ, usando a regra de retenção de ordem zero Ā = exp(ΔA). Cada etapa calcula então h_t = Ā h_(t-1) + B̄ x_t e y_t = C h_t. Quando A, B e C permanecem fixos ao longo das etapas temporais, desenrolar a recorrência mostra que a saída é uma convolução da entrada com o kernel (CB̄, CĀB̄, CĀ²B̄, ...). Por isso, a mesma camada pode ser executada como uma convolução durante o treino, processando toda a sequência em paralelo numa GPU, e como recorrência durante a geração, com uma atualização de estado de tamanho constante por token. Em PyTorch, as duas perspetivas produzem as mesmas saídas:

import torch

torch.manual_seed(0)
L, N = 64, 16  # comprimento da sequência, tamanho do estado
x = torch.randn(L)

# SSM contínuo diagonal, discretizado com passo delta (retenção de ordem 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

# Perspetiva recorrente: uma atualização de estado por etapa, como na geração
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)

# Perspetiva convolucional: pré-calcular o kernel K_k = C . (A_bar^k * B_bar), como no treino
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))

Imprime Same outputs: True. A recorrência mantém apenas 16 números no estado em cada etapa, enquanto o kernel pré-calculado permite ao treino calcular os 64 resultados em paralelo.

Do S4 ao Mamba#

Os primeiros SSMs profundos eram demasiado lentos e exigiam demasiada memória para serem treinados em grande escala. O S4, de Albert Gu, Karan Goel e Christopher Ré, deu à matriz de estado uma estrutura que torna barato o cálculo do kernel de convolução e uma inicialização concebida para memorizar históricos longos. Estabeleceu o estado da arte em todas as tarefas do benchmark Long Range Arena, tornou-se o primeiro modelo a resolver a tarefa Path-X, com sequências de 16K etapas, e gerava texto 60 vezes mais depressa do que um Transformer.

Um kernel fixo trata todas as entradas da mesma forma, pelo que não consegue decidir ignorar conteúdo supérfluo e reter um facto importante. O Mamba, de Albert Gu e Tri Dao, torna Δ, B e C funções do token atual, permitindo ao modelo escolher o que escreve no seu estado e o que esquece. Estes parâmetros seletivos quebram a perspetiva convolucional, pelo que o Mamba calcula a recorrência com uma varredura paralela adaptada ao hardware. O Mamba atingiu cinco vezes o débito de inferência dos Transformers, e o Mamba-3B superou Transformers do mesmo tamanho e igualou os que tinham o dobro do seu tamanho. Cada bloco Mamba envolve o SSM seletivo com uma convolução curta e uma porta SiLU:

O Mamba-2 mostrou que os SSMs seletivos e uma forma de atenção linear são perspetivas duais do mesmo cálculo, uma estrutura a que os autores chamam dualidade de espaço de estados, e usou-a para tornar a camada principal 2 a 8 vezes mais rápida.

SSMs vs. arquiteturas relacionadas#

• Transformers: A autoatenção compara cada token com todos os outros, pelo que pode copiar exatamente qualquer token anterior, com um custo que cresce de forma quadrática e uma cache KV que cresce com a janela de contexto. Um SSM comprime o histórico num estado fixo, o que é mais barato, mas implica perda de informação, e o artigo de 2024 Repete depois de mim mostrou que os Transformers copiam e recuperam melhor informação de entradas longas do que os SSMs de estado fixo.
• Redes neuronais recorrentes (RNNs): Uma RNN ou LSTM aplica uma atualização não linear, pelo que o treino tem de percorrer a sequência um token de cada vez. As atualizações dos SSMs são lineares no estado, o que permite usar uma convolução ou uma varredura paralela.
• Modelos ocultos de Markov (HMMs): Um HMM tem um conjunto finito de estados discretos com probabilidades de transição. O modelo de espaço de estados linear gaussiano subjacente ao filtro de Kalman é o seu equivalente contínuo, e os SSMs profundos mantêm esse vetor de estado contínuo.

Onde se usam os SSMs#

A maioria das utilizações em produção ocorre em grandes modelos de linguagem (LLMs) híbridos, que mantêm algumas camadas de atenção para a recuperação exata de informação. O Jamba da AI21 intercala camadas Transformer e Mamba com camadas de mistura de especialistas (MoE), cabe numa única GPU de 80 GB e apresenta resultados sólidos com contextos de até 256K tokens; os modelos híbridos Granite 4.0-H da IBM combinam camadas Mamba-2 e Transformer numa proporção de 9:1. O Vision Mamba aplica varreduras seletivas a patches de imagem, e os SSMs também modelam áudio bruto e sequências de ADN demasiado longas para a atenção total. A forma clássica é usada no rastreamento do Ultralytics YOLO26: o rastreador TrackTrack predefinido, tal como o ByteTrack e o BoT-SORT, prevê a próxima caixa de cada objeto com um filtro de Kalman cujo estado de 8 dimensões contém o centro e o tamanho da caixa, bem como as respetivas velocidades, atualizados com uma matriz de transição de velocidade constante e uma matriz de observação linear que lê as quatro coordenadas da caixa.