Resumo Técnico: Sparse Delta Memory (SDM)

Declaração do Problema

As Redes Neurais Recorrentes (RNNs) Lineares, como a Gated DeltaNet (GDN) e a Mamba2, oferecem a vantagem de memória e computação fixas por token, tornando-as teoricamente adequadas para contextos arbitrariamente longos. No entanto, seu desempenho prático em recuperação de contexto longo e aprendizado em contexto é limitado pelo tamanho reduzido de seu estado oculto. Aumentar o tamanho do estado em RNNs lineares padrão leva a um aumento linear nos FLOPs devido às atualizações de matriz densa (multiplicações de matriz-vetor), criando um gargalo computacional proibitivo. Por outro lado, as arquiteturas Transformer padrão com atenção softmax mantêm uma alta capacidade de recuperação, mas sofrem com custos de memória e computação do cache Key-Value (KV) que crescem linearmente conforme o comprimento da sequência aumenta.

O desafio central abordado é como escalar a capacidade do estado oculto das RNNs lineares em ordens de magnitude para melhorar o desempenho em contextos longos sem aumentar o custo computacional (FLOPs) por token.

Metodologia: Sparse Delta Memory (SDM)

Os autores introduzem a Sparse Delta Memory (SDM), uma arquitetura que estende a Gated DeltaNet ao substituir o produto externo denso de chave-valor por um esquema de endereçamento esparso baseado em Product-Key Memory (PKM).

Mecanismo Central

1. Memória Esparsa Explícita: Em vez de manter uma matriz densa $M_t \in \mathbb{R}^{d_{qk} \times d_v}$, a SDM mantém uma tabela de memória explícita $M_t \in \mathbb{R}^{N \times d_v}$ com $N$ slots, onde $N$ pode ser ordens de magnitude maior que a dimensão oculta (ex: $N \approx 10^6$).
2. Seleção de Chave Esparsa: Em cada passo de tempo, as chaves projetadas ($k'$) e as consultas ($q'$) são divididas em duas metades. Sua soma externa gera $N$ pontuações para os slots de memória. Usando uma estratégia de seleção top-$W$ (escrita) e top-$R$ (leitura), o modelo identifica slots de memória específicos para atualizar ou ler. Essa seleção é realizada eficientemente sem materializar a matriz de pontuação completa de tamanho $N$, aproveitando a propriedade de que $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$.
3. Atualizações Delta com Portão (Gated Delta Updates): Para os $W$ slots de escrita selecionados, o modelo aplica a regra de atualização Gated DeltaNet:
   • Portão de Esquecimento (Forget Gate): $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • Atualização Delta: $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • Slots não selecionados permanecem inalterados.
4. Leitura Esparsa: A saída é uma soma ponderada dos $R$ slots de leitura selecionados.

Design IsoFLOP

Para garantir uma comparação justa, os autores projetam a SDM para corresponder à contagem de parâmetros e FLOPs de uma linha de base GDN densa:

• FLOPs: O custo por token é $O((W+R) \times d_v)$, independente do tamanho da memória $N$. Ao definir $W$ e $R$ iguais à dimensão de consulta da linha de base GDN, o custo computacional permanece idêntico.
• Parâmetros: As projeções lineares são dimensionadas identicamente à linha de base. O tamanho da memória $N$ é controlado através do número de cabeças $H$ e das dimensões de projeção, permitindo que o tamanho do estado escale como $O(1/H^2)$ sem afetar os FLOPs.

Estado Inicial Aprendido

Ao contrário da GDN, que tipicamente inicializa seu pequeno estado como zero, a SDM trata o estado de memória inicial $M_0$ como um parâmetro aprendível. Isso permite que o modelo armazene o conhecimento de pré-treinamento diretamente na estrutura de memória, atuando como uma memória paramétrica que persiste durante a inferência.

Principais Contribuições

