Resumo Técnico: MOMENTKV

Definição do Problema

A decodificação autorregressiva em modelos de linguagem baseados em Transformer depende do cache de Chave-Valor (KV), que armazena projeções computadas anteriormente para reutilizá-las em etapas de decodificação. À medida que os comprimentos de sequência escalam para dezenas ou centenas de milhares de tokens, a pegada de memória do cache KV cresce linearmente, tornando-se o principal gargalo para a inferência de contexto longo.

A evasão (eviction) do cache KV aborda isso retendo um subconjunto de tamanho fixo de pares chave-valor e descartando o restante. Embora métodos existentes tenham melhorado as estratégias de seleção de tokens (por exemplo, com base em pesos de atenção cumulativos ou pontuação consciente de valor), eles compartilam uma limitação fundamental: a inferência pós-evasão renormaliza a atenção exclusivamente sobre os tokens retidos, descartando completamente o conjunto evadido.

Os autores identificam que a principal fonte de degradação da saída não é meramente a massa de atenção residual nos tokens evadidos (que os métodos existentes já minimizam), mas sim um desajuste direcional entre o conjunto de tokens retidos e o evadido. A análise empírica revela que, na prática, os tokens evadidos são frequentemente quase ortogonais aos retidos. Consequentemente, mesmo uma pequena quantidade de massa de atenção evadida pode produzir um erro direcional substancial que a renormalização padrão não consegue recuperar. Isso cria um "gap direcional" onde a melhoria apenas na seleção de tokens gera retornos decrescentes, pois a retenção mais seletiva concentra o conjunto retido em direções específicas alinhadas à consulta (query), ampliando ainda mais o gap com o conjunto evadido.

Metodologia: MOMENTKV

Para abordar este gap direcional, os autores propõem o MOMENTKV, um framework que mantém estatísticas de momentos compactas e de pequeno tamanho sobre o conjunto de tokens evadidos. Essas estatísticas servem a um duplo propósito: guiar as decisões de evasão e corrigir a saída da atenção durante a inferência.

1. Estatísticas de Momentos Compactas

Em vez de armazenar os tokens evadidos, o MOMENTKV mantém quatro somas acumuladas por cabeça de atenção, atualizadas incrementalmente a cada evasão:

• Contagem ($n_e$): O número de tokens evadidos.
• Soma de Chaves ($s_k$): $\sum_{i \in E} k_i$.
• Soma de Valores ($s_v$): $\sum_{i \in E} v_i$.
• Soma de Produto Externo ($S$): $\sum_{i \in E} v_i k_i^\top$.

Essas estatísticas permitem o cálculo da média de chave ($\bar{k}$), da média de valor ($\bar{v}$) e da covariância valor-chave ($\tilde{S}$) sem revisitar os tokens evadidos. O overhead de armazenamento total é $O(d^2)$ por cabeça, independente do comprimento do contexto.

2. Aproximação de Primeira Ordem da Saída Evadida

Durante a inferência, o MOMENTKV deriva uma aproximação de forma fechada de primeira ordem da saída de atenção evadida, $\hat{f}_E(q)$. Ao recentralizar os logits em torno da média da chave e aplicar uma expansão de Taylor de primeira ordem à função softmax, os autores derivam:
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• O termo $\bar{v}$ atua como uma linha de base independente da consulta, recuperando a direção média da sub-saída evadida.
• O termo $\tilde{S} q$ atua como uma correção adaptativa à consulta, capturando a variação perpendicular (conteúdo direcional) que a linha de base sozinha não captura.

Esta aproximação permite que o modelo reconstrua uma estimativa da contribuição perdida dos tokens evadidos, restaurando efetivamente a informação direcional que a evasão padrão oblitera.

3. Evasão Informada por Momentos

As estatísticas de momentos também guiam o próprio processo de evasão. Para cada token retido $j$, o algoritmo calcula um resíduo de momento:
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
Este resíduo mede quanta informação direcional nova o token $j$ carrega além do que o resumo acumulado já captura. A pontuação de evasão é definida como o produto do peso de atenção e a norma do resíduo: $\text{score}(j) = \alpha_j \cdot \|r_j\|$.

