MFA Atenção Fatorizada Multi-Matriz para LLMs Eficientes

Introdução à Necessidade de Mecanismos de Atenção Eficientes

O uso crescente de modelos de linguagem grandes (LLMs) e o surgimento de novos paradigmas na inferência trouxeram à tona o desafio da inferência eficiente em larga escala. Um gargalo significativo é o cache Key-Value (KV) dentro dos mecanismos de atenção tradicionais. Este cache se expande linearmente com o tamanho do lote e o comprimento da sequência, tornando-se um 'consumidor de memória' que dificulta a escalabilidade e expansão dos LLMs.

Embora variantes como MQA, GQA e MLA tenham surgido para abordar este problema, elas muitas vezes lutam para manter o desempenho sob restrições de memória rigorosas ou introduzem complexidades que apresentam desafios de engenharia e problemas de compatibilidade.

Atenção por Fatorização Multi-Matriz (MFA): Uma Nova Abordagem

Um artigo recente de pesquisadores da Stepes, Universidade de Tsinghua e outras instituições apresenta uma nova arquitetura de mecanismo de atenção: Multi-matrix Factorization Attention (MFA) e sua variante MFA-Key-Reuse (MFA-KR). Este mecanismo reduz significativamente o custo da inferência de modelos de linguagem, ao mesmo tempo que melhora o desempenho.

MFA e MFA-KR não só superam o MLA em desempenho, mas também igualam o desempenho tradicional do MHA, reduzindo o uso do cache KV em até 93,7%. O MFA foi projetado para simplicidade, fácil reprodução, baixa sensibilidade a hiperparâmetros e compatibilidade com vários métodos de Pos-embedding.

Abordagem e Análise do MFA

A equipe de pesquisa analisou o design geral e a capacidade dos mecanismos de atenção, identificando duas dimensões críticas relacionadas à capacidade. Esta análise levou ao desenvolvimento de novos métodos analíticos e princípios de design.

Eles introduziram o conceito de Generalized Multi-Head Attention (GMHA) como uma estrutura unificadora para entender diferentes variantes do MHA. A equipe também explorou a computação e o armazenamento de key-values de uma perspectiva de inferência e examinou a capacidade do modelo de uma perspectiva de decomposição.

Fully Parameterized Bilinear Attention (FPBA) foi estabelecido como o limite superior teórico de desempenho. Eles descobriram que o MHA e suas variantes são decomposições de baixo rank do FPBA.

Comparação com MQA e MLA

A análise focou em dois esquemas de melhoria representativos: Multi-Query Attention (MQA) e Multi-Head Latent Attention (MLA).

O MQA usa uma estratégia de compartilhamento de parâmetros mais agressiva, onde todas as heads de atenção compartilham o mesmo conjunto de parâmetros de key-value. Isso reduz o uso de memória, mas pode afetar a expressividade do modelo.
O MLA introduz um espaço latente compartilhado para compressão de parâmetros, mas o poder expressivo real é limitado pela menor dimensão, o que significa que aumentar as dimensões intermediárias não melhora significativamente o desempenho.

Principais Inovações do MFA

O desenvolvimento do MFA foi impulsionado pelo objetivo de criar um mecanismo de atenção que minimize o consumo de recursos, aproximando-se dos limites teóricos de desempenho. O design do MFA incorpora três principais inovações:

Aumento significativo do número e dimensão das heads de atenção para maximizar a capacidade do modelo.
Emprego de uma estratégia agressiva de decomposição de baixo rank para manter a eficiência dos parâmetros, expandindo o número e as dimensões das heads de atenção.
Utilização de um design de uma única head de key-value para manter o consumo de memória mínimo, mesmo com o aumento da complexidade do modelo.

Medição e Comparação de Capacidade

Para analisar mais profundamente o MFA e outros mecanismos de atenção, a equipe introduziu duas métricas principais:

Total Effective Rank (TER): O produto do número de heads de atenção e o rank de fatorização por head (FRH).
Shared Latent Subspace Dimension (SLSD): A dimensão do espaço oculto compartilhado por todas as heads de atenção.

O MFA alcança um SLSD e TER mais elevados em comparação com o MQA. Em comparação com o MLA, o MFA alcança um tamanho de cache KV menor e um TER mais elevado com orçamentos de parâmetros semelhantes, mantendo um SLSD comparável. Em comparação com o MHA tradicional, o MFA tem um TER mais elevado, embora seu SLSD seja menor.

Resultados Experimentais

Experimentos extensivos foram conduzidos para avaliar o desempenho da nova arquitetura em escalas maiores, testando modelos que variam de 1B a 7B parâmetros e dados de treinamento de 10B a 1T.

O MFA demonstrou capacidades de escalabilidade comparáveis ao MHA tradicional, mantendo um excelente desempenho mesmo em escalas maiores. Enquanto o MFA-KR exibiu um desempenho ligeiramente inferior, sua tendência de escalabilidade alinhou-se com o MHA.

As vantagens de economia de memória do MFA e MFA-KR continuaram a se expandir com o tamanho do modelo, com o MFA alcançando 87,5% de economia de memória e o MFA-KR reduzindo o uso de memória para 6,25% na maior escala.

Estudos de Ablação

Estudos de ablação validaram a eficácia do MFA e MFA-KR. Suas vantagens de desempenho também foram confirmadas em vários métodos de codificação posicional convencionais.

Perspectivas Futuras

O MFA oferece melhorias significativas com um design simples, abordando efetivamente o gargalo de memória na inferência de LLM sem adicionar complexidade extra de engenharia. Ele se integra perfeitamente ao ecossistema Transformer existente, acelerando a aplicação de LLMs em vários cenários.