Resumo Técnico: RATEQUANT

Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) exigem o armazenamento em cache de todos os pares chave-valor (KV) previamente computados durante a geração, criando um gargalo de memória que escala linearmente com o comprimento da sequência, o tamanho do lote e a profundidade do modelo. Embora a quantização do cache KV reduza essa pegada, os métodos existentes atribuem uniformemente a mesma largura de bits a cada cabeça de atenção. Essa abordagem ignora a heterogeneidade bem documentada na importância das cabeças.

Tentativas recentes de alocação de precisão mista no nível da camada ou do canal dependem de regras heurísticas vinculadas a quantizadores específicos, carecendo de garantias teóricas e transferibilidade. Os autores identificam um obstáculo mais profundo, anteriormente não reconhecido: mismatch do modelo de distorção. Diferentes quantizadores seguem curvas de distorção-taxa distintas, modeladas como $D(b) = \alpha \beta^{-b}$. A taxa de decaimento $\beta$ varia significativamente entre designs (por exemplo, de 3,6 no TURBOQUANT a 5,3 no QuaRot). Aplicar ingenuamente uma estratégia de alocação de precisão mista derivada do modelo de distorção de um quantizador a outro inverte a ordem de ganho marginal das cabeças de atenção, frequentemente resultando em desempenho inferior à quantização uniforme.

Metodologia

O RATEQUANT aborda esses desafios formalizando a alocação de bits do cache KV por cabeça como um problema de otimização de taxa-distorção. O framework consiste em quatro fases:

1. Formulação do Problema e Alocação Ótima

Os autores modelam o Erro Quadrático Médio (MSE) da quantização por cabeça como uma função exponencial da largura de bits: $D(b) = \alpha \beta^{-b}$. O objetivo é minimizar a soma ponderada das distorções através de $N$ grupos de quantização (cabeças) sujeito a um orçamento total de bits $B$:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
onde $w_i$ representa a sensibilidade da cabeça $i$.

Usando a estacionaridade lagrangiana, os autores derivam uma solução de forma fechada via enchimento reverso de água (reverse waterfilling):
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
onde $\bar{b}$ é a largura de bits média e $\bar{w}$ é a média geométrica das sensibilidades. Esta solução aloca mais bits para cabeças com maior sensibilidade. Teoricamente, a redução de distorção alcançável é limitada pela razão entre a média aritmética e a média geométrica (AM/GM) das sensibilidades das cabeças.

2. Estimativa de Sensibilidade

Para determinar os pesos $w_i$, o RATEQUANT emprega sensibilidade baseada em gradiente. Ele calcula as normas quadradas dos gradientes das saídas de projeção KV em relação à perda de modelagem de linguagem causal:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
Os autores demonstram que proxies baseados em gradiente são superiores a métricas baseadas em ativação (por exemplo, magnitude do canal), que medem a amplificação de erro em vez do impacto direto na perda. A sensibilidade baseada em gradiente revela um padrão em forma de U de importância (alta sensibilidade nas camadas iniciais e finais), enquanto métricas baseadas em ativação mostram um aumento monotônico com a profundidade.

3. Alocação Inteira e Extensões

Como as larguras de bits devem ser inteiros, a solução contínua é refinada usando um algoritmo de ganho marginal ganancioso (Algoritmo 1). Este algoritmo atribui iterativamente bits à cabeça que oferece a maior redução na distorção ponderada até que o orçamento seja esgotado.

Duas extensões críticas tornam o RATEQUANT agnóstico ao quantizador:

• Calibração de Distorção: Em vez de assumir um $\beta$ universal, o RATEQUANT ajusta um modelo de distorção por quantizador ($\alpha, \beta$) usando um pequeno conjunto de calibração (16 sequências). Isso previne o modo de falha de "mismatch do modelo de distorção".
• Separação K/V: Reconhecendo que chaves e valores frequentemente usam esquemas de quantização diferentes (por exemplo, simétrica por canal para chaves versus assimétrica por token para valores no KIVI), o framework trata-os como $2N$ componentes independentes com curvas de distorção distintas. Isso permite que o algoritmo descubra orçamentos de bits assimétricos (por exemplo, alocando mais bits para chaves do que para valores).

4. Pipeline

O processo envolve custos offline únicos:

1. Calibração de Sensibilidade: Passadas direta e inversa em um pequeno conjunto de dados (~1,6s para um modelo de 8B).
2. Modelagem de Distorção: Ajuste de $\alpha$ e $\beta$ em 5 larguras de bits (<0,1s).
3. Alocação: Execução do algoritmo ganancioso (<0,01s).
4. Inferência: Aplicação das larguras de bits estáticas com zero sobrecarga em tempo de execução.

Principais Resultados

Os experimentos foram conduzidos em modelos Qwen3 (4B, 8B, 32B) usando perplexidade (PPL) do WikiText-2 e tarefas downstream (ARC-Challenge, HellaSwag, etc.).

• Resolução do Mismatch: Aplicar um modelo de distorção incompatível (por exemplo, usar o $\beta$ do TURBOQUANT para o KIVI) degradou o PPL do KIVI de 49,3 para 87,0 em 2,5 bits médios. O RATEQUANT calibrado resolveu isso, reduzindo o PPL para 14,9 (uma redução de 70%).
• Ganhos de Desempenho:
  • KIVI (2,5 bits): O RATEQUANT melhorou o PPL de 49,3 para 14,9.
  • QuaRot (2,5 bits): Melhorou o PPL em 6,6 pontos.
  • TURBOQUANT (4,0 bits): Recuperou 66% da degradação induzida pela quantização para o Qwen3-8B.
• Separação K/V: Para o KIVI em 2,5 bits, a alocação ótima descoberta pelo RATEQUANT foi de 2,85 bits para chaves e 2,15 bits para valores, superando significativamente a alocação conjunta.
• Eficiência: O método não adiciona sobrecarga em tempo de execução. A calibração leva menos de 2 segundos em uma única GPU.
• Precisão Downstream: Em 4,0 bits, o RATEQUANT alcançou 64,2% de precisão média em tarefas downstream, quase igualando a linha de base FP16 (64,4%) e recuperando 89,8% da lacuna de precisão.

Significado e Alegações

O artigo alega que o RATEQUANT é o primeiro framework a aplicar a teoria taxa-distorção à quantização do cache KV na granularidade por cabeça, ao mesmo tempo que aborda as diferenças estruturais entre chaves e valores.

A descoberta central é que o mismatch do modelo de distorção é um modo de falha crítico na alocação de precisão mista. Sem calibrar a curva de distorção ($\beta$) para o quantizador base específico, estratégias de precisão mista podem ser prejudiciais. Ao combinar:

1. Sensibilidade baseada em gradiente (superior a proxies baseados em ativação),
2. Calibração empírica de distorção (para corresponder ao quantizador), e
3. Separação K/V (para lidar com características de erro assimétricas),

o RATEQUANT transforma a quantização KV de precisão mista em uma camada de alocação geral e agnóstica ao quantizador. Os autores postulam que a razão AM/GM das sensibilidades das cabeças serve como um preditor prático para quando a alocação de precisão mista produzirá ganhos significativos. O trabalho sugere que a perspectiva taxa-distorção pode generalizar para alocação de recursos em outros componentes heterogêneos de redes neurais além dos caches KV.