Resumo Técnico: COMPILERKV

1. Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) enfrentam um gargalo de memória devido ao crescimento linear do cache Chave-Valor (KV) com o comprimento do contexto. A compressão de KV apenas no pré-preenchimento (prefill-only) emergiu como uma solução amigável para implantação, onde um subconjunto de tokens é selecionado no final da etapa de pré-preenchimento e congelado para toda a passagem de decodificação. Essa abordagem converte o uso de memória em tempo de execução em um custo fixo, com zero intervenção no tempo de decodificação.

No entanto, o artigo identifica um descompasso estatístico fundamental nos métodos existentes de apenas pré-preenchimento (por exemplo, SnapKV, PyramidKV, HeadKV). Esses métodos tentam estimar sinais corretivos — especificamente confiabilidade por cabeça (quais cabeças de atenção são críticas para a recuperação) e sensibilidade de compressão no nível do prompt (quanto um prompt específico pode ser comprimido) — online, usando estatísticas de um único prompt ruidoso. Os autores argumentam que esta é a unidade estatística errada:

• E1 (Ruído de prompt único): As pontuações de atenção por prompt são de cauda pesada e enviesadas por escala; grandes magnitudes não garantem importância semântica, e pequenos erros de estimativa causam erros de seleção irreversíveis.
• E2 (Heterogeneidade de cabeças): As cabeças de atenção são funcionalmente especializadas. Algumas são "críticas para recuperação", enquanto outras são consistentemente ruidosas. Agregação ingênua permite que cabeças ruidosas dominem as pontuações.
• E3 (Variância de risco no nível do prompt): Diferentes prompts têm limites de compressão diferentes. Uma regra de retenção fixa falha em prompts de alta entropia ou alta perplexidade, levando a uma perda catastrófica de informações.

Métodos existentes reestimam esses sinais online para cada prompt, falhando em aproveitar o fato de que esses sinais exibem alta regularidade entre prompts em vez de estabilidade dentro do prompt.

2. Metodologia: COMPILERKV

Os autores propõem o COMPILERKV, uma política de retenção de KV que compila tabelas corretivas offline a partir de um corpus de calibração, em vez de estimá-las online. A ideia central é que as classificações de confiabilidade das cabeças e o mapeamento de características de risco para limites seguros são priors arquitetônicos estáveis que podem ser aprendidos uma vez e reutilizados.

O framework consiste em três componentes acoplados operando durante a etapa de pré-preenchimento:

A. Utilidade de Token Estabilizada (Abordando E1)

Para remover picos transitórios de atenção e viés de escala entre cabeças sem aprendizado, o método calcula uma pontuação de utilidade livre de parâmetros:

• Atenção agregada por janela: Suaviza picos transitórios agregando a atenção sobre uma janela de observação $\Omega$ na cauda do prompt.
• Normalização relativa à cabeça: Normaliza as pontuações dos tokens pela norma média de valor da cabeça, removendo o viés de escala.
• Fórmula: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, onde $\alpha_t$ é a massa da janela e $\rho$ é a norma de valor normalizada.

B. Tabela de Heterogeneidade de Cabeças (Abordando E2)

Em vez de tratar todas as cabeças igualmente ou reestimar a importância online, o COMPILERKV compila uma tabela congelada $W_{head}[l, h]$ de pesos de confiabilidade contínua.

• Compilação Offline: Aprendida via Q-Learning Conservador (CQL) em um corpus de calibração. O estado é $(l, h, B_l)$, e a ação é um peso de confiabilidade discretizado.
• Mecanismo: Durante a inferência, a utilidade de um token é reescalada pela confiabilidade de sua cabeça. Crucialmente, a agregação usa max-pooling ponderado ($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$). Isso concede "poder de veto" às cabeças críticas para recuperação: se uma única cabeça crítica atender fortemente a um token, ele é retido, independentemente do consenso de cabeças ruidosas.

C. Porta de Limite Adaptativa ao Risco (Abordando E3)

Para lidar com prompts com limites de compressão variados, uma Porta de Limite Adaptativa ao Risco mapeia características de risco no nível do prompt para um limite de retenção.

