Grouped-query attention이란 무엇인가?

최근 발표된 성능 좋은 오픈소스 모델은 대부분 Attention을 개선하여 성능을 개선시키고 있는 것 같습니다. 이번 포스트에서는 Meta의 Llama2와 Google의 gemma에서 사용된 GQA와 MQA에 대해서 알아보려고 합니다.

https://www.blossominkyung.com/deeplearning/transformer-mha

Transformer의 멀티-헤드 어텐션(Multi-Head Attention)은 self-attention을 다양한 각도에서 head 수 만큼 수행하고, 각자의 시각에서 문장을 분석하게 됩니다.

https://wikidocs.net/31379

이러한 계산 방법은 이전 순환 신경망에 비해 병렬처리도 가능하고 훈련이 빠르고 성능이 매우 뛰어났습니다. 하지만 Incremental Inference(Auto-regressive하게 이전 단어를 활용해 새로운 단어를 생성하도록 순환적인 추론)부분에서는 병렬처리가 불가능하기 때문에 “키(keys)”와 “값(values)”이라 불리는 큰 tensor를 반복해서 로딩해야 하는 메모리 대역폭 비용 때문에 병목현상이 발생합니다.

이 문제를 해결하기 위해 2019년 구글 소속인 Noam Shazeer가 MQA(Multi-query attention라는 것을 제안합니다. 논문 제목은 “Fast Transformer Decoding: One Write-Head is All You Need”. 추론 속도를 높이기 위해서는 한개의 머리만 있으면 된다는 의미입니다.

MQA는 멀티헤드 어텐션에서 각 어텐션 헤드가 독립적인 쿼리(Query), 키(Key), 값(Value) 세트를 사용하는 대신 모든 head가 동일한 키와 값 세트를 공유하는 방식입니다.

MHA와 MQA를 수식으로 표현하면 다음과 같습니다.(논문 사진)

MHA는 각 계산마다 head에 따라 다른 K, V를 사용하여 4차원 tensor 연산인 반면 MQA 모든 head가 동일한 K, V를 사용하기 때문에 3차원 tensor가 됩니다. 이러한 구조의 디코더는 연산량과 tensor를 불러오는 속도를 최소화하기 때문에 추론 속도를 대폭 향상시킵니다. 하지만 이러한 구조는 성능이 떨어질 수 있다는 문제가 있습니다.

이를 해결하기 위해서 등장한 것이 Grouped-query attention입니다.

기존 MHA, MQA의 중간 버전이라고 생각하시면 됩니다.

K, V를 중간정도만 공유하는 GQA가 단일 K, V를 사용하는 MHA의 속도, 성능 면에서 모두 우수한 결과가 나왔다는 것을 실험 결과를 통해서 보여주고 있습니다. 어중간하지 않고 오히려 좋은 결과를 보여주는 것이 신기합니다. (역시 구글..)

입력 길이가 2048이고 출력 길이가 512일 때, GQA의 적절한 그룹 크기는 8개라고 합니다.

또한 GQA는 기존에 학습된 모델을 추가로 조금만 학습하여 적용이 가능합니다. MHA를 MQA로 변환하는 과정은 위 그림과 같습니다.

우선 멀티헤드 모델의 모든 헤드에서 사용된 키와 값에 대한 투영 행렬을 평균 풀링하여 단일 헤드의 투영 행렬을 생성합니다. 이는 여러 헤드에 걸쳐 분산된 정보를 통합하여 단일 헤드가 접근할 수 있게 합니다. 이후 그룹화된 head를 추가로 학습하여 모델이 이 구조에 적응할 수 있도록 합니다.

오른쪽 그림을 보면 추가로 10% 정도 학습했을 때 원래 MHA 만큼의 추론 성능을 보여주는 그림입니다. (추가 학습을 아예 하지 않아도 성능 저하가 심하지 않습니다.)

얼마 전까지 오픈소스 SOTA였던 라마2 논문을 자세히 들여다보면 위와 같은 그래프가 나와있습니다. MQA와 GQA모델은 더 큰 배치 크기에서 더 많은 처리량(throughput)을 가능하게 하며, 작은 배치 크기에서는 MHA와 유사한 지연시간(latency)을 보입니다. MHA은 256 토큰 컨텍스트에서 배치 크기 1024, 2k 컨텍스트에서 배치 크기 128에서 메모리 부족 오류를 일으켰습니다. 반면 MQA와 GQA는 이러한 설정에서 성공적으로 실행되었습니다.

따라서 34B와 70B Llama 2 모델에서는 MQA대신 GQA를 사용했습니다. 어텐션 레이어의 축소를 보상하기 위해 피드포워드 네트워크(FFN)의 차원을 1.3배 정도 증가시켰다고 합니다. GQA는 특히 대규모 모델과 큰 배치 크기에서 추론 속도를 향상시키는 데 유리한 것으로 나타났습니다. 이후 많은 오픈소스 모델들이 이 모델을 fine-tunning해서 제작했습니다.

그런데 신기하게도 구글의 Gemma는 7B에서는 기존의 MHA를 사용했고, 2B에서는 1 kv head를 사용하는 GQA를 사용했습니다. 실험을 통해 이렇게 정했다고 했지만, 정확한 근거는 나와있지 않았습니다.

결론적으로는 구글이 Llama 2 13B 모델보다도 더 좋은 성능의 오픈 소스를 출시했습니다. 이 외에도 몇 가지 llama2와 다른 점이 있긴 하지만, 토큰을 transformation 하는 것이 핵심인 모델인 만큼 성능에 가장 많은 영향을 끼쳤을 것이라고 생각합니다. (추론 속도는 더 느릴 것으로 생각됨.)

Self-Attention의 응용이 LLM 성능 개선의 핵심이기 때문에, 위와 같은 연구들이 더 필요할 것이라고 생각합니다.