[논문 리뷰] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Joshua Ainslie, James Lee-Thorp외 4명|arXiv (Cornell University)|2023. 05. 22.

Topic Modeling인용 수 27

한 줄 요약

이 논문은 uptraining을 통해 multi-head attention을 multi-query로 변환하고 일반화된 그룹화된 쿼리 주의(GQA)를 도입하여, 원래의 사전 학습 컴퓨트의 아주 작은 부분만 사용하면서도 빠른 추론과 경쟁력 있는 품질을 가능하게 한다.

ABSTRACT

Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference. We (1) propose a recipe for uptraining existing multi-head language model checkpoints into models with MQA using 5% of original pre-training compute, and (2) introduce grouped-query attention (GQA), a generalization of multi-query attention which uses an intermediate (more than one, less than number of query heads) number of key-value heads. We show that uptrained GQA achieves quality close to multi-head attention with comparable speed to MQA.

연구 동기 및 목표

• autoregressive transformers에서 메모리 대역폭으로 인한 디코더 병목 현상을 동기 부여하고 품질을 손실하지 않으면서 더 빠른 추론을 모색한다.
• 기존의 multi-head 체크포인트를 멀티-쿼리 구성으로 변환하기 위한 비용 효과적인 uptraining 레시피를 제안한다.
• multi-head와 multi-query attention 사이의 보간으로서 grouped-query attention (GQA)를 도입한다.
• uptrained GQA가 다양한 태스크에서 multi-head attention에 비해 품질에 가까우면서도 속도는 거의 multi-query attention에 근접하다는 것을 입증한다.

제안 방법

• multi-head attention 체크포인트를 key 및 value 투영 행렬의 평균 풀링을 통해 MQA를 위해 하나의 헤드로 형성하거나 각 GQA 그룹에 대해 형성한다.
• 원본 사전 학습 단계의 소량 비율α의 동일한 데이터와 레시피로 변환된 모델을 uptrain한다.
• GQA를 G개의 쿼리 헤드 그룹으로 정의하고, 그룹당 하나의 KV 헤드를 공유하며, MQA (G=1)와 MHA (G=H) 사이를 보간한다.
• decoder self-attention과 cross-attention에 MQA와 GQA를 적용하되 encoder self-attention은 다루지 않는다.
• 요약, 번역, QA 벤치마크에서 추론 시간당 샘플 및 개발 세트 성능을 보고한다.

실험 결과

연구 질문

• RQ1 다중-head 체크포인트가 제한된 추가 계산으로 효과적으로 빠른 multi-query 형태로 uptraining될 수 있는가?
• RQ2 그룹화된 쿼리 주의가 속도와 품질 간의 tunable trade-off를 제공하여 MQA를 능가하고 MHA에 근접하는가?
• RQ3 uptraining이 다양한 태스크와 모델 규모에서의 안정성과 성능에 어떤 영향을 미치는가?

주요 결과

• uptrained MQA는 MHA-Large보다 더 빠른 추론과 더 높은 품질을 제공하지만 일부 경우 XXL-MHA보다 뒤처진다.
• uptrained GQA는 MHA-XXL에 거의 근접한 품질과 MQA에 가까운 속도를 달성하여 유리한 트레이드오프를 제공한다.
• 체크포인트 변환에서의 평균 풀링은 헤드를 선택하거나 무작위 초기화하는 것보다 정보를 더 잘 보존한다.
• 5% uptraining에서 성능이 향상되며, MQA와 GQA 모두 5–10%를 넘으면 수익 감소가 나타난다.
• GQA 성능 이점은 대형 모델에서 대역폭 다이나믹과 캐시 고려로 인해 더 뚜렷하다.
• GQA 그룹 수를 1(MQA)에서 더 많이 늘리면 증가하는 비용에 비해 속도 증가가 작아지며, 8개 그룹이 바람직한 중간 지점으로 확인된다.