O que é a atenção de consulta agrupada (GQA)?

Ainslie et al conceberam a atenção de consulta agrupada como uma otimização da atenção multi-título (MHA), o inovador algoritmo de autoatenção introduzido no seminal artigo de 2017 “Attention is All You Need” que estabeleceu as redes neurais transformadora. Mais especificamente, foi proposta como uma generalização e aplicação mais contida da atenção multi-query (MQA), uma otimização anterior da MHA.

Embora a atenção multi-cabeça padrão tenha catalisado um salto evolutivo em aprendizado de máquina, processamento de linguagem natural (NLP) e IA generativa, ela é extremamente exigente em termos de recursos computacionais e largura de banda de memória. À medida que os LLMs ficaram maiores e mais sofisticados, os requisitos de uso de memória se tornaram um obstáculo para o avanço, especialmente para os LLMs baseados apenas em decodificadores autoregressivos usados em geração de texto, resumo e outras tarefas de IA generativa.

Pesquisas subsequentes concentraram-se em técnicas para aprimorar ou simplificar a atenção multi-cabeça. Técnicas como flash attention e ring attention melhoram a maneira como as GPUs responsáveis pelo treinamento e execução dos modelos gerenciam cálculos e uso de memória. Outras, como a GQA e a MQA, exploraram mudanças na forma como as arquiteturas de transformação processam tokens.

A atenção de consulta agrupada tem como objetivo equilibrar as compensações entre a atenção multi-título padrão e a atenção multi-query. A primeira maximiza a precisão ao custo de um maior overhead de largura de banda de memória e velocidade reduzida, enquanto a segunda maximiza a velocidade e a eficiência em detrimento da precisão.

Para entender como a atenção de consulta agrupada otimiza os modelos transformer, é importante primeiro compreender como a atenção multi-cabeça funciona de maneira geral. Tanto a GQA quanto a MQA simplesmente refinam, em vez de substituírem, a metodologia central da MHA.

A força motriz por trás dos LLMs e de outros modelos que utilizam a arquitetura de transformação é a autoatenção, um framework matemático para compreender as relações entre os diferentes tokens em uma sequência. A autoatenção permite que um LLM interprete dados textuais não apenas com base em definições estáticas, mas também considerando o contexto fornecido por outras palavras e expressões.

Nos LLMs autoregressivos utilizados para geração de texto, o mecanismo de atenção ajuda o modelo a prever o próximo token em uma sequência ao determinar quais tokens anteriores merecem mais “atenção” naquele momento. Informações dos tokens considerados mais relevantes recebem pesos de atenção maiores, enquanto informações dos tokens considerados irrelevantes recebem pesos de atenção próximos de 0.

O mecanismo de atenção multi-título que impulsiona os modelos transformer gera informações contextuais ricas ao calcular a autoatenção diversas vezes em paralelo, dividindo as camadas de atenção em múltiplas títulos de atenção.

Os autores de “Attention is All You Need” descreveram seu mecanismo de atenção utilizando a terminologia de um banco de dados relacional: consultas, chaves e valores. Bancos de dados relacionais foram projetados para simplificar o armazenamento e a recuperação de dados relevantes: eles atribuem um identificador único (“chave”) a cada dado, e cada chave está associada a um valor correspondente.O objetivo de um banco de dados relacional é associar cada consulta à chave apropriada.

Atenção de consulta agrupada versus atenção multi-query versus atenção multi-título.

• Um vetor de consulta, Q, que representa as informações que o token está "buscando". Por exemplo, o vetor de consulta de um substantivo pode representar uma busca por adjetivos que o descrevam.
• Um vetor chave, K, que representa as informações que o token contém. Pontuações de alinhamento, que representam a relevância do vetor de chave de cada token em relação ao vetor de consulta dos outros tokens, são utilizadas para calcular os pesos de atenção.
• Um vetor de valor, V, representando a informação contextual que será atualizada pelas contribuições ponderadas por atenção dos vetores de chave de outros tokens.

As interações matemáticas entre esses 3 vetores, mediadas pelo mecanismo de atenção, são a forma como um modelo ajusta sua compreensão contextual específica de cada token.
 

Para gerar cada um desses 3 vetores para um determinado token, o modelo parte da representação vetorial original do token, incorporação de vetor: uma codificação numérica na qual cada dimensão do vetor corresponde a algum elemento abstrato do significado semântico do token. O número de dimensões desses vetores é um hiperparâmetro predeterminado.

Os vetores Q, K e V de cada token são gerados ao se passar a representação vetorial original do token por uma camada linear que precede a primeira camada de atenção. Essa camada linear é particionada em 3 matrizes únicas de pesos do modelo: W_Q, W_K e W_V. Os valores específicos desses pesos são aprendidos por meio de pré-treinamento auto supervisionado em um conjunto massivo de exemplos textuais.

