Resumo Técnico: xKV - Compressão de Cache KV entre Camadas via Extração de Vetores Singulares Alinhados

1. Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) de contexto longo habilitam aplicações avançadas, como recuperação de informações em larga escala e depuração de bases de código, mas enfrentam um gargalo crítico: o consumo de memória dos caches de Chave-Valor (KV) durante a inferência. À medida que os comprimentos de contexto se estendem para centenas de milhares ou milhões de tokens, a ocupação de memória do Cache KV torna-se proibitiva, limitando severamente as solicitações de inferência concorrentes e a taxa de transferência.

Estratégias de mitigação existentes geralmente se enquadram em duas categorias:

• Métodos intra-camada: Técnicas como quantização, evicção de tokens e decomposição de baixo posto de camada única comprimem o Cache KV de cada camada independentemente. Embora eficazes até certo ponto, elas falham em explorar redundâncias potenciais entre as camadas.
• Métodos entre camadas: Abordagens como Atenção entre Camadas (CLA) e YOCO exigem modificações arquitetônicas caras e pré-treinamento. Outras, como MiniCache, baseiam-se na fusão post-hoc de camadas adjacentes com base na similaridade cosseno token a token. No entanto, a análise empírica mostra que a similaridade token a token entre camadas é frequentemente muito fraca para suportar compressão robusta sem degradação significativa da precisão.

2. Metodologia: xKV

Os autores propõem o xKV, um método de compressão pós-treinamento, plug-and-play, que explora redundância estrutural profunda entre camadas sem exigir re-treinamento do modelo ou alterações arquitetônicas.

2.1. Observação Central: Vetores Singulares Alinhados

O artigo desafia a suposição de que as camadas são distintas com base na similaridade cosseno token a token. Utilizando Alinhamento de Kernel Centralizado (CKA), os autores demonstram que, embora os embeddings de tokens individuais possam divergir, os vetores singulares dominantes (a base que define as variações principais no espaço de tokens) estão altamente alinhados entre camadas adjacentes. Essa alta pontuação CKA indica que a estrutura geométrica dos subespaços do Cache KV é compartilhada, mesmo que as similaridades ponto a ponto sejam baixas.

2.2. Fatoração entre Camadas (CLF)

Aproveitando esse alinhamento, o xKV particiona as $N$ camadas do modelo em janelas contíguas de tamanho $W$. Para cada janela, ele realiza uma Fatoração entre Camadas:

1. Concatenação: Os Caches KV de todas as camadas dentro de uma janela são concatenados horizontalmente.
2. SVD: Uma Decomposição em Valores Singulares (SVD) de posto-$r$ é computada no cache concatenado.
3. Decomposição: O resultado produz uma base de tokens compartilhada ($A_k$) para a janela e matrizes de reconstrução específicas da camada ($B_\ell$).
4. Armazenamento: Em vez de armazenar $W$ caches de posto completo, o sistema armazena uma base compartilhada e $W$ matrizes de reconstrução menores, reduzindo a complexidade de memória de $O(W \cdot L \cdot d)$ para $O(L \cdot r + W \cdot r \cdot d)$.

2.3. Reconstrução Seletiva (SR)

Para abordar o gargalo computacional de reconstruir todos os tokens durante a fase de decodificação (que escala linearmente com o comprimento da sequência $L$), o xKV introduz a Reconstrução Seletiva (SR).

• Mecanismo: Em vez de reconstrução densa, a SR identifica um subconjunto esparsos de tokens ($S$) mais relevantes para a consulta atual usando seleção de blocos guiada por marcos (semelhante ao ShadowKV).
• Execução: Apenas as linhas correspondentes a esses tokens críticos são reconstruídas usando a base compartilhada e as matrizes específicas da camada.
• Benefício: Isso reduz o custo de reconstrução para ser constante em relação ao comprimento da sequência, permitindo acelerações de ponta a ponta.

3. Contribuições Principais

1. Descoberta de Redundância entre Camadas: O artigo revela que os Caches KV compartilham vetores singulares dominantes altamente alinhados entre camadas, uma forma de redundância obscurecida pela baixa similaridade token a token.
2. Algoritmo xKV: Um método sem treinamento que extrai uma base compartilhada via Fatoração entre Camadas (CLF), permitindo compartilhamento eficaz de informações entre $W$ camadas ($W > 2$).
3. Reconstrução Seletiva: Uma otimização no momento da decodificação que reconstrói apenas tokens relevantes para a consulta, eliminando a sobrecarga computacional linear da reconstrução densa.
4. Compatibilidade Plug-and-Play: O método não exige ajuste fino ou modificações arquitetônicas, tornando-o aplicável a modelos pré-treinados existentes (por exemplo, Llama-3.1, Qwen2.5, DeepSeek-V2).

4. Resultados Experimentais

Os autores avaliaram o xKV em Llama-3.1-8B-Instruct, Qwen2.5-7B/14B-Instruct e DeepSeek-Coder-V2-Lite usando os benchmarks RULER, LongBench e NIAH de múltiplas voltas.

• Compressão vs. Precisão:
  • O xKV alcança até 8× de compressão de Cache KV com perda mínima de precisão (≤3% em média) em comparação com baselines de atenção completa.
  • No Llama-3.1-8B, o xKV com compressão 8× alcança 88,50% de precisão média no RULER, superando significativamente o MiniCache (45,04% em 1,3×) e igualando ou superando métodos de evicção de tokens como SnapKV e PyramidKV.
  • Em tarefas NIAH de múltiplas voltas, o xKV mantém desempenho estável ao longo das voltas da conversa, enquanto métodos de evicção de tokens sofrem declínios acentuados após a primeira volta.
• Eficiência e Taxa de Transferência:
  • Memória: O xKV permite tamanhos de lote maiores sob o mesmo orçamento de memória de GPU.
  • Latência: A reconstrução densa reduz a latência de atenção em até 3,6×.
  • Taxa de Transferência: Combinado com Reconstrução Seletiva (xKV-SR), o sistema alcança até 4,23× de aceleração de ponta a ponta sobre baselines de atenção completa em GPUs NVIDIA A100.
  • Comparação: O xKV-SR supera o ShadowKV em 30% na taxa de transferência sob níveis de precisão semelhantes, mantendo tanto Chaves quanto Valores na GPU, evitando gargalos de transferência PCIe.
• Robustez: O método permanece eficaz em diferentes arquiteturas de modelos, incluindo Atenção com Consulta Agrupada (GQA) e Atenção Latente Multi-cabeça (MLA) no DeepSeek-V2.

5. Significado e Alegações

O artigo afirma que o xKV fornece uma abordagem plug-and-play prática para reduzir simultaneamente a ocupação de memória e a latência para inferência de LLMs de contexto longo. Ao deslocar o foco da similaridade token a token para o alinhamento de vetores singulares, o xKV desbloqueia redundância inter-camada não explorada que métodos anteriores perderam.

Os autores enfatizam que sua abordagem:

• Evita o alto custo de pré-treinamento exigido por alterações arquitetônicas.
• Supera as limitações de precisão da fusão baseada em similaridade (MiniCache).
• Resolve o gargalo computacional da reconstrução de baixo posto densa por meio da Reconstrução Seletiva.
• Permite inferência totalmente no dispositivo (GPU) para caches comprimidos, contornando os limites de largura de banda PCIe que restringem métodos baseados em descarregamento, como o ShadowKV.

O trabalho é apresentado como um passo fundamental em direção à inferência eficiente de contexto longo, com código disponibilizado publicamente para facilitar a adoção.