DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

1. Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) enfrentam um gargalo crítico de memória durante a inferência de longo contexto. Embora a geração autoregressiva dependa do cache de Chave-Valor (KV) para evitar a recomputação de tokens passados, a pegada de memória desse cache cresce linearmente com o comprimento da sequência ($O(L \cdot N \cdot d)$, onde $L$ é o número de camadas, $N$ é o comprimento da sequência e $d$ é a dimensão oculta).

As soluções existentes de pós-treinamento utilizam principalmente o pruning de cache KV, que descarta tokens com baixas pontuações de atenção para caber dentro de um orçamento de memória fixo. No entanto, uma limitação fundamental dos métodos atuais (por exemplo, H2O, StreamingLLM) é a suposição de pruning uniforme: eles aplicam a mesma taxa de pruning em todas as camadas do transformador. Os autores argumentam que isso é subótimo, pois assume implicitamente que todas as camadas contribuem igualmente para o desempenho do modelo, ignorando a sensibilidade heterogênea de diferentes camadas à remoção de tokens.

2. Metodologia: DepthKV

O framework proposto, DepthKV, desafia a suposição de pruning uniforme ao introduzir alocação de cache KV dependente da camada. Em vez de descartar uma porcentagem fixa de tokens por camada, o DepthKV realoca um orçamento global de KV fixo entre as camadas com base em sua sensibilidade específica ao pruning.

A. Pré-estudo: Análise de Sensibilidade da Camada

Antes de projetar a estratégia de alocação, os autores conduziram um estudo de ablação camada por camada:

• Método: Eles realizaram o pruning de tokens em apenas uma camada específica por vez, mantendo as outras intactas, medindo a queda de desempenho resultante (ROUGE-1).
• Descoberta: A degradação de desempenho é altamente não uniforme. Certas camadas "críticas" causam quedas massivas quando submetidas ao pruning, enquanto outras mostram impacto negligenciável ou até redundância.
• Amplificação de Conteúdo: Eles identificaram "camadas de amplificação de conteúdo". O pruning dessas camadas não apenas reduz a precisão, mas suprime significativamente o comprimento e a informatividade da saída (medido pelo YapScore), indicando que essas camadas são cruciais para sustentar a geração.
• Validação Estatística: Testes de permutação confirmaram que a variação na importância das camadas é estatisticamente significativa entre modelos (Gemma, LLaMA, Qwen) e conjuntos de dados.

B. Métricas de Representação para Importância

Para automatizar a alocação sem ablação exaustiva, os autores analisaram propriedades de representação para encontrar um proxy para a importância da camada.

• Seleção de Métricas: Eles avaliaram seis métricas (Entropia Espectral, Rank Efetivo, Curvatura, DiME, LiDAR e InfoNCE).
• Descoberta Chave: InfoNCE (uma perda contrastiva que mede a robustez da representação a perturbações de entrada) mostrou a correlação negativa mais forte com a degradação de desempenho.
  • Interpretação: Camadas com menores valores de InfoNCE (representações mais robustas) são mais críticas. O pruning delas causa quedas maiores de desempenho.
  • Uso: O DepthKV utiliza pontuações de InfoNCE para classificar as camadas e determinar a alocação de orçamento.

C. Estratégias de Alocação

O DepthKV opera sob uma taxa global de pruning fixa ($\rho$), mas a distribui de forma não uniforme. Três estratégias são propostas:

1. Proteção de Camada Média (MLP): Uma abordagem heurística que preserva um subconjunto fixo de camadas médias (por exemplo, camadas $\lfloor L/2 \rfloor$ e $\lfloor L/2 \rfloor + 1$) enquanto realiza pruning uniforme nas demais.
2. Alocação Guiada por Métrica (MGA): O método principal proposto. Ele calcula uma pontuação de importância $s^{(l)}$ baseada no InfoNCE. Camadas com maior robustez (menor InfoNCE) recebem um orçamento de KV maior (menor taxa de pruning), enquanto camadas menos críticas são submetidas a um pruning mais agressivo.
3. Alocação de Métrica de Camada Média (MLMA): Uma abordagem híbrida que protege estruturalmente as camadas médias, enquanto distribui o orçamento restante entre as outras camadas com base nas pontuações de InfoNCE.

