Resumo Técnico: Ablação, Inferência Estatística e Validação para Compressão de KV-Cache

Declaração do Problema

A inferência de Transformers em escala é limitada pela largura de banda de memória, com o acesso ao cache de Chave-Valor (KV) dominando o tráfego total de memória durante a geração de contextos longos. Embora a quantização de chaves e valores para 2–4 bits por elemento seja uma abordagem padrão para reduzir o tamanho do cache, avaliar esses esquemas de compressão apresenta um desafio significativo. As diferenças de qualidade observadas são frequentemente confundidas por propriedades distributivas do tráfego real de Large Language Models (LLMs), efeitos de hardware e suposições algorítmicas, tornando difícil atribuir ganhos ou falhas de desempenho a escolhas de design específicas.

Metodologia

Os autores introduzem um framework de avaliação controlado consistindo em seis regimes estatísticos sintéticos projetados para isolar suposições estruturais específicas do pipeline de compressão (ex: forma da distribuição, alinhamento do subespaço K/V, decaimento de autovalores). Este framework é pareado com uma metodologia de validação estatística para distinguir diferenças algorítmicas sistemáticas de variância induzida pela implementação.

Configuração Experimental

• Esquemas Avaliados: Duas famílias primárias de compressão são comparadas:
  1. TurboQuant (TQ): Uma abordagem data-oblivious (independente dos dados) usando uma Transformada de Walsh-Hadamard (WHT) aleatória para espalhar a energia dos tokens uniformemente, seguida por um código de livro (codebook) Beta Lloyd-Max data-oblivious. Opcionalmente utiliza um esboço residual de 1-bit Johnson-Lindenstrauss (QJL) no caminho da chave (K) ou do valor (V).
  2. SpectralQuant (SQ): Uma abordagem data-adaptive (adaptável aos dados) que calibra uma base de autovetores por cabeça a partir de tokens representativos, aloca bits via water-filling para dimensões de alta variância e utiliza um esboço QJL seletivo nas dimensões efetivas superiores.
• Métricas: O estudo utiliza uma métrica de erro 2D ($d^2$) combinando erro quadrático médio relativo (relMSE) e erro de direção (similaridade de cosseno). Também emprega métricas de isolamento (relMSE_K, relMSE_V) e uma geometria de erro 6D para decompor erros no cache K, cache V e na saída final da atenção.
• Validação: Para separar o desempenho algorítmico do ruído de implementação, os autores comparam implementações Python (oráculo) e HIP/GPU (produção) usando diferentes geradores de números aleatórios (PCG64 vs. mt19937). O teste de Kolmogorov-Smirnov (KS) é aplicado para identificar diferenças sistemáticas versus variância dependente de semente (seed).

Contribuições Principais

1. Ablação Sistemática de Esquemas de Compressão

Através de uma ablação completa de oito variantes de TQ e múltiplas configurações de SQ, o estudo identifica três esquemas não dominados:

• Quantização Escalar (Plain-KV): Sem rotação, eficaz em dados isotrópicos.
• TQ-KV: Rotação WHT com codebook Beta Lloyd-Max.
• TQ-QKV: TQ-KV aumentado com um esboço residual QJL de 1-bit no caminho da chave (K).
  O estudo elimina variantes de QJL no caminho V, demonstrando que aplicar QJL aos valores (que são somados sobre os tokens) é uma heurística que consistentemente degrada o desempenho em comparação com correções no caminho K.

2. Validação Estatística e Desigualdade de Jensen

O artigo introduz uma metodologia para separar diferenças de codec de ruído de implementação. Uma descoberta crítica é a identificação da desigualdade de Jensen como um mecanismo de amplificação de erro.

• A não-linearidade softmax é convexa; portanto, a variância de um erro de correção de média zero ($\epsilon$) adicionado às pontuações de atenção é exponencialmente amplificada ($E[\exp(s+\epsilon)] \approx \exp(s) \cdot \exp(\sigma^2/2)$).
• A correção QJL no caminho K introduz tal variância. O estudo mostra que a variância do gap Python-HIP para TQ-QKV é significativamente maior do que para TQ-KV, confirmando que o caminho K é uma assinatura direta da desigualdade de Jensen amplificando a variância da pontuação. Inversamente, correções no caminho V entram linearmente na saída e não disparam essa amplificação.

