Resumo Técnico: A3: uma Framework Analítica de Aproximação de Baixo RANK para Atenção

Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) alcançaram desempenho notável, mas sofrem com contagens massivas de parâmetros, resultando em altos custos computacionais, requisitos de memória e implantação dispendiosa. Embora a aproximação de baixo rank ofereça uma estratégia de compressão promissora ao decompor matrizes de pesos ($W \approx A_r B_r$), as abordagens existentes enfrentam duas limitações críticas:

1. Cegueira Arquitetural: Elas tipicamente minimizam o erro de saída de camadas lineares individuais sem considerar as características funcionais específicas dos componentes do Transformer (por exemplo, pontuações de atenção versus saídas de atenção).
2. Sobrecarga de Tempo de Execução: Elas decompõem uma única matriz de pesos grande em duas matrizes menores, o que substitui uma multiplicação de matriz geral (GEMM) por duas GEMMs menores. Isso introduz sobrecargas de tempo de execução, incluindo lançamentos extras de kernel e operações de memória, frequentemente resultando em desempenho inferior ao da poda ou quantização.

Consequentemente, os métodos atuais de baixo rank frequentemente falham em igualar a eficiência e o desempenho de outras técnicas de compressão, como poda e quantização.

Metodologia: A Framework A3

Os autores propõem o A3, uma framework de aproximação analítica de baixo rank pós-treinamento que aborda essas limitações ao decompor a camada do Transformer em três componentes funcionais: QK (Query-Key), OV (Output-Value) e MLP (Perceptron Multicamada). Em vez de minimizar erros de pesos ou saídas genéricas de camadas, o A3 minimiza a perda funcional específica de cada componente, derivando soluções analíticas de forma fechada que reduzem diretamente as dimensões ocultas.

1. A3-QK (Componente Query-Key)

• Objetivo: Minimizar o erro das pontuações de atenção pré-softmax ($A_i = Q_i K_i^T$).
• Método: Os autores formulam o problema como a minimização da norma de Frobenius da diferença entre as pontuações de atenção originais e aproximadas em um conjunto de calibração.
• Solução: Eles derivam uma solução de forma fechada (Teorema 3.2) envolvendo as matrizes de autocorrelação das entradas de query e key/value ($R_{X_q X_q}$ e $R_{X_{kv} X_{kv}}$). A aproximação de baixo rank ótima é obtida por meio de uma Decomposição em Valores Singulares (SVD) truncada ponderada da matriz de pesos fundida.
• Extensão: Para Atenção com Query em Grupo (GQA), uma SVD conjunta é aplicada entre cabeças que compartilham a mesma key/value. Para Incorporações de Posição Rotacionais (RoPE), uma decomposição CUR é utilizada para lidar com rotações em pares, ordenando as dimensões pela soma da importância de pares adjacentes.

2. A3-OV (Componente Output-Value)