• Sinais de Risco: Usa Entropia de Atenção (risco estrutural) e Perplexidade Local (risco semântico) calculadas sobre a janela de observação.
• Compilação Offline: Uma tabela de consulta (LUT) mapeia bins de risco discretizados $(b_{ent}, b_{ppl})$ para um limite de retenção $\tau^{(l)}$. Isso também é compilado via CQL.
• Viés Conservador: O objetivo do CQL penaliza limites agressivos em estados raros e de alto risco (alta entropia/perplexidade), enviesando a tabela para uma retenção mais segura onde heurísticas online normalmente falham.
• Seleção Elástica: A seleção final retém todos os tokens com utilidade acima do limite $\tau^{(l)}$. Se o conjunto de candidatos exceder o orçamento $B_l$, é limitado aos top-$B_l$ tokens. Se estiver abaixo, todos os candidatos são mantidos (sub-retenção elástica).

D. Fundamentação Teórica

O artigo fornece uma análise formal mostrando que o erro de aproximação é controlado pela massa de cauda descartada (Proposição 4.1). Argumenta-se que o erro de calibração de tabela finita é controlado no nível do prompt (através de prompts independentes) em vez de assumir independência dentro das linhas de atenção de um único prompt, reduzindo assim a variância.

3. Contribuições Principais

1. Reenquadramento do Problema: Identifica que a unidade estatística para correção na compressão apenas no pré-preenchimento deve ser a escala entre prompts, e não o prompt único.
2. Compilação Offline: Introduz um método para compilar a confiabilidade das cabeças e limites adaptativos ao risco offline, transformando-os em priors arquitetônicos portáteis.
3. Política Acoplada: Propõe uma otimização conjunta onde os pesos de confiabilidade das cabeças e os limites de risco são calibrados contra a mesma recompensa a jusante, em vez de serem módulos independentes.
4. Portabilidade: Demonstra que tabelas compiladas transferem-se entre corpora disjuntos e até entre diferentes famílias de modelos (por exemplo, LLaMA para Mistral) com degradação mínima de desempenho.

4. Resultados Experimentais

O método foi avaliado em quatro backbones de 7B–8B (LLaMA-3, Mistral, Qwen2, InternLM) nos benchmarks LongBench, RULER e Needle-in-a-Haystack (NIAH).

• LongBench (orçamento de 512 tokens): O COMPILERKV alcança SOTA comprimido em todos os quatro backbones, superando a melhor linha de base (por exemplo, ChunkKV, HeadKV) em uma média de +1,67 pontos (IC 95% [+1,08, +2,37]). Em alguns backbones, supera ligeiramente o desempenho do FullKV, o que os autores atribuem à remoção de ruído e não ao ganho de informação.
• Regimes de Pressão:
  • RULER (contexto de 128k): Sob uma razão fixa de cache 512/32k, o COMPILERKV mantém uma pontuação de 73, superando significativamente o SnapKV (38) e o KVZip (66,5), permanecendo abaixo do FullKV (79).
  • NIAH (32k): Alcança 0,89 de precisão contra 0,42 do SnapKV.
  • Agrupamento (Batching): Em entrada de 32k com tamanho de lote de 16, o COMPILERKV permanece viável, enquanto o FullKV esgota a memória (OOM).
• Transferibilidade:
  • Entre Corpora: As classificações de confiabilidade das cabeças mostram alta correlação de Spearman ($\bar{\rho} \approx 0,90$) entre corpora de calibração disjuntos.
  • Entre Modelos: A transferência direta de tabela entre modelos (por exemplo, LLaMA-3 para Mistral) incorre em uma perda de apenas 0,4–0,8 pontos no LongBench comparado à compilação específica do alvo.
• Eficiência: A sobrecarga online é mínima: $O(1)$ consultas de tabela por camada e uma única ordenação sobre o pool de candidatos. Alcança até 2,45× de vazão de decodificação comparado ao FullKV em contexto de 32k.

5. Significado e Afirmações

O artigo afirma que o COMPILERKV fecha a lacuna para o FullKV em orçamentos fixos ao corrigir o descompasso estatístico na forma como as políticas de retenção são computadas. Ao deslocar o cálculo dos sinais corretivos de uma estimativa online, ruidosa e de prompt único para um prior offline, compilado e entre prompts, o método estabiliza o processo de seleção contra ruído e heterogeneidade.

Os autores enfatizam que a novidade não é meramente "consciência de cabeças" (que métodos anteriores como HeadKV possuem), mas a compilação offline conjunta de confiabilidade contínua das cabeças e limites de risco de prompt. Isso transforma a confiabilidade das cabeças de um truque de ajuste em um objeto científico reutilizável que pode ser compilado uma vez e implantado em modelos e contextos com sobrecarga online de $O(1)$. O trabalho sugere que, para compressão apenas no pré-preenchimento, os sinais "certos" são propriedades arquitetônicas do modelo, e não características transitórias do prompt de entrada.