Grouped Query Attention (GQA)

그룹 쿼리 어텐션(GQA)은 트랜스포머를 위한 멀티헤드 어텐션의 한 형태로, 쿼리 헤드 그룹이 키 헤드 하나와 값 헤드 하나를 공유합니다. 예를 들어 Llama 3.1 8B의 각 레이어에서는 쿼리 헤드 32개를 키-값 헤드 8개와 짝지어 사용합니다. GQA는 각 쿼리 헤드가 자체 키와 값을 갖는 멀티헤드 어텐션(MHA)과 모든 쿼리 헤드가 단일 키 헤드와 값 헤드를 공유하는 멀티쿼리 어텐션(MQA)의 중간에 해당합니다. Google Research의 Joshua Ainslie와 동료들은 EMNLP 2023에서 GQA: 멀티헤드 체크포인트에서 일반화된 멀티쿼리 트랜스포머 모델 학습을 통해 GQA를 소개하고, GQA 모델이 MQA에 가까운 디코딩 속도로 MHA에 근접한 품질을 달성한다고 보였습니다. Llama 2 70B, Llama 3 계열, Mistral 7B, Gemma 2와 같은 공개 대규모 언어 모델(LLM)이 GQA를 사용합니다. Gemma 2를 기반으로 한 Google의 PaliGemma 2와 같이 이 모델을 기반으로 구축한 비전 언어 모델(VLM)은 언어 디코더에 GQA를 상속합니다.

GQA의 작동 방식#

셀프 어텐션에서는 각 토큰이 각 헤드에 대해 쿼리, 키, 값으로 투영되며 각 헤드는 쿼리를 키와 비교합니다. GQA는 쿼리 헤드 전체를 유지하되 키와 값 헤드 수를 줄이고, 쿼리 헤드 그룹마다 해당 그룹이 공유하는 키와 값으로 어텐션을 수행합니다. 쿼리 헤드가 8개이고 키-값 헤드가 2개라면 키-값 헤드 하나가 쿼리 헤드 4개를 담당합니다.

각 헤드 내부의 어텐션 계산은 변경되지 않는 스케일드 닷 프로덕트 어텐션입니다. 그룹이 하나면 레이어는 MQA가 되고, 그룹 수가 쿼리 헤드 수와 같으면 MHA가 됩니다. PyTorch는 scaled_dot_product_attention의 enable_gqa 인수를 통해 이 레이아웃을 직접 지원합니다.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Output (1, 8, 16, 64), same as copied heads: True과 256 with GQA, 1024 with MHA을 출력합니다. 쿼리 헤드 8개 모두 MHA가 저장하는 키와 값의 4분의 1만 사용해 출력을 생성합니다.

GQA가 KV 캐시를 줄이는 방식#

텍스트 생성 중 모델은 이전 모든 토큰의 키와 값을 KV 캐시에 보관하고 새 토큰마다 캐시 전체를 읽습니다. 따라서 디코딩 속도는 연산량보다 메모리 대역폭에 더 큰 영향을 받습니다. 캐시 크기는 키-값 헤드 수에 비례합니다. Llama 3.1 8B에는 32개 레이어가 있으며, 각 레이어는 차원 128의 키-값 헤드 8개를 갖습니다. 따라서 16비트 정밀도에서 토큰당 128KiB를 저장하며, 전체 131,072토큰 컨텍스트 윈도우에서는 16GiB가 필요합니다. MHA를 사용하면 같은 모델은 토큰당 512KiB, 최대 길이에서 64GiB가 필요합니다. 확보된 메모리는 더 긴 컨텍스트, 연속 배칭을 위한 더 큰 배치, 토큰당 더 낮은 추론 지연 시간에 사용할 수 있습니다. Ultralytics LLM 인터페이스는 자체 호스팅 Llama 3.1 8B와 같은 OpenAI 호환 서버에 사용할 base_url을 받습니다. 이 서버에서는 키-값 헤드가 8개이므로 32개 헤드 MHA 레이아웃보다 동일한 캐시 메모리에 대화 토큰을 4배 더 많이 보관할 수 있습니다.

모델을 GQA로 변환하기#

GQA는 처음부터 학습할 필요가 없습니다. 저자들은 각 그룹의 키와 값 헤드를 평균 풀링해 기존 MHA 체크포인트를 변환한 뒤, 원래 연산량의 약 5%를 사용해 사전 학습을 이어갔습니다. 이를 업트레이닝이라고 합니다. 공유 헤드 하나만 사용하는 극단적인 형태인 MQA는 2019년 Noam Shazeer가 소개했으며 가장 빠른 레이아웃입니다. 그러나 품질이 저하되고 학습이 불안정해질 수 있으며, 중간 수의 그룹을 사용하면 품질의 대부분을 회복할 수 있습니다.

GQA와 관련 기법 비교#

• 멀티헤드 잠재 어텐션(MLA): DeepSeek-V2와 함께 소개되었으며, 각 토큰의 키와 값을 작은 공유 잠재 벡터로 압축합니다.
• PagedAttention: KV 캐시를 고정 크기 블록에 할당해 단편화를 방지하며, 캐시에 포함되는 헤드 수와는 무관합니다.
• 슬라이딩 윈도우 어텐션: 레이어당 캐시되는 토큰 수를 제한하며, Mistral 7B는 이를 GQA와 결합합니다.
• Flash attention: 정확한 어텐션을 더 빠르게 계산하고 GQA 레이아웃을 기본적으로 지원하는 GPU 커널입니다.