Resumo Técnico: WUSH: Transformadas Adaptativas Próximas do Ótimo para Quantização de LLM

1. Definição do Problema

Quantizar pesos e ativações de Grandes Modelos de Linguagem (LLMs) é uma abordagem padrão para implantação eficiente. No entanto, uma dificuldade central surge dos "outliers" — alguns valores extremos em pesos e ativações que esticam a faixa dinâmica. Esses outliers forçam escalas de quantização a serem grandes, o que degrada a resolução efetiva das representações de poucos bits e amplifica os erros de quantização.

As estratégias de mitigação anteriores frequentemente utilizam transformadas lineares (ex: rotações) antes da quantização para espalhar a variância de forma mais uniforme entre os canais. Embora as rotações de Hadamard tenham se mostrado empiricamente eficazes, elas são tipicamente fixas e agnósticas aos dados. Elas não se adaptam às estatísticas específicas dos pesos ou ativações subjacentes. Isso levanta uma questão fundamental: se a transformada de Hadamard não é consciente dos dados, em que sentido ela pode ser considerada ótima para a quantização? Métodos existentes que aprendem transformadas (ex: SpinQuant, FlatQuant) dependem de otimização iterativa (descida de gradiente) em dados de calibração, o que pode ser dispendioso e heurístico.

2. Metodologia

Os autores propõem o WUSH, uma transformada linear por blocos, de forma fechada e consciente dos dados, projetada para a quantização conjunta de pesos e ativações. O método é derivado para minimizar a perda de saída $L_2$ sob quantizadores de bloco com escala AbsMax padrão de arredondamento para o mais próximo (RTN), cobrindo tanto formatos de ponto flutuante (FP) quanto inteiros (INT).

Derivação Teórica

O problema é formulado como a minimização do erro de quantização esperado para um bloco de pesos $W$ e ativações $X$. Os autores modelam o erro de quantização probabilisticamente:

• Para Ponto Flutuante (FP): O erro é modelado como proporcional ao valor absoluto da entrada ($\varepsilon(\alpha) = \text{diag}(\eta)\alpha$).
• Para Inteiro (INT): O erro é modelado como proporcional ao valor absoluto máximo dentro de um grupo ($\varepsilon(\alpha) = \|\alpha\|_\infty \eta$).

Ao analisar as estatísticas de segunda ordem (covariância) dos blocos de peso e ativação, os autores derivam uma transformada ótima de forma fechada. Sejam $W'$ e $X'$ matrizes derivadas da decomposição de Cholesky dos momentos de segunda ordem normalizados de pesos e ativações, respectivamente. Sejam $U, S, V$ a Decomposição de Valor Singular (SVD) de $W'^\top X'$. As transformadas ótimas $T_W$ e $T_X$ são construídas como:
$$T_X = H S^{-1/2} V^\top X'^\top$$
$$T_W = H S^{-1/2} U^\top W'^\top$$
onde $H$ é uma matriz de Hadamard normalizada.

Propriedades fundamentais desta construção:

1. Consciente dos Dados: Ao contrário das transformadas de Hadamard fixas, o WUSH incorpora $S^{-1/2}$ e os componentes da SVD ($U, V$), que dependem das estatísticas de segunda ordem específicas dos dados.
2. Não Ortogonal: A transformada resultante é geralmente não ortogonal, permit o equilíbrio de escalas entre pesos e ativações de uma forma que transformadas ortogonais não conseguem.
3. Otimalidade: O artigo prova que esta construção é provadamente ótima para quantizadores de bloco FP e assintoticamente ótima para quantizadores de bloco INT (dentro de um fator de $d$ ou $d^{o(1)}$ dependendo da distribuição).

Integração com GPTQ

O método integra-se naturalmente com o GPTQ (Greedy Per-Block Quantization). Como o GPTQ utiliza o Hessiano (informação de segunda ordem) para propagação de erro, o WUSH pode utilizar as mesmas estatísticas para construir a transformada. Os autores propõem um cronograma computacional intercalado onde as atualizações de peso e a construção da transformada são acopladas, permitindo que o WUSH seja aplicado junto ao GPTQ com sobrecarga adicional negligenciável.

