ReSpinQuant: Efficient Layer-Wise LLM Quantization via Subspace Residual Rotation Approximation

Título: ReSpinQuant: Quantização Eficiente de Camada por Camada para LLMs via Aproximação de Rotação de Subespaço Residual

1. O Problema

A implementação de Grandes Modelos de Linguagem (LLMs) em dispositivos com recursos limitados enfrenta desafios significativos devido ao grande footprint de memória e custos computacionais. A Quantização Pós-Treinamento (PTQ) é uma solução padrão para comprimir pesos e ativações para baixa precisão (ex: 4-bit ou 3-bit). No entanto, a presença de outliers extremos nos canais de ativação expande desproporcionalmente o intervalo dinâmico, causando erros severos de quantização.

As abordagens existentes para mitigar outliers baseadas em rotação dividem-se em dois paradigmas, ambos com limitações:

1. Rotação Global: Aplica uma única matriz de rotação aprendida em todo o modelo.
   • Vantagem: Permite fusão offline das rotações nos pesos, mantendo a sobrecarga de inferência mínima.
   • Desvantagem: Baixa expressividade, pois todas as camadas são forçadas a compartilhar a mesma base, incapaz de lidar com padrões de outliers específicos de cada camada.
2. Transformação por Camada (Layer-wise): Atribui uma matriz de rotação única e aprendível a cada camada.
   • Vantagem: Alta expressividade e adaptação local, resultando em maior precisão.
   • Desvantagem: Impede a fusão offline das rotações nos pesos, exigindo computações online adicionais (transformação de ativações) que geram uma sobrecarga computacional significativa e degradam o throughput.

O objetivo do ReSpinQuant é resolver esse trade-off, alcançando a alta expressividade das transformações por camada com a eficiência de inferência das rotações globais.

2. Metodologia

O ReSpinQuant propõe um framework que combina rotações por camada de tamanho completo (full-size) com uma aproximação de rotação residual em subespaço.

A. Fusão Offline de Pesos e Rotações por Camada

Diferente dos métodos globais, o ReSpinQuant atribui matrizes ortogonais aprendíveis distintas ($R_i$) para diferentes estágios de cada camada (entrada do MHSA, saída do FFN, etc.).

• Fusão: As matrizes de rotação são matematicamente absorvidas nas matrizes de peso das camadas anteriores (ex: $W_{merged} = R^T W R$). Isso permite que a computação principal durante a inferência permaneça idêntica a uma camada linear padrão, sem necessidade de transformar as ativações online.

B. O Problema do Desalinhamento de Base (Basis Mismatch)

Ao usar rotações diferentes para a entrada e saída de um bloco residual ($x_{out} = x_{in} + Block(x_{in})$), surge um desalinhamento de base. A conexão residual na coordenada rotacionada torna-se:
$$\tilde{x}_{out} = R_{out} R_{in}^T \tilde{x}_{in} + \dots$$
Onde $T = R_{out} R_{in}^T$ é a matriz de transição. Em métodos globais, $T=I$. No ReSpinQuant, $T$ não é identidade, exigindo correção.

C. Aproximação de Rotação Residual em Subespaço

A inovação central do papel é a observação empírica de que, ao inicializar as rotações com uma matriz de Hadamard e otimizá-las via otimizador de Cayley, as matrizes aprendidas não se desviam significativamente de sua inicialização. Consequentemente, a matriz de transição residual $T$ permanece diagonalmente dominante e próxima da identidade ($T \approx I$).

Para resolver o desalinhamento sem o custo $O(D^2)$ de uma multiplicação de matriz densa $D \times D$:

1. Identificação do Subespaço Principal: Realiza-se uma Decomposição em Valores Singulares (SVD) na diferença $\Delta T = T - I$. Os vetores singulares principais (que capturam a maior parte do desalinhamento) formam uma base de subespaço de baixa dimensão $Q \in \mathbb{R}^{D \times r}$, onde $r \ll D$.
2. Projeção e Rotação: Em vez de aplicar $T$ diretamente, o método projeta o vetor de entrada no subespaço de baixa dimensão, aplica uma rotação densa aprendível apenas nesse subespaço ($r \times r$) e reprojeta de volta.
3. Complexidade: Isso reduz a complexidade computacional da correção residual de $O(D^2)$ para $O(D)$ (linear), mantendo a expressividade da rotação densa.

3. Principais Contribuições

• Framework ReSpinQuant: Introduz um método que utiliza rotações por camada de tamanho completo (maximizando a capacidade de representação) que são totalmente fundidas nos pesos offline, eliminando a sobrecarga de transformação de ativações.
• Aproximação de Subespaço Residual: Resolve o problema de desalinhamento de base em conexões residuais através de uma projeção de baixa dimensão, garantindo eficiência computacional sem sacrificar a precisão.
• Análise Empírica: Demonstra que as matrizes de rotação otimizadas permanecem próximas à inicialização (Hadamard), validando a premissa de que o desalinhamento é de baixo rank.
• Eficiência e Desempenho: Valida que o método alcança o estado da arte (SOTA) em configurações extremas (W4A4 e W3A3) com uma sobrecarga de inferência quase nula.

4. Resultados Experimentais

Os experimentos foram conduzidos nas famílias de modelos LLaMA-2, LLaMA-3 e LLaMA-3.2.

• Desempenho (W4A4 e W3A3):

  • O ReSpinQuant superou consistentemente métodos de rotação global (como SpinQuant e QuaRot) e métodos por camada estruturados (como FlatQuant e OSTQuant).
  • No cenário desafiador W3A3 (3-bit para pesos e ativações), o ReSpinQuant reduziu a Perplexidade (PPL) em mais de 19 pontos no LLaMA-3.2 1B em comparação com o SpinQuant, demonstrando robustez onde outros métodos falham.
  • Em configurações W4A4, o modelo quantizado com ReSpinQuant (3B) superou modelos de precisão completa de menor tamanho (1B), deslocando a fronteira de Pareto.

• Eficiência Computacional:

  • Parâmetros: Embora o treinamento utilize $O(L \cdot D^2)$ parâmetros (1091M para LLaMA-3 8B), a maioria é fundida. A footprint online é de apenas 8.4M parâmetros (0.1% do modelo).
  • Custo Computacional: A sobrecarga adicional de MACs (Multiplicações e Acumulamentos) é de apenas ~0.2% (32.3M vs 15.37T MACs totais).
  • Latência: A latência de inferência (TTIT) aumentou apenas marginalmente (~1.7%) em comparação com o SpinQuant global, mantendo velocidades comparáveis em GPUs H100.

5. Significado e Impacto

O ReSpinQuant representa um avanço significativo na democratização da IA e na eficiência de modelos grandes:

• Viabilidade de Implantação: Permite a execução de LLMs de alta qualidade em hardware de consumo e dispositivos de borda, reduzindo drasticamente os requisitos de memória e largura de banda.
• Green AI: Ao reduzir a sobrecarga de inferência e permitir a execução de modelos maiores em hardware menor, contribui para a redução da pegada de carbono associada à implantação de modelos.
• Equilíbrio Ideal: O trabalho demonstra que é possível superar a barreira entre expressividade (adaptação por camada) e eficiência (fusão de pesos), oferecendo uma solução prática para a quantização extrema de LLMs sem a necessidade de kernels de hardware especializados complexos.

Em suma, o ReSpinQuant estabelece um novo padrão para a quantização pós-treinamento, provando que a alta precisão em baixos bits não precisa custar desempenho de inferência.