Resumo Técnico: FPTQuant: Transformações de Preservação de Função para Quantização de LLMs

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) incorrem em custos computacionais e de energia significativos durante a inferência, limitando sua implantação em dispositivos de borda com recursos limitados. Embora a quantização (mapeamento de valores contínuos para uma grade inteira uniforme) seja a técnica mais eficaz para melhorar a eficiência, a quantização ingênua de LLMs frequentemente leva a uma severa degradação de desempenho. Isso se deve principalmente aos "outliers" — ativações e pesos com magnitudes elevadas que perturbam a distribuição uniforme necessária para a quantização de poucos bits.

O desafio fundamental reside no compromisso entre amplitude e precisão:

1. Aumentar a amplitude para capturar outliers resulta em uma perda de precisão para a maioria dos valores agrupados em torno de zero.
2. Manter a precisão ao realizar o clipping de outliers leva à perda de informação e ao colapso do modelo.

Trabalhos anteriores exploraram Transformações de Preservação de Função (FPTs) — operações como escalonamento ou rotações que suavizam os outliers sem alterar a saída do modelo na ausência de quantização. No entanto, os métodos existentes frequentemente enfrentam um compromisso entre expressividade (a capacidade de reduzir efetivamente o erro de quantização) e overhead de computação (custo de inferência). Muitas transformações expressivas não são fundíveis (non-mergeable), exigindo novas operações no grafo computacional e incorrendo em latência significativa.

Metodologia: FPTQuant

O artigo apresenta o FPTQuant, um framework que utiliza três FPTs inovadores, leves e expressivos, projetados especificamente para transformers. A filosofia central é aproveitar as equivariâncias e independências inerentes às operações canônicas de transformers para projetar transformações que sejam tanto fundíveis (fundindo-se nos pesos existentes) quanto expressivas.

As três transformações primárias são:

1. Transformação Pre-RoPE Fundível para Queries e Keys ($T_k, \bar{T}_k$)

• Contexto: LLMs modernos utilizam Rotary Positional Embeddings (RoPE), o que impede a fusão ingênua de transformações nas matrizes de query ($W_q$) e key ($W_k$), pois o RoPE não comuta com transformações lineares arbitrárias.
• Mecanismo: Os autores introduzem um par de transformações aplicadas antes do RoPE. $T_k$ é aplicada às keys, e $\bar{T}_k$ (uma aproximação inversa) é aplicada às queries. Estas consistem em matrizes de rotação de $2 \times 2$ escalonadas.
• Propriedade: Devido à estrutura de bloco diagonal do RoPE, essas matrizes de rotação/escalonamento específicas comutam com a operação RoPE. Isso permite que as transformações sejam fundidas diretamente nos pesos $W_q$ e $W_k$ após o treinamento, incorrendo em zero overhead de inferência.
• Aplicabilidade: Válido para modelos com RoPE (ex: série Llama) e modelos sem embeddings posicionais.

2. Transformação de Valor Multihead Fundível ($T_v, \bar{T}_v$)

• Contexto: No mecanismo de atenção, a multiplicação de matrizes em lote (BMM) opera sobre os valores ($V$) e a projeção de saída ($W_o$). A dimensão $d$ (dimensão do valor) é independente do cálculo da probabilidade de atenção.
• Mecanismo: Uma transformação invertível $T_v$ é aplicada aos pesos de projeção de valor ($W_v$), e sua inversa $\bar{T}_v$ é aplicada aos pesos de projeção de saída ($W_o$).
• Propriedade: Como as diferentes cabeças de atenção são independentes, uma matriz invertível única pode ser otimizada para cada cabeça. Esta transformação é fundida em $W_v$ e $W_o$.
• Vantagem: Oferece alta expressividade (matrizes invertíveis completas por cabeça) enquanto permanece fundível, ao contrário de métodos baseados em rotação anteriores que frequentemente compartilhavam uma única rotação entre as cabeças ou não eram fundíveis.

3. Escalonamento Residual Pseudo-Dinâmico ($S_n$)

• Contexto: Conexões residuais em transformers frequentemente ignoram camadas de normalização, levando a tokens com escalas vastamente diferentes. Isso causa outliers nos blocos subsequentes (ex: a saída do SwiGLU).
• Mecanismo: Os autores propõem mover o RMSNorm para aplicar também ao fluxo residual. Um fator de escalonamento por token $S_n$ (baseado na raiz quadrada da média do quadrado do residual) é aplicado à entrada residual. Para preservar a saída, as saídas dos blocos de atenção e MLP ($Y_n$) são reescalonadas pelo mesmo $S_n$ antes de serem adicionadas ao residual.
• Propriedade: Multiplicação de matrizes e camadas lineares sem viés (bias) comutam com o escalonamento por token. Isso permite que o escalonamento seja aplicado profundamente dentro dos blocos de atenção e MLP para reduzir os outliers intra-bloco.
• Implementação: As escalas são computadas recursivamente usando a saída normalizada do bloco anterior. Na prática, a divisão final pela escala é desnecessária, pois o LM head começa com um RMSNorm.