3. Contribuições Principais

1. Insight Teórico: Demonstrou que as camadas do transformador exibem sensibilidade estatisticamente significativa e não uniforme ao pruning de cache KV, invalidando a suposição de pruning uniforme.
2. Descoberta de Amplificação de Conteúdo: Identificou camadas específicas cujo pruning suprime a geração de conteúdo, ligando a robustez da representação diretamente à qualidade e ao comprimento da saída.
3. Proposta de Framework: Introduziu o DepthKV, um framework de alocação dependente da camada que redistribui um orçamento de memória fixo com base na importância da camada (especificamente usando InfoNCE).
4. Validação Empírica: Mostrou que o DepthKV supera consistentemente o pruning uniforme em diversas tarefas e modelos sem exigir treinamento adicional.

4. Resultados Experimentais

Os autores avaliaram o DepthKV em três famílias de modelos (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) em sete conjuntos de dados:

• Tarefas: Resumo de documentos longos (arXiv, PubMed, GovReport, LegalCase), QA fundamentada em documentos (HotpotQA, Qasper) e Raciocínio Matemático (GSM-$\infty$).
• Configuração: Redução global fixa de KV de 60% (retenção de 40%).
• Desempenho:
  • Resumo: O DepthKV (especificamente MGA) melhorou significativamente o ROUGE-1 e a similaridade SBERT em comparação com o pruning uniforme (por exemplo, +3.0 ROUGE-1 no arXiv para Gemma).
  • QA e Raciocínio: Variantes do DepthKV alcançaram pontuações mais altas de Correspondência Exata (EM) e F1. Por exemplo, no Qasper, o MLMA-6L superou todas as linhas de base.
  • Robustez: Os benefícios estenderam-se às tarefas de raciocínio (GSM-$\infty$), onde todas as variantes do DepthKV superaram o pruning uniforme.
  • LLM como Juiz: Avaliações usando o Prometheus confirmaram que as saídas do DepthKV foram mais corretas, completas e concisas do que as linhas de base de pruning uniforme.
• Ablação: O estudo confirmou que como o orçamento é distribuído (alocação por camada) importa mais do que o estimador específico de importância do token (por exemplo, H2O com V vs. sem V).

5. Significado e Limitações

Significado:

• Eficiência: O DepthKV oferece uma solução simples de pós-treinamento para o gargalo de memória de longo contexto, permitindo janelas de contexto mais longas em hardware existente sem retreinamento.
• Generalizabilidade: O método é agnóstico à arquitetura e eficaz em diferentes famílias de modelos e tipos de tarefas.
• Mudança de Paradigma: Move o campo do pruning em nível de "token" para alocação de recursos em nível de "camada", reconhecendo a heterogeneidade funcional das camadas do transformador.

Limitações:

• Não Consciente da Consulta (Non-Query-Aware): O DepthKV estima a importância do token durante a fase de preenchimento (prefill) com base na atenção histórica. Ele não se ajusta dinamicamente à consulta de decodificação específica, potencialmente perdendo tokens que se tornam críticos apenas para consultas futuras específicas.
• Granularidade em Nível de Cabeça (Head-Level): O framework atual aloca orçamentos no nível de camada, mas agrega a atenção através das cabeças. Ele não leva em conta a variabilidade intra-camada onde cabeças de atenção específicas podem ser mais críticas do que outras.

Conclusão

O DepthKV estabelece que nem todas as camadas do transformador são iguais na inferência de longo contexto. Ao realocar orçamentos de memória para preservar as camadas mais sensíveis e robustas (identificadas via InfoNCE), o método alcança trade-offs de desempenho por memória superiores em comparação com o pruning uniforme, oferecendo um caminho prático para a escalabilidade das janelas de contexto de LLMs.