QK-Normed MLA: QK normalization without full key caching

Resumo Técnico: QK-Normed MLA

Problema
Os modelos de linguagem de grande escala modernos enfrentam duas pressões conflitantes: a necessidade de decodificação eficiente em termos de memória em cenários de contexto longo e a necessidade de estabilidade no treinamento.

• Eficiência de Memória: O Multi-Head Attention (MHA) padrão armazena vetores completos de chave (key) e valor (value) para cada token, criando altos custos de memória e largura de banda. O Multi-Head Latent Attention (MLA) aborda isso armazenando estados latentes de baixa dimensão ($c_t$) em vez de chaves completas, recuperando o efeito total da chave via uma up-projeção ($W^K_g$) durante a decodificação.
• Estabilidade de Treinamento: Escalas descontroladas em queries e keys podem levar à saturação da atenção e instabilidades no treinamento. A normalização de Query-Key (QK) (especificamente o RMSNorm) estabiliza o treinamento ao normalizar os vetores antes do produto escalar. No entanto, o RMSNorm de QK pós-projeção padrão parece incompatível com o MLA porque aparentemente exige a chave totalmente projetada ($k_{t,g} = c_t W^K_g$) para cada token armazenado em cache para calcular a estatística de normalização. Uma implementação ingênua exigiria a materialização de chaves completas, negando as vantagens de memória do MLA.

Metodologia
O artigo propõe o QK-Normed MLA, uma formulação que integra o QK RMSNorm de pós-up-projeção no caminho de decodificação latente do MLA sem materializar chaves completas. A ideia central é a decomposição algébrica do RMSNorm:
$$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$$
onde $\gamma$ é um peso afim estático e $S(z)$ é uma estatística RMS dinâmica escalar.

Os autores derivam uma formulação de "absorção" que separa esses componentes:

1. Absorção de Peso Estático: O peso estático do lado da chave $\gamma_k$ é absorvido na matriz de projeção do lado da query. Em vez de projetar a query $q$ e então normalizar a chave, a query é pré-multiplicada por $\gamma_k$ (e pelo peso estático $\gamma_q$) antes do produto escalar latente.
2. Armazenamento de Escalar Dinâmico: A estatística dinâmica $S(z)$ reduz-se a um único escalar de inverso-RMS ($\alpha_{t,g} = 1/S_{t,g}$) por token e por grupo de KV.
3. Layout do Cache: O cache de decodificação retém o estado latente original de baixa dimensão $c_t$ e adiciona um pequeno cache escalar $A^k_{cache}$ contendo os valores de inverso-RMS. Ele não armazena as chaves normalizadas completas.
4. Normalização por Blocos: Para manter a eficiência, o método normaliza os componentes de "conteúdo" (content) e "RoPE" separadamente. O caminho de conteúdo usa a formulação latente absorvida, enquanto o caminho RoPE usa vetores materializados padrão. Isso evita pontos de sincronização que complicariam as implementações de kernels fundidos.

Durante a decodificação, os logits de atenção são computados como:
$$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$$
onde $\tilde{q}$ inclui os pesos absorvidos, e os fatores escalares são aplicados como uma etapa de escala linha-coluna antes do softmax.

Principais Contribuições

1. Formulação Exata: O artigo deriva uma formulação algebricamente exata de QK RMSNorm de pós-up-projeção que é compatível com o caminho de decodificação latente do MLA.
2. Design de Cache Eficiente: Propõe um layout de cache que preserva o cache KV latente como o estado dominante, adicionando apenas um escalar de estatística RMS de chave por token e grupo de KV (por exemplo, aumentando o tamanho do cache em ~1,56% nas configurações testadas).
3. Validação Empírica: Os autores demonstram que esta abordagem alcança um erro de treinamento (loss) menor e melhor precisão em tarefas downstream comparado ao QK Clipping, introduzindo um overhead de latência negligenciável.

Resultos
Experimentos foram conduzidos em modelos MLA de 400 milhões de parâmetros treinados por até 100 bilhões de tokens em GPUs H800:

• Estabilidade de Treinamento: O QK-Normed MLA alcançou consistentemente um erro de treinamento menor que o QK-Clip ao longo da execução de 100 bilhões de tokens. Em testes de estresse de alta taxa de aprendizado ($2 \times 10^{-2}$), o QK-Clip divergiu com NaNs de gradiente, enquanto o QK-Normed MLA permaneceu estável.
• Desempenho Downstream: Em um conjunto de tarefas downstream, o QK-Normed MLA melhorou a precisão em 7 de 8 tarefas, elevando a pontuação média de 44,75 para 46,33, e reduziu a perplexidade do LAMBADA de 16,28 para 14,18.
• Eficiência de Decodificação: Benchmarks em contextos variando de 4k a 256k tokens mostraram que o método adiciona menos de 2% de overhead de latência (média de 1,40%). O cache escalar aumentou o tamanho total do cache em apenas ~2,8% por camada.

Significância e Alegações
O artigo alega que a aparente incompatibilidade entre a normalização QK e o MLA é um artefato de implementação, não uma restrição arquitetônica. Ao mostrar que os pesos estáticos podem ser absorvidos e as estatísticas dinâmicas reduzidas a escalares, os autores permitem a normalização QK como um mecanismo prático de estabilização para modelos MLA sem sacrificar a eficiência de memória da decodificação latente.

Os autores mantêm modéstia quanto ao escopo de suas alegações:

• Os resultados baseiam-se em modelos de 400 milhões de parâmetros; eles não alegam estabelecer comportamento de escala de fronteira para modelos de bilhões de parâmetros, embora a equivalência algébrica se mantenha independentemente da escala.
• A análise da escala de logit é apresentada como um diagnóstico de nível de mecanismo, não como uma teoria de otimização completa.
• A velocidade prática depende da integração de kernels (fundindo fatores escalares no caminho de atenção) para evitar tráfego de memória desnecessário.