Resumo Técnico: SPECTRA

Declaração do Problema

Na era dos modelos de linguagem de grande escala (LLMs) agentes, a inferência é cada vez mais dominada por entradas de contexto longo, que variam de bases de código inteiras a diálogos de múltipls turnos. À medida que o comprimento do contexto cresce, o cache Key-Value (KV) — o armazenamento contínuo de chaves e valores de atenção — torna-se o principal gargalo, frequentemente excedendo o tamanho dos próprios pesos do modelo. Este cache deve residir em memória rápida e ser relido em cada etapa de geração, limitando tanto o comprimento máximo do contexto quanto o tamanho do lote (batch size) que uma GPU pode suportar.

Os métodos de compressão existentes geralmente se dividem em três categorias: eliminação de tokens (descarte de tokens não importantes), projeção de baixo rank (redução de dimensionalidade) e quantização (redução de precisão). Embora a quantização seja amplamente utilizada, ela atinge um "abismo de 2 bits" onde o desempenho degrada-se drasticamente. A 2 bits por valor (quatro níveis), a gama dinâmica limitada não consegue acomodar a distribuição rica em outliers do cache; entradas grandes consomem os níveis disponíveis, colapsando valores menores e críticos em ruído. Além disso, a quantização padrão aplica uma largura de bits uniforme em todos os canais, ignorando o fato de que os canais brutos do cache KV são fortemente correlacionados, o que torna difícil identificar quais canais específicos carregam mais informação.

Metodologia: SPECTRA

Os autores propõem o SPECTRA, um codec drop-in livre de treinamento que reformula a compressão do cache KV como um problema de taxa-distorção (rate-distortion). Em vez de fixar um orçamento antecipadamente (por exemplo, um rank ou largura de bits específico), o SPECTRA aloca dinamicamente bits para as partes do cache que mais importam. O método baseia-se em três observações principais derivadas da sondagem de LLMs pré-treinados:

1. Concentração em uma Base Correta: Na base de canais brutos, os canais do cache KV são altamente correlacionados. No entanto, quando rotacionados para uma base computada a partir das próprias estatísticas do cache (especificamente, a matriz de Gram de ativação), os canais tornam-se não correlacionados e a energia se concentra fortemente: uma pequena fração dos canais carrega a vasta maioria do sinal.
2. Alocação Ponderada pela Importância: Alocar bits com base na importância do canal (energia) supera significativamente a precisão uniforme, particularmente no regime abaixo de 2 bits. Canais com zero bits são efetivamente descartados, unificando quantização e redução de rank em um único mecanismo de controle.
3. Importância Ponderada pela Ativação: A importância de um canal é determinada pela força com que as ativações de entrada o impulsionam, não meramente pela magnitude de suas projeções de peso. Portanto, a transformação deve ser ponderada pelas estatísticas de ativação ($G = E[h^\top h]$) em vez de apenas pelas matrizes de peso.

O Pipeline SPECTRA

O SPECTRA implementa estas observações através de um processo de três estágios:

1. Transformação Ponderada por G (Cache Latente):
   O método fatoriza as projeções de chave e valor ($W_K, W_V$) usando uma Decomposição de Valor Singular (SVD) ponderada pela matriz de Gram de ativação $G$. Ao branquear os pesos com $G^{1/2}$, o método encontra uma fatoração de baixo rank que minimiza o erro na métrica das ativações reais produzidas pelo modelo, em vez de no espaço de pesos bruto. Isso gera uma representação latente compacta e ordenada por energia ($C$), onde as coordenadas são ortogonais e ordenadas por variância.

   • Refinamento: Uma correção de forma fechada e livre de treinamento é aplicada à projeção de saída ($W_O$) para compensar o resíduo estruturado introduzido pela truncagem de baixo rank.

2. Quantização Taxa-Distorção (Water-Filling):
   Uma vez que o cache é transformado neste espaço latente ordenado por energia, o SPECTRA aplica um algoritmo de "preenchimento de água reverso" (reverse water-filling) para alocar bits. Com base nas variâncias marginais medidas das coordenadas latentes, o algoritmo atribui mais bits aos canais de alta variância (alta energia) e menos bits (até zero) aos canais de baixa variância.

   • Agrupamento: Para manter uma estrutura de cache de adição (append-only) adequada para inferência de fluxo (streaming), coordenadas adjacentes são agrupadas para compartilhar uma única largura de bits.

