Resumo Técnico: Kalman Linear Attention (KLA)

Problema

Os atuais modelos de estado-espaço (SSMs), como Mamba e Gated Linear Attention (GLA), oferecem alternativas de complexidade linear e paralelizáveis ao softmax attention, abordando o gargalo de escalonamento quadrático dos Transformers. No entanto, esses modelos dependem de atualizações de estado oculto lineares ou afins, o que limita sua expressividade em comparação com as interações não lineares induzidas pela normalização softmax. Além disso, os SSMs e RNNs lineares existentes carecem de mecanismos explícitos para representar a incerteza do estado, tratando os estados ocultos como determinísticos. Este artigo aborda a lacuna entre a eficiência das atualizações lineares e a expressividade/robustez da filtragem não linear, questionando se as atualizações exatas do filtro de Kalman podem ser implementadas eficientemente como um mixer de sequências.

Metodologia: Kalman Linear Attention (KLA)

Os autores propõem o Kalman Linear Attention (KLA), um mixer de sequências probabilístico que formula a modelagem de sequências como inferência bayesiana exata. A inovação central reside na reparametrização do filtro de Kalman para a forma de informação (usando precisões e parâmetros naturais) para permitir o computação paralela.

1. Formulação Probabilística

O KLA modela a mistura de sequências como inferência em um modelo de estado-espaço linear-gaussiano:

• Prior de Dinâmica: A evolução do estado latente é modelada como um processo de Ornstein-Uhlenbeck (OU), discretizado exatamente. Isso introduz um termo de ruído de processo ($\bar{p}_t$) que captura a incerteza nas transições de estado, distinguindo-se das dinâmicas determinísticas dos SSMs padrão.
• Modelo de Observação: Os tokens de entrada são tratados como observações ruidosas do estado latente, parametrizadas por um operador de observação ($k_t$) e uma precisão de valor ($\Lambda^v_t$) que representa a confiança na evidência do token.
• Leitura (Readout): A saída é derivada da média posterior do estado latente, projetada via um operador de consulta ($q_t$).

2. Paralelização via Transformações de Möbius

O filtragem de Kalman clássica é inerentemente sequencial. O KLA supera isso explorando a forma de informação do filtro:

• Recursão de Precisão: A atualização da precisão posterior ($\lambda_t$) é demonstrada como uma transformação de Möbius (fracionária-linear) da precisão anterior ($\lambda_{t-1}$).
  $$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$$
  Crucialmente, as transformações de Möbius compõem-se de forma associativa. Isso permite que as atualizações de precisão sejam computadas via scans de prefixo paralelos (scans associativos) com trabalho de $O(T)$ e profundidade de $O(\log T)$, igualando a eficiência paralela dos modernos SSMs como o Mamba.
• Recursão de Média: Uma vez que o caminho da precisão é computado em paralelo, as atualizações da média posterior (média de informação) seguem uma transformação afim que também é paralelizável.

3. Gating Não Linear a partir da Incerteza

Ao contrário dos SSMs lineares, onde o gating é frequentemente linear ou dependente do token, mas determinístico, o gating do KLA emerge naturalmente da razão de precisões. O "portão de esquecimento" (forget gate) e o "portão de escrita" (write gate) são dependentes do histórico e não lineares, impulsionados pela incerteza relativa entre o estado anterior e a nova observação. Isso permite que o modelo incorpore informações seletivamente com base na confiança, sem exigir parâmetros de gating aprendidos explicitamente para cada token.

4. Arquitetura

O KLA é projetado como um substituto direto (drop-in replacement) para camadas padrão de SSM ou atenção. Ele segue uma arquitetção de MLP fundido similar ao Mamba, incorporando convoluções causais 1D e ativações SiLU. Os autores implementam kernels CUDA e Triton eficientes para o scan de Möbius para evitar a materialização de estados expandidos em memória de alta largura de banda (HBM).

Principais Contribuições