Implementação

Para enfrentar o desafio de aplicar uma transformada distinta por bloco durante a inferência, os autores implementam um kernel de GPU fundido (fused GPU kernel). As matrizes de transformação WUSH são armazenadas em um layout transposto $(G, G, C)$ para maximizar a utilização da largura de banda de memória. O kernel funde a transformação por bloco com a quantização de ativação, alcançando um throughput comparável aos kernels de Hadamard otimizados.

3. Principais Contribuições

1. Transformadas Ótimas de Forma Fechada: A derivação das primeiras transformadas lineares de forma fechada e conscientes dos dados para a quantização conjunta de peso-ativação, que são provadamente próximas do ótimo para formatos tanto FP quanto INT.
2. Insight Teórico: O artigo esclarece o papel da transformada de Hadamard. Mostra que, embora a Hadamard seja uma transformada ortogonal agnóstica aos dados eficaz, a solução ótima requer um componente não ortogonal ($S^{-1/2}$) derivado das estatísticas dos dados para realmente minimizar o erro de quantização.
3. Implementação Eficiente: Um kernel de GPU especializado e fundido que lida com transformadas adaptativas por bloco sem sacrificar a velocidade de inferência, igualando o desempenho de baselines de transformada fixa.
4. Validação Empírica: Experimentos extensos demonstrando melhorias significativas de precisão sobre baselines de estado da arte (MR-GPTQ baseado em Hadamard, SmoothQuant, QuaRot, SpinQuant) através de múltiplas famílias de modelos (Llama, Qwen) e formatos de quantização (MXFP4, NVFP4, INT4).

4. Resultados Experimentais

O artigo avalia o WUSH no Llama-3.1-8B-Instruct, Llama-3.2-3B-Instruct e modelos Qwen3 (8B, 14B, 32B) usando o LM Evaluation Harness e Platinum Benchmarks.

• Perda por Camada: O WUSH atinge consistentemente a menor perda de quantização RTN por camada nos formatos MXFP4, NVFP4 e INT4 em comparação com variantes Identidade, Rotação Aleatória, Hadamard e não adaptativas.
• Precisão End-to-End:
  • No Llama-3.1-8B-Instruct com MXFP4 (W4A4), o WUSH melhora a precisão em +2.8 pontos médios sobre o baseline baseado em Hadamard (RTN) e +0.7 pontos (GPTQ).
  • Ele reduz substancialmente a lacuna de desempenho entre os formatos NVFP e MXFP. Por exemplo, no Qwen3-14B, o WUSH reduz a lacuna média NVFP4-MXFP4 para 0.36 pontos.
  • O WUSH combinado com GPTQ alcança os melhores resultados na maioria dos benchmarks, frequentemente recuperando mais de 96% do desempenho BF16 em configurações W4A4.
• Throughput: A implementação do kernel fundido entrega até 5.8× de throughput por camada sobre BF16 em GPUs RTX 5090, com apenas uma diferença de 1.3% de throughput em relação ao baseline otimizado de Hadamard.
• Sobrecarga: Os custos de pré-processamento offline são comparáveis ao GPTQ padrão. A sobrecarga de armazenamento para as transformações do lado da ativação é negligenciável (ex: ~1.4% para MXFP4 em Qwen3-8B).

5. Significância e Alegações

O artigo afirma que o WUSH esclarece o sucesso empírico das rotações de Hadamard ao mostrar que elas são um componente específico e agnóstico aos dados de uma solução ótima mais ampla. Ao combinar uma estrutura de base de Hadamard fixa com um componente de segundo momento dependente dos dados, o WUSH alcança uma transformada não ortogonal provadamente próxima do ótimo.

Os autores enfatizam que sua abordagem é fundamentada e consciente dos dados, indo além de métodos de aprendizado heurísticos ou iterativos. Eles afirmam que o WUSH torna o formato MXFP competitivo com o NVFP, abordando uma limitação chave na inferência de baixa precisão atual. O trabalho demonstra que transformadas por bloco baseadas em princípios e conscientes dos dados podem gerar melhorias significativas de precisão sem comprometer a eficiência necessária para a implantação de LLMs em larga escala.