High-Rate Quantized Matrix Multiplication II

Resumo Técnico: Multiplicação de Matrizes Quantizadas de Alta Taxa II

Declaração do Problema
Este artigo investiga os limites teóricos e práticos da quantização apenas de pesos para Modelos de Linguagem de Grande Escala (LLMs). O cenário específico assume que a matriz de ativação $X$ está disponível em precisão total durante a inferência, enquanto a matriz de pesos $W$ deve ser quantizada para uma representação de baixa resolução $\hat{W}$ a fim de reduzir a largura de banda de memória e os requisitos de armazenamento. O objetivo é minimizar a distorção esperada do erro quadrático médio ponderado (WMSE):
$$D = \frac{1}{n} \mathbb{E} \| X^\top (W - \hat{W}) \|_F^2 = \frac{1}{n} \text{tr}((W - \hat{W})^\top \Sigma_X (W - \hat{W}))$$
onde $\Sigma_X = \mathbb{E}[XX^\top]$ é a matriz de covariância das ativações, que é assumida como conhecida (via calibração), mas o decodificador (motor de inferência) não tem acesso a $\Sigma_X$ ou sua base, apenas aos bits quantizados.

Metodologia e Estrutura Teórica
Os autores abordam o problema através de duas lentes: limites teóricos da informação e design algorítmico prático.

1. Limites Teóricos da Informação (Regime de Alta Taxa):

   • Codificação Isotrópica vs. Anisotrópica: O artigo contrasta a quantização "isotrópica" (usando um livro de códigos padrão alheio a $\Sigma_X$) com a quantização "anisotrópica" (adaptando-se a $\Sigma_X$).
   • Solução de Enchimento de Água (Waterfilling): Baseando-se na teoria clássica de codificação de fontes, os autores derivam que a distorção ótima $D^*$ no regime de alta taxa é alcançada por uma estratégia de enchimento de água reverso. Isso aloca taxas de quantização de forma desigual entre as coordenadas com base nos autovalores de $\Sigma_X$, resultando em uma distorção proporcional à média geométrica dos autovalores: $D^* \propto |\Sigma_X|^{1/n} 2^{-2R}$.
   • A Lacuna: A codificação isotrópica padrão produz uma distorção proporcional à média aritmética dos autovalores. A lacuna entre as soluções ótima e isotrópica é governada pela desigualdade da Média Aritmética-Média Geométrica (AM-GM).
   • Decodificador Alheio: Crucialmente, o artigo demonstra que, mesmo que o decodificador seja alheio a $\Sigma_X$ (não podendo aplicar a base SVD), um codificador adaptativo usando um livro de códigos isotrópico pode alcançar assintoticamente o limite de distorção de enchimento de água no regime de alta taxa, desde que o arredondamento para o livro de códigos seja ótimo.

2. Algoritmos Práticos (Cancelamento Sucessivo de Interferência):

   • O artigo analisa o algoritmo GPTQ (e LDLQ equivalente), que usa uma taxa fixa por entrada e uma abordagem de cancelamento sucessivo de interferência (SIC) (algoritmo de Babai) baseada na decomposição de Cholesky $\Sigma_X = U^\top U$.
   • Limite Teórico do GPTQ: Os autores mostram que o GPTQ padrão (com espaçamento igual) alcança uma distorção proporcional à média aritmética dos elementos diagonais ao quadrado de $U$ ($U_{i,i}^2$). Este desempenho é sensível à base de $\Sigma_X$; especificamente, ele performa pior na base PCA (onde $U$ é diagonal) e melhor em bases aleatórias.
   • Algoritmo WaterSIC: Motivado pela teoria de enchimento de água, os autores propõem o WaterSIC. Este algoritmo ajusta o tamanho do passo de quantização (espaçamento) para cada coordenada com base nos elementos diagonais do fator de Cholesky $U$. Especificamente, ele aloca mais bits para coordenadas com maiores valores de $U_{i,i}$.
   • Desempenho: O WaterSIC alcança uma distorção proporcional à média geométrica de $U_{i,i}^2$. Como o produto de $U_{i,i}$ é igual a $\sqrt{|\Sigma_X|}$, o WaterSIC atinge uma distorção dentro de um fator multiplicativo de $\frac{2\pi e}{12}$ (aprox. 1,423, ou 0,25 bits/entrada) do limite inferior teórico da informação.

Resultados Principais

• Quase-Optimalidade do WaterSIC: O algoritmo WaterSIC é mostrado como "alheio à base" (seu desempenho depende apenas do determinante de $\Sigma_X$, tornando-o imune a rotações aleatórias do espaço de entrada) e está dentro de 0,25 bits/entrada do ótimo teórico no regime de alta taxa.
• Desempenho do GPTQ: Embora o GPTQ padrão seja subótimo em bases específicas, os autores descobrem que, para camadas do Llama-3-8B, o GPTQ com uma rotação aleatória dos pesos performa surpreendentemente bem, frequentemente dentro de 0,1 bits do WaterSIC. Isso sugere que a base padrão usada em LLMs (provavelmente influenciada pela dinâmica do otimizador Adam) não é o pior cenário para o GPTQ, e a rotação aleatória pode mitigar a sensibilidade à base.
• Validação Empírica: Experimentos em camadas do Llama-3-8B usando dados de calibração Wikitext-2 confirmam que o WaterSIC com codificação de entropia rastreia de perto o limite teórico de enchimento de água, enquanto o GPTQ padrão mostra desempenho variável dependendo da camada e da base.

Significado e Alegações
O artigo afirma preencher a lacuna entre a teoria da informação abstrata e a quantização prática de LLMs:

1. Justificativa Teórica: Fornece uma justificativa rigorosa da teoria da informação para por que a alocação de taxa por coordenada (enchimento de água) é superior à alocação uniforme na quantização apenas de pesos.
2. Melhoria Algorítmica: Introduz o WaterSIC como um algoritmo prático e de baixa complexidade que realiza esses ganhos teóricos, oferecendo um caminho provável para distorção quase ótima.
3. Insight sobre Métodos Existentes: Esclarece o comportamento do GPTQ, explicando sua sensibilidade à base de $\Sigma_X$ e demonstrando que rotações aleatórias podem torná-lo quase ótimo, validando assim seu uso generalizado enquanto destaca o potencial para ganhos adicionais via WaterSIC.
4. Limitações e Trabalho Futuro: Os autores notam modestamente que suas aproximações de alta taxa assumem alta resolução. Eles identificam que, em taxas baixas (por exemplo, < 2 bits), outros fatores como redução de dimensionalidade e estruturas específicas de pesos tornam-se críticos. Além disso, destacam que, embora a codificação de entropia (EC) seja teoricamente ótima, seu custo computacional em GPUs pode ser proibitivo, e métodos de conformação eficientes para GPUs permanecem um desafio aberto. Eles também notam que a lacuna de 0,25 bits do WaterSIC decorre da suboptimalidade do reticulado inteiro $\mathbb{Z}^n$, que poderia ser reduzida usando reticulados de dimensão superior, embora isso complique a codificação de entropia.