Resumo Técnico: KVpop — Compressão de Cache Chave-Valor com Poda Online Preditiva

1. Declaração do Problema

A decodificação autorregressiva em modelos de linguagem de grande escala (LLMs) baseados em Transformer depende de um cache de Chave-Valor (KV) para evitar a recomputação de representações de tokens passados. No entanto, o cache KV cresce linearmente com o comprimento do contexto, criando um gargalo significativo de memória e largura de banda durante a inferência de contexto longo.

As soluções existentes para reduzir esse custo geralmente se dividem em duas categorias:

• Evicção Estática/Heurística: Métodos como Janela Deslizante (Sliding Window) ou preservação de Tokens de Sumidouro (Sink-Token) (ex: StreamingLLM) dependem de regras fixas, enquanto métodos baseados em pontuação usam proxies online (ex: atenção cumulativa). Estes frequentemente falham em rastrear a utilidade futura dos tokens, levando a uma evicção frágil quando a relevância muda.
• Evicção Aprendida: Abordagens como Compressão de Memória Dinâmica (DMC) ou Esparsidade de Memória Dinâmica (DMS) treinam políticas para fundir ou descartar tokens. No entanto, muitos desses métodos tomam decisões de evicção no momento da inserção do token ou dependem de relaxamentos diferenciáveis que não supervisionam diretamente a decisão específica de "manter ou descartar" baseada na utilidade futura. Além disso, métodos com janelas protegidas atrasam a remoção de tokens, mas frequentemente decidem seu destino imediatamente após a entrada, ignorando evidências que se acumulam enquanto o token permanece na janela protegida.

O desafio central é prever quais tokens serão úteis para consultas futuras, pois tokens localmente salientes podem se tornar irrelevantes, enquanto tokens com baixa atenção imediata podem se tornar críticos mais tarde (ex: em traços de raciocínio).

2. Metodologia: KVpop

O KVpop (compressão de KV com Poda Online Preditiva) é um retrofit de atenção esparsa projetado para impor um orçamento fixo por cabeça de KV sem alterar a arquitetura do modelo base. Ele aborda as limitações de trabalhos anteriores através de dois mecanismos primários: Supervisão de Atenção Futura e Pontuação de Estado Atrasada.

2.1 Arquitetura e Orçamento

O KVpop mantém um cache limitado por cabeça de KV definido pelo orçamento $B = s + w + k$:

• $s$ Tokens de Sumidouro (Sink Tokens): Os primeiros tokens são sempre retidos.
• $w$ Janela Protegida: Os $w$ tokens mais recentes são protegidos da evicção.
• $k$ Top-K de Longo Alcance: Tokens mais antigos competem por um orçamento fixo de $k$ slots baseados em pontuações de importância aprendidas.

2.2 Supervisão de Alvo de Atenção Futura

Ao contrário de métodos que pontuam tokens com base em atenção passada ou estatísticas locais, o KVpop supervisiona o pontuador usando a massa de atenção futura que um token recebe após sair da janela protegida.

• Definição do Alvo: Para um token $t$, o alvo é a massa de log agregada de atenção que ele recebe de consultas futuras (posições $d \ge t + w$).
• Computação Eficiente: Para evitar o custo $O(S^2)$ de materializar mapas de atenção densos durante o treinamento, o KVpop utiliza uma passagem de atenção transposta. Ao trocar os papéis de queries e keys e reutilizar os normalizadores log-sum-exp (LSE) retornados pelos kernels de atenção esparsa, o sistema computa o alvo de atenção futura de forma eficiente sem a materialização de matrizes densas.
• Função de Perda: O objetivo de treinamento é uma perda logística pareada consciente de fronteira (boundary-aware). Ela compara a pontuação prevista de um novo token elegível ($t_{new}$) contra a pontuação do token de corte do professor ($t_{bnd}$). A perda força o aluno a classificar corretamente $t_{new}$ em relação à fronteira, garantindo que a seleção top-$k$ corresponda ao ranking de utilidade futura do professor.

2.3 Pontuação de Estado com Decisão Atrasada

O KVpop introduz um pontuador baseado em memória com atraso (implementado via um mLSTM) que adia a decisão de pontuação até que o token atinja a fronteira de evicção.

• Mecanismo: Em vez de pontuar um token imediatamente após sua inserção, o pontuador atualiza seu estado interno com a representação do token e o contexto subsequente. Quando o token sai da janela protegida ($q - w$), o pontuador lê a memória acumulada para gerar a pontuação de importância.
• Vantagem: Isso permite que o pontuador incorpore o contexto de "quase-futuro" (evidência que se acumula enquanto o token está na janela protegida) antes de tomar a decisão de manter ou descartar.

2.4 Eficiência de Inferência

Durante a inferência, o KVpop impõe o orçamento fixo $B$. A seleção dos $k$ tokens de longo alcance é otimizada usando um limiar contínuo sobre prioridades estáticas de tokens e uma árvore Fenwick, permitindo que a máscara esparsa seja gerada on-the-fly dentro do kernel de atenção (ex: FlexAttention) sem armazenar máscaras densas $S \times S$.

