Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs

논문 원본: https://arxiv.org/abs/2502.14837

Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs

---

Introduction

이 논문은 대형 언어모델의 추론 비용에서 가장 큰 병목 중 하나가 KV 캐시 메모리라는 문제의식에서 출발한다. MHA 기반 LLM은 시퀀스 길이와 모델 크기에 비례해 KV 캐시가 선형으로 커지기 때문에, 긴 문맥 추론 시 메모리와 대역폭 비용이 매우 커진다. 이를 줄이기 위해 GQA, MQA 같은 변형들이 제안되었지만, 이들은 KV 캐시뿐 아니라 attention 파라미터 수도 같이 줄여서 성능 저하를 유발한다. DeepSeek이 제안한 MLA는 키·값을 저차원 잠재벡터로 공동 압축하는 구조를 통해 MHA보다 더 작은 KV 캐시로도 좋은 성능을 내는 아키텍처이지만, 기존에 MHA로 잘 학습된 Llama 계열 모델들을 MLA 구조로 옮기려면 원래는 처음부터 다시 pretrain 해야 하는 문제가 있다.

 

이 논문은 “이미 학습된 MHA 기반 LLM을, 적은 데이터로 MLA 기반으로 갈아타게 만드는” 첫 번째 데이터 효율적 파인튜닝 프레임워크 MHA2MLA를 제안한다. 핵심 아이디어는 (1) attention score에 거의 기여하지 않는 차원에서 RoPE를 제거하는 partial-RoPE, (2) 기존 key·value projection 파라미터를 공동 SVD로 저차원 잠재공간에 매끄럽게 사상하는 joint low-rank approximation이다. 이 조합을 통해 전체 데이터의 0.6~1% 정도만으로도 성능을 대부분 회복하면서, Llama2-7B 기준 KV 캐시를 92.19%까지 줄이고 LongBench 성능 감소를 약 1% 수준으로 유지했다고 보고한다.

Preliminary

저자들은 먼저 MHA, MLA, 그리고 RoPE를 중심으로 현재 구조를 정리한다. MHA에서는 각 토큰의 쿼리·키·값이 head마다 분리된 공간에서 계산되고, 전체 KV 캐시는 모든 레이어·모든 헤드·모든 토큰의 K,V 벡터를 그대로 저장한다. 이 구조에서는 시퀀스 길이가 길어질수록 저장해야 하는 텐서가 폭발적으로 증가해, 메모리와 메모리 대역폭이 추론 속도를 제약하게 된다.

 

DeepSeek MLA는 이 구조를 바꿔서, K,V를 바로 저장하는 대신 하나의 저차원 잠재 공간으로 압축해 latent 벡터 형태로 KV 캐시를 저장하고, attention 연산 시에는 이 latent로부터 필요한 head-specific K,V를 복원하는 식의 low-rank joint compression을 사용한다. 개념적으로는 “모든 헤드가 공유하는 작은 잠재 공간” 위에 여러 헤드가 얹혀 있는 구조라서, KV 캐시의 공간 복잡도가 크게 줄어든다.

 

또한 최근 Llama 계열 모델들이 사용하는 RoPE(회전 위치 임베딩)은 쿼리·키 차원에 위치 정보를 곱해주는 방식인데, MLA처럼 저차원 압축을 하려면 “어떤 차원에 RoPE가 적용돼 있는지, 어떤 차원은 순수한 내용 표현인지”가 구조적으로 중요해진다. 이 논문은 바로 이 점을 이용해서, RoPE 적용을 재구성한 뒤에 공동 저차원화를 설계한다.

 

MHA2MLA

MHA2MLA 프레임워크는 크게 두 가지 구성 요소로 설명된다. 첫 번째는 partial-RoPE이다. 기존 MHA에서는 q,k 전 차원에 RoPE를 적용하지만, 모든 차원이 attention score에 똑같이 중요하게 기여하는 것은 아니다. 저자들은 q·k 차원별로 attention score 기여도를 분석한 후, 상대적으로 기여도가 낮은 차원에서는 RoPE를 제거하고 내용 전용(non-RoPE) 차원으로 두며, 중요한 상위 r차원에만 RoPE를 유지하는 구조를 만든다. 이렇게 하면 q,k가 “RoPE가 걸린 위치 의존 차원”과 “순수 내용 표현 차원”으로 분해되고, 나중에 저차원 잠재 벡터로 압축할 때 RoPE 부분과 non-RoPE 부분을 구분해 더 안정적으로 압축할 수 있다. 두 번째는 low-rank approximation을 위한 joint SVD 초기화이다.

 

