Resumo Técnico: Quantize o que Importa: Mais para Chaves, Menos para Valores

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) enfrentam gargalos severos de memória durante a inferência, impulsionados pelo cache Chave-Valor (KV), que escala linearmente com o tamanho do modelo e o comprimento do contexto. Embora a quantização do cache KV (redução da precisão de BF16 para INT4 ou inferior) ofereça um caminho para economia de memória, os métodos existentes frequentemente dependem de ajuste heurístico e ad hoc para a alocação de bits entre chaves e valores. As abordagens atuais tratam tipicamente a divisão chave-valor como um hiperparâmetro a ser buscado em grade ou derivado de estatísticas de ativação em tempo de inferência, carecendo de uma fundamentação teórica enraizada em propriedades intrínsecas do modelo. Isso levanta uma questão fundamental: como a alocação de bits pode ser determinada de maneira principial e generalizável, minimizando o erro de quantização enquanto preserva a precisão do modelo?

Metodologia

O artigo propõe uma abordagem guiada pela geometria para quantização KV de precisão mista, fundamentada nas normas espectrais e de Frobenius intrínsecas das matrizes de pesos do Transformer. A metodologia procede em duas etapas teóricas e empíricas:

1. Análise Teórica da Disparidade de Normas: Os autores analisam a dinâmica de treinamento das matrizes de projeção de chave ($W_K$) e valor ($W_V$). Eles demonstram que $W_K$ sistematicamente adquire normas espectrais e de Frobenius maiores do que $W_V$ devido à assimetria arquitetural nos mecanismos de atenção. Especificamente, $W_K$ molda o mapa de atenção e interage multiplicativamente com as consultas, levando a sinais de gradiente amplificados durante o treinamento, enquanto $W_V$ influencia apenas as representações pós-atenção. Isso resulta em $W_K$ possuindo maior densidade de informação e magnitude.
2. Derivação de Teoremas de Quantização: Com base na observação de que o erro de quantização escala com a energia do sinal (magnitude da norma), os autores derivam dois teoremas:
   • Teorema da Disparidade de Normas Chave-Valor: Prova que as normas espectrais e de Frobenius esperadas de $W_K$ excedem as de $W_V$ em famílias proeminentes de LLMs.
   • Teorema de Quantização Prioritária para Chaves: Estabelece que, para qualquer orçamento de memória fixo, alocar maior precisão para chaves do que para valores reduz estritamente o erro total de quantização. Isso ocorre porque o erro introduzido no caminho de alta norma da chave é amplificado de forma mais significativa através da rede do que os erros no caminho de valor.

Os autores validam esses teoremas em modelos diversos (Llama-3, Mistral, Phi, Qwen, DeepSeek), conjuntos de dados (C4, MMLU, GSM8K, CoQA, EQ-Bench, LongBench) e backends de quantização (Optimum Quanto e HQQ). Eles também investigam a integração dessa estratégia com métodos de redistribuição de outliers baseados em rotação (por exemplo, QuaRot).

Contribuições Principais

• Fundamentação Teórica: O artigo fornece a primeira justificação teórica para alocação assimétrica de bits em caches KV, avançando além do ajuste heurístico para um princípio de design baseado nas propriedades algébricas lineares dos pesos do Transformer.
• Teorema da Disparidade de Normas: Prova que as projeções de chave exibem consistentemente normas maiores do que as projeções de valor, uma propriedade preservada nos caches resultantes.
• Teorema de Quantização Prioritária para Chaves: Demonstra formalmente que priorizar a precisão para chaves em vez de valores minimiza o erro de quantização e preserva a precisão da inferência mais efetivamente do que alocações simétricas ou invertidas.
• Validação Empírica: Experimentos extensivos mostram que alocações favorecendo chaves (por exemplo, 4 bits para chaves, 2 bits para valores) retêm até 98,3% da precisão de alocações uniformes de 4 bits (K4V4) enquanto reduzem a memória do cache KV em 25%.
• Ortogonalidade e Integração: O estudo mostra que essa estratégia guiada pela geometria é ortogonal às técnicas existentes de redistribuição de outliers. Combinar quantização prioritária para chaves com rotação apenas de chaves gera ganhos sinérgicos, superando as bases uniformes em 4,4–18% de precisão em diversas tarefas.

Resultados

• Erro de Reconstrução: Em sete modelos e múltiplos conjuntos de dados, os caches de chave consistentemente incorrem em erros de reconstrução maiores do que os caches de valor em larguras de bits idênticas (por exemplo, K2 vs. V2), confirmando sua maior sensibilidade à quantização.
• Precisão a Jusante:
  • Nas avaliações do GSM8K usando Optimum Quanto, a configuração K4V2 (4 bits para chaves, 2 bits para valores) supera significativamente a K2V4, recuperando aproximadamente 94% da precisão da linha de base de precisão total, enquanto alocações favorecendo valores incorrem em perdas de até 30 pontos percentuais.
  • Usando HQQ em uma faixa mais ampla de larguras de bits (1–8 bits) e modelos (0,6B–32B), alocar até mesmo um bit extra para chaves em precisões ultra-baixas gera ganhos massivos (por exemplo, +48 pontos percentuais para Qwen3-8B no GSM8K ao passar de 1 bit para 2 bits para chaves).
  • Em média, K4V2 retém 98,3% da precisão de K4V4 em CoQA, EQ-Bench e GSM8K.
• Integração com Rotação: Aplicar rotação apenas às chaves em uma configuração K4V2 alcança uma precisão quase igual à linha de base completa K4V4, indicando que o benefício primário da rotação decorre da mitigação de outliers de chaves. Rotacionar apenas valores fornece retornos marginais ou negativos.
• Sensibilidade ao Tamanho do Grupo: Tamanhos de grupo mais finos (por exemplo, 32) para chaves melhoram a precisão devido à maior densidade de informação, enquanto tamanhos de grupo mais grossos (64 ou 128) para valores são suficientes com perda de precisão negligenciável.

Significado

O artigo afirma transformar a alocação de bits de uma escolha de design ad hoc em um princípio de design teoricamente fundamentado e guiado pela geometria para inferência eficiente de LLMs. Ao ancorar estratégias de precisão mista nas normas intrínsecas dos pesos do modelo, os autores fornecem uma regra generalizável: "Mais para chaves, menos para valores". Essa abordagem oferece um custo analítico leve e único por modelo, sem exigir introspecção em tempo de inferência ou re-treinamento. As descobertas sugerem que futuros frameworks de quantização KV e o co-design hardware-algoritmo devem priorizar a precisão das chaves para alcançar máxima eficiência de memória com degradação mínima de precisão. O trabalho também destaca que essa estratégia complementa as técnicas existentes de redistribuição de outliers, permitindo ganhos adicionais de desempenho sem aumentar a complexidade.