Resumo Técnico: ForesightKV

Definição do Problema

Modelos de Linguagem de Grande Escala (LLMs), particularmente modelos de raciocínio, demonstraram capacidades excepcionais na geração de longos rastros de raciocínio. No entanto, à medida que o comprimento da sequência aumenta, o cache de Chave-Valor (KV) expande-se linearmente, levando a uma sobrecarga significativa de memória e ao aumento da latência devido à decodificação limitada pela memória. Embora os métodos existentes de expulsão de cache KV tentem mitigar isso descartando pares KV menos importantes com base em regras heurísticas (por exemplo, pontuações de atenção, posições ou características), eles frequentemente falam em capturar dependências complexas e dinâmicas entre diferentes cabeças de atenção. Especificamente, esses métodos livres de treinamento têm dificuldade com padrões "semanticamente dependentes", onde a importância de um par KV muda dinamicamente com base no contexto em evolução. Além disso, a expulsão de pares KV críticos pode causar aumentos abruptos na perda para tokens de baixa entropia (partes determinísticas do raciocínio), levando a erros factuais que distorcem etapas subsequentes de raciocínio.

Metodologia

Os autores propõem o ForesightKV, um framework baseado em treinamento projetado para aprender a contribuição de longo prazo dos pares KV e expulsar adaptativamente os menos importantes durante a geração. O framework opera sob um orçamento de memória fixo e emprega um paradigma de treinamento em dois estágios para otimizar um modelo de pontuação leve ($\pi_\theta$) que prevê a importância de cada par KV. Crucialmente, os parâmetros do LLM permanecem congelados durante todo o processo; apenas o modelo de pontuação é atualizado.

1. Processo de Expulsão Dinâmica

A expulsão é acionada periodicamente (a cada $L$ tokens). Quando o tamanho do cache atinge $B + L$ (onde $B$ é o orçamento alvo), o sistema retém os $L$ tokens mais recentes e usa o modelo de pontuação para podar o restante do cache de volta para $B - L$. O modelo de pontuação recebe como entrada a concatenação da chave ($k_n$), valor ($v_n$) e características de atenção ($a_n$) para prever uma pontuação de importância ($\phi_n$). As decisões de expulsão são feitas via uma estratégia de Top-K seguida de amostragem multinomial para introduzir estocasticidade.

2. Estágio Um: Aprendizado Supervisionado com "Expulsão de Ouro" (Golden Eviction)

Para treinar o modelo de pontuação, os autores introduzem o algoritmo Golden Eviction para construir rótulos de supervisão.

• Mecanismo: Para um rastro de raciocínio completo, a matriz de atenção é particionada em blocos. O algoritmo calcula "pontuações de bloco" agregando pontuações de atenção dentro dos blocos e entre as cabeças de atenção.
• Consciência Futura: Para cada etapa de expulsão, o algoritmo identifica a pontuação máxima de bloco em todas as etapas futuras para cada par KV. Esta "pontuação futura" representa a estratégia de retenção ideal que minimiza o impacto nas computações de atenção futuras.
• Treinamento: O modelo de pontuação é treinado usando uma Perda de Ranking Pareado (Pairwise Ranking Loss) para garantir que as pontuações previstas para quaisquer dois pares KV se alinhem com o ranking de suas pontuações futuras derivadas do algoritmo Golden Eviction.

3. Estágio Dois: Aprendizado por Reforço (RL)

O treinamento supervisionado sozinho pode não considerar o deslocamento de distribuição causado pela expulsão durante a inferência. Para preencher essa lacuna, os autores formulam a expulsão de cache KV como um Processo de Decisão de Markov (MDP) e aplicam o algoritmo GRPO (Group Relative Policy Optimization).

• Estado: O cache KV restante atual.
• Ação: Selecionar um subconjunto de pares KV para reter.
• Recompensa: A função de recompensa é projetada para mitigar a degradação de desempenho em tokens de baixa entropia. Especificamente, ela foca em tokens que originalmente tinham baixa entropia, mas experimentaram um aumento significativo na perda após a expulsão. A recompensa é definida como o Erro Quadrático Médio (MSE) negativo dos aumentos de perda para este subconjunto específico de tokens.
• Otimização: O modelo de política é otimizado para maximizar a recompensa cumulativa ao longo da sequência, aprendendo efetivamente a preservar o contexto crítico necessário para etapas de raciocínio determinísticas.

Principais Contribuições

1. Framework ForesightKV: Um novo método de treinamento em dois estágios (Aprendizado Supervisionado + RL) que aprende a prever a contribuição de longo prazo dos pares KV, indo além das regras heurísticas estáticas.
2. Algoritmo Golden Eviction: Um método para gerar rastros de expulsão ideais usando informações de atenção futura, servindo como supervisão de alta qualidade para o modelo de pontuação.
3. Formulação de MDP para Expulsão: Modelar o processo de expulsão como um MDP e utilizar GRPO para visar especificamente a mitigação de picos de perda em tokens de baixa entropia, que são críticos para a precisão do raciocínio.
4. Eficiência: O framework treina apenas um modelo de pontuação leve sem exigir operações de retropropagação ou atualizações de parâmetros no LLM base.

Resultados Experimentais

Os autores avaliaram o ForesightKV em três modelos de raciocínio (Qwen3-1.7B, Qwen3-4B e DeepSeek-R1-Distill-Qwen-7B) utilizando os benchmarks matemáticos AIME2024 e AIME2025.

• Desempenho: O ForesightKV superou consistentemente as linhas de base existentes (SnapKV, H2O, R-KV). Notavelmente, alcançou desempenho comparável ou superior aos baselines de cache total enquanto utilizava apenas metade do orçamento de cache KV. Por exemplo, no AIME2024 com Qwen3-4B, o ForesightKV com um orçamento de 1K superou o R-KV com um orçamento de 2K (54.5 vs. 44.8 Pass@1).
• Throughput: O método melhorou significamente o throughput. Com um orçamento de 1K em uma sequência de 32K, o ForesightKV alcançou até 9.79× ganho de throughput em comparação com o cache KV total.
• Generalização: Apesar de ter sido treinado principalmente em tarefas matemáticas, o ForesightKV demonstrou forte generalização nos benchmarks GPQA (ciência) e LiveCodeBench (codificação), superando as linhas de base com metade do orçamento de cache.
• Estudos de Ablação: Os resultados confirmaram que tanto o estágio de treinamento supervisionado (fornecendo capacidade preditiva inicial) quanto o estágio de aprendizado por reforço (refinando a política para preservação de tokens de baixa entropia) são cruciais para o sucesso. O design específico de recompensa focando em tokens de baixa entropia com grandes aumentos de perda foi considerado o mais eficaz.

Significância e Alegações

O artigo afirma que o ForesightKV oferece uma nova direção promissora para o gerenciamento adaptativo de cache KV em contextos longos de raciocínio. Ao mudar de heurísticas baseadas em regras para uma abordagem baseada em aprendizado que modela explicitamente a contribuição de longo prazo dos pares KV e visa a preservação de tokens de baixa entropia, o método alcança um equilíbrio superior entre eficiência de memória e qualidade de geração. Os autores enfatizam que sua abordagem captura a importância intrínseca das entradas do cache KV em vez de sofrer overfitting para restrições de orçamento específicas, permitindo que generalize bem através de diferentes configurações de orçamento e domínios de tarefas.