"Attention Is All You Need" 의 대항마 : Multi-Head Latent Attention

대규모 언어 모델(LLM)의 Self-Attention 메커니즘은 모델의 성능을 좌우하는 가장 핵심적인 요소로, 문맥을 전역적으로 이해하고 복잡한 언어의 패턴을 학습하는 데 필수적이다. 그러나, Multi-Head Attention(이하 MHA)은 높은 계산 비용과 메모리 사용량으로 인해 대규모 모델의 효율적인 학습과 추론에 제약적이다. 이를 해결하기 위해 제안된 Grouped-Query Attention(이하 GQA)가 제안되었으며, 이는 쿼리 헤드를 적은 수의 키-값(KV)헤드로 그룹화하여 메모리 사용량을 줄이고 추론 속도와 효율성을 높였다. 하지만, GQA의 방식은 모델에 따라 품질 저하가 발생하는 단점이 있다. 이러한 한계를 극복하기 위해 DeepSeek에서 제안하였던  Multi-Head Latent Attention (이하 MLA)이 큰 주목을 받았다. 따라서 본 게시물에서는 MLA의 작동 원리에 대해 알아보고자 하며, 추후 게시물에서는 "TransMLA: Multi-Head Latent Attention Is All You Need" 논문을 살펴보고자 합니다. 

Multi-Head Attention (MHA) 이해하기

 

Multi-Head Attention (MHA)는 트랜스포머 모델의 핵심 요소로, 입력된 문장에서 단어(토큰) 간의 관계를 파악하는 데 매우 효과적이다.  각 단어는 쿼리(Query, Q), 키(Key, K), 값(Value, V)이라는 세 가지 벡터로 변환되는데, 여기에서 Q는 특정 단어가 다른 단어들에게 보내는 '질문'이고, K는 각 단어가 가진 '고유한 식별표'이며, V는 그 단어가 가진 '의미'라고 생각할 수 있습니다. Q가 모든 K와 비교(Attention Score 계산)를 통해 자신과 가장 관련 깊은 단어를 찾고, 그 단어의 V(의미)를 더 많이 가져오는 방식이죠. 하지만 MHA는 자기회귀(Autoregressive) 방식으로 다음 단어를 예측할 때, 새로운 단어 하나를 생성하기 위해 이전에 계산했던 모든 단어와의 Attention Score를 다시 계산해야 하는 비효율성이 존재한다. 

 

이러한 문제점을 해결하기 위해 KV Cache 기법이 고안되었다. 이름 그대로, 이전에 계산했던 모든 토큰의 K와 V를 메모리에 저장해두고, 다음 단어를 생성할 때 이 캐시를 재사용하는 방식이다. 이를 통해 중복적인 계산을 방지하고 추론 속도를 크게 향상시킬 수 있다는 큰 장점이 있다. 이에 대한 예시는 다음과 같다. 서점에 가서 점원에게 "혹시 머신러닝 관련된 책이 있을까요?" 라고 물어보면, 점원은 책의 제목과 내용을 확인한 후, "핸즈온 머신러닝 이라는 책이 있습니다" 라고 이야기를 해줄 것이다. 그렇다면, 다음에 방문해서 이전과 조금 다르게 "딥러닝 관련된 책이 있을까요?" 라고 물어보면 직원은 다시 서점에 있는 책의 제목과 내용을 확인해야한다. 이처럼 질문이 바뀔 때 마다 모든 책을 다시 검색하고 확인하는 작업이 비효율적이므로, 검색했던 내용을 메모하는 것이 KV Cache인 셈이다. 

 

그럼 간단한 예시를 통해 알아보고자 한다. 만약 손님이 서점에 방문해서 점원에게 "머신러닝 관련된 책 있나요?" 라고 물어본다면, 점원은 모든 책의 제목(Key)과 내용(Value)을 훑어본 후 "핸즈온 머신러닝이라는 책이 있습니다" 라고 답을 하게된다. 여기에서 핵심은, 손님의 다음 질문, 예를 들어 "혹시 그 책의 저자에 대해서 알 수 있을까요?" 라는 질문에 답하기 위해 점원이 방금 훑어봤던 책들의 정보(K,V)를 기억하고 있다는 점이다. 이것이 바로 KV Cache의 원리이다. 

 

하지만, 대화가 길어지면 길어질수록 점원이 기억해야 할 책의 정보(KV Cache)는 계속 쌓이게 된다. 문장이 1,000개의 토큰으로 이루어져 있다면, 1,001번째 토큰을 생성하기 위해 1,000개 토큰의 K,V를 모두 메모리에 올려둬야 하는 상황이 필요하다. 이처럼 LLM이 처러하는 토큰의 양이 기하급수적으로 증가하게 되면서, KV Cache가 차지하는 메모리 또한 엄청나게 증가하는 문제가 발생하게 된다. 

 

이러한 KV Cache 메모리 사용량 증가 문제를 해결하기 위해 Multi-Query Attention (MQA)과 Grouped-Query Attention (GQA) 기법이 고안되었다. 이 두 개의 방식은 모두 MHA의 구조를 살짝 변형하여 K, V 헤드의 수를 줄임으로써 메모리 사용량을 최적화한다. 

---