1. Reparametrização Associativa: Os autores demonstram que a filtragem de Kalman linear-gaussiana diagonal em forma de informação admite uma estrutura de Möbius, permitindo scans de prefixo paralelos exatos para filtragem bayesiana.
2. Não Linearidade Impulsionada por Incerteza: O KLA introduz gating não linear derivado de razões de precisão, oferecendo uma expressividade estritamente maior do que as atualizações de SSM linear/afim, enquanto mantém a estrutura de scan de tempo linear.
3. Camada KLA: Um mixer de sequência prático e substituível que produz incerteza de estado de crença explícita e pode ser empilhado em redes profundas.
4. Escalonamento e Validação: O artigo apresenta a primeira primitiva de filtragem bayesiana empilhada treinada na escala de bilhões de tokens, validando sua estabilidade e desempenho.

Resultados Experimentais

Os autores avaliam o KLA através de tarefas sintéticas, benchmarks de rastreamento de estado e modelagem de linguagem em larga escala.

• Tarefas Sintéticas (MAD-Lab): O KLA supera o GDN, GLA, Mamba e mLSTM em tarefas que exigem seleção sob incerteza, especificamente Fuzzy Recall e Selective Copy. Ele alcança resultados de estado da arte em Compressão, sugerando que sua filtragem em forma de informação atua como um compressor de contexto ótimo.
• Rastreamento de Estado (A5 Permutation): Na tarefa de composição de permutação A5 (um benchmark para rastreamento de estado NC1-completo), o KLA resolve a tarefa com profundidade constante (1–2 blocos). Em contraste, SSMs lineares e modelos baseados em atenção requerem profundidade que cresce com o comprimento da sequência. Isso demonstra que as atualizações fracionárias-lineares do KLA quebram os limites de expressividade dos modelos recorrentes lineares, mantendo-se paralelizáveis.
• Recall Associativo de Longo Contexto (MQAR): O KLA atinge uma precisão quase perfeita (>95%) na tarefa de Multi-Query Associative Recall com comprimento de sequência de 2048, superando o Mamba e excedendo significativamente o GLA, que falha em aprender a tarefa neste cenário.
• Pré-treinamento em Larga Escala:
  • Standalone: Com 45M e 180M de parâmetros, o KLA standalone iguala ou melhora o desempenho do GPT, Mamba e GDN em benchmarks de senso comum zero-shot (ex: HellaSwag, PIQA).
  • Híbrido: Substituir a última camada de atenção de um GPT pelo KLA (GPT+KLA) produz o desempenho mais forte entre os híbridos, superando GPT+Mamba e GPT+GDN.
• Ablação: Os autores mostem que o prior de Ornstein-Uhlenbeck (OU) é crítico para o empilhamento profundo estável. Sem a dinâmica estocástica de retorno à média, o treinamento de pilhas profundas de KLA torna-se instável.

Significância e Alegações

O artigo afirma que o KLA consegue unir com sucesso a eficiência dos modelos de estado-espaço lineares com a expressividade dos mecanismos de atenção não lineares. Ao tratar a mistura de sequências como filtragem bayesiana exata, o KLA fornece:

• Incerteza Explícita: Diferente de mixers determinísticos, o KLA carrega uma incerteza de estado de crença explícita que controla diretamente o gating de informação.
• Expressividade Estritamente Superior: As atualizações de Möbius permitem que o KLA resolva tarefas de rastreamento de estado (como A5) que SSMs lineares e atenção não conseguem, sem sacrificar as capacidades de treinamento paralelo.
• Escalabilidade: O trabalho demonstra que primitivas de filtragem bayesiana exata podem ser empilhadas e treinadas na escala de bilhões de tokens, desafiando a noção de que a filtragem probabilística é inerentemente sequencial ou computacionalmente proibitiva para a modelagem de linguagem em larga escala.

Os autores posicionam o KLA não apenas como uma curiosidade teórica, mas como uma primitiva prática e substituível que aumenta a robustez e as capacidades de rastreamento de estado dos modelos de linguagem modernos, mantendo a eficiência computacional necessária para o processamento de contexto longo.