Resumo Técnico: Evicção de Cache KV Consciente de Epifania Sem a Matriz de Atenção

Declaração do Problema

À medida que modelos de raciocínio (ex: DeepSeek-R1) geram cadeias de pensamento que abrangem dezenas de milhares de tokens, o cache Key-Value (KV) cresce linearmente, tornando-se o principal gargalo de memória para implantação. As estratégias atuais de evicção de cache KV tipicamente classificam os tokens com base em pesos de atenção. Esta abordagem apresenta duas limitações críticas:

1. Procurador de Importância Ruidoso: Os pesos de atenção são indicadores pouco confiáveis da importância dos tokens em longos rastros de raciocínio. "Sumidouros de atenção" (attention sinks) absorvem peso independentemente do conteúdo, e tokens de preenchimento atraem peso durante a geração, mas raramente são referenciados posteriormente.
2. Incompatibilidade Arquitetural: Ler pesos de atenção requer a materialização da matriz de atenção $n \times n$. Isso força o uso de kernels "eager", desabilitando kernels fundidos como o FlashAttention. Consequentemente, o uso de memória escala quadraticamente ($O(n^2)$), causando erros de Out-of-Memory (OOM) em comprimentos de contexto (ex: 8192 tokens) que são muito menores do que os rastros de raciocínio típicos.

Métodos existentes que tentam mitigar isso (ex: ThinKV, RaaS) ainda dependem de sinais derivados da atenção, enquanto outros (ex: LagKV) são genéricos, mas carecem de otimização específica para raciocínio.

Metodologia: EPIKV

Os autores propõem o EPIKV (Epiphany-Aware KV Cache), uma política de evicção livre de treinamento que pontua tokens com base em mudanças na representação interna do modelo, em vez de pesos de atenção.

1. O Score de Epifania

Em vez da atenção, o EPIKV mede a mudança no fluxo residual (estado oculto) entre passos de decodificação consecutivos. Para um token $t$ na camada $l$, o sinal é definido pela diferença L2:
$$g_l(t) = \|h_l(t) - h_l(t-1)\|_2$$
Grandes mudanças indicam "tokens de epifania" — pontos de transição onde o modelo muda de um raciocínio exploratório para um convergente, conclui um passo ou se compromete com um insight.

2. Anatomia de Duas Bandas de Camadas

Através de experimentos de oclusão contrafactual, os autores identificam uma estrutura de camada específica em modelos de raciocínio de 32 camadas:

• Banda A (Camadas 7–13): Mudanças elevadas no estado oculto aqui correlacionam-se positivamente com a importância do token. Estas camadas alinham-se com locais conhecidos de recuperação de fatos e roteamento de características.
• Banda B (Camadas 18–25): Mudanças elevadas no estado oculto aqui correlacionam-se negativamente com a importância. Estas camadas superiores intermediárias são ativas mesmo para continuações fluentes e de baixa surpresa; portanto, grandes mudanças aqui sinalizam tokens previsíveis e dispensáveis.

O score final combina estas bandas:
$$s(t) = \bar{g}_{10}(t) - \bar{g}_{21}(t)$$
onde $\bar{g}$ representa uma média móvel sobre uma janela causal.

3. Correção de Tendência Temporal

O sinal bruto exibe uma tendência posicional monotônica (rastreando a posição tanto quanto o conteúdo). Para corrigir isso, os autores aplicam um z-score de rolagem causal:
$$z_l(t) = \frac{g_l(t) - \mu_l(t)}{\sigma_l(t) + \epsilon}$$
O score de evicção final é $z_{10}(t) - z_{21}(t)$. Esta normalização remove o viés posicional, garantindo que o score reflita a desviação de conteúdo local em vez da magnitude absoluta.

4. Compatibilidade com FlashAttention

Crucialmente, o EPIKV calcula os scores usando apenas:

1. Keys e Values em cache (já presentes na memória de alta largura de banda).
2. Estados ocultos por camada expostos por passagens de forward padrão.
   Ele nunca solicita output_attentions=True e nunca materializa a matriz de atenção. Isso permite que o EPIKV rode dentro de stacks de inferência FlashAttention padrão sem forçar um fallback para kernels eager.

Principais Contribuições

1. Identificação da Anatomia de Duas Bandas: Descoberta de que mudanças no estado oculto em camadas médias (7–13) correlacionam-se positivamente com a importância, enquanto mudanças em camadas superiores intermediárias (18–25) correlacionam-se negativamente, superando sinais baseados em atenção.
2. Pontuação Independente de Posição: Demonstração de que um z-score de rolagem causal remove tendências posicionais inerentes à variância do estado oculto, recuperando uma pontuação de evicção de alta qualidade.
3. Paridade de Desempenho sem Atenção: O EPIKV iguala ou excede os baselines mais fortes baseados em atenção (ThinKV, H2O) no MATH-500 e AIME-2024 com orçamentos de cache implantáveis, sem nunca materializar a matriz de atenção.
4. Eficiência de Engenharia: O método roda até 2.8× mais rápido que a evicção baseada em atenção em orçamentos iguais e evita o muro de memória que limita o comprimento de contexto em pontuações baseadas em atenção.

Resultados

Experimentos foram conduzidos no DeepSeek-R1-Distill-LLaMA-8B utilizando os datasets MATH-500 e AIME-2024.

• Acurácia:
  • Com um cache de 4096 tokens no MATH-500, o EPIKV atinge 72% de acurácia, igualando o ThinKV (71%) e superando o H2O (67%).
  • Com um cache de 8192 tokens no AIME-2024, a variante de EPIKV com normalização de lag atinge 37%, superando o melhor baseline baseado em atenção (33%).
• Velocidade:
  • No AIME-2024 (8192 tokens), o EPIKV processa problemas em 440 segundos, comparado a 721 segundos para o baseline de atenção mais rápido (ThinKV) e 1239 segundos para o RaaS. Isso representa um aceleração de 1.6× a 2.8×.
• Escalabilidade de Memória:
  • Enquanto a pontuação baseada em atenção esgota uma GPU A100 de 80 GB em aproximadamente ~8192 tokens devido à matriz de atenção $O(n^2)$, o EPIKV (via FlashAttention) escala para 65.536 tokens (um aumento de 16× no comprimento de contexto factível) utilizando apenas 48 GB de memória.
• Modos de Falha de Baselines: Métodos baseados em atenção como o H2O colapsam em rastros de raciocínio, produzindo saídas vazias ou texto não estruturado em orçamentos apertados (ex: 1024 tokens), em vez de degradar graciosamente.

Significância e Alegações

O artigo afirma que o EPIKV torna a evicção de cache KV implantável em stacks de serviço FlashAttention padrão. Ao desacoplar a pontuação de importância da matriz de atenção, o método:

• Elimina o custo arquitetural de materializar pesos de atenção.
• Fornece um sinal mais robusto para rastros de raciocínio onde os pesos de atenção são ruidosos.
• Permite janelas de contexto significativamente mais longas e maiores batches em hardware existente.

Os autores observam explicitamente que a anatomia de duas bandas é específica para matemática competitiva e muda com a dificuldade da tarefa (ex: movendo-se para camadas anteriores no GSM8K), sugerindo que o sinal rastreia a consolidação do raciocínio em vez de índices de camadas fixos. Eles liberam rótulos de importância contrafactual como um recurso de validação para trabalhos futuros.