3. Principais Contribuições

1. Política de Evicção Supervisionada: Os autores formulam a evicção de KV com orçamento fixo como um problema de previsão supervisionada na fronteira de evicção, usando um alvo de atenção futura para treinar decisões top-$k$ de longo alcance.
2. Computação de Alvo Eficiente: Eles introduzem uma passagem de atenção transposta exclusiva para treinamento para computar alvos de atenção futura sem materializar mapas de atenção densos, garantindo que não haja overhead no tempo de inferência.
3. Pontuação Atrasada: Eles demonstram que um pontuador de estado pode explorar o contexto de quase-futuro ao adiar a decisão de pontuação até que o token se torne evictível, superando métodos de pontuação imediata.
4. Validação Empírica: O método é avaliado em raciocínio matemático (AIME, HMMT) e tarefas fora de domínio (GPQA-D, LiveCodeBench), mostrando que políticas de evicção aprendidas podem manter alto desempenho sob compressão agressiva.

4. Resultados Experimentais

Os autores avaliaram o KVpop nos modelos Qwen3-4B e Qwen3-8B, treinados no dataset Nemotron-Math v2 com comprimento de sequência de 16.384.

4.1 Raciocínio Matemático (AIME & HMMT)

• Taxa de Compressão 75% (Orçamento 2048):
  • Qwen3-4B: O KVpop reteve 98% do desempenho do professor de atenção total.
  • Qwen3-8B: O KVpop reteve 100% (desempenho próximo ao do professor).
  • Comparação: O KVpop superou significativamente os baselines livres de treinamento (StreamingLLM, TOVA) e o baseline aprendido DMS. Por exemplo, com 75% de compressão, o StreamingLLM caiu para ~76% de desempenho relativo no Qwen3-4B, enquanto o KVpop manteve 98%.
• Taxa de Compressão 88% (Orçamento 4096):
  • Qwen3-4B: O KVpop reteve 97% do desempenho do professor.
  • Qwen3-8B: O KVpop reteve 100% do desempenho do professor.
  • Comparação: Métodos livres de treinamento colapsaram significativamente (ex: StreamingLLM caiu para ~47% de desempenho relativo), enquanto o KVpop permaneceu robusto.

4.2 Generalização Fora de Domínio

Avaliado em GPQA-Diamond (raciocínio científico) e LiveCodeBench (geração de código):

• Modelos KVpop treinados apenas em dados de raciocínio matemático generalizaram bem para esses domínios.
• Com 75% de compressão, o KVpop no Qwen3-4B alcançou precisão comparável ao professor denso no GPQA-D (0.59 vs 0.59) e resultados competitivos no LCB.

4.3 Eficiência de Inferência

• Memória: Diferente da atenção densa, onde o uso de VRAM escala linearmente com o comprimento da geração (atingindo 36GB em 131k tokens), o KVpop e o DMS mantiveram uma pegada de memória quase constante (19GB), independentemente do comprimento da sequência.
• Latência: O KVpop demonstrou menor latência de decodificação end-to-end em comparação ao DMS e à atenção densa em comprimentos de geração longos (até 131k tokens). Os autores atribuem isso ao orçamento fixo por cabeça do KVpop, que evita as estruturas de cache "irregulares" do DMS que complicam a compilação e execução.

4.4 Estudos de Ablação

• Pontuação Atrasada: Remover o atraso (pontuando imediatamente após a inserção) resultou em uma queda de 0.2 pontos na precisão do token, confirmando o benefício de integrar o contexto de quase-futuro.
• Stateless vs. Stateful: A variante stateful (KVpop) geralmente igualou ou superou ligeiramente a variante stateless (KVpopmlp), particularmente em taxas de compressão mais altas.

5. Significância e Alegações

O artigo afirma que supervisionar a evicção com sinais de atenção futura é uma estratégia superior para a compressão de cache KV do que heurísticas estáticas ou scores de proxy. Ao treinar diretamente uma política para prever quais tokens importarão para consultas futuras, o KVpop alcança:

1. Alta Fidelidade: Mantém o desempenho próximo ao do professor mesmo em altas taxas de compressão (até 88%), um regime onde outros métodos falham.
2. Robustez: A política aprendida transfere-se efetivamente para tarefas fora de domínio, sugerindo que o modelo aprende princípios gerais de utilidade de tokens em vez de fazer overfitting a distribuições específicas de treinamento.
3. Eficiência: Impõe um limite rígido de memória, permitindo decodificação de contexto longo estável com uso de memória constante e melhor throughput em comparação com métodos de gating dinâmico como o DMS.

Os autores concluem que o KVpop demonstra que políticas de evicção aprendidas, quando adequadamente supervisionadas pela utilidade futura e permitem o uso de contexto de quase-futuro, podem resolver efetivamente o gargalo de memória na geração autorregressiva sem sacrificar a qualidade do modelo.