TransMLA: Многоголовое скрытое внимание - все, что вам нужно.

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Авторы: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Аннотация

Современные большие языковые модели (LLM) часто сталкиваются с узкими местами в области коммуникаций на текущем оборудовании, а не исключительно с вычислительными ограничениями. Метод многоголового латентного внимания (MLA) решает эту проблему, используя матрицы низкого ранга в слоях ключ-значение (KV), что позволяет кэшировать сжатые латентные состояния KV. Этот подход значительно сокращает размер кэша KV по сравнению с традиционным многоголовым вниманием, что приводит к более быстрой инференции. Более того, MLA использует матрицу повышения проекции для увеличения выразительности, обменивая дополнительные вычисления на сокращение коммуникационной нагрузки. Хотя MLA продемонстрировал эффективность и эффективность в Deepseek V2/V3/R1, многие крупные поставщики моделей до сих пор полагаются на групповое внимание к запросу (GQA) и не объявили о планах принять MLA. В этой статье мы показываем, что GQA всегда может быть представлено с помощью MLA, сохраняя ту же самую нагрузку кэша KV, но обратное не верно. Для поощрения более широкого использования MLA мы представляем **TransMLA**, метод посттренировочного преобразования, который преобразует широко используемые предварительно обученные модели на основе GQA (например, LLaMA, Qwen, Mixtral) в модели на основе MLA. После преобразования модель может пройти дополнительное обучение для увеличения выразительности без увеличения размера кэша KV. Кроме того, мы планируем разработать техники ускорения вывода, специфичные для MLA, чтобы сохранить низкую задержку в преобразованных моделях, обеспечивая более эффективное дистиллирование 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.