MLA Analysis

MLA (Multi-head Latent Attention) 심층 분석

Low-Rank KV Joint Compression · Decoupled RoPE · Weight Absorption · DeepSeek-V2/V3

MLA(Multi-head Latent Attention)는 DeepSeek-V2가 도입한 어텐션 기법으로, Key/Value를 저차원 latent(잠재) 벡터로 압축해 KV 캐시를 크게 줄입니다. 긴 문맥에서는 KV 캐시가 메모리와 대역폭 병목이 되기 쉬운데, MLA는 저장 대상을 작게 만들어 이 병목을 직접 줄입니다.

GQA/MQA가 'KV 헤드 수'를 줄이는 방식이라면, MLA는 '저장하는 내용 자체'를 저차원으로 압축합니다. 그 결과 캐시는 작게 유지하면서도 품질 저하를 줄이는 방향으로 설계되며, decoupled RoPE와 weight absorption을 통해 위치 정보 처리와 추론 경로도 함께 정리합니다.

핵심 개념

• latent 압축 — K,V를 작은 latent c_KV로 압축해 캐시에 저장합니다.
• low-rank 분해 — 큰 투영 행렬을 down-projection과 up-projection으로 나눕니다.
• decoupled RoPE — 내용 정보와 위치 정보를 분리해 RoPE를 별도로 처리합니다.
• weight absorption — up-projection을 Q/O 쪽에 흡수해 복원 단계를 줄입니다.

1. 핵심 아이디어 — K,V를 저차원 latent로 압축

MLA의 출발점은 GQA/MQA와 같은 문제(KV 캐시가 메모리·대역폭 병목)이지만, 접근이 다릅니다. 헤드를 공유해 표현력을 희생하는 대신, Key/Value를 저차원 잠재 벡터로 압축해 저장하고 사용할 때 복원합니다.

그림 1. 기존 MHA/GQA는 헤드별 K,V를 저장, MLA는 작은 latent c_KV만 저장

GQA와의 근본적 차이

• GQA/MQA — 'KV 헤드 수'를 줄여 캐시를 줄입니다(헤드 공유). 표현력을 일부 희생하는 트레이드오프가 있습니다.

• MLA — 저차원 행렬 분해(low-rank)로 K,V를 압축합니다. 큰 투영 행렬 W_KV를 down-projection(W_DKV)과 up-projection(W_UK, W_UV) 두 단계로 분해해, 입력을 작은 latent c_KV(d_c 차원, 예: 512)로 압축 저장하고, 사용 시 헤드별 full K,V로 복원합니다.

• KV 캐시에는 작은 c_KV만 저장 → 헤드 수가 아니라 'latent 차원'만큼만 메모리를 씁니다.

• 핵심 주장(DeepSeek): MLA는 MHA보다 KV 캐시는 훨씬 작으면서 품질은 오히려 더 좋습니다. 즉 '헤드를 줄이면 품질 손실'이라는 GQA의 트레이드오프를 차원 압축으로 우회합니다.

2. 데이터 흐름 — 압축 경로와 decoupled RoPE

MLA는 두 갈래의 경로로 작동합니다. 하나는 K,V의 '내용'을 압축하는 경로이고, 다른 하나는 '위치 정보'(RoPE)를 따로 처리하는 경로입니다. 이 둘을 분리하는 것이 MLA 구현의 핵심 난점이자 해법입니다.

그림 2. 압축 경로(down→latent→up)와 RoPE 경로(별도 rotary key)의 분리

두 경로

• 압축 경로(내용) — 입력 hₜ를 down-projection으로 latent c_KV로 압축 → 이것만 캐시에 저장 → 어텐션 시 up-projection으로 K^C, V를 복원.

• RoPE 경로(위치) — 별도의 작은 투영으로 위치 부분을 만들고 RoPE를 적용해, 모든 헤드가 공유하는 작은 rotary key(k^R)로 저장. 캐시에는 c_KV와 작은 k^R 둘만 저장됩니다.

왜 RoPE를 분리(decouple)하나

문제는 저차원 압축(low-rank)과 RoPE(회전 위치 인코딩)가 호환되지 않는다는 점입니다. RoPE는 위치에 따라 K를 회전시키는데, 이 회전이 압축→복원 과정과 교환(commute)되지 않아 위치 정보를 latent에 함께 흡수할 수 없습니다. 해결책이 decoupled RoPE입니다. K를 내용 부분(K^C, latent로 압축·위치 없음)과 위치 부분(K^R, 별도 차원에 RoPE 적용·모든 헤드가 공유)으로 나누고, 어텐션 시 둘을 concat해 사용합니다. 이렇게 하면 압축의 이점을 지키면서도 위치 정보를 정확히 반영할 수 있습니다.

3. Weight Absorption — 복원 단계를 없애는 추론 트릭

MLA를 순진하게 구현하면 디코딩마다 latent에서 K,V를 복원하는 연산이 추가됩니다. 그러나 행렬곱의 결합법칙을 이용하면 이 복원 단계를 아예 없앨 수 있습니다. 이것이 weight absorption(행렬 흡수)입니다.

그림 3. 순진한 복원 방식 vs up-projection 행렬을 Q·O에 흡수하는 방식

weight absorption의 원리

• Q·Kᵀ에서 K=W_UK·c_KV이므로, W_UK를 Query 투영에 미리 곱해두면(흡수) 복원 없이 latent c_KV로 바로 어텐션 점수를 계산할 수 있습니다.

• 마찬가지로 출력 쪽의 W_UV는 출력 투영 W_O에 흡수됩니다. 둘 다 학습이 끝난 뒤 가중치를 미리 합쳐두는 것이라 수학적으로 완전히 동등합니다.

