VQKV: High-Fidelity and High-Ratio Cache Compression via Vector-Quantization

Resumo Técnico: VQKV

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) enfrentam um desafio crítico à medida que o comprimento do contexto aumenta: o crescimento contínuo do Cache de Chave-Valor (KV Cache). Este cache consome uma quantidade massiva de memória, limitando a viabilidade de implantar LLMs em ambientes com recursos restritos (como GPUs de consumidor) e restringindo o comprimento máximo de geração.

As abordagens existentes para compressão de KV Cache geralmente apresentam um compromisso (trade-off) entre taxa de compressão e fidelidade da reconstrução:

• Evicção de Tokens: Remove tokens menos importantes. Embora alcance altas taxas de compressão, causa perda de informação irreversível, degradando o desempenho em tarefas que exigem recuperação de contexto de longo alcance.
• Redução de Dimensão de Recursos (Low-rank): Projeta vetores em espaços de menor dimensão. Frequentemente requer treinamento adicional ou sacrifica a fidelidade em taxas de compressão agressivas.
• Quantização Escalar: Reduz a precisão de cada valor individualmente. Falha em explorar correlações estruturais dentro dos vetores, resultando em erros de quantização significativos em taxas de compressão muito altas.

A necessidade é de um método sem treinamento (training-free) que atinja simultaneamente altas taxas de compressão e alta fidelidade de reconstrução, sem alterar os parâmetros do modelo original.

2. Metodologia: VQKV

O VQKV propõe uma nova abordagem baseada em Quantização Vetorial (Vector Quantization - VQ) para comprimir o cache KV. A ideia central é mapear vetores de cache de alta dimensão (ponto flutuante) para um conjunto compacto de códigos inteiros, preservando a estrutura de dados intrínseca.

Componentes Principais:

• Quantização Vetorial Residual (RSimVQ):

  • O método utiliza uma variante da Quantização Vetorial Simples com Resíduos (RSimVQ).
  • Em vez de quantizar vetores inteiros de uma só vez, o processo é iterativo e residual. O primeiro códigobook captura a maior parte da informação; o resíduo (erro) é então passado para o próximo códigobook, e assim por diante.
  • Isso permite que múltiplos códigobooks (conjuntos de vetores de referência) capturem subespaços independentes dos dados originais, permitindo uma reconstrução de alta fidelidade com poucos índices inteiros.
  • Uma matriz de projeção é aplicada em cada códigobook para adaptar o espaço de representação.

• Treinamento Sem Modificação do Modelo (Training-Free):

  • Os códigobooks são pré-treinados independentemente usando uma pequena amostra de dados (0,1% do OpenWebText) antes da inferência.
  • Durante a inferência, o modelo LLM original não é alterado. O VQKV atua como uma camada de compressão/descompressão transparente.

• Integração nas Etapas de Inferência:

  • Etapa de Preenchimento (Prefilling): Os vetores KV do contexto inicial são comprimidos em índices inteiros e armazenados.
  • Etapa de Decodificação: Para novos tokens, o cache é comprimido e adicionado aos códigos existentes.
  • Reconstrução Sob Demanda: Ao contrário de métodos que mantêm o cache inteiro em memória, o VQKV reconstroi apenas os vetores KV necessários para o passo de atenção atual (usando FlashAttention customizado), reduzindo ainda mais o uso de memória.
  • Janela Deslizante Local: Mantém uma janela local de tokens recentes (não comprimidos ou com alta fidelidade) e comprime o histórico antigo, garantindo que informações críticas recentes não sejam perdidas.

• Lidando com RoPE (Rotary Positional Embeddings):

  • O design residual do VQKV distribui as variações introduzidas pela transformação RoPE através de múltiplos códigobooks, permitindo uma reconstrução robusta mesmo após a rotação de frequências heterogêneas.

3. Principais Contribuições

1. Alta Taxa de Compressão com Alta Fidelidade: O VQKV consegue comprimir o cache KV em até 82,8% mantendo 98,6% do desempenho do modelo base (LLaMA3.1-8B) em benchmarks de longo contexto.
2. Método Sem Treinamento: Não requer fine-tuning ou re-treinamento do LLM, sendo aplicável imediatamente a qualquer checkpoint existente.
3. Extensão de Contexto Dramática: No mesmo hardware, o VQKV permite gerações 4,3 vezes mais longas do que o modelo com cache completo.
4. Superioridade sobre o Estado da Arte: Supera métodos de evicção (SnapKV), redução de dimensão (ASVD, Palu) e quantização escalar (KIVI) em tarefas de recuperação de longo alcance e geração.

4. Resultados Experimentais

Os experimentos foram conduzidos nos modelos LLaMA3.1-8B e LLaMA3.2-3B nos benchmarks LongBench, Needle-In-A-Haystack (NIAH) e RULER.

• LongBench: O VQKV obteve pontuações médias muito próximas ao baseline de cache completo (98,6% da performance), superando consistentemente todos os outros métodos de compressão em taxas comparáveis. Em algumas tarefas, superou até mesmo o baseline não comprimido.
• NIAH (Agulha no Palheiro): O método manteve uma pontuação perfeita de 100, igual ao modelo completo, enquanto outros métodos sofreram degradação significativa na capacidade de recuperar informações profundas no contexto.
• RULER (Contexto Longo): Em contextos de até 32K tokens, o VQKV preservou a capacidade de longo contexto, superando largamente concorrentes como ASVD e Palu.
• Eficiência de Memória:
  • Em uma GPU NVIDIA RTX 4090 (48 GB), o modelo base (LLaMA3.1-8B) atingiu o limite de memória (OOM) em aproximadamente 190k tokens.
  • Com VQKV, o mesmo hardware suportou sequências de mais de 824k tokens (um aumento de mais de 4x).
  • O uso de memória de pico foi reduzido para quase metade do baseline em comprimentos de geração equivalentes.

5. Significado e Conclusão

O VQKV representa um avanço significativo na otimização de LLMs para inferência de longo contexto. Ao demonstrar que a Quantização Vetorial pode ser aplicada eficazmente ao cache KV sem treinamento, o trabalho resolve o dilema tradicional entre compressão agressiva e perda de qualidade.

A capacidade de gerar textos com contextos extremamente longos (centenas de milhares de tokens) em hardware de consumo abre novas portas para aplicações práticas de LLMs que exigem memória de longo alcance, como análise de documentos extensos, livros inteiros ou sessões de diálogo prolongadas, sem a necessidade de clusters de servidores caros. O trabalho sugere que o design diferenciado de códigobooks para chaves e valores é crucial, já que os vetores de valor são mais sensíveis à configuração de compressão.

Limitações: A principal limitação identificada é a eficiência de decodificação. Embora kernels customizados (Triton) tenham otimizado o tempo, ainda existe um custo computacional adicional devido à operação de quantização que pode ser alvo de otimizações futuras.