TransMLA: Multi-head Latente Aandacht Is Alles Wat Je Nodig Hebt

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Auteurs: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Samenvatting

Moderne grote taalmodellen (LLM's) ondervinden vaak communicatieknelpunten op huidige hardware, eerder dan louter computationele beperkingen. Multi-head Latente Aandacht (MLA) pakt deze uitdaging aan door het gebruik van lage-rang matrices in de sleutel-waarde (KV) lagen, waardoor gecomprimeerde latente KV-staten gecachet kunnen worden. Deze aanpak vermindert aanzienlijk de KV-cachegrootte ten opzichte van traditionele multi-head aandacht, wat resulteert in snellere inferentie. Bovendien maakt MLA gebruik van een up-projectiematrix om de expressiviteit te vergroten, waarbij extra berekeningen worden verruild voor verminderde communicatie-overhead. Hoewel MLA efficiëntie en effectiviteit heeft aangetoond in Deepseek V2/V3/R1, vertrouwen veel belangrijke modelaanbieders nog steeds op Groepsvraagaandacht (GQA) en hebben zij geen plannen aangekondigd om MLA over te nemen. In dit artikel tonen we aan dat GQA altijd kan worden voorgesteld door MLA met behoud van dezelfde KV-cache-overhead, maar andersom niet geldt. Om het bredere gebruik van MLA aan te moedigen, introduceren we **TransMLA**, een post-trainingmethode die veelgebruikte op GQA gebaseerde vooraf getrainde modellen (bijv. LLaMA, Qwen, Mixtral) omzet in MLA-gebaseerde modellen. Na conversie kan het model extra training ondergaan om de expressiviteit te verhogen zonder de KV-cachegrootte te vergroten. Bovendien zijn we van plan MLA-specifieke inferentieversnellingsmethoden te ontwikkelen om een lage latentie te behouden in getransformeerde modellen, waardoor een efficiëntere distillatie van Deepseek R1 mogelijk wordt.

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.