Estratégia de Otimização

O FPTQuant emprega um processo de otimização de dois estágios:

1. Otimização Local: As transformações são otimizadas local e independentemente para minimizar a norma $L_p$ (especificamente $L_4$) dos pesos fundidos. Isso reduz os outliers e fornece uma inicialização estável.
2. Treinamento End-to-End (Student-Teacher): Um modelo estudante quantizado (com transformações) é treinado para aproximar um modelo professor de precisão total usando a perda de Divergência de Jensen-Shannon (JSD).
   • Distinção Crucial: Ao contrário de métodos que utilizam perda de predição do próximo token, a abordagem student-teacher evita o overfitting ao conjunto de dados de treinamento (ex: Wikitext) ao igualar a distribuição de probabilidade completa do professor. Isso melhora a generalização para tarefas downstream.

Principais Contribuições

1. Novas FPTs: Introdução de três transformações específicas (Pre-RoPE, Value e Residual Scaling) que foram projetadas para serem altamente expressivas e fundíveis, eliminando o overhead de inferência.
2. Quantização INT4 Estática: Demonstração de que o FPTQuant permite a quantização estática INT4 (pesos, ativações e KV cache) com overhead mínimo, alcançando até 3.9× de aceleração sobre FP16 em tarefas de prefill.
3. Compromisso entre Precisão e Velocidade: O FPTQuant atinge um desempenho equivalente ou superior aos métodos de estado da arte (SOTA) como FlatQuant e SpinQuant, sendo significativamente mais rápido (até 29% mais rápido que o FlatQuant) e não exigindo kernels customizados.

Resultados Experimentais

Os autores avaliaram o FPTQuant em vários modelos (Llama-2/3, Ministral, Qwen) através de múltiplas configurações de quantização (W4A4KV4, W4A8KV8, etc.).

• Velocidade: O FPTQuant alcança consistentemente um aumento de velocidade de prefill de 2.8×–3.9× em relação ao FP16. É mais rápido que o SpinQuant e consistentemente mais rápido que o FlatQuant (aceleração relativa de 15–29%).
• Precisão:
  • Em configurações de implantação "difíceis" (quantizando todas as ativações intermediárias), o FPTQuant supera os baselines como SpinQuant e FlatQuant.
  • Em configurações padrão (Linear+KV), o FPTQuant performa comparavelmente ao FlatQuant, que é atualmente o SOTA para precisão.
  • No desafiador cenário W4A4KV4, o FPTQuant mostra perplexidade competitiva e pontuações de raciocínio zero-shot, frequentemente superando o SpinQuant e o QuaRot.
• Estabilidade: A análise de sensibilidade mostra que o FPTQuant é preservador de função mesmo sob atualizações de parâmetros ruidosas, enquanto métodos como FlatQuant (que dependem de inversões de ponto flutuante) e OSTQuant (não preservadores de função) apresentam instabilidade ou degradação com ruído.
• Custo de Treinamento: O treinamento do FPTQuant é eficiente. Embora seja ligeiramente mais lento que o RTN-opt, é geralmente mais rápido que o FlatQuant porque suas transformações são fundíveis, evitando as passagens de forward online exigidas por transformações não fundíveis.

Significância e Alegações

O artigo afirma que o FPTQuant representa um passo significativo para tornar os LLMs viáveis para implantação em borda (edge deployment). Ao explorar explicitamente as propriedades matemáticas das operações de transformers (equivariâncias), os autores criaram transformações que não sacrificam a velocidade de inferência pela precisão.

• Impacto Prático: O FPTQuant permite o uso de larguras de bits menores (INT4) sem o pesado custo computacional da quantização dinâmica ou de transformações não fundíveis. Isso reduz a pegada de memória e o consumo de energia dos LLMs, potencialmente democratizando o acesso à IA em smartphones e dispositivos de borda.
• Contribuição Metodológica: O trabalho destaca que entender as propriedades de comutação das operações existentes é a chave para projetar estratégias de quantização eficazes. Ele argumenta contra o uso de perda de predição do próximo token para o treinamento de FPTs devido aos riscos de overfitting, defendendo o treinamento student-teacher com correspondência de distribuição.
• Limitações: Os autores observam modestamente que, embora a abordagem seja robusta em diferentes famílias de modelos, não há garantia de que se traduza igualmente para todos os LLMs (ex: Mixture-of-Experts). Além disso, as limitações atuais de hardware (falta de suporte nativo para INT4 em muitos GPUs) significam que o benchmarking depende de emulação, e os benefícios totais podem ser plenamente realizados apenas em hardware futuro ou aceleradores específicos de borda.