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

개요

대규모 언어 모델의 추론 시 가장 큰 병목 중 하나는 KV 캐시(Key-Value Cache)입니다. 시퀀스 길이가 길어질수록 KV 캐시의 메모리 사용량이 선형으로 증가하고, 이를 HBM(High Bandwidth Memory)에서 반복적으로 로드하는 비용이 추론 속도를 크게 저하시킵니다. GQA(Grouped-Query Attention)는 이 문제를 효과적으로 해결하는 어텐션 변형 기법으로, MHA(Multi-Head Attention)와 MQA(Multi-Query Attention) 사이의 균형점을 제공합니다.

배경 및 문제

멀티헤드 어텐션 (MHA)

표준 트랜스포머의 MHA는 $H$ 개의 헤드가 각각 독립적인 쿼리, 키, 값 프로젝션을 가집니다.

$Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V$

헤드 수가 $H$ 개일 때, KV 캐시의 크기는 레이어당 시퀀스 길이 × $2 \times H \times d_{k}$ 에 비례합니다. 배치 처리와 긴 컨텍스트 환경에서 이 메모리 비용이 주요 병목이 됩니다.

멀티쿼리 어텐션 (MQA)

Shazeer(2019)가 제안한 MQA는 모든 쿼리 헤드가 단 하나의 키-값 헤드를 공유합니다. KV 캐시가 $H$ 배 감소하여 메모리와 대역폭 효율이 크게 향상되지만, 표현력이 줄어들어 품질 저하가 발생할 수 있습니다. 또한 기존 MHA 체크포인트에서 MQA로 전환하려면 처음부터 다시 학습해야 하는 비용이 있었습니다.

핵심 아이디어

GQA는 $H$ 개의 쿼리 헤드를 $G$ 개의 그룹으로 나누고, 각 그룹 내 쿼리 헤드들이 하나의 키-값 헤드를 공유하는 구조를 도입합니다.

$G = 1$ : MQA와 동일 (모든 쿼리가 단일 KV 공유)
$G = H$ : MHA와 동일 (각 쿼리가 독립적인 KV 보유)
$1 < G < H$ : GQA (중간 설정, 품질-효율 균형)

그룹 수 $G$ 를 조절함으로써 KV 캐시 크기를 $H / G$ 배 줄이면서도 MHA에 근접한 품질을 유지할 수 있습니다.

방법론

업트레이닝 (Uptraining)

GQA의 핵심 기여 중 하나는 기존 MHA 체크포인트를 GQA로 효율적으로 변환하는 업트레이닝 방법입니다.

변환 절차:

기존 MHA 모델의 $H$ 개 키-값 헤드를 $G$ 개 그룹으로 묶음
각 그룹 내 키-값 헤드들을 평균 풀링(mean pooling)하여 하나의 대표 헤드 생성
원본 학습 스텝의 **5%**만 추가 학습(업트레이닝)하여 변환 완료

이 방법을 통해 T5-XXL(11B 파라미터) 규모 모델을 기준으로 전체 재학습 없이 GQA 모델로 전환할 수 있음을 보였습니다.

그룹 수에 따른 KV 캐시 절감

설정	KV 헤드 수	KV 캐시 비율	품질
MHA	$H$	100%	최고
GQA-G	$G$	$G / H \times 100%$	중간
MQA	$1$	$1/ H \times 100%$	최저

예를 들어 $H = 32$ 이고 $G = 8$ 이면 KV 캐시가 기존 대비 25%로 줄어듭니다.

추론 속도 분석

KV 캐시의 메모리 접근 비용이 줄면 메모리 대역폭 병목이 완화되어 토큰 생성 속도가 향상됩니다. T5 모델 기준 실험에서:

GQA-8: MHA 대비 약 1.5~2x 속도 향상
MQA: MHA 대비 약 2~3x 속도 향상
GQA-8의 품질은 MQA보다 MHA에 훨씬 가까움

실험 결과

T5 기반 실험 (요약 태스크)

모델	XSum ROUGE-2	CNN/DM ROUGE-2	추론 속도
T5-XXL (MHA)	21.7	21.3	1.0x
T5-XXL (MQA)	21.2	20.9	2.8x
T5-XXL (GQA-8)	21.6	21.2	1.9x

GQA-8은 MHA와 ROUGE-2 차이가 0.1 이내를 유지하면서 MQA에 근접한 속도 향상을 달성했습니다.

업트레이닝 효율성

원본 학습의 5% 스텝만으로도 MQA 및 GQA 변환 모델이 처음부터 학습한 모델과 동등한 성능에 도달했습니다. 이는 기존에 학습된 대형 모델을 재사용할 수 있는 실용적 경로를 제시합니다.

의의 및 한계

의의

실용성: LLaMA 2(34B, 70B), Mistral-7B, Falcon 등 주요 오픈소스 모델이 GQA를 채택하여 산업 표준 기법으로 정착
유연성: 그룹 수 $G$ 를 조절하여 메모리-품질 트레이드오프를 태스크별로 조정 가능
업트레이닝: 기존 MHA 체크포인트를 5% 추가 학습으로 재활용 가능하여 재학습 비용 절감
KV 캐시 감소: 긴 컨텍스트와 대용량 배치 환경에서 특히 유리

한계

그룹 내 키-값을 공유하므로 MHA보다 표현력이 떨어지며, 매우 세밀한 어텐션 패턴이 필요한 태스크에서 품질 손실 가능성 존재
최적의 그룹 수 $G$ 는 모델 아키텍처와 태스크에 따라 다르므로 별도 탐색이 필요
FlashAttention, 양자화 등 다른 효율화 기법과의 조합 효과는 별도 검증 필요

GQA는 추론 효율과 모델 품질 사이의 실용적인 균형점을 제공하며, 현재 대부분의 최신 LLM 아키텍처에서 기본 구성 요소로 채택되고 있습니다.

도형준 지식 볼트

탐색기