Multiplicar a representação vetorial original do token por W_Q, W_K e W_V resulta, respectivamente, em seu vetor de consulta, vetor de chave e vetor de valor correspondentes. O número de dimensões d que cada vetor contém é determinado pelo tamanho de cada matriz de pesos. Q e K terão o mesmo número de dimensões, _dk.

Esses 3 vetores são então encaminhados para a camada de atenção.

Na camada de atenção, os vetores Q, K e V são utilizados para calcular uma pontuação de alinhamento entre cada token em cada posição em uma sequência. Essas pontuações de alinhamento são então normalizadas em pesos de atenção usando uma função softmax.

Para cada token x em uma sequência, as pontuações de alinhamento são calculadas computando o produto escalar do vetor de consulta do token Q_x com o vetor de chave K de cada um dos outros tokens: ou seja, multiplicando-os. Se uma relação significativa entre dois tokens se reflete na semelhança entre seus respectivos vetores, multiplicá-los resultará em um valor elevado. Se os dois vetores não estiverem alinhados, multiplicá-los resultará em um valor pequeno ou negativo. A maioria dos modelos transformadores utilizam uma variante chamada atenção por produto escalar escalada, na qual QK é escalada, ou seja, multiplicada por $\frac{1}{\sqrt{d_K}}$  para melhorar a estabilidade de treinamento.

Essas pontuações de alinhamento entre consulta e chave são então inseridas em uma função softmax. O softmax normaliza todas as entradas para um valor entre 0 e 1, de forma que a soma total seja 1. Os resultados da função softmax são os pesos de atenção, cada um representando a parcela (de 1) da atenção do token x a ser direcionada a cada um dos outros tokens. Se o peso de atenção de um token estiver próximo de 0, ele será ignorado. Um peso de atenção de 1 significa que um token recebe toda a atenção de x e os demais serão ignorados.

Por fim, o vetor de valor de cada token é multiplicado por seu peso de atenção. Essas contribuições ponderadas pela atenção de cada token anterior são calculadas em média e somadas à representação vetorial original do token x. Com isso, a incorporação vetorial do token x é atualizada para refletir o contexto fornecido pelos outros tokens da sequência que lhe são relevantes.

A incorporação vetorial atualizada é então enviada para outra camada linear, com sua própria matriz de pesos W_Z, onde o vetor com contexto atualizado é normalizado para um número consistente de dimensões e, em seguida, encaminhado para a próxima camada de atenção. Cada camada de atenção progressiva captura maiores nuances contextuais.
 

Utilizar as médias das contribuições ponderadas por atenção de outros tokens, em vez de considerar individualmente cada parte do contexto ponderado, é matematicamente eficiente, mas resulta em perda de detalhes.

Para compensar, as redes transformador dividem a representação vetorial do token de input original em h pedaços de tamanho igual. Da mesma forma, eles dividem W_Q, W_K e W_V em h subconjuntos denominados títulos de consulta, títulos de chave e títulos de valor, respectivamente. Cada título de consulta, de chave e de valor recebe uma parte da representação vetorial original do token. Os vetores produzidos por cada um desses trios paralelos de cabeças de consulta, de chave e de valor são encaminhados para a respectiva cabeça de atenção. Eventualmente, as saídas desses circuitos paralelos de h são concatenadas novamente para atualizar a incorporação vetorial completa do token.

Durante o treinamento, cada circuito aprende pesos distintos que capturam um aspecto separado dos significados semânticos. Isso, por sua vez, ajuda o modelo a processar as diversas maneiras pelas quais as implicações de uma palavra podem ser influenciadas pelo contexto das outras palavras ao seu redor.

A desvantagem da atenção multi-título padrão não reside tanto em algum defeito crucial, mas na ausência de qualquer otimização. A MHA foi o primeiro algoritmo do gênero e representa a execução mais complexa do seu mecanismo geral de cálculo de atenção.

A maior parte da ineficiência da MHA decorre da abundância de cálculos e parâmetros do modelo. Na MHA padrão, cada cabeça de consulta, de chave e de valor em cada bloco de atenção possui sua própria matriz de pesos. Assim, por exemplo, um modelo com 8 cabeças de atenção em cada camada, muito menos do que a maioria dos modernos LLMs, exigiria 24 matrizes de pesos únicas apenas para as cabeças Q, K e V da camada, implicando um enorme número de cálculos intermediários em cada camada.

Uma consequência dessa configuração é que ela é computacionalmente cara. Os requisitos computacionais para a MHA aumentam de forma quadrática em relação ao comprimento da sequência: dobrar o número de tokens em uma sequência de entrada exige quatro vezes a complexidade. Isso impõe limites práticos rígidos ao tamanho das janelas de contexto.

A MHA também impõe uma grande demanda à memória do sistema. As GPUs não dispõem de muita memória interna para armazenar as saídas da enorme quantidade de cálculos intermediários que precisam ser acessados a cada etapa subsequente do processamento. Esses resultados intermediários são, em vez disso, armazenados em memória de alta largura de banda (HBM), que não está localizada no próprio chip da GPU. Isso implica uma pequena latência sempre que as chaves e os valores precisam ser lidos da memória. À medida que os modelos transformer passaram a escalar para muitos bilhões de parâmetros, o tempo e o poder de computação necessários para treinar e executar a inferência tornaram-se um gargalo no desempenho do modelo.

