[논문 리뷰] TransMLA: Multi-Head Latent Attention Is All You Need

Fanxu Meng, Tang, Pingzhi외 4명|ArXiv.org|2025. 02. 11.

Topic Modeling인용 수 5

한 줄 요약

논문은 같은 KV 캐시에서 MLA가 GQA보다 더 큰 표현력을 가지며, GQA 기반의 사전학습 모델을 MLA 기반 모델로 변환하는 사후 학습 방법 TransMLA를 도입하고, KV 캐시 크기를 늘리지 않으면서 다운스트림 성능을 향상시키기 위한 파인튜닝을 제시한다.

ABSTRACT

In this paper, we present TransMLA, a framework that seamlessly converts any GQA-based pre-trained model into an MLA-based model. Our approach enables direct compatibility with DeepSeek's codebase, allowing these models to fully leverage DeepSeek-specific optimizations such as vLLM and SGlang. By compressing 93% of the KV cache in LLaMA-2-7B, TransMLA achieves a 10.6x inference speedup at an 8K context length while preserving meaningful output quality. Additionally, the model requires only 6 billion tokens for fine-tuning to regain performance on par with the original across multiple benchmarks. TransMLA offers a practical solution for migrating GQA-based models to the MLA structure. When combined with DeepSeek's advanced features, such as FP8 quantization and Multi-Token Prediction, even greater inference acceleration can be realized.

연구 동기 및 목표

• 대형 언어 모델에서 자체 주의에서 KV 캐시 병목 현상을 줄이려는 동기를 제시한다.
• KV 캐시 비용이 같을 때 MLA의 이론적 표현력 우위를 GQA와 비교하여 확립한다.
• 인기 있는 GQA 기반 모델을 MLA 기반 모델로 변환하는 실용적 변환 워크플로우(TransMLA)를 제공한다.
• 변환 후 미세조정을 통해 KV 캐시 오버헤드를 최소화하면서도 다운스트림 작업에서 GQA 대비 향상된 성능을 보이는 TransMLA 모델을 보여준다.

제안 방법

• KV 캐시와 표현력 측면에서 MHA, GQA, MQA, MLA를 정의하고 비교한다.
• 복제와 저랭크 인수분해를 통한 논거로 GQA를 동일한 KV 캐시로 MLA로 변환할 수 있음을 보인다.
• 잠재 KV 표현으로 GQA를 MLA로 표현하기 위한 SVD 기반의 저랭크 분해를 밝힌다.
• W_K^a, W_K^b, W_V^a, W_V^b 매개변화를 통해 표현력을 확장하면서 KV 캐시를 고정한 채로 후학습 변환인 TransMLA를 도입한다.
• 추론 시 특정 행렬을 흡수하는 연산(absorb)으로 잠재 KV 차원을 고정적으로 유지한다.
• Qwen2.5 및 이와 유사한 모델을 GQA에서 MLA로 변환하고 SmolTalk에서 수학/코딩 작업에 대해 파인튜닝하는 실험적 구성를 제시한다.

(a) Group Query Attention (GQA)

실험 결과

연구 질문

• RQ1 MLA가 동일 KV 캐시 크기에서 GQA 성능과 같거나 더 높게 도달할 수 있는가?
• RQ2 KV 캐시 오버헤드를 늘리지 않고 GQA 기반 사전학습 모델을 MLA 기반 모델로 변환하는 것이 가능한가?
• RQ3 전환 후 TransMLA를 파인튜닝하면 수학 및 코드 벤치마크와 같은 다운스트림 작업에서 측정 가능한 이득이 있는가?
• RQ4 정교한 직교 분해가 TransMLA의 성능 향상에 어떤 역할을 하는가?

주요 결과

• KV 캐시 크기가 고정된 상태에서 MLA가 GQA보다 더 표현력이 있으며 모든 GQA 구성이 MLA 표현으로 변환될 수 있다.
• 실용적이고 저오버헤드의 변환(TransMLA)이 LLaMA, Qwen, Mixtral 등 인기 있는 GQA 기반 모델을 KV 캐시 크기를 늘리지 않고 MLA로 변환할 수 있다.
• 세밀하게 튜닝된 TransMLA 모델은 특히 수학 및 코딩 작업에서 GQA 대비 다운스트림 작업 성능이 향상된다.
• 성능 증가는 잠재 KV 표현으로부터의 확장된 표현력과 인수분해의 직교 분해 덕분이며, 단순한 매개변수 수 증가 때문이 아니다.
• 직교 분해 없이 아이덴티티 맵형 차원 확장은 미미한 이득만을 야기하므로 잠재 분해 접근의 중요성을 강조한다.

(b) Multi-Head Attention (MHA)