MLA에서는 키·값을 잠재 공간으로 투영하는 여러 projection 행렬이 등장하는데, 이 논문은 단순히 랜덤 초기화하는 대신 기존 MHA의 Wk, Wv 파라미터를 합쳐 공동 SVD를 수행해, 주요 특이벡터 방향을 latent 방향으로 삼는다. 이 과정을 통해 MLA의 Wdkv, Wuk, Wuv, Wuq 같은 projection 행렬을 “기존 MHA 파라미터의 저차원 근사”로 초기화할 수 있고, 이는 곧 학습 초기에 MHA와 MLA 간 함수 동작이 최대한 맞춰지도록 하는 역할을 한다.

 

요약하면, MHA2MLA는 (1) RoPE 구조를 부분적으로 해제해서 압축 친화적인 표현 구조를 만들고, (2) 원래 MHA 파라미터에서 최대한 정보를 가져와 저차원 latent space를 초기화한 뒤, (3) 적은 양의 데이터로 전체 모델을 파인튜닝해서 MLA로 부드럽게 이행시키는 절차라고 볼 수 있다.

Experiment

실험에서는 Llama2-7B 등 대표적인 MHA 기반 LLM을 대상으로, MHA2MLA를 적용했을 때의 성능 유지와 비용 절감 효과를 평가한다. 사용 데이터는 전체 프리트레이닝 코퍼스의 0.6~1% 정도에 해당하는 소량으로 제한해 “데이터 효율성”을 강조했고, 그 상태에서 LongBench와 같은 긴 문맥 벤치마크에서 base 모델 대비 성능이 얼마나 유지되는지를 본다. 그 결과 Llama2-7B의 경우, MLA 구조로 전환하면서 KV 캐시 크기를 92.19%까지 줄였음에도 LongBench 평균 성능 하락은 약 1% 수준에 그친다고 보고한다.

 

또 하나의 포인트는 MHA2MLA가 KV 캐시 양자화와도 잘 결합된다는 점이다. MLA 자체가 이미 low-rank 압축으로 KV 캐시를 줄인 상태이지만, 여기에 추가로 KV 캐시를 낮은 비트 수로 양자화했을 때도 성능 저하가 크지 않으며, 전체 메모리와 대역폭을 한 번 더 줄일 수 있다는 것을 실험으로 보여준다. 즉, “MHA → MLA 구조 전환 + KV quantization”을 동시에 적용해도 합성적인 성능이 잘 유지된다는 메시지를 강조한다.

 

Related Work

관련 연구 정리는 크게 세 축으로 서술된다. 첫째, KV 캐시 비용을 줄이기 위한 기존 attention 변형들(GQA, MQA 등)은 헤드 수와 파라미터 수 자체를 줄이는 방식이라 모델 용량이 감소하고 성능 저하를 동반하는 한계가 있다. 둘째, KV 캐시를 그대로 두되 양자화나 eviction(부분 삭제), 슬라이딩 윈도우 등으로 메모리 사용량을 제어하는 다양한 캐시 압축 기법들이 있지만, 이들은 기본적인 MHA 구조 자체는 변경하지 않기 때문에 구조적으로 얻을 수 있는 추가 이득이 제한적이다. 셋째, DeepSeek MLA와 같이 latent space를 이용해 attention을 재설계하는 연구들이 등장하고 있지만, 기존에 MHA로 학습된 LLM을 MLA로 옮기는 구체적인 fine-tuning 경로를 제시한 연구는 거의 없었다. 이 논문은 바로 이 세 번째 축에서, “pretrained MHA LLM → MLA 구조로의 data-efficient adaptation”이라는 새로운 연구 방향을 제시한다고 자기 위치를 잡는다.

 

Conclusion

정리하면, 이 논문은 KV 캐시 메모리가 병목이 되는 MHA 기반 LLM 추론 환경에서, 이미 잘 학습된 모델을 처음부터 다시 학습하지 않고 MLA 구조로 옮겨 효율을 극대화하는 방법을 제안한다. 핵심 구성은 RoPE를 차원별 중요도에 따라 부분적으로만 유지하는 partial-RoPE와, 기존 Wk/Wv 파라미터를 공동 SVD로 저차원 잠재 공간에 투영하는 joint low-rank 초기화이며, 이를 통해 극히 적은 비율의 데이터만으로도 MLA 성능을 회복하게 만든다. 실험적으로는 Llama2-7B 기준 KV 캐시를 90% 이상 줄이면서도 긴 문맥 벤치마크 성능 감소를 약 1% 수준에 묶어두고, 동시에 KV quantization과도 자연스럽게 결합된다는 점을 보여준다. 전체적으로 “KV 캐시 양자화 → MLA 기반 구조적 압축 → 이를 기존 MHA LLM에 이식하는 MHA2MLA”라는 흐름을 통해, 경제적인 LLM 추론을 위한 하나의 설계 패턴을 제시한 논문이라고 이해하면 될 것 같다.