Progressos adicionais exigiram métodos para reduzir o número de etapas computacionais sem diminuir a capacidade dos transformadores de aprender e reproduzir padrões linguísticos intrincadamente complexos. Foi nesse contexto que a MQA e, posteriormente, a GQA foram introduzidas.

A atenção multi-query (MQA) é um mecanismo de atenção mais eficiente computacionalmente que simplifica a atenção multi-cabeça para reduzir o uso de memória e os cálculos intermediários. Em vez de treinar uma cabeça de chave e uma cabeça de valor únicas para cada cabeça de atenção, a MQA utiliza uma única cabeça de chave e uma única cabeça de valor em cada camada. Portanto, os vetores de chave e de valor são calculados apenas uma vez; esse conjunto único de vetores de chave e de valor é então compartilhado entre todas as h cabeças de atenção.

Essa simplificação reduz consideravelmente o número de projeções lineares que o modelo precisa calcular e armazenar em memória de alta largura de banda. De acordo com o artigo de 2019 que introduziu a MQA , a MQA permite um armazenamento de pares chave-valor de 10 a 100 vezes menor (ou cache de KV ) e uma inferência de decodificador 12 vezes mais rápida. A redução no uso de memória proporcionada pela MQA também acelera significativamente o treinamento, permitindo um tamanho de lote maior.

Apesar de seus benefícios, a MQA apresenta várias desvantagens inevitáveis.

• Degradação de desempenho: Sem surpresa, reduzir o número de parâmetros únicos e treináveis do modelo diminui sua capacidade de conhecimento e nuance. A MQA implica uma queda significativa na precisão em comparação à MHA padrão, tornando-a inadequada para certas situações e casos de uso.
  
  
• Deve ser treinado do zero: Um modelo treinado com a MHA padrão não pode ser simplesmente adaptado para a MQA, mas deve ser treinado com a MQA desde o início. Isso significa que a MQA não pode ser utilizada para otimizar modelos existentes e acarreta um custo de oportunidade considerável ao se experimentar a MQA em novos modelos.
  
  
• Redundâncias no paralelismo tensor: Um dos principais benefícios de treinar modelos transformador em GPUs é a capacidade de executar múltiplas operações complexas com tensores em paralelo. Os valores K e V devem estar presentes em cada nó do cluster de GPUs que executa essas operações, o que significa que, na prática, eles precisam ser replicados para cada nó. Isso não representa o uso ideal dos recursos computacionais, apesar de ainda ser mais eficiente que a MHA padrão.

A atenção de consulta agrupada é uma formulação mais geral e flexível da atenção multi-query que particiona as cabeças de consulta em múltiplos grupos, cada um compartilhando um conjunto de chaves e valores, em vez de compartilhar um único conjunto de chaves e valores entre todas as cabeças de consulta.

Após a publicação de “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints” em maio de 2023, muitos LLMs adotaram rapidamente o GQA. Por exemplo, a Meta adotou inicialmente o GQA para seus modelos Llama 2 em julho de 2023 e manteve o GQA nos modelos Llama 3 lançados em 2024. Mistral IA usou o GQA no modelo Mistral 7B, lançado em setembro de 2023. Da mesma forma, os modelos IBM’s Granite 3.0 empregam GQA para inferência rápida.

Em teoria, a GQA pode ser considerada uma generalização do espectro entre a MHA padrão e a MQA completa. A GQA com o mesmo número de grupos de cabeças chave-valor que cabeças de atenção equivale à MHA padrão; a GQA com 1 grupo de cabeças equivale à MQA.

Na prática, a GQA quase sempre implica uma abordagem intermediária, na qual o número de grupos é, por si só, um hiperparâmetro importante.

A atenção de consulta agrupada oferece várias vantagens que levaram à sua adoção relativamente ampla em LLMs de ponta.

• Uso eficiente da GPU: A distribuição de pares chave-valor na GQA aproveita o paralelismo de tensores, reduzindo a quantidade de computação “desperdiçada” pela replicação de valores redundantes.
  
  
• Compromisso efetivo: A GQA oferece um equilíbrio ideal entre a velocidade de inferência do decodificador e a precisão de desempenho, sendo quase tão precisa quanto a MHA e quase tão rápida quanto a MQA.
  
  
• Redução do overhead de largura de banda de memória: Assim como a MQA, a GQA reduz significativamente o número de cálculos intermediários que devem ser realizados, armazenados e recuperados no momento da inferência.
  
  
• Treinamento flexível: ao contrário do MQA, a atenção à consulta em grupo não exige que os modelos sejam treinados do zero usando a abordagem. Modelos pré-treinados com a MHA padrão podem ser adaptados para usar a GQA por meio de um processo de ajuste fino chamado “uptraining.”