Resumo Técnico: Um JOLT para o KV Cache

Declaração do Problema

Na inferência de transformadores autorregressivos, o cache Key-Value (KV) representa o custo de memória dominante, escalando como $O(\text{batch} \times T \times L \times n_h \times d_h)$. Para contextos longos ou grandes lotes (batches), o cache KV, em vez dos pesos do modelo, torna-se o principal gargalo de memória e limita o throughput. Os métodos de compressão existentes dividem-se em duas famílias, ambas falhando em explorar totalmente a redundância estrutural do cache:

1. Métodos de baixo posto (low-rank): Estes fatorizam fatias bidimensionais (ex: matrizes por cabeça ou blocos entre camadas), mas comprometem-se com um único eixo de redundância, ignorando os outros.
2. Métodos de quantização: Estes reduzem a largura de bits de cada entrada, mas são limitados por taxas de bits fixas (ex: 4 bits), criando um "piso de compressão" que impede o alcance de razões de compressão intermediárias (2–3×) onde a qualidade é frequentemente preservada.

A observação central que motiva este trabalho é que o cache KV em uma camada é um tensor de terceira ordem ($n_h \times T \times d_h$) onde os eixos carregam quantidades vastamente diferentes de redundância. Especificamente, os eixos de cabeça (head) e camada (layer) são amplamente incompressíveis, enquanto os eixos de token e característica (feature) detêm a maior parte da estrutura compressível. Além disso, empiricamente, os valores (values) são 2–3× mais difíceis de comprimir do que as chaves (keys) devido às diferenças em suas propriedades espectrais.

Metodologia: JoLT

O método proposto, JoLT (Joint Tucker and JL-residual allocation), aborda esses desafios tratando o cache como um tensor de terceira ordem e aplicando uma estratégia de compressão direcionada.

1. Decomposição Tucker Parcial

Ao contrário das fatorizações de tensor completas que comprimem todos os eixos, o JoLT aplica uma decomposição Tucker parcial. Ele trunca apenas os eixos de token e característica (que carregam redundância), mantendo os eixos de cabeça e camada intactos (que são incompressíveis). Isso é alcançado usando SVD de ordem superior sequencialmente truncada (ST-HOSVD).

• Motivação: Comprimir os eixos de cabeça ou camada acarreta erro significativo, enquanto truncar os eixos de token e característica captura a maior parte da redundância.

2. Resíduo Rotacionado por JL

Um backbone de baixo posto puro é insuficiente para uma fidelidade próxima da sem-perda (near-lossless), particularmente para os valores, que possuem um espectro plano onde o truncamento descarta energia significativa. O JoLT restaura essa energia descartada usando um resíduo de baixa taxa rotacionado por Johnson-Lindenstrauss (JL):

• O resíduo de truncamento $R = X - \tilde{X}$ é computado.
• Uma rotação JL aleatória é aplicada ao resíduo.
• O resíduo rotacionado é quantizado uniformemente a $b \in \{0, 2, 4, 8\}$ bits.
• A decodificação envolve inverter a rotação. Este resíduo atua como uma alavanca para converter um backbone com perda em um quase sem-perda, sendo especialmente valioso para os valores.

3. Alocação Lagrangiana Conjunta

O método utiliza um único Lagrangiano dual para alocar recursos conjuntamente entre as duas etapas (ranks de Tucker e larguras de bits de resíduo) e entre diferentes alvos (Keys vs. Values, Grupos de Camadas).

• Objetivo: Minimizar o erro de reconstrução total sujeito a um orçamento global de bytes.
• Mecanismo: A variável dual permite que o alocador desloque bits entre ranks e resíduos, e entre chaves e valores, explorando dinamicamente a assimetria na compressibilidade (ex: financiando pesadamente o resíduo de valores enquanto reduz os bits do resíduo de chaves).
• Granularidade: A alocação é realizada por grupo de camada e separadamente para Keys e Values.

4. FlashJoLT (Variante Rápida)

Para lidar com o custo computacional do SVD exato em contextos longos, o FlashJoLT substitui o SVD exato no modo de token por um SVD de baixo posto randomizado.

