Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Resumo Técnico: Tensor Cache

Declaração do Problema
A inferência autoregressiva de Transformers depende do armazenamento em cache de pares Chave-Valor (KV) para evitar a recomputação de prefixos. No entanto, o estado KV retido cresce linearmente com o comprimento do contexto, a profundidade e a concorrência, criando um gargalo de memória. O armazenamento em cache com janela deslizante limita o uso de memória retendo apenas os $W$ tokens mais recentes, mas essa abordagem sofre de "esquecimento rígido": uma vez que um token é expulso da janela, sua informação é totalmente descartada, tornando evidências relevantes fora da janela inacessíveis. Métodos existentes de estado limitado ou descartam tokens expulsos inteiramente ou tentam comprimir cada token em uma média móvel, o que pode introduzir ruído ou falhar em distinguir efetivamente entre contexto recente e distante.

Metodologia: Tensor Cache (TC)
O Tensor Cache introduz uma arquitetura de estado limitado de dois níveis que transforma a expulsão de KV de um evento de exclusão em uma operação de escrita de memória associativa.

1. Arquitetura de Dois Níveis:

   • L1 (Cache de Primeiro Nível): Um buffer de anel KV de tamanho fixo que mantém os $W$ tokens mais recentes. Isso suporta atenção softmax exata para contexto local.
   • L2 (Cache de Segundo Nível): Uma matriz de produto externo $A$ de tamanho fixo (memória de pesos rápidos) que armazena informações exclusivamente de tokens expulsos da janela L1.
   • Mecanismo de Escrita: Quando o buffer de anel sobrescreve uma entrada, o par $(k, v)$ deslocado não é descartado. Em vez disso, é escrito na matriz L2 $A$ por meio de uma atualização de produto externo: $A \leftarrow \lambda A + \eta (k \otimes v)$. Aqui, $\lambda$ (decaimento) e $\eta$ (taxa de escrita) são escalares aprendidos por cabeça.
   • Mecanismo de Leitura: Consultas futuras $q_t$ acessam o passado comprimido calculando $r_t = q_t A$. Aproveitando a identidade de atenção linear $q(k \otimes v) = \langle q, k \rangle v$, essa operação recupera efetivamente informações de todos os tokens expulsos sem armazená-los individualmente.
   • Fusão: Um portão escalar aprendido $g$ funde as saídas da atenção local (L1) e da leitura de memória (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Otimização de Treinamento (Varredura Paralela):
   O artigo identifica uma questão crítica com atalhos padrão de treinamento em blocos. Uma abordagem ingênua resume um bloco por sua média ($\bar{k} \otimes \bar{v}$), o que introduz $C^2 - C$ produtos cruzados de tokens espúrios ($k_i \otimes v_j$ onde $i \neq j$) que não existem durante a inferência em streaming.

   • Solução: Os autores implementam uma varredura ponderada paralela que é matematicamente equivalente a escritas sequenciais por token. Isso fecha a lacuna entre treinamento e inferência dentro da epsilon float32 (erro relativo < $10^{-7}$), garantindo que o modelo aprenda as dinâmicas associativas corretas sem o ruído de interferência entre tokens.

Principais Contribuições

• Memória Associativa Condicionada à Expulsão: Diferentemente de memórias de produto externo anteriores (como Infini-attention, mLSTM) que ingerem cada token, o TC usa a matriz de produto externo estritamente como um cache L2 alimentado apenas por expulsões de janela deslizante. Isso particiona limpiamente os tokens em "residentes no L1" e "comprimidos no L2".
• Correção de Artefatos de Treinamento: O artigo identifica e resolve um atalho de treinamento previamente negligenciado (média de blocos) que introduz produtos cruzados de tokens espúrios. A varredura ponderada paralela proposta elimina essa incompatibilidade com o mesmo custo computacional.
• Eficiência de Estado Limitado: O TC fornece um mecanismo onde o contexto mais antigo influencia previsões sem reter entradas individuais expulsas, oferecendo uma alternativa prática à retenção total de KV.

Resultados Experimentais
Os autores avaliaram o TC em escalas de sistemas, recuperação associativa sintética e modelagem de linguagem de contexto longo com texto real (OpenWebText e Shakespeare).

• Recuperação Associativa Sintética: Em tarefas que exigem recuperação de informações além da janela deslizante, o TC alcançou 100% de precisão, igualando o desempenho de KV Completo. Reduziu o estado de inferência retido em 72–84% em comparação com KV Completo, enquanto superou significativamente o Infini-attention (que degradou conforme as lacunas aumentavam) e as bases de janela deslizante (que performaram em níveis de acaso).
• Modelagem de Linguagem de Contexto Longo: No OpenWebText, o TC alcançou a menor Negativa Log-Verossimilhança (NLL) média em cada comprimento de contexto de avaliação de 1.024 a 32.768 tokens, com a única exceção de $L=2.048$, onde foi estatisticamente indistinguível do Infini-attention. Em $L=32.768$, o TC atingiu uma NLL de 5,14, superando KV Completo (6,00) e Infini-attention (5,42), enquanto usava 2,4$\times$ menos memória de pico da GPU que KV Completo.
• Vazão: O TC incorreu em um sobrecarga de vazão de ~30% em comparação com métodos puros de janela deslizante (devido à leitura/atualização L2), mas permaneceu consistentemente mais rápido que o Infini-attention em todos os comprimentos de contexto testados.
• Diagnósticos de Capacidade: Experimentos na capacidade bruta de memória associativa mostraram que o desempenho de recuperação degrada conforme mais associações não relacionadas são comprimidas no estado fixo, confirmando a natureza limitada da memória.

Significado e Alegações
O artigo afirma que o Tensor Cache conecta com sucesso a lacuna entre a eficiência de memória da atenção com janela deslizante e as capacidades de recuperação da atenção de contexto completo. Ao tratar a expulsão de KV como uma operação de memória estruturada em vez de uma exclusão destrutiva, o TC permite que Transformers acessem evidências relevantes fora da janela imediata sem o custo de memória linear de caches KV completos. Os autores enfatizam que sua principal contribuição não é a invenção da memória de produto externo em si, mas sua aplicação específica como um cache L2 condicionado à expulsão e a correção rigorosa de artefatos de treinamento que anteriormente impediam tais métodos. O TC é apresentado como uma alternativa prática de estado limitado que melhora a fronteira de qualidade-memória para inferência de contexto longo.