TransMLA: Atenção Latente Multi-cabeça É Tudo o Que Você Precisa

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Autores: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) modernos frequentemente enfrentam gargalos de comunicação no hardware atual, em vez de limitações puramente computacionais. A Atenção Latente Multi-Cabeça (MLA, do inglês *Multi-head Latent Attention*) aborda esse desafio utilizando matrizes de baixo posto nas camadas de chave-valor (KV, do inglês *key-value*), permitindo que estados latentes KV comprimidos sejam armazenados em cache. Essa abordagem reduz significativamente o tamanho do cache KV em comparação com a atenção multi-cabeça tradicional, resultando em inferência mais rápida. Além disso, a MLA emprega uma matriz de projeção ascendente para aumentar a expressividade, trocando computação adicional por redução na sobrecarga de comunicação. Embora a MLA tenha demonstrado eficiência e eficácia nos modelos Deepseek V2/V3/R1, muitos dos principais provedores de modelos ainda dependem da Atenção de Consulta em Grupo (GQA, do inglês *Group Query Attention*) e não anunciaram planos para adotar a MLA. Neste artigo, mostramos que a GQA pode sempre ser representada pela MLA mantendo a mesma sobrecarga de cache KV, mas o inverso não é verdadeiro. Para incentivar o uso mais amplo da MLA, introduzimos o **TransMLA**, um método pós-treinamento que converte modelos pré-treinados amplamente utilizados baseados em GQA (por exemplo, LLaMA, Qwen, Mixtral) em modelos baseados em MLA. Após a conversão, o modelo pode passar por treinamento adicional para aumentar a expressividade sem aumentar o tamanho do cache KV. Além disso, planejamos desenvolver técnicas de aceleração de inferência específicas para MLA, a fim de preservar a baixa latência em modelos transformados, permitindo assim uma destilação mais eficiente do Deepseek R1.

English

Modern large language models (LLMs) often encounter communication bottlenecks on current hardware, rather than purely computational constraints. Multi-head Latent Attention (MLA) tackles this challenge by using low-rank matrices in the key-value (KV) layers, thereby allowing compressed latent KV states to be cached. This approach significantly reduces the KV cache size relative to traditional multi-head attention, leading to faster inference. Moreover, MLA employs an up-projection matrix to increase expressiveness, trading additional computation for reduced communication overhead. Although MLA has demonstrated efficiency and effectiveness in Deepseek V2/V3/R1, many major model providers still rely on Group Query Attention (GQA) and have not announced any plans to adopt MLA. In this paper, we show that GQA can always be represented by MLA while maintaining the same KV cache overhead, but the converse does not hold. To encourage broader use of MLA, we introduce **TransMLA**, a post-training method that converts widely used GQA-based pre-trained models (e.g., LLaMA, Qwen, Mixtral) into MLA-based models. After conversion, the model can undergo additional training to boost expressiveness without increasing the KV cache size. Furthermore, we plan to develop MLA-specific inference acceleration techniques to preserve low latency in transformed models, thus enabling more efficient distillation of Deepseek R1.