• Objetivo: Minimizar o erro da saída de atenção ($O = \sum A'_i V_i W_o$).
• Método: Em vez de minimizar o erro global da saída de atenção (o que incorre em altos custos computacionais), a framework relaxa o objetivo para minimizar o erro da saída de cada cabeça independentemente.
• Solução: A solução ótima (Teorema 3.3) é derivada de forma semelhante ao A3-QK, mas foca na autocorrelação das pontuações de atenção pós-softmax multiplicadas pelas entradas de key/value ($P_i$). Isso reduz a dimensão da cabeça compartilhada pelas projeções de value e output.

3. A3-MLP (Perceptron Multicamada)

• Objetivo: Minimizar o erro da saída da camada MLP.
• Desafio: A presença de funções de ativação não lineares (por exemplo, SiLU, ReLU) impede a aplicação direta de SVD.
• Método: A framework formula o problema como a seleção de um subconjunto de colunas na matriz de pesos de projeção para baixo ($W_d$) a serem mantidas, reduzindo efetivamente a dimensão intermediária ($d_{inter}$).
• Solução: Usando princípios de decomposição CUR, os autores selecionam as $r$ principais colunas com base no produto das normas $L_2$ das linhas correspondentes nos pesos de projeção para cima e nas colunas da matriz de autocorrelação das ativações intermediárias. Esta é uma heurística simples e eficaz para a seleção de colunas NP-difícil.

Característica de Design Chave: Eficiência de Tempo de Execução

Diferentemente dos métodos tradicionais de baixo rank que substituem $XW$ por $(XA)B$ (introduzindo uma GEMM extra), o A3 reduz as dimensões ocultas ($d_{qk}$, $d_{vo}$, $d_{inter}$) diretamente. Isso resulta no mesmo número de operações GEMM que o modelo original, mas com tamanhos de matriz menores. Consequentemente, o A3 reduz o tamanho do modelo, o tamanho do cache KV e os FLOPs sem introduzir qualquer sobrecarga de tempo de execução (sem lançamentos extras de kernel).

Contribuições Principais

1. Decomposição de Componentes Funcionais: Uma configuração de baixo rank em três partes para Atenção Multi-Cabeça (MHA) que formula a otimização como a minimização de perdas funcionais para pontuações QK, saídas OV e saídas MLP, em vez de saídas genéricas de camadas.
2. Soluções Analíticas de Forma Fechada: Derivação de soluções ótimas para os três objetivos que reduzem naturalmente as dimensões ocultas, evitando a sobrecarga de tempo de execução de matrizes decompostas.
3. Adaptabilidade Arquitetural: Extensão da framework para suportar variantes modernas do Transformer, incluindo Atenção com Query em Grupo (GQA) e Incorporações de Posição Rotacionais (RoPE), superando a limitação de métodos existentes restritos ao MHA vanilla.
4. Aplicações Versáteis: Demonstração da compatibilidade do A3 com compressão de cache KV, quantização, ajuste fino e atribuições de rank misto.

Resultados Experimentais

Os autores conduziram experimentos extensos em vários LLMs (LLaMA 1/2/3.1, MPT, Phi 3, Mistral) em tarefas de pré-treinamento (WikiText-2, C4, SlimPajama) e benchmarks downstream (ARC, BoolQ, MMLU, etc.).

• Superioridade de Desempenho: O A3 supera significativamente os métodos de baixo rank mais avançados (SoTA) (por exemplo, SVD-LLM, ASVD, FWSVD).
  • Exemplo: No LLaMA 3.1-70B com uma taxa de compressão de 10%, o A3 alcança uma perplexidade de 4,69 no WikiText-2, superando o SoTA anterior (SVD-LLM) que obteve 7,87 (uma redução de 3,18 pontos de perplexidade).
  • No LLaMA-2-7B, o A3 mantém consistentemente uma perplexidade menor em taxas de compressão de 10%, 20% e 40% em comparação com SVD, FWSVD e SVD-LLM.
• Ganhos de Vazão: Como o A3 reduz os tamanhos dos problemas sem lançamentos extras de kernel, ele alcança acelerações consistentes em Tokens por Segundo (TPS). Em contraste, métodos como SVD-LLM frequentemente sofrem penalidades de velocidade ou mostram ganhos apenas em níveis de compressão agressivos.
• Eficiência do Cache KV: O A3 fornece compressão de cache KV mais forte em comparação com métodos como Palu e CLOVER, mantendo perplexidade menor mesmo em altas taxas de compressão (por exemplo, 80% no MPT-30B).
• Compatibilidade: O A3 combina-se efetivamente com quantização (por exemplo, HQQ), criando um espectro contínuo de níveis de compressão e alcançando uma fronteira de Pareto melhor do que a quantização isolada em regimes sub-4-bit.

Significado e Alegações

O artigo afirma que o A3 representa um passo significativo para frente na compressão de modelos ao deslocar o foco da otimização de erros genéricos de camadas lineares para erros de componentes funcionais. Ao derivar soluções analíticas que reduzem diretamente as dimensões ocultas, o A3 alcança:

• Precisão Superior: Manter o desempenho do modelo de ponta a ponta significativamente melhor do que os métodos existentes de aproximação de baixo rank.
• Eficiência de Hardware: Redução do tamanho do modelo, footprint de memória (incluindo cache KV) e FLOPs sem as penalidades de tempo de execução associadas à multiplicação de matrizes decompostas.
• Ampla Aplicabilidade: Adaptação bem-sucedida a arquiteturas modernas (GQA, RoPE) que métodos anteriores de baixo rank não conseguiam lidar efetivamente.

Os autores enfatizam que o A3 é um método pós-treinamento, sem treinamento, que oferece uma alternativa prática e de alto desempenho à poda e quantização, particularmente para cenários onde a sobrecarga de tempo de execução deve ser minimizada.