• Tokens com baixos resíduos são bem previstos pelo modelo de momentos e são priorizados para evasão.
• Tokens com altos resíduos (carregando informação direcional única) são retidos mesmo que seus pesos de atenção sejam menores.

Isso cria um loop de reforço mútuo: a evasão informada por momentos mantém o conjunto evadido "geometricamente regular" (baixa dispersão de logit $\sigma$), o que estreita o limite de aproximação para a correção de inferência. Por sua vez, a correção precisa fornece melhores resíduos para futuras decisões de evasão.

4. Inferência Corrigida por Normalização

A saída de atenção final é reconstruída como uma combinação convexa da saída retida ($f_R$) e da saída evadida aproximada ($\hat{f}_E$):
$$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
O peso de mistura $\hat{w}_R$ é estimado usando um limite inferior na função de partição evadida derivado pela desigualdade de Jensen. Isso garante que o modelo considere a magnitude da massa evadida enquanto corrige sua direção.

Principais Contribuições

1. Identificação do Gap Direcional: O artigo estabelece que a principal fonte de erro na evasão do cache KV é a quase ortogonalidade entre as sub-saídas retidas e evadidas, não apenas a magnitude da massa de atenção evadida.
2. Estatísticas de Momentos Compactas: Introdução de um método eficiente em termos de memória para rastrear as propriedades geométricas (média e covariância) dos tokens evadidos usando apenas $O(d^2)$ de armazenamento por cabeça.
3. Correção de Forma Fechada: Derivação de uma aproximação de primeira ordem para a saída de atenção evadida que recupera tanto a direção média quanto as variações dependentes da consulta.
4. Mecanismo de Papel Duplo: Um framework unificado onde as estatísticas de momentos otimizam simultaneamente a seleção de evasão (identificando redundâncias) e corrigem as saídas de inferência (restaurando a informação direcional perdida).

Resultados Experimentais

Os autores avaliaram o MOMENTKV nos benchmarks LongBench e RULER utilizando LLaMA-3.1-8B-Instruct e Qwen3-4B-Instruct.

• Desempenho: O MOMENTKV superou consistentemente todos os baselines (H2O, SnapKV, PyramidKV, Ada-KV) em todos os orçamentos de cache ($L \in \{128, 256, 512, 1024\}$).
• Compressão Agressiva: Os maiores ganhos foram observados sob compressão agressiva. No LongBench com $L=128$ e LLaMA-3.1-8B, o MOMENTKV alcançou uma melhoria média de +1.35 no escore em relação ao baseline mais forte (Ada-KV), retendo 94.1% do desempenho do cache total.
• Tarefas de Recuperação: No RULER, que testa a recuperação precisa de tokens, o MOMENTKV mostrou ganhos proporcionais de +3.4 e +3.5 sobre o Ada-KV para LLaMA e Qwen, respectivamente, validando sua eficácia em configurações intensivas de recuperação.
• Estudos de Ablação:
  • A Inferência Corrigida por Normalização (NC) contribuiu significativamente mais (+2.41) do que a Evasão Informada por Momentos (MI) sozinha (+0.57), confirmando que recuperar a direção evadida é o principal motor de desempenho.
  • A combinação de ambos os mecanismos gerou um efeito sinérgico (+3.17 de ganho total), excedendo a soma dos ganhos individuais.
• Eficiência: O método adiciona apenas 1.6 ms de latência por token e requer 4.1 MB de memória auxiliar (6.4% do cache retido), enquanto reduz a latência por token em mais de 65% em comparação com um cache total.

Significância

O artigo argumenta que o caminho para melhorias futuras na inferência de contexto longo não reside apenas em uma melhor seleção de tokens, mas em recuperar a informação direcional que as estratégias de evasão atuais obliteram. O MOMENTKV demonstra que, ao manter estatísticas compactas sobre o conjunto evadido, é possível fechar o gap direcional, transformando o "custo irreduzível" da evasão em um erro de aproximação corrigível. Esta abordagem oferece uma solução prática e escalável para implantar modelos de linguagem de contexto longo sem sacrificar a eficiência dos caches KV de orçamento fixo.