When Local Variance Optimality Is Not Enough: RoPE-Aligned Q/K Rotations for Dynamic 4-Bit Quantisation

Resumo Técnico: Rotações Q/K Alinhadas ao RoPE para Quantização Dinâmica de 4 Bits

Enunciado do Problema
A quantização pós-treinamento (PTQ) de Grandes Modelos de Linguagem (LLMs) em larguras de bits baixas (ex: W4A4KV4) é fortemente influenciada por outliers de ativação. Métodos baseados em rotação abordam isso aplicando transformações ortogonais para redistribuir a energia de ativação enquanto preservam a função não quantizada da rede. Uma escolha de design crítica nesses métodos é o "suporte de mistura" (mixing support): se misturam todos os canais dentro de uma cabeça de atenção ou se restringem a subespaços estruturais específicos.

O Rotary Position Embedding (RoPE) particiona cada cabeça de atenção em pares de frequências bidimensionais fixos. Essa estrutura levanta uma questão específica: Será que um transform que respeita essa decomposição (mistura por pares) pode melhorar a precisão da quantização em comparação com a mistura de cabeça total (full-head mixing)? Embora trabalhos anteriores (ex: FPTQuant) tenham utilizado rotações por par que comutam com o RoPE, permanece incerto se restringir o suporte de mistura a esses pares, embora teoricamente ideal para certos objetivos locais, traduz-se em melhor desempenho de ponta a ponta sob quantização dinâmica.

Metodologia
Os autores investigam essa questão comparando três configurações de transformação online aplicadas às queries (Q) e chaves (K) após o RoPE:

1. Hadamard de cabeça total (Full-head Hadamard): O método de linha de base que mistura todos os $d_h$ canais em uma cabeça.
2. Apenas por pares (Pairwise-only): Uma rotação compartilhada pela cabeça, restrita aos dois canais de cada par de frequência do RoPE, comutando com o RoPE.
3. Par por par + Hadamard (Pairwise + Hadamard): A composição da rotação por pares com o Hadamard de cabeça total.

O estudo foca na quantização dinâmica W4A4KV4 através de quatro checkpoints de modelos (Llama-3.2-1B, Llama-3.2-3B, Llama-3.1-8B e Mistral-7B-v0.3).

Caracterização Teórica:
O artigo estabelece um resultado de conversa teórica: para uma única cabeça com frequências RoPE distintas, os únicos mapas ortogonais que comutam com o RoPE são rotações planares independentes dentro dos pares de frequência. Isso define a "família que comuta com o Rope".

Objetivo de Otimização:
Para a parametrização compartilhada por cabeça avaliada (onde um ângulo é compartilhado entre todas as cabeças em uma camada para um par de frequência específico), os autores derivam um ângulo de rotação de forma fechada $\phi^*_k$. Este ângulo minimiza a variância máxima do canal sob um surrogate de covariância agrupada e média por posição. A implementação é verificada para atingir esse mínimo analítico.

Estratégias de Estimativa:
O estudo avalia diferentes estimadores para o ângulo compartilhado:

• Padrão (Default): Agrupa acumuladores de momentos brutos de Q e K entre as cabeças antes do centramento.
• Apenas K (K-only): Constrói a covariância a partir de observações de K apenas (já que apenas K é quantizado pós-transformação).
• Equilibrado por fluxo (Stream-balanced): Atribui pesos iguais aos momentos não centrados de Q e K.

Principais Contribuições

1. Benchmark Local Exato: O artigo fornece uma caracterização de conversa provando que, para frequências distintas, a família ortogonal que comuta com o RoPE consiste exatamente em rotações independentes por pares. Ele também deriva o ângulo de variância mínima minimax para um surrogate de covariância agrupada e verifica sua obtenção.
2. Resultado Negativo Escopado: Sob quantização dinâmica W4A4KV4, a configuração por pares compartilhada por cabeça resulta em perplexidade mais alta do que a mistura Hadamard de cabeça total em todos os checkpoints e comprimentos de contexto testados (até 128K). Mesmo usando o optimum do surrogate verificado $\phi^*_k$, a configuração apenas por pares apresenta desempenho inferior ao método de linha de base.
3. Evidência sobre Alinhamento de Proxy e Suporte de Mistura: O artigo identifica dois fatores que explicam a discrepância entre a otimalidade local e o desempenho de quantização:
   • Descompasso de Objetivo (Objective Mismatch): A rotação de forma fechada minimiza um segundo momento médio por posição (covariância agrupada), enquanto o quantizador dinâmico define sua escala com base em um intervalo de grupo por token que abrange múltiplos pares de frequência. Otimizar o primeiro não garante o controle sobre o segundo.
   • Suporte de Mistura (Mixing Support): As rotações por pares têm um suporte de mistura de $b=2$, enquanto o Hadamard de cabeça total tem $b=d_h$. Um experimento de interpolação variando o suporte de bloco-Hadamard de $b=2$ para $b=d_h$ mostra que o intervalo de K, o erro de quantização relativo e a degradação da perplexidade diminuem monotonicamente à medida que o suporte aumenta.

Resultos

• Perplexidade: Substituir o Hadamard de cabeça total pela configuração por pares compartilhada por cabeça aumenta a perplexidade em cada comparação de contexto curto (ex: +1.33 PPL no Llama-3.2-1B, +0.48 PPL no Llama-3.2-3B). Esta ordem mantém-se para contextos longos (até 128K).
• Composição: Compor a rotação por pares com o Hadamard (Pairwise + Hadamard) satisfaz o critério de intervalo de $\pm 0.05$-PPL em relação à linha de base, indicando que o componente por pares não degrada a mistura de cabeça total quando adicionado.
• Impacto do Estimador: Estimar o ângulo compartilhado a partir de K sozinho melhora a configuração por pares em todos os checkpoints em comparação com o estimador padrão, mas não elimina a lacuna de desempenho em relação à mistura de cabeça total.
• Interpolação de Suporte: À medida que o suporte de mistura aumenta de 2 canais (por pares) para a cabeça total, o intervalo de K, o erro de quantização relativo e a lacuna de perplexidade diminuem, aproximando-se do desempenho da linha de base de cabeça total em $b=d_h$.

Significância e Alegações
O artigo conclui que o alinhamento estrutural isolado (comutar com o RoPE) é insuficiente para garantir que uma rotação reduza o erro de quantização. A otimalidade de uma transformação para um surrogate estruturado (covariância agrupada) não reduz necessariamente o erro de quantização se o surrogate e o suporte de mistura da transformação estiverem desalinhados com a estatística de definição de escala do quantizador (intervalo por grupo por token).

Os autores alegam modestamente que seus resultados separam três conceitos distintos:

1. Otimalidade para um surrogate de covariância agrupada.
2. Melhoria na estatística de intervalo por token do quantizador.
3. Qualidade de ponta a ponta do modelo (perplexidade).

Os achados sugerem que, para quantização dinâmica, um suporte de mistura mais amplo (cabeça total) é mais eficaz do que a mistura por pares estruturalmente restrita, principalmente porque o mecanismo de definição de escala do quantizador opera em uma estatística (intervalo por grupo por token) que se beneficia de um tamanho de suporte maior, o qual a transformação por pares não consegue prover. O artigo não afirma ter encontrado um método superior, mas sim fornece um resultado negativo que esclarece as limitações da mistura por pares alinhada ao RoPE neste regime específico de quantização dinâmica.