GQA (Grouped Query Attention): 성능과 속도의 절묘한 줄타기, MHA와 MQA를 넘어서

GQA (Grouped Query Attention): 성능과 속도의 절묘한 줄타기, MHA와 MQA를 넘어서

거대 언어 모델(LLM)의 심장에는 '어텐션(Attention)'이라는 핵심 메커니즘이 자리 잡고 있습니다. 이는 모델이 문장의 어떤 부분에 더 집중해야 할지 결정하는 능력으로, 인간의 독해 방식과 유사합니다. 초기 트랜스포머 모델의 멀티 헤드 어텐션(MHA)은 뛰어난 성능을 보였지만, 모델이 커지면서 막대한 계산 비용과 메모리 대역폭이라는 치명적인 병목 현상에 부딪혔습니다. 이러한 문제를 해결하기 위한 다양한 시도 끝에, 성능과 효율성 사이의 완벽한 균형점을 찾은 기술이 등장하게 되었습니다. GQA(Grouped Query Attention) 입니다. 이 글에서는 MHA(Mulit Head Attention)에서 MQA(Mulit Query Attention)를 거쳐 GQA에 이르기까지의 기술적 진화 과정을 '도서관'이라는 비유를 통해 살펴보며, GQA가 어떻게 두 선행 기술의 장점만을 적용하였는지 알아보도록 하겠습니다.

---

 

◇ 멀티 헤드 어텐션(MHA): 여러 질문, 여러 전문 도서관

MHA(Mulit Head Attention)는 하나의 정보를 여러 관점에서 동시에 바라보기 위해 고안된 독창적인 방식입니다. 모델은 입력된 정보에 대해 여러 개의 독립적인 '질문(Query)', '색인(Key)', '정보(Value)' 세트를 만들어냅니다.

• 개념: MHA는 입력 데이터로부터 다수의 Query, Key, Value(QKV) 헤드(Head)를 생성합니다. 각 헤드는 독립적으로 어텐션 계산을 수행하여 정보의 다양한 측면을 포착합니다.
• 비유: 여러 전문 도서관에 문의하기
  어떤 주제(예: "사과")에 대해 깊이 있는 정보를 얻고 싶다고 상상해봅시다. 여러분은 이 주제를 "사과의 영양학적 가치는?", "사과의 문화적 의미는?", "사과 재배 방법은?"과 같이 여러 관점의 질문(Multi-Query)으로 세분화합니다. 이 질문들을 각각 영양학, 문화, 농업 등 특정 분야에 특화된 여러 전문 도서관(Multi Key-Value Space)에 보냅니다. 각 도서관의 사서(Head)는 질문(Query)에 가장 부합하는 색인(Key)을 찾아 해당 책(Value)을 꺼내주고, 이 책들을 종합하여 하나의 완성된 보고서를 만듭니다.
• 한계: 이 방식은 매우 풍부하고 다각적인 정보를 제공하지만, 수많은 전문 도서관을 모두 운영하고 관리하려면 엄청난 비용(GPU 메모리)과 시간(계산 속도), 즉 메모리 대역폭 오버헤드가 발생합니다.

---

◇ 멀티 쿼리 어텐션(MQA): 여러 질문, 단 하나의 종합 도서관

MHA의 비효율 문제를 해결하기 위한 첫 시도로 MQA(Multi-Query Attention)가 제안되었습니다. "굳이 비싼 도서관을 여러 개 둘 필요가 있을까? 모든 질문을 하나의 거대한 종합 도서관에 던져보면 어떨까?"라는 아이디어에서 출발했습니다.

• 개념: MQA는 여러 개의 Query 헤드를 사용하지만, Key와 Value 헤드는 단 하나만 공유하여 계산량을 줄입니다.
• 비유: 하나의 거대한 종합 도서관
  여러 전문 도서관을 없애는 대신, 모든 분야의 책을 망라한 하나의 거대한 국립중앙도서관(Single Key-Value Space)만 남겨둡니다. 이제 우리는 모든 다양한 질문(Multi-Query)을 이 한 곳의 도서관에 문의합니다. 박학다식한 사서 한 명이 모든 질문에 대한 답을 찾아주는 셈입니다.