3. Desempenho Específico de Regime e Sensibilidade de Calibração

• Dados de Cauda Pesada (Heavy-Tailed): Para dados com caudas pesadas (ex: distribuições Student-t), métodos baseados em autovetores (SQ) falham catastroficamente. As covariâncias amostrais são desestabilizadas por outliers, levando a uma base calibrada sistematicamente desalinhada. Nenhum aumento no orçamento de bits recupera essa perda. O TQ-KV domina nesses regimes porque a WHT espalha os outliers de cauda pesada, tornando as marginais projetadas aproximadamente Gaussianas.
• Dados Estruturados: Em regimes de baixo rank e estruturados, o SQ supera o TQ quando K e V são calibrados em conjuntos representativos separados. No entanto, o estudo descobre que a alocação de bits via water-filling fornece benefício negligenciável sobre a alocação uniforme em todos os regimes testados, pois as razões de autovalores dentro do subespaço efetivo raramente excedem o limiar necessário para justificar a alocação não uniforme.
• Dimensão Semântica Efetiva ($d_{eff}$): O artigo caracteriza $d_{eff}$ como auto-calibrável ao orçamento de calibração disponível, em vez de recuperar o verdadeiro rank dos dados. Um único autovalor espúrio pode causar falha catastrófica (saída anti-correlacionada), necessitando de um critério de gap espectral para podar a dimensão estimada.

Resultados e Recomendações

Desempenho Comparativo

• Orçamento $b=2$: Nenhum esquema único domina. Variantes de SQ vencem em regimes estruturados (baixo rank), TQ-KV vence em dados de cauda pesada, e Plain-KV é competitivo em dados isotrópicos.
• Orçamento $b \ge 4$: O TQ-KV afirma dominância na maioria dos regimes. Variantes de SQ apenas mantêm uma vantagem em regimes altamente estruturados onde a base de autovetores calibrada pode concentrar bits eficientemente, mas elas falham em recuperar do desalinhamento da base de autovetores em dados de cauda pesada.
• Armazenamento vs. Qualidade: Embora o SQ possa alcançar custos de armazenamento menores em orçamentos mais altos devido à sua capacidade de ignorar dimensões de cauda, o "piso da cauda" (bits fixos para dimensões não modeladas) torna-se uma restrição vinculante. Para sequências longas, erros de cauda acumulados corroem a vantagem semântica do SQ, tornando o TQ-KV a opção mais segura para tarefas de geração.

Recomendações Práticas

Os autores concluem com diretrizes específicas para implantação:

1. Use TQ-KV quando a distribuição dos dados for desconhecida, de cauda pesada ou sujeita a mudança de distribuição, e para qualquer orçamento acima de 2 bits.
2. Use SQU-KV (SpectralQuant com alocação uniforme) apenas quando os dados forem conhecidos como estruturados (subespaço KV de baixo rank), o orçamento for baixo (2–3 bits), a sequência for curta (workloads de prefill) e conjuntos de calibração separados para K e V estiverem disponíveis.
3. Descarte o Water-Filling: Em dados sintéticos, a alocação uniforme performa de forma equivalente ao water-filling.
4. Evite QJL no K para Caudas Pesadas: A correção de pontuação QJL no caminho das chaves amplifica a variância através do softmax (efeito Jensen) em dados de cauda pesada e deve ser evitada nesses regimes.

Significância

A significância deste artigo reside em sua metodologia rigorosa e controlada para avaliar a compressão de KV-cache. Ao isolar suposições estruturais através de regimes sintéticos e aplicar inferência estatística para separar o sinal algorítmico do ruído de implementação, os autores fornecem um framework de diagnóstico claro. Este trabalho vai além do simples relato de acurácia para explicar por que métodos funcionam ou falham, revelando limitações fundamentais de métodos baseados em autovetores em dados de cauda pesada e o comportamento não-monotônico da dimensionalidade efetiva sob orçamentos de calibração finitos. As descobertas oferecem um guia definitivo para a seleção de esquemas de compressão baseados em características de dados e restrições operacionais.