• Ele computa apenas as $q_{cap}$ direções do modo de token.
• Utiliza a contabilização de massa de cauda (tail-mass accounting), adicionando a cauda espectral omitida como um valor singular de resíduo sintético para evitar que o alocador subestime a energia capturada por um determinado rank.
• Esta variante alcança uma aceleração de 5–13× no tempo de compressão com perda de qualidade negligenciável na "zona livre".

Contribuições Principais

1. Análise Espectral: Um estudo empírico identificando que os eixos de cabeça e camada são incompressíveis, enquanto os eixos de token e característica detêm redundância, com os valores sendo significativamente mais difíceis de comprimir do que as chaves.
2. Algoritmo JoLT: Um novo compressor que combina decomposição Tucker parcial com um resíduo rotacionado por JL, acoplados via um único dual Lagrangiano para alocação conjunta de ranks e bits.
3. FlashJoLT: Uma variante de SVD randomizado que reduz drasticamente o tempo de compressão enquanto preserva a qualidade.
4. Validação Empírica: Um estudo abrangente nos modelos GQA (Mistral-7B) e MHA (LLaMA-2-13B) estabelecendo uma "zona livre" de compressão quase sem-perda.

Resultados

A "Zona Livre" (Compressão de 2–3×)

O artigo identifica uma razão de compressão de 2–3× como uma "zona livre" onde a qualidade do modelo é preservada dentro do ruído estatístico da linha de base sem compressão:

• Perplexidade: Tanto no Mistral-7B (GQA) quanto no LLaMA-2-13B (MHA), o desvio de perplexidade é insignificante a 2–3× para comprimentos de contexto de até 1024.
• Erro de Reconstrução: A 2× de compressão, o JoLT atinge um erro de Frobenius relativo de 0,009 (Keys) e 0,006 (Values). Isso é aproximadamente uma ordem de magnitude inferior à quantização de 4 bits e às linhas de base de SVD entre camadas.
• Tarefas Downstream:
  • GSM8K: Modelos comprimidos (2–3×) permanecem dentro do intervalo de confiança de 95% da linha de base.
  • RULER (Needle-in-a-Haystack): O Mistral mantém 100% de precisão de recuperação a 2–3×; o LLaMA iguala o teto de sua linha de base.

Divergência de Arquitetura

Além da zona livre de 2–3×, as duas arquiteturas divergem:

• Mistral (GQA): Degrada graciosamente, perdendo cerca de 4% de perplexidade por passo de razão inteira além de 4×.
• LLaMA (MHA): Degrada bruscamente entre 4× e 5×, com a perplexidade aumentando em ~27% a 4× e ~68% a 5×. Tarefas de raciocínio (GSM8K) para o LLaMA colapsam significativamente em razões mais altas.

Eficiência

• Memória: A pegada persistente do KV diminui linearmente com a razão de compressão (ex: de 268,4 MB para 134,6 MB a 2× para Mistral em ctx-2048). A memória de processo de pico permanece limitada pelos pesos/ativações.
• Velocidade: O FlashJoLT oferece uma aceleração de 5–13× no tempo de compressão em comparação com o backbone exato.

Significância e Alegações

O artigo afirma que o JoLT fornece uma compressão quase sem-perda de 2–3× do cache KV, removendo efetivamente o gargalo de memória para inferência de contexto longo sem sacrificar a qualidade do modelo. Ao tratar o cache como um tensor e explorar a assimetria espectral específica entre chaves e valores, o JoLT supera os métodos existentes de baixo posto e quantização no regime de compressão intermediária.

Os autores posicionam o método como uma alavanca direta para reduzir o custo de servir contextos longos. Eles observam dois limites principais:

1. Limites de Arquitetura: O método é robusto para modelos GQA, mas enfrenta degradação acentuada para modelos MHA além de 4×, sugerindo que a faixa de 2–3× é a recomendação segura e agnóstica ao modelo.
2. Custo de Implementação: Embora o armazenamento seja reduzido, a decodificação atualmente requer a reconstrução do cache a partir dos fatores em cada etapa. Os autores identificam o desenvolvimento de um kernel de reconstrução e atenção fundido como o próximo passo crítico para alcançar paridade de inferência e realizar plenamente o potencial do método.