TransMLA: L'attenzione latente a più teste è tutto ciò di cui hai bisogno

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

Autori: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

Abstract

I moderni modelli di linguaggio di grandi dimensioni (LLM) spesso si trovano di fronte a vincoli di comunicazione sull'hardware attuale, piuttosto che a vincoli puramente computazionali. Il Multi-head Latent Attention (MLA) affronta questa sfida utilizzando matrici a rango ridotto nei livelli chiave-valore (KV), consentendo così di memorizzare stati KV latenti compressi. Questo approccio riduce significativamente la dimensione della cache KV rispetto alla tradizionale attenzione multi-head, portando a un'infereza più veloce. Inoltre, MLA utilizza una matrice di up-projection per aumentare l'espressività, scambiando calcoli aggiuntivi per ridurre il carico di comunicazione. Anche se MLA ha dimostrato efficienza ed efficacia in Deepseek V2/V3/R1, molti dei principali fornitori di modelli continuano a fare affidamento sull'Attention a Query di Gruppo (GQA) e non hanno annunciato piani per adottare MLA. In questo articolo, dimostriamo che GQA può sempre essere rappresentato da MLA mantenendo lo stesso overhead di cache KV, ma il contrario non è vero. Per promuovere un uso più ampio di MLA, presentiamo **TransMLA**, un metodo di post-training che converte modelli pre-addestrati ampiamente utilizzati basati su GQA (ad esempio, LLaMA, Qwen, Mixtral) in modelli basati su MLA. Dopo la conversione, il modello può essere sottoposto a ulteriore addestramento per aumentare l'espressività senza aumentare la dimensione della cache KV. Inoltre, prevediamo di sviluppare tecniche di accelerazione specifiche per MLA per preservare bassa latenza nei modelli trasformati, consentendo così una distillazione più efficiente di 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.