• 장점과 단점: 도서관 유지 비용, 즉 KV 캐시(Key-Value Cache)의 메모리 사용량이 획기적으로 줄고 정보 처리(추론) 속도가 매우 빨라집니다. 하지만 성능 저하와 훈련 불안정성이라는 대가를 치러야 했습니다. 하나의 도서관이 모든 전문 분야의 깊이를 감당하기는 어렵기 때문입니다.

---

◇ 그룹 쿼리 어텐션(GQA): 여러 질문, 소수의 권역별 대표 도서관

MHA의 높은 성능과 MQA의 빠른 속도 사이의 트레이드오프를 극복하기 위해 GQA(Grouped-Query Attention)가 등장했습니다. "모든 전문 도서관은 너무 비싸고, 종합 도서관 하나는 정보가 부족하니, 질문들을 몇 개의 그룹으로 묶어 권역별 대표 도서관에 물어보자"는 영리한 절충안입니다.

• 개념: GQA는 전체 Query 헤드를 몇 개의 그룹으로 나눈 뒤, 각 그룹이 하나의 Key와 Value 헤드 세트를 공유하는 방식입니다. 예를 들어 32개의 Query 헤드를 4개 그룹으로 묶어, 총 4개의 Key-Value 헤드만 사용하는 식입니다.
• 비유: 권역별 대표 도서관 운영하기
  수십 개의 동네 도서관(MHA)을 모두 운영하는 대신, 서울, 경기, 강원 등 권역별로 대표 도서관(Grouped Key-Value Space) 몇 개를 전략적으로 운영합니다. 이제 각 지역의 질문(Query 그룹)들은 해당 권역의 대표 도서관에 문의합니다. 이렇게 하면 전국 모든 도서관을 운영하는 것(MHA)보다 비용이 훨씬 저렴하면서, 단 하나의 중앙 도서관(MQA)에 의존하는 것보다 훨씬 전문적이고 안정적인 정보 수집이 가능해집니다.
• 일반화된 아키텍처: GQA는 그룹 수를 어떻게 설정하느냐에 따라 MHA나 MQA가 될 수 있는 일반화된 구조입니다. 그룹 수가 전체 헤드 수와 같으면 MHA, 그룹이 단 하나뿐이면 MQA가 되는 유연성을 가집니다.

---

GQA의 성능과 실제 적용 사례

실험 결과는 매우 인상적입니다. GQA는 MQA에 필적하는 빠른 추론 속도를 보이면서도 MHA에 근접하는 높은 성능을 달성하여 '두 마리 토끼를 다 잡은' 해결책임을 증명했습니다. 이러한 강력한 장점 덕분에 GQA는 Llama, Mistral 등 여러 최신 LLM에 빠르게 도입되었습니다.

• Llama 2: 70B와 같은 대규모 모델에 GQA를 적용하여 성능 저하를 최소화하면서 추론 속도를 크게 향상시켰습니다.
• Llama 3: GQA의 효과를 확신하고, 모델 크기에 상관없이 모든 모델에 GQA를 전면 적용했습니다. 흥미롭게도 Llama 3의 모든 모델은 K-V 헤드를 8개로 통일했는데, 이는 실험을 통해 8개가 성능과 속도의 '최적점(sweet spot)'임을 찾아냈음을 시사합니다.

---

마무리하며

GQA는 MHA의 높은 성능과 MQA의 빠른 속도 사이에서 최적의 균형점을 찾아낸 혁신적인 어텐션 메커니즘입니다. 단순히 두 기술을 절충한 것을 넘어, 메모리 사용량과 연산 비용을 획기적으로 줄이면서도 모델의 품질을 유지하는 새로운 패러다임을 제시했습니다. GQA의 등장은 제한된 컴퓨팅 자원으로도 고성능 AI 모델을 더 효율적으로 개발하고 배포할 수 있는 길을 열어주었습니다.

---

https://arxiv.org/abs/2305.13245

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

https://arxiv.org/abs/1911.02150

Fast Transformer Decoding: One Write-Head is All You Need

https://arxiv.org/abs/2407.21783

The Llama 3 Herd of Models