Multi-Query Attention (MQA) 및 Grouped Query Attention (GQA) 이해 하기 

이미지 출처: medium.com

 

• MQA (Multi-Query Attention): MQA는 MHA의 여러 Q 헤드를 그대로 유지하되, 단 하나의 K,V 헤드만을 모든 Q 헤드가 공유하는 방식이다. 조금 더 직관적으로 설명을 해보자면, 여러 명의 손님(다수의 Q 헤드)이 각자 다른 질문을 하지만, 이들의 질문은 모두 단 한명의 점원(하나의 K,V 헤드)이 전담하여 처리하게 된다. 한명의 점원이 모든 질문에 대해서 처리하기 때문에, KV Cache에 저장할 데이터가 극적으로 줄어들어 메모리 효율이 가장 좋지만, 모든 쿼리가 동일한 K,V를 참조하기 떄문에 표현력이 감소하며 이는 MHA보다 성능이 저하될 수 있는 단점을 가지고 있다. 
• GQA (Grouped-Query Attention): GQA 같은 경우, MHA와 MQA의 중간 형태이다. 이는 여러 개의 Q 헤드를 몇 개의 그룹으로 나누고, 각 그룹이 하나의 K,V 헤드를 공유하도록 한다. 만약에 여러 손님(Q 헤드)들을 몇 개의 그룹으로 묶고, 각 그룹마다 분야를 배정하여 전담 직원을 한명 씩 배치하는 것과 같다. 같은 그룹 내(책의 분야)내의 손님들은 전담 점원에게 추천을 받지만, 다른 그룹은 다른 점원의 추천을 받아 활용하게 된다. 이러한 특징으로 인하여 MHA 표현력 저하를 어느정도 완하하면서도 메모리 사용량을 줄일 수 있다.

이처럼 MQA와 GQA는 MHA의 높은 KV Cache 메모리 요구량을 줄여 Inference(추론) 시의 효율성을 극대화하기 위해 다양한 LLM 모델에서 채택되었다. 하지만 이는 KV Cache 메모리 절감이라는 이점을 얻는 대신, MHA에 비해 Downstream Task에서의 정확도(표현력 저하로 인한 성능 하락) 측면에서 Trade-off 관계가 존재한다. 즉, 메모리를 적게 쓰는 대신 성능의 일부를 희생하게 되는데 그렇다면 이에 대해 해결할 수 있는 방법이 없을까? 

 

---

한계를 뛰어넘는 새로운 패러다임: Multi-Head Latent Attention (MLA)

GQA가 제시한 메모리 효율성과 MHA가 가진 높은 정확도(표현력), 이 두 가지를 모두 가질 수 없을까? 이러한 "Trade-Off" 관계를 깨기 위해 등장한 것이 바로 Multi-Head Latent Attention(MLA) 이다.

이미지 출처: G3LU

 

기존의 Multi-Head Attention(MHA)은 모든 토큰 쌍(Key-Value Pair)을 저장하기 위해 막대한 양의 KV Cache 메모리를 필요로 한다. Multi-Head Latent Attention(MLA)은 Key-Value Pair를 그대로 저장하는 대신 저차원 행렬(Low-Rank Matrices)을 사용해 이를 압축된 버전으로 저장하게 된다. 이를 통해 모델의 핵심적인 성능을 희생하지 않으면서도 메모리 사용량을 획기적으로 줄일 수 있다.

 

물론 GQA(Grouped-Query Attention) 역시 여러 쿼리가 하나의 KV 헤드를 공유하는 방식으로 메모리 문제를 완화시지만, 이러한 구조는 토큰 간의 미세하고 복잡한 관계를 포착하는 데 한계가 있다. 반면 MLA는 단순히 정보를 공유하는 것을 넘어, 입력을 더 작은 잠재 공간(Latent Space)으로 압축했다가 어텐션 계산 시 다시 복원하는 접근 방식을 사용한다. 이러한 과정에서 모델은 불필요한 정보를 걸러내고 데이터의 핵심만 학습하게 되어, 메모리 효율성과 높은 표현력을 모두 달성할 수 있게 되는 것이다.

 

MLA에 대해 조금 더 직관적인 예시를 들어보면 다음과 같다. 여러 고객(Query)이 책에 대한 추천을 점원에게 요청을 하게 된다면, 굉장히 많은 양의 정보를 파악해야한다. 이에 기존에는 단순히 모든 질문에 대해 기억해서 알려주는 MHA 혹은 각 분야의 점원이 추천해주는 방식인 GQA이 아니라, 모든 고객의 요청을 종합적으로 분석하여 "요약본"을 만드는 베테량 점원과 같다. 이 베터랑 점원은 개별 추천 기록을 일일이 저장하는 대신, "인문학 베스트셀러" 등과 같은 고도로 압축된 지식(Latent Cache)을 생성한다. 이에 고객이 질문을 하면, 이 요약된 지식을 바탕으로 깊이 있는 추천을 해줄 수 있게 되는 것이다. 

 

이러한 방식은 DeepSeek-R1 혹은 하위 모델의 연구를 통해 이런 매핑(압축)하는 과정이 모델의 불필요한 정보를 버리면서 일반화 성능을 이끌어냈으며, 결과적으로 MHA 보다 더 뛰어난 성능을 입증하였다.