Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space

Resumo Técnico: Compressed Convolutional Attention (CCA)

1. O Problema

Os modelos de transformadores baseados em Atenção Multi-Cabeça (MHA) enfrentam dois gargalos principais, especialmente em contextos longos:

1. Complexidade Computacional Quadrática: O custo de computação da atenção cresce quadraticamente com o tamanho da sequência ($S$) e a dimensão oculta ($E$), tornando o prefill (processamento de entrada) e o treinamento caros.
2. Crescimento Linear do KV-Cache: Durante a geração autoregressiva, o cache de Chaves e Valores (KV-cache) cresce linearmente com a sequência, consumindo largura de banda de memória e limitando a velocidade de decodificação.

Limitações das Soluções Atuais:

• GQA (Grouped Query Attention): Compartilha parâmetros entre cabeças de K e V para reduzir o tamanho do KV-cache, mas não reduz a contagem de FLOPs (operações de ponto flutuante) para treinamento ou prefill.
• MLA (Multi-Latent Attention): Comprime K e V em um espaço latente compartilhado. Embora reduza o KV-cache, exige projeções de "up-projection" (expansão) antes da atenção, mantendo o custo computacional alto (semelhante ao MHA) e adicionando complexidade para integrar posicionamentos como RoPE.

2. Metodologia: Compressed Convolutional Attention (CCA)

Os autores propõem o CCA, um método que realiza a operação de atenção inteiramente dentro de um espaço latente comprimido. A ideia central é projetar down (reduzir) as consultas (Q), chaves (K) e valores (V) para uma dimensão menor $\tilde{e} = E/C$ (onde $C$ é o fator de compressão) e executar toda a atenção nesse espaço reduzido.

Inovações Chave do CCA:
Para compensar a perda de expressividade devido à compressão e superar o desempenho do MLA, o CCA introduz três operações adicionais antes da atenção padrão:

1. Mistura Convolucional (Convolutional Mixing):
   • Aplica convoluções nas dimensões de sequência e canal (dentro de cada cabeça) para as projeções comprimidas de Q e K.
   • Isso adiciona expressividade ao espaço latente, permitindo uma transferência de informação mais eficiente, similar a como convoluções causais melhoram modelos SSM (como Mamba).
2. Operação QK-Mean:
   • Adiciona a média dos valores de Q e K (pré e pós-convolução) aos valores convoluídos.
   • Atua como uma conexão de skip que permite ao modelo interpolar a força da convolução e aumenta a esparsidade da diagonal da matriz de atenção quando combinado com normalização.
3. Value-Shift:
   • Gera vetores de valor a partir da entrada atual e da entrada anterior da sequência.
   • Introduz um viés indutivo onde metade das cabeças não "vê" o token presente, melhorando a estabilidade e a perplexidade (semelhante ao token-shift do RWKV).

Integração com RoPE:
Diferente do MLA, que precisa de cabeças separadas para RoPE (Rotary Positional Embedding) devido à compressão, o CCA aplica o RoPE diretamente no espaço latente comprimido, simplificando a arquitetura e eliminando a necessidade de caches separados.

Compressed Convolutional Grouped Query Attention (CCGQA):
Os autores combinam o CCA com o compartilhamento de parâmetros do GQA. O CCGQA aplica o compartilhamento de cabeças de K e V dentro do espaço latente comprimido. Isso permite:

• Redução adicional do KV-cache (até 8x ou mais).
• Desacoplamento das taxas de compressão para Q e para KV (ex: comprimir Q menos que KV).
• Uma fronteira de Pareto suave entre limites de computação (FLOPs) e memória.

3. Resultados Principais

Os experimentos foram realizados em modelos densos (1B parâmetros) e MoE (Mixture of Experts, 350M ativos/1.5B totais) treinados no conjunto de dados Zyda2.

• Desempenho vs. MHA/GQA/MLA:
  • O CCGQA supera consistentemente o MHA, GQA e MLA em qualidade do modelo (perda e benchmarks como HellaSwag, ARC) com o mesmo tamanho de KV-cache comprimido.
  • Em modelos MoE, o CCGQA com metade do KV-cache do GQA/MLA supera todos os outros métodos, alcançando 8x de compressão de KV-cache sem queda de desempenho em relação ao MHA padrão.
• Eficiência Computacional (FLOPs):
  • O CCA reduz drasticamente o custo de FLOPs para atenção (fator de $1/C$).
  • Em modelos densos, o CCA supera o MLA com 16x menos FLOPs de decodificação e 8x menos no CCGQA.
• Latência e Velocidade (H100 GPU):
  • Implementação de kernel fundido (fused kernel) em BF16.
  • Redução de latência de prefill de ~1.7x em relação ao MHA em sequências de 16k tokens.
  • Aceleração de ~1.3x no backward pass (treinamento).
  • A aceleração cresce com o tamanho da sequência, pois os custos fixos são amortizados.

4. Contribuições Chave

1. Atenção Totalmente Comprimida: É o primeiro método a realizar a atenção completa (treino e inferência) no espaço latente comprimido, eliminando a necessidade de projeções de "up" caras.
2. Ortogonalidade de Estratégias: Demonstra que compressão de parâmetros (CCA) e compartilhamento de parâmetros (GQA) são estratégias ortogonais e podem ser combinadas (CCGQA) para otimizar a fronteira de eficiência.
3. Integração Nativa de RoPE: Resolve o problema complexo de aplicar RoPE em caches comprimidos (um desafio no MLA) de forma simples e eficiente.
4. Novos Primitivos de Mistura: Introduz convoluções sequenciais e operações de shift dentro do bloco de atenção, provando que operações adicionais baratas no espaço latente podem recuperar a expressividade perdida pela compressão.

5. Significância e Impacto

O trabalho oferece uma solução prática para escalar modelos de linguagem para contextos extremamente longos (ex: raciocínio complexo, análise de documentos longos) sem os custos proibitivos de memória e computação atuais.

• Viabilidade de Longo Contexto: Ao reduzir o custo quadrático por um fator de compressão, permite processar sequências muito mais longas com o mesmo orçamento de FLOPs.
• Flexibilidade de Implantação: O CCGQA permite que os engenheiros ajustem a compressão para atender a limitações de hardware específicas (limitadas por memória ou por computação) sem sacrificar a qualidade do modelo.
• Aceleração de Treinamento: A redução significativa de FLOPs no prefill e no backward torna o treinamento de modelos de longo contexto mais rápido e barato, democratizando o acesso a modelos de alta capacidade.

Em resumo, o CCA e CCGQA representam um avanço fundamental na arquitetura de transformadores, superando as limitações de métodos anteriores ao unir compressão de espaço latente, compartilhamento de parâmetros e operações convolucionais eficientes.