TransMLA : L'attention latente à têtes multiples est tout ce dont vous avez besoin

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Auteurs: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Résumé

Les modèles de langage larges modernes (LLMs) rencontrent souvent des goulots d'étranglement de communication sur le matériel actuel, plutôt que des contraintes purement computationnelles. L'Attention Latente Multi-Tête (MLA) relève ce défi en utilisant des matrices de rang faible dans les couches clés-valeurs (KV), permettant ainsi de mettre en cache des états KV latents compressés. Cette approche réduit significativement la taille du cache KV par rapport à l'attention multi-tête traditionnelle, ce qui accélère l'inférence. De plus, le MLA utilise une matrice de sur-projection pour augmenter l'expressivité, échangeant des calculs supplémentaires contre une réduction des frais de communication. Bien que le MLA ait démontré son efficacité dans Deepseek V2/V3/R1, de nombreux grands fournisseurs de modèles continuent de s'appuyer sur l'Attention de Groupe par Requête (GQA) et n'ont pas annoncé de plans pour adopter le MLA. Dans cet article, nous montrons que le GQA peut toujours être représenté par le MLA tout en maintenant le même surcoût de cache KV, mais l'inverse n'est pas vrai. Pour encourager une utilisation plus large du MLA, nous introduisons **TransMLA**, une méthode de post-entraînement qui convertit des modèles pré-entraînés largement utilisés basés sur le GQA (par exemple, LLaMA, Qwen, Mixtral) en modèles basés sur le MLA. Après conversion, le modèle peut subir un entraînement supplémentaire pour renforcer l'expressivité sans augmenter la taille du cache KV. De plus, nous prévoyons de développer des techniques d'accélération spécifiques au MLA pour préserver une faible latence dans les modèles transformés, permettant ainsi une distillation plus efficace de 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.