TransMLA: Atención Latente Multi-cabeza es Todo lo que Necesitas

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Autores: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Resumen

Los modernos modelos de lenguaje de gran tamaño (LLMs) a menudo se enfrentan a cuellos de botella en la comunicación en el hardware actual, en lugar de limitaciones puramente computacionales. La Atención Latente Multi-Cabeza (MLA) aborda este desafío mediante el uso de matrices de rango bajo en las capas clave-valor (KV), lo que permite almacenar estados KV latentes comprimidos en caché. Este enfoque reduce significativamente el tamaño de la caché KV en comparación con la atención multi-cabeza tradicional, lo que conduce a una inferencia más rápida. Además, MLA emplea una matriz de proyección ascendente para aumentar la expresividad, intercambiando cálculos adicionales por una reducción en la sobrecarga de comunicación. Aunque MLA ha demostrado eficiencia y efectividad en Deepseek V2/V3/R1, muchos proveedores de modelos importantes aún confían en la Atención de Consulta de Grupo (GQA) y no han anunciado planes para adoptar MLA. En este documento, mostramos que GQA siempre se puede representar mediante MLA manteniendo la misma sobrecarga de caché KV, pero lo contrario no es cierto. Para fomentar un uso más amplio de MLA, presentamos **TransMLA**, un método de post-entrenamiento que convierte modelos pre-entrenados ampliamente utilizados basados en GQA (por ejemplo, LLaMA, Qwen, Mixtral) en modelos basados en MLA. Después de la conversión, el modelo puede someterse a entrenamiento adicional para aumentar la expresividad sin aumentar el tamaño de la caché KV. Además, planeamos desarrollar técnicas de aceleración de inferencia específicas de MLA para preservar una baja latencia en los modelos transformados, lo que permite una destilación más eficiente 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.