Resumo Técnico: DistillCache

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) baseados em Transformer dependem de caches de Chave-Valor (KV) para evitar a recomputação de representações de tokens durante a decodificação autorregressiva. No entanto, o tamanho do cache KV cresce linearmente com o comprimento da sequência, criando um gargalo severo de memória que limita a janela de contexto e o rendimento (throughput) de inferências de contexto longo.

As soluções existentes baseiam-se principalmente em estratégias de evasão heurísticas (ex: H2O, SnapKV) que retêm tokens com base em estatísticas de atenção estáticas ou sinais posicionais. O autor argumenta que essas heurísticas frequentemente falham em capturar a influência preditiva futura de um token, particularmente em cenários envolvendo dependências de longo alcance ou raciocínio estruturado, onde a importância de um token é contingente às consequências da geração subsequente, em vez de apenas aos padrões de atenção passados.

Metodologia: DistillCache

O artigo propõe o DistillCache, um framework de aprendizado por reforço (RL) que formula a evasão do cache KV como um problema de tomada de decisão sequencial sob orçamentos de memória rigorosos. Diferente dos métodos heurísticos, o DistлillCache aprende uma política para reter um subconjunto de tokens que preserva a distribuição de saída de um modelo com cache completo.

Componentes Principais

1. Espaço de Estados (Representação de Características):
   Para cada token no cache, a política observa um vetor de características de 5 dimensões derivado de sinais internos do modelo disponíveis durante a decodificação:

   • $\alpha_i$: Massa de atenção média recebida pelo token em todas as cabeças e camadas.
   • $\|v_i\|_2$: A norma $\ell_2$ do vetor de valor do token.
   • $\sigma_i$: Variância da massa de atenção recebida entre as cabeças.
   • $H_i$: Entropia média ponderada pela atenção das consultas (queries) recentes que atendem ao token (capturando a importância sensível ao contexto).
   • $pos_i$: Posição sequencial normalizada.

2. Espaço de Ações:
   Uma rede de política MLP (Multi-Layer Perceptron) leve atribui uma pontuação escalar a cada token no cache. Durante o treinamento, um ranking estocástico é induzido pela adição de ruído Gumbel a essas pontuações, permitindo que a política amostre um subconjunto de top-$B$ tokens (onde $B$ é o orçamento) a partir da distribuição Plackett–Luce. Na inferência, os top-$B$ tokens são selecionados deterministicamente.

3. Arquitetura de Treinamento (Configuração Professor-Aluno):
   Para treinar a política sem rótulos de verdade fundamental (ground-truth), o autor emprega uma abordagem de destilação professor-aluno:

   • Cache Professor: Mantém o histórico completo e não podado de tokens para computar a distribuição "ouro" do próximo token ($P_t$).
   • Cache Aluno: Mantém o histórico podado de acordo com as decisões da política para computar a distribuição podada ($\hat{P}_t$).
   • Sinal de Recompensa: A recompensa em cada etapa é definida como a divergência de Kullback-Leibler (KL) negativa entre as distribuições professor e aluno: $R_t = -D_{KL}(P_t \parallel \hat{P}_t)$. Este objetivo otimiza diretamente a fidelidade distributiva em cada etapa de decodificação, em vez de depender de sinais substitutos (proxies) como pontuações de atenção.

4. Otimização:
   A política é otimizada usando o algoritmo REINFORCE com regularização de entropia e um baseline de redução de variância (Leave-One-Out). O processo de treinamento é realizado offline, dobrando temporariamente a pegada de memória, mas a inferência requer apenas uma passagem direta (forward pass) com overhead negligenciável para extração de características e pontuação da política.

Contribuições Principais

• Framework de RL para Evasão de Cache: Introdução do DistillCache, um framework de RL de propósito geral que otimiza a retenção de tokens para fidelidade distributiva subsequente sob restrições de memória, indo além da pontuação heurística estática.
• Objetivo de Preservação de Distribuição: Demonstração de que a otimização da divergência KL por etapa produz ganhos mensuráveis sobre alternativas heurísticas e supervisionadas ao aprender explicitamente compensações (trade-offs) de longo horizonte.
• Avaliação Abrangente: Avaliação extensa no Mistral-7B-Instruct-v0.3 e transferência zero-shot para Llama-3-8B-Instruct em LongBench, benchmarks de raciocínio (GSM8K, BBH) e métricas de rendimento.

Resultados Experimentais

Avaliado no Mistral-7B-Instruct-v0.3 com um orçamento de cache KV de 25%:

• Desempenho no LongBench: O DistillCache retém 94,2% da precisão do cache total, superando fortes baselines heurísticos (H2O, SnapKV) em até 2,7 pontos absolutos e métodos de RL concorrentes (ForesightKV, RLKV) em até 1,4 ponto.
• Tarefas de Raciocínio: Em benchmarks de raciocínio, o DistillCache é competitivo com os melhores métodos concorrentes e os supera sob compressão agressiva (orçamento de 12,5%).
• Transferência entre Modelos: A política treinada no Mistral-7B transfere-se via zero-shot para o Llama-3-8B-Instruct, mantendo a competitividade com baselines heurísticos sem exigir o retreinamento por modelo (um requisito para os baselines de RL concorrentes).
• Eficiência: O método entrega até 2,1× mais rendimento (throughput) do que o cache total. Embora introduza um pequeno overhead (~1,4 ms/token) para extração de características e avaliação da política, permanece comparável a outros métodos aprendidos e é significancialmente mais rápido que a inferência de cache total.
• Análise de Erro: O DistillCache exibe as menores taxas de erros repetitivos e respostas incorretas em comparação com heurísticas e outros métodos de RL, sugerindo uma superior preservação da coerência da sequência e precisão factual.

Significância e Alegações

O artigo postula que o DistillCache demonstra a eficácia de políticas aprendidas e conscientes da distribuição para a inferência de LLMs de contexto longo com eficiência de memória. Ao otimizar a decisão de retenção de tokens contra a fidelidade distributiva futura, em vez de estatísticas de atenção retrospectivas, o framework aprende políticas limitadas por orçamento que levam em conta as consequências da geração subsequente.

O autor enfatiza que o DistillCache não é meramente um substituto aprendido para métricas heurísticas, mas uma abordagem distinta que otimiza diretamente a preservação da distribuição de saída do modelo. Ao reconhecer limitações como custos de treinamento (aprox. 130 horas de GPU) e uma pequena lacuna de desempenho com o RLKV em tarefas de raciocínio com orçamentos moderados, o artigo conclui que a evasão adaptativa de preservação de distribuição é uma direção promissora para permitir a inferência de contexto longo em hardware com recursos limitados. O trabalho sugere que combinar a fidelidade distributiva ao nível do token com estratégias de controle (gating) ao nível da cabeça (conforme explorado em experimentos híbridos) oferece um caminho complementar para otimizações futuras.