1. Arquitetura: A proposta da SDM, que escala a capacidade de estado de RNNs lineares em três ordens de magnitude (ex: de kilobytes para gigabytes) mantendo FLOPs constantes por token.
2. Validação IsoFLOP: Demonstração de que aumentar a capacidade de estado isoladamente (sem aumentar o computo) melhora significativamente o desempenho em tarefas de contexto longo.
3. Memória Paramétrica: A introdução de um estado inicial $M_0$ aprendido para RNNs lineares, mostrando que este mecanismo armazena efetivamente o conhecimento comum e melhora tarefas de raciocínio.
4. Leis de Escala (Scaling Laws): Validação empírica de que a SDM segue leis de escala previsíveis, superando a GDN densa e, em escalas maiores (8B parâmetros), superando modelos de Atenção Total (Full Attention) em perda de treinamento.

Resultados Experimentais

Os autores treinaram modelos variando de 1.4B a 8B parâmetros em mais de 1 trilhão de tokens.

• Leis de Escala: Em uma escada de escala, a SDM consistentemente alcançou menor perda de treinamento do que a GDN em todos os níveis de computo. Na escala de 8B, a SDM alcançou uma perda menor do que tanto a GDN quanto um modelo de Atenção Total.
• Recuperação de Contexto Longo (Benchmark RULER): A SDM superou significativamente a GDN em tarefas de recuperação de contexto longo. Nas escalas de 1.4B e 8B, a SDM alcançou pontuações de 31.2 e 50.2 respectivamente, comparadas aos 20.0 e 34.2 da GDN. Notavelmente, a SDM igualou ou excedeu o desempenho da Atenção Total em várias sub-tarefas do RULER, apesar de usar uma representação de memória fixa.
• Contexto Curto e Raciocínio: A SDM alcançou menor perplexidade em dados de código e maior precisão média em benchmarks de raciocínio (ex: MMLU, GSM8K, HumanEval) comparado à GDN. Na escala de 8B, a SDM superou a Atenção Total em precisão média.
• Estudos de Ablação:
  • Tamanho do Estado: O desempenho degradou monotonicamente conforme o tamanho da memória era reduzido, confirmando que o grande estado é o principal motor dos ganhos.
  • Inicialização Aprendida: Aprender $M_0$ melhorou o desempenho para a SDM, mas proporcionou ganhos negligenciáveis para a GDN, sugerindo que a grande capacidade é necessária para utilizar efetivamente a inicialização aprendida.
  • Eficiência: Embora o throughput de treinamento da SDM tenha sido 1.49x mais lento que o da GDN na escala de 8B (devido ao acesso de memória HBM vs SRAM), foi 6x mais rápido que a Atenção Total durante a inferência.

Significância e Alegações

O artigo afirma que a SDM consegue unir com sucesso a eficiência das RNNs lineares com as capacidades de contexto longo dos Transformers. Ao desacoplar a capacidade de memória do custo computacional através da esparsidade, a SDM permite:

• Espaço/Computo Constante: Processar sequências arbitrariamente longas (até 1 milhão de tokens demonstrados) sem o crescimento ilimitado de memória dos caches KV.
• Desempenho Superior em Contexto Longo: Alcançar resultados de estado da arte entre as RNNs lineares e resultados competitivos contra a Atenção Total em recuperação de contexto longo e aprendizado em contexto.
• Armazenamento de Conhecimento Paramétrico: A capacidade de aprender e armazenar conhecimento de pré-treinamento dentro do próprio estado de memória, melhorando o raciocínio geral e tarefas de conhecimento comum.

Os autores reconhecem limitações, especificamente que a grande pegada de memória (comparável aos parâmetros do modelo) exige recursos significativos de HBM, e que as implementações atuais de kernel são menos eficientes do que os kernels otimizados da GDN densa. No entanto, eles posicionam a SDM como uma abertura de novas possibilidades para o desenvolvimento de agentes com memória de longo prazo aprimorada e para o processamento de vídeos longos ou traços de raciocínio estendidos.