• 결과: 저장은 작은 latent로(메모리↓), 계산도 복원 없이(연산↓). 다만 latent에 직접 계산하면 헤드 간 공유 구조가 생겨 커널 최적화가 까다롭습니다.

• 실행 방식 선택 — latent 재사용(연산↓·메모리 접근↑)과 복원 재계산(메모리↓·연산↑) 사이의 트레이드오프가 있어, 하드웨어·문맥 길이에 따라 유리한 쪽이 달라집니다.

4. 비교/분석

그림 4. KV 캐시 크기·DeepSeek-V2 수치·품질 vs KV 산점도·채택/마이그레이션

네 방식의 KV 캐시 비교

DeepSeek-V2 보고 수치 (vs DeepSeek 67B)

• KV 캐시 93.3% 감소 — latent 압축으로 토큰당 캐시가 크게 줄어듭니다.

• 생성 처리량 5.76배 — 8×H800 노드에서 50K tokens/s를 초과하는 최대 생성 처리량.

• 학습 비용 42.5% 절감 — MLA와 DeepSeekMoE를 함께 적용한 결과.

• 모델 규모 — 총 236B 파라미터 중 토큰당 21B만 활성화(MoE), 128K 문맥 길이 지원.

품질-KV 평면에서 MLA는 '작은 KV + 높은 품질'의 좌상단에 위치합니다. MQA는 KV가 작지만 품질이 낮고, MHA는 품질이 높지만 KV가 크며, GQA는 그 중간인 반면, MLA는 두 목표를 동시에 달성한다는 것이 핵심 주장입니다. 실제 배포 시에는 FP8과 KV 캐시 양자화까지 결합해 메모리를 더 줄일 수 있습니다 — 즉 MLA는 양자화와 직교적으로 함께 쓰입니다.

장단점

5. 관련 기술

기존 MHA 모델을 MLA로 — MHA2MLA

MLA의 강력함 때문에, 이미 학습된 MHA/GQA 모델을 MLA로 전환하려는 연구가 활발합니다. MHA2MLA는 두 가지 기법으로 이를 데이터 효율적으로 수행합니다: (1) partial-RoPE — 어텐션 점수에 덜 기여하는 차원에서 RoPE를 제거하고, (2) joint SVD — 사전학습된 Key/Value 파라미터를 특이값 분해(SVD)로 저차원 근사합니다. 이 방법은 데이터의 0.3~0.6%만으로 성능을 회복하며, 예를 들어 Llama2-7B의 KV 캐시를 92.19% 줄이면서 LongBench 성능은 0.5%만 하락한다고 보고합니다. KV 캐시 양자화 같은 압축 기법과도 매끄럽게 결합됩니다.

KV 절감 기법들 속 MLA의 위치

MLA는 LLM 추론 메모리 최적화의 여러 축 중 '어텐션 아키텍처' 축에서 GQA의 다음 진화에 해당하며, 다른 축의 기법들과 직교적입니다.

정리하면 'KV를 몇 개 헤드로/얼마나 압축할지(MLA·GQA), 몇 비트로(양자화), 어떻게 담고(PagedAttention), 어느 토큰까지(축출)'가 함께 작동하는 설계 공간이며, MLA는 그중 가장 침습적이지만 강력한 아키텍처 수준 압축을 담당합니다.

참고 문헌

• DeepSeek-AI et al., DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (arXiv:2405.04434)
• Ji et al., Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (arXiv:2502.14837)

핵심 정리

핵심 정리 — MLA는 Key/Value를 저차원 latent 벡터로 압축해 KV 캐시를 줄이는 어텐션 기법으로, GQA가 헤드를 공유하는 것과 달리 '저장 내용 자체'를 압축한다. down-projection으로 작은 c_KV만 캐시에 저장하고 up-projection으로 복원하되, 압축과 호환되지 않는 RoPE는 decoupled RoPE로 분리(내용 K^C + 위치 K^R)해 처리한다. 추론 시에는 up-projection 행렬을 Q·O에 흡수(weight absorption)해 복원 단계를 없앨 수 있어 메모리와 연산을 함께 절감한다. DeepSeek-V2는 MLA로 67B 대비 KV 캐시 93.3% 감소·생성 처리량 5.76배를 보고했고, MLA는 양자화·PagedAttention과 직교적이며, MHA2MLA 같은 마이그레이션으로 기존 모델에도 적용할 수 있다(Llama2-7B KV 92.19% 감소).

주의 — 본문 수치는 원논문(DeepSeek-V2 arXiv 2405.04434)과 후속 연구(MHA2MLA arXiv 2502.14837)의 보고값이다. 'KV 93.3% 감소·5.76배 처리량'은 DeepSeek-V2 대 DeepSeek 67B의 특정 조건(8×H800, 특정 프롬프트/생성 길이 분포) 결과이며 d_c=512 등 구성은 모델마다 다르다. 'Llama2-7B KV 92.19% 감소·LongBench 0.5% 하락'은 별도 연구(MHA2MLA)의 보고값으로, 압축률·정확도 균형은 모델·설정에 따라 달라진다. MLA는 V3/R1에도 사용되나 구성 세부는 버전마다 차이가 있다.

# 어텐션 점수: Q · K^t (K^t 는 K의 전치)
# 그런데 MLA에서 K = W_UK · c_KV (up-projection × latent)
#
#   Q · K^t = Q · (W_UK · c_KV)^t
#           = Q · c_KV^t · W_UK^t
#           = (Q · W_UK^t) · c_KV^t
#             └────┬────┘
#         W_UK 를 Query 쪽에 흡수 → 미리 합쳐둘 수 있음
#
# 마찬가지로 출력의 W_UV 는 출력 투영 W_O 에 흡수 가능
# 결과: 캐시된 latent c_KV 로 '바로' 어텐션 → K,V 복원 단계 제거