3. Integração ao Sistema:
   As matrizes de transformação ($W_{down}, W_{up}$) e a correção $W_O$ são incorporadas de volta aos pesos do modelo, tornando o SPECTRA um substituto direto (drop-in replacement) que não requer ajuste fino (fine-tuning). Durante a inferência, os tokens são projetados para baixo, quantizados por token e anexados ao cache latente. As chaves e valores completos são reconstruídos on-the-fly para a atenção padrão.

Principais Contribuições

• Visão Unificada de Compressão: O artigo recasualiza a compressão do cache KV como um problema único de gasto de um orçamento fixo de bits entre canais, unificando a projeção de baixo rank (descarte de canais) e a quantização (redução de bits) como dois extremos do mesmo espectro de controle.
• Transformação Ponderada por Ativação: Identifica que minimizar o erro de reconstrução na métrica das ativações do modelo (via uma transformação ponderada por $G$) é superior à fatoração no espaço de pesos, gerando um latente que é simultaneamente ortogonal de segundo momento e ordenado por energia.
• Codec SPECTRA: Uma implementação prática e livre de treinamento que concentra o orçamento de bits nos canais que carregam sinal, ultrapassando o abismo de 2 bits.

Resultados Experimentais

Os autores avaliaram o SPECTRA no Llama-3.1-8B, Mistral-7B e Qwen2.5-7B utilizando os benchmarks LongBench e RULER (Needle-in-a-Haystack).

• Desempenho vs. Compressão:
  • Com 4x de compressão, o SPECTRA é quase sem perdas (near-lossless), igualando ou superando o desempenho do modelo denso (fp16) e superando baselines como KIVI, OTT e PolarQuant.
  • Com 8x de compressão, onde métodos de quantização uniforme colapsam (caindo significativamente abaixo do baseline denso), o SPECTRA permanece competitivo. Para o Llama-3.1-8B, o SPECTRA atinge uma pontuação média no LongBench de 51.89 com 8.84x de compressão, enquanto baselines como TurboQuant e RotateKV caem para ~47 e ~46, respectivamente.
  • O SPECTRA permanece utilizável até 12x de compressão (ex: pontuação de 48.81 para 11.12x no Llama), um regime onde outros métodos falham totalmente.
• Capacidade de Memória:
  • Em uma GPU H200 com 24GB de memória, um modelo Llama-3.1-8B denso suporta um contexto de ~65k tokens. O SPECTRA com 12x de compressão estende isso para 793k tokens, multiplicando efetivamente a capacidade de contexto por mais de 12x sem alterar o hardware.
• Recuperação (Retrieval):
  • Em tarefas de Needle-in-a-Haystack, o SPECTRA mantém uma precisão de recuperação quase perfeita até 12x de compressão, enquanto baselines de 2 bits degradam-se bruscamente em contextos longos.

Significância e Alegações

O artigo alega que o principal alavancador para a compressão do cache KV não é meramente quão finamente quantizar, mas o que armazenar. Ao rotacionar o cache para um sistema de coordenadas derivado das próprias estatísticas do modelo, o SPECTRA revela que a informação é altamente concentrada. Isso permite que o método gaste bits onde eles mais importam, efetivamente contornando o "abismo de 2 bits" que assola a quantização uniforme.

Os autores enfatizam que o SPECTRA é uma solução drop-in livre de treinamento que não requer pré-treinamento, ajuste fino ou mudanças arquiteturais. Embora a implementação atual foque na economia de memória (reconstruindo chaves/valores on-the-fly), o artigo observa que a estrutura do espaço latente abre teoricamente um caminho para economia de computação via atenção direta no espaço latente, o que é deixado como trabalho futuro. A contribuição central é demonstrar que empurrar a compressão além do abismo de 2 bits é possível ao tratar o cache como um sinal a ser codificado via princípios de codificação de transformada, em vez de um tensor a ser trimado uniformemente.