Resumo Técnico: StiefAttention

Definição do Problema

A inferência de contexto longo em Grandes Modelos de Linguagem (LLMs) é cada vez mais limitada pela capacidade de memória e pelos requisitos de largura de banda do cache Key-Value (KV). Durante a decodificação autorregressiva, o cache KV cresce linearmente com o comprimento da sequência e o tamanho do lote (batch size), consumindo frequentemente uma memória comparável aos pesos do próprio modelo (ex: ~16 GB para o Llama3-8B em um contexto de 32k).

Os métodos de compressão de pós-treinamento existentes tipicamente empregam decomposição de matrizes de baixo posto (como a Decomposição de Valor Singular, SVD) para projetar chaves (keys) e valores (values) em dimensões inferiores. No entanto, esses métodos otimizam objetivos de proxy, como a minimização do erro de reconstrução de Frobenius dos vetores de chave/valor brutos ou dos escores de atenção pré-softmax ($QK^\top$). Os autores argumentam que esses objetivos de proxy falham em considerar como os erros de reconstrução interagem com a softmax de atenção, a mistura de valores (value mixing) e as transformações subsequentes das camadas do decodificador (normalização, caminhos residuais, MLPs). Consequentemente, uma projeção que é ótima para reconstruir tensores intermediários ainda pode induzir erros significativos na saída final da camada do decodificador, degradando a qualidade da geração de ponta a ponta.

Metodologia: StiefAttention

O artigo introduz o StiefAttention, um método de compressão de cache KV de pós-treinamento que minimiza diretamente o erro de reconstrução na saída da camada do decodificador, em vez de em quantidades de atenção intermediárias.

Mecanismo Central

1. Função Objetivo: Em vez de minimizar $\|K - \tilde{K}\|_F$, o StiefAttention minimiza o erro relativo entre a saída original da camada do decodificador $f_\ell(x)$ e a saída $\tilde{f}_\ell(x)$ produzida pelo mecanismo de atenção comprimido. Este objetivo captura o efeito cumulativo da softmax, da mistura de valores e das não-linearidades subsequentes.
2. Projeção Ortonormal na Variedade de Stiefel: O método aprende bases de projeção ortonormais ($P_K, P_V$) que residem na variedade de Stiefel (matrizes com colunas ortonormais).
3. Predição de Base Baseada em Gradiente:
   • Em vez de computar SVD no conjunto de calibração completo (o que é intensivo em termos de memória), o StiefelAttention utiliza um preditor treinável leve (um MLP) para gerar bases de projeção a partir de estatísticas simples de ativação.
   • Características de Entrada: O preditor recebe a média ($\mu$) e a variância ($\sigma^2$) por dimensão dos vetores de chave e valor como características de entrada.
   • Treinamento: O MLP é treinado via gradiente descendente para mapear essas estatísticas em matrizes $A_K$ e $A_V$.
   • Ortonormalização: As matrizes de saída são ortonormalizadas via decomposição QR ($A = QR$). O fator ortogonal $Q$ serve como a base de projeção, pois o fator triangular $R$ representa apenas uma mudança de coordenadas dentro do subespaço.
4. Alocação de Posto (Rank Allocation): O método constrói perfis de posto de erro por camada. Ele emprega uma estratégia de alocação de orçamento sequencial onde os postos são selecionados camada por camada para minimizar o erro cumulativo da saída da camada do decodificador sob um orçamento total de cache KV especificado pelo usuário. Isso leva em conta a propagação de erro entre as camadas.
5. Inferência: Durante a inferência, as chaves e valores são armazenados em forma comprimida ($K\downarrow = K P_K$). As dimensões completas são reconstruídas on-the-fly ($\tilde{K} = K\downarrow P_K^\top$) antes da computação da atenção. As bases de projeção podem ser incorporadas (folded) nos pesos lineares de atenção, não incorrendo em sobrecarga adicional de tempo de execução em comparação com os métodos padrão de baixo posto.

Principais Contribuições

