Resumo Técnico: KVzap - Poda de KV Cache Rápida, Adaptativa e Fiel

1. Declaração do Problema

À medida que os comprimentos de contexto em modelos de linguagem baseados em transformers se expandem para dezenas ou centenas de milhares de tokens, o cache de Chave-Valor (KV) tornou-se um gargalo dominante para a eficiência da inferência de Grandes Modelos de Linguagem (LLMs). O cache KV, que armazena estados de atenção intermediários, consome uma quantidade significativa de memória GPU (por exemplo, 335 GB para um modelo de 65B com comprimento de sequência de 128k) e aumenta o tempo para o primeiro token (time-to-first-token), reduzindo o throughput de decodificação.

Embora modificações arquiteturais (ex: Grouped Query Attention, Multi-head Latent Attention) e mecanismos de atenção esparsa tenham abordado a compressão ao longo dos eixos de cabeça ($H$), dimensão ($D$) e camada ($L$), o eixo de comprimento de sequência ($T$) permanece amplamente não abordado por métodos amplamente adotados. Os métodos ad-hoc existentes frequentemente falham em atender a quatro critérios críticos para implantação em produção:

1. Rápido e leve: Sobrecarga (overhead) desprezível.
2. Independente de fase: Aplicável tanto ao preenchimento (prefilling - contexto longo) quanto à decodificação (raciocínio).
3. Amigável à otimização: Compatível com kernels como FlashAttention2 e PagedAttention.
4. Fiel: Degradação mínima de precisão.

Métodos de estado da arte atuais, como o KVzip, alcançam alta compressão, mas exigem uma passagem direta dupla (double forward pass) computacionalmente cara em prompts estendidos, tornando-os inadequados para inferência online e tarefas de decodificação.

2. Metodologia

Os autores introduzem o KVzap, uma aproximação rápida e adaptativa ao input do KVzip que opera durante o preenchimento e a decodificação. A metodologia consiste em três componentes principais:

2.1. Política de Pontuação Aprimorada (KVzip+)

Os autores primeiro refinam o mecanismo de pontuação do KVzip. O KVzip pontua tokens com base na "reconstrução de contexto", medindo o peso máximo de atenção que um token recebe quando o prompt é repetido. Os autores aprimoram isso incorporando um termo de normalização inspirado em Devoto et al. (2025), que contabiliza a contribuição do vetor de valor de um token para o fluxo residual (residual stream).
A pontuação KVzip+ ($s^+_i$) é definida como:
$$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$$
onde $a_{ji}$ é o peso de atenção, $W_O$ é a matriz de projeção de saída, $v_i$ é o vetor de valor e $h_j$ é o estado oculto (hidden state).

2.2. Modelagem Substituta (Surrogate Modeling)

Para evitar o custo proibitivo da passagem direta dupla do KVzip, o KVzap treina modelos substitutos leves (uma camada linear ou um MLP de duas camadas) para prever os log-scores ($\log(s^+)$) diretamente a partir dos estados ocultos de entrada ($h$) do modelo.

• Input: Estados ocultos $h_t \in \mathbb{R}^{D_h}$.
• Output: Pontuações para cada cabeça KV em $\mathbb{R}^H$.
• Treinamento: Os modelos são treinados em 1,2 milhão de pares de (estado oculto, log-score) amostrados do Nemotron-Pretraining-Dataset. A variante MLP utiliza uma largura oculta de $D_h/8$ com ativação GELU.
• Desempenho: Os substitutos alcançam uma correlação média $R^2$ de 0,60–0,80 com o oráculo KVzip+, demonstrando que scores caros podem ser aproximadamente previstos a partir de estados ocultos.

2.3. Poda e Expulsão Adaptativa

Ao contrário do KVzip, que impõe um orçamento fixo (ex: manter exatamente 50% dos tokens), o KVzap emprega poda baseada em limiar (threshold-based pruning):

• Limiarização (Thresholding): Tokens com pontuações previstas abaixo de um limiar fixo $\tau$ são descartados. Isso permite que a razão de compressão se adapte dinamicamente à densidade de informação do prompt de entrada.
• Janela Deslizante (Sliding Window): Para preservar o contexto local, os últimos $w=128$ tokens são sempre mantidos, seguindo a abordagem do StreamingLLM.
• Implementação: A poda ocorre após a operação de atenção. O processo é computacionalmente eficiente, adicionando uma sobrecarga desprezível (limitada a 1,1% para modelos MLP e 0,02% para modelos Lineares).

3. Principais Contribuições

1. KVzip+: Uma métrica de pontuação melhorada para a importância do cache KV que incorpora a normalização do estado oculto.
2. Aproximação Substituta: Uma demonstração de que os scores do KVzip+ podem ser efetivamente aproximados por modelos leves treinados em estados ocultos, eliminando a necessidade de passagens diretas duplas.
3. KVzap: Uma nova técnica de poda que aplica esses substitutos com expulsão baseada em limiar e janelas deslizantes, alcançando desempenho de estado da arte enquanto atende a todos os quatro critérios práticos (velocidade, independência de fase, amigabilidade à otimização e fidelidade).

4. Resultados Experimentais

Experimentos foram conduzidos no Qwen3-8B, Llama-3.1-8B-Instruct e Qwen3-32B através de três benchmarks: RULER (recuperação/agregação de contexto longo), LongBench (multitarefa de contexto longo) e AIME25 (raciocínio matemático).

• Compressão vs. Precisão: O KVzap alcança 2–4× de compressão de cache KV com perda de precisão insignificante.
  • No RULER 4k, o KVzap iguala ou supera o desempenho do oráculo KVzip+ e supera outros 15 métodos de poda (incluindo Expected Attention e Duo Attention).
  • No LongBench, o KVzap mantém precisão quase perfeita até 2–3× de compressão.
  • No AIME25, o KVzap-MLP preserva a precisão do raciocínio mesmo ao descartar mais de 50% do cache KV.
• Adaptabilidade: A abordagem baseada em limiar ajusta automaticamente as razões de compressão com base na dificuldade da tarefa. Por exemplo, ele alcança maior compressão em dados sintéticos do RULER comparado aos dados mais densos e reais do LongBench.
• Eficiência: O método adiciona sobrecarga de computação e memória desprezíveis. Durante a decodificação, os FLOPs adicionais utilizam ciclos da GPU que de outra forma estariam ociosos devido às limitações de largura de banda de memória.

5. Significância e Alegações

O artigo afirma que o KVzap preenche a lacuna entre a pesquisa acadêmica de poda e os motores de inferência do mundo real. Sua significância reside em:

• Praticidade: É o primeiro método a alcançar compressão de estado da arte sem sacrificar a velocidade ou exigir mudanças arquiteturais, tornando-o compatível com os principais motores de inferência (vLLM, SGLang, TRT-LLM).
• Fidelidade: Demonstra que os LLMs não exploram totalmente o cache KV e que pares não utilizados podem ser identificados via estados ocultos sem o retreinamento end-to-end do modelo base.
• Escalabilidade: A abordagem é validada em modelos variando de 8B a 32B de parâmetros e é projetada para ser compatível com futuras arquiteturas de atenção esparsa.

Os autores reconhecem limitações, observando que o KVzap é uma adição post-hoc e não um objetivo treinado end-to-end, e que mais engenharia é necessária para traduzir as razões de compressão em ganhos de velocidade de tempo real via kernels otimizados (ex: PagedAttention). No entanto, eles postulam que a combinação de simplicidade, alta compressão e desempenho robusto do KVzap o torna um forte candidato para implantação em produção.