Resumo Técnico: GQA-µP: A Atualização de Parametrização Máxima para Atenção com Consultas Agrupadas

Declaração do Problema

A transferência de hiperparâmetros entre arquiteturas de modelos, especificamente a Parametrização de Atualização Máxima (µP), permite a determinação de hiperparâmetros ótimos para modelos de grande escala através da varredura de modelos proxy menores. Embora o µP tenha sido aplicado com sucesso à transferência de taxa de aprendizado entre larguras de modelos, sua aplicação a arquiteturas inovadoras e a outros hiperparâmetros (como decaimento de peso e profundidade) permanece desafiadora.

Existe uma lacuna específica para a Atenção com Consultas Agrupadas (GQA), um mecanismo amplamente utilizado em Modelos de Linguagem de Grande Porte (LLMs) modernos para reduzir o uso de memória ao compartilhar cabeças de chave e valor entre grupos de cabeças de consulta. Tentativas anteriores de aplicar o µP padrão à GQA revelaram uma falha crítica: embora a implementação tenha passado nos "testes de coordenada" padrão (validando normas de ativação), ela falhou em transferir taxas de aprendizado efetivamente entre diferentes números de cabeças KV. Além disso, os frameworks existentes de µP lutaram para derivar escalas principiais para decaimento de peso e profundidade sem depender de suposições de "aprendizado preguiçoso". A dificuldade teórica central surge porque as matrizes de peso da GQA são frequentemente degeneradas em rank (de baixo rank), fazendo com que a norma espectral padrão falhe em capturar com precisão a deformação esperada das entradas durante o treinamento.

Metodologia

Os autores baseiam-se na visão de aprendizado de características espectrais de Yang et al. (2023a) para derivar novas leis de escala. A metodologia envolve três avanços teóricos e analíticos primários:

1. Redefinição do Aprendizado de Características via Normas Espectrais:
   O artigo eleva as condições de norma espectral sobre matrizes de peso de uma heurística para a definição formal de aprendizado de características. Os autores argumentam que garantir que a norma espectral das atualizações de peso escale corretamente é uma condição mais forte e confiável para o aprendizado de características do que analisar apenas as normas de ativação. Essa mudança permite a derivação de escalas de profundidade e decaimento de peso Complete-P sem recorrer a suposições de aprendizado preguiçoso.

2. Introdução da Norma Operadora Esperada:
   Para abordar a degeneração de rank das matrizes GQA, os autores introduzem uma norma modificada: a norma operadora esperada ($||A||_E$).

   • As normas espectrais padrão ($||A||$) medem a deformação máxima de um vetor de entrada. Para matrizes de rank completo, isso é atingido assintoticamente.
   • Para matrizes degeneradas em rank (como a GQA), os vetores que causam a deformação máxima ocorrem com probabilidade zero durante o treinamento (já que as entradas são i.i.d.). Consequentemente, a norma espectral superestima a deformação real.
   • A norma operadora esperada é definida como a expectativa da razão entre as normas de saída e entrada sobre a distribuição de entrada: $||A||_{E,\Omega,p} := \mathbb{E}_{x \sim \Omega} [\frac{||Ax||_p}{||x||_p}]$.
   • Esta norma reflete corretamente a deformação real encontrada durante o treinamento, restaurando a validade das leis de escala para a GQA.

3. Derivação das Escalas GQA-µP:
   Usando a norma operadora esperada, os autores derivam os fatores de escala específicos para camadas GQA.

   • Para uma camada GQA com $r$ repetições de cabeças de chave-valor, a taxa de aprendizado ($\eta$) e o decaimento de peso ($\lambda$) devem ser escalados de forma diferente do que na atenção padrão.
   • A escala de taxa de aprendizado derivada é $\eta = \Theta(\frac{1+\sqrt{r}}{n})$, onde $n$ é a dimensão de incorporação.
   • A escala de decaimento de peso é ajustada para $\lambda = \Theta(\frac{2n}{1+\sqrt{r}})$ para manter o equilíbrio entre regularização e atualizações de gradiente.
   • Essas derivações são resumidas na Tabela 1 do artigo, fornecendo uma parametrização unificada para Transformers com GQA.

4. Escala de Profundidade (Complete-P):
   O framework é aplicado a redes residuais para derivar escalas de profundidade. Ao aplicar a condição espectral à composição de blocos residuais, os autores mostram que o fator de escala $\beta$ para ramos residuais deve ser $\Theta(L^{-1})$ (onde $L$ é a profundidade) para evitar explosão ou desaparecimento do sinal, recuperando os mesmos resultados de Dey et al. (2025) sem exigir a suposição de "sem aprendizado preguiçoso".