1. Otimização Direta da Camada do Decodificador: Os autores propõem otimizar a compressão do cache KV diretamente contra o erro de saída da camada do decodificador, indo além dos objetivos de proxy intermediários, como reconstrução de tensores ou aproximação de escores de atenção.
2. Framework StiefAttention: Um método de pós-treinamento que aprende bases ortonormais de estatísticas de ativação usando um preditor leve, permitindo compressão eficiente sem retreinar todo o modelo.
3. Alocação Sequencial de Posto: Uma estratégia para alocar postos através das camadas com base em superfícies de erro locais avaliadas em entradas que já incluem os efeitos das compressões das camadas anteriores.
4. Validação Empírica: Avaliação abrangente no Llama3-8B e Qwen3-8B mostrando que o StiefAttention supera os baselines de estado da arte (especificamente o EigenAttention) sob orçamentos de cache KV equivalentes.

Resultados Experimentais

Os autores avaliaram o StiefAttention no Llama3-8B e Qwen3-8B, comparando-o com o baseline FP16 não comprimido e o EigenAttention (um método baseado em SVD líder que incorpora as queries na projeção).

• Ganhos de Desempenho: Sob condições de iso-compressão (mesmo orçamento de cache KV), o StiefAttention superou significativamente o EigenAttention:
  • Llama3-8B: Alcançou uma melhora de 4,2 pontos na perplexidade C4 e um aumento de 8,9 pontos na precisão de zero-shot MMLU em comparação ao EigenAttention.
  • Qwen3-8B: Demonstrou tendências semelhantes, com o StiefAttention mantendo menor perplexidade e maior precisão, particularmente conforme as taxas de compressão aumentavam (compressão mais forte).
• Robustez: O método mostrou uma degradação suave de desempenho à medida que a razão do cache KV diminuía, sem colapso abrupto.
• Compatibilidade de Quantização: O StiefAttention mostrou-se substancialmente mais robusto à quantização de cache KV de baixa precisão (ex: 4-bit e 8-bit) do que o EigenAttention, mantendo menor perplexidade em tamanhos de cache comparáveis.

Análise de Significância

O artigo fornece uma análise diagnóstica explicando por que o StiefAttention melhora o desempenho de ponta a ponta:

• Consistência Direcional: Embora o EigenAttention reconstrua melhor as saídas de atenção intermediárias (reduzindo o erro de saída de atenção em ~29,6% em relação ao StiefAttention), o StiefAttention alcança uma melhor preservação da direção da saída da camada do decodificador (maior similaridade de cosseno, +3,3%) e um erro de saída de camada total menor (melhora de 5,2%).
• Sensibilidade das Camadas Iniciais: A análise revela que a lacuna de similaridade de cosseno é mais pronunciada nas camadas iniciais. Os autores postulam que perturbações nas camadas iniciais são amplificadas pelos blocos não lineares subsequentes (softmax, caminhos residuais, MLP). Métodos do tipo SVD, que otimizam para reconstrução intermediária, frequentemente falham em preservar a informação direcional específica necessária para esses blocos downstream. Ao otimizar diretamente para a saída do decodificador, o StiefAttention alinha as bases de projeção com o "subespaço relevante para a saída", levando a uma melhor qualidade de geração final.

Limitações e Trabalhos Futuros

Os autores reconhecem várias limitações:

• Escopo: Os experimentos foram realizados principalmente em uma sequência de 2048 devido a restrições de memória de GPU; uma avaliação mais ampla em contextos mais longos é necessária.
• Simplificação de Otimização: A seleção de posto usa uma trajetória sequencial gulosa (greedy) em vez de uma otimização global exaustiva.
• Representação de Características: O preditor atualmente utiliza apenas média e variância diagonal; resumos de ativação mais ricos (ex: sketches de covariância) ficam para trabalhos futuros.
• Avaliação de Sistema: Uma avaliação completa do sistema, incluindo latência de prefill/decode e throughput sob condições reais de serviço, está pendente.

O artigo conclui que otimizar diretamente para a saída final da camada do decodificador é uma estratégia mais eficaz para a compressão do cache KV do que otimizar para objetivos de proxy intermediários, oferecendo um caminho prático para reduzir a pegada de memória sem sacrificar a qualidade da geração.