[용어] GQA(Grouped-Query Attention)

(개념) 어텐션의 쿼리 헤드 여러 개가 하나의 KV 세 트를 공유하도록 묶는 방식

 

트랜스포머 모델에서 처음 제시된 MHA(MultiHead Attention)은 헤드 각각 Q, K, V의 가중치 행렬을 가지기 때문에 다양한 관점에서 어텐션 값 추출이 가능하여 추론 성능을 극대화하는 반면 헤드별로 어텐션 연산을 수행하기 때문에 연산량이 높다. 이를 해결하기 위해 초기에 제안된 방안은 헤드를 하나만 가지며 쿼리만 여러 개로 나누는 MQA(Multi-Query Attention) 방안이 제안되었는데, 속도는 향상되는 방면 다양성이 지나치게 감소하여 추론 성능이 저하되었다. GQA는 둘 방안 사이의 보완책으로 [그림 1]과 같이 헤드를 몇 개의 그룹으로 묶고, 쿼리를 각 그룹별로 분배하는 방식이다. 이를 통해 연산량은 MHA 대비 감소하는 대신 성능은 MHA에 비해 근소하게 열화되는 성과를 달성했다. 현재는 생성형 AI 모델 전반적으로 GQA를 매우 널리 적용하고 있다.

 

 

출처: Joshua Ainslie, et al., “Gqa: Training generalized multi-query transformer models from multihead checkpoints”, 2023.

 

https://arxiv.org/abs/2305.13245