Contribuições Principais

1. Primeira Derivação µP para GQA: O artigo fornece, ao conhecimento dos autores, a primeira derivação principial de escalas µP para Atenção com Consultas Agrupadas, resolvendo a discrepância onde implementações padrão de µP falharam em transferir taxas de aprendizado.
2. Extensão Teórica do µP Espectral: O trabalho estende a teoria do µP espectral de Yang et al. (2023a) promovendo condições de norma espectral à definição de aprendizado de características e introduzindo a norma operadora esperada para lidar com pesos degenerados em rank. Isso permite a derivação de escalas para decaimento de peso, blocos de recursão e GQA.
3. Validação Empírica da Transferência: Os autores demonstram que suas escalas derivadas permitem a transferência de taxas de aprendizado entre diferentes contagens de repetição GQA ($r$) e de decaimento de peso entre tamanhos de modelo. Eles mostram que a implementação Adam-µP padrão falha nos testes de coordenada espectral para GQA, enquanto sua escala proposta passa neles.

Resultados Experimentais

Os autores validam suas derivações teóricas através de experimentos extensos:

• Transferência de Taxa de Aprendizado: Experimentos variando o número de cabeças KV (de 1 a 12) mostram que a parametrização padrão e o Adam-µP padrão falham em convergir para uma taxa de aprendizado ótima consistente. Em contraste, a escala GQA-µP proposta estreita significativamente a região de taxa de aprendizado ótima, permitindo uma transferência eficaz.
• Testes de Coordenada: A análise das normas de atualização de ativação ($||\Delta h||$) sugeriu que a implementação padrão estava correta, mas a análise das normas espectrais das matrizes de peso ($||\Delta W||$) revelou uma dependência não linear do número de cabeças KV, explicando a falha na transferência. A escala GQA-µP proposta remove essa dependência.
• Transferência de Decaimento de Peso e $\tau_{\text{epoch}}$: O artigo demonstra que, com as escalas corretas, tanto o decaimento de peso quanto a constante de tempo de treinamento $\tau_{\text{epoch}}$ (introduzida por Wang & Aitchison, 2024) são transferíveis entre tamanhos de modelo. A variância nos hiperparâmetros ótimos é significativamente menor para a implementação proposta em comparação com as linhas de base.
• Ruído na GQA: Os autores observam que, embora a transferência de taxa de aprendizado seja possível, o treinamento com GQA introduz ruído inerente que se torna mais pronunciado à medida que o número de cabeças KV diminui. Eles notam uma compensação: alcançar o aprendizado de características através da escala derivada leva a dinâmicas mais ruidosas à medida que $r$ diminui, enquanto restringir a variância para estabilizar o treinamento desloca a taxa de aprendizado ótima.

Significado e Alegações

O artigo afirma fechar uma lacuna significativa na aplicabilidade do µP a arquiteturas modernas de LLMs de nível de produção. Ao identificar as limitações da norma espectral padrão para matrizes de baixo rank e propor a norma operadora esperada, os autores fornecem uma base teórica robusta para derivar escalas µP para arquiteturas complexas como a GQA.

O significado deste trabalho reside em:

• Utilidade Prática: Permite que praticantes determinem hiperparâmetros ótimos para grandes modelos baseados em GQA (por exemplo, LLaMA 3, Mistral 7B) usando modelos proxy menores, reduzindo o custo computacional do ajuste de hiperparâmetros.
• Rigor Teórico: Refina a definição de aprendizado de características, indo além de verificações baseadas em ativação para condições espectrais baseadas em matrizes de peso, o que prova ser mais eficaz para validar implementações.
• Ampla Aplicabilidade: O framework recupera com sucesso escalas existentes para decaimento de peso e profundidade (Complete-P) enquanto as estende para GQA, sugerindo uma abordagem unificada para futuras inovações arquitetônicas.

Os autores notam modestamente que, embora sua análise assintótica seja sólida, o intervalo prático do parâmetro de repetição $r$ é finito e tipicamente pequeno (por exemplo, $r \le 16$). O objetivo principal da escala em relação a $r$ é empírico: prevenir o desvio da taxa de aprendizado ao alterar o número de repetições, em vez de modelar um limite infinito. Eles também destacam que o treinamento com GQA permanece inerentemente ruidoso, sugerindo cautela ao transferir taxas de aprendizado entre configurações de GQA drasticamente diferentes.