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

Paper: Joshua Ainslie; James Lee-Thorp; Michiel de Jong; Yury Zemlyanskiy; Federico Lebrón; Sumit Sanghai (2023). "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints". arXiv:2305.13245v3. 검토 PDF · arXiv:2305.13245v3 · EMNLP 2023 출판본 · DOI. 2023년 12월 공개된 v3를 기준으로 설명합니다.

언어 모델이 다음 단어를 만들 때는 앞에서 읽은 내용도 참고해야 합니다. 매번 긴 문장을 처음부터 다시 계산하지 않도록, 앞선 단어에서 만든 Key와 Value를 저장해 두는데 이 공간을 KV 캐시라고 부릅니다. 질문하는 쪽은 여러 개로 유지하면서, 찾아볼 Key와 가져올 Value는 함께 쓸 수 있을까요?

GQA(Grouped-Query Attention)는 그 절충을 시도합니다. 기존처럼 query head는 다양하게 두되 K/V head를 몇 개의 그룹으로 줄입니다. 논문은 여기에 기존 모델의 K/V 가중치를 평균해 구조를 바꾸고, 추가 사전학습으로 적응시키는 방법을 결합합니다. 따라서 이미 만든 KV 캐시를 손실 없이 압축하는 기술과는 다릅니다. 상세 읽기에서는 변환과 실험 수치의 조건을 더 자세히 살펴봅니다.

1. KV 캐시는 다음 토큰이 다시 볼 내용을 저장합니다

Query는 지금 무엇을 찾을지 나타내고, Key는 어떤 정보가 관련 있는지 비교하는 기준이며, Value는 실제로 가져올 내용입니다. Query와 Key를 비교해 관련도를 얻고, 그 비중만큼 Value를 모읍니다. 다음 토큰을 생성할 때 앞선 토큰의 Key와 Value를 다시 쓸 수 있도록 저장하는 것이 KV 캐시입니다.

다중 헤드 어텐션(MHA)은 각 query head마다 별도의 Key·Value 표현을 둡니다. 서로 다른 관점의 질문을 유지하는 대신 저장할 K/V도 많아집니다. 반대편의 다중 쿼리 어텐션(MQA)은 여러 query가 K/V 한 쌍을 함께 씁니다. 캐시는 작아지지만 공유가 너무 많아지면 표현의 선택지가 줄 수 있습니다. GQA는 그 중간에서 K/V를 여러 그룹으로 나눕니다.

2. GQA는 query는 유지하고 K/V head를 묶습니다

Figure 2. Query head 수는 유지되며, 가운데 도식은 query 8개가 K/V 4쌍을 공유하는 구조 예시입니다. 이는 GQA-8의 실험 구성을 그대로 나타내지 않습니다. 출처: Ainslie et al. (2023), v3, p. 2, Figure 2.

예를 들어 query head가 32개이고 K/V 그룹이 8개인 설명용 설정을 생각해 보겠습니다. 그룹을 균등하게 나누면 네 query head가 K/V 한 쌍을 함께 사용합니다. 이는 공유되는 표현을 줄여도 query마다 다른 질문을 유지할 수 있음을 보여 주는 예입니다. 32 query·8 K/V 설정은 이해를 위한 예시이며 논문의 T5 실제 head 설정이 아닙니다.

GQA-8에서 숫자 8은 query 수가 아니라 K/V 그룹 수입니다. 그림처럼 MHA는 각 query가 자기 K/V를 쓰고, MQA는 모든 query가 한 쌍을 나눠 씁니다. GQA는 그 사이에 공유 그룹을 둡니다. 같은 K/V를 쓰더라도 query는 서로 다르므로 어떤 정보를 중요하게 보는지는 달라질 수 있습니다.

3. 평균으로 구조를 바꾼 뒤 추가 학습으로 적응합니다

이미 학습된 MHA 모델을 GQA로 바꿀 때 저자들은 한 그룹에 들어갈 기존 K/V projection 가중치 행렬을 평균합니다. Query 가중치나 이미 만들어 둔 KV 캐시를 평균하는 것이 아닙니다. 가중치 평균은 여러 head를 하나의 공유 표현으로 만드는 출발점이지만, 원래 모델의 계산을 그대로 보존하지는 않습니다. 그래서 변환 뒤에도 학습을 이어 갑니다.

대표 설정은 원래 사전학습 step 수의 5%만큼 추가 학습하는 것입니다. 저자들은 이를 약 600 TPUv3 chip-days로 보고합니다. Chip-days는 여러 가속기를 합산한 자원량이지, 작업을 600일 동안 기다렸다는 뜻은 아닙니다. 추가 학습을 전혀 하지 않는 0% 설정도 이후 task별 fine-tuning까지 생략했다는 뜻은 아닙니다.

Figure 5. T5-XXL의 MQA와 GQA-8을 CNN/Daily Mail·MultiNews·TriviaQA 세 과제에서 비교합니다. 가로축의 α=0은 추가 사전학습이 없다는 뜻이며, 과제별 미세조정까지 생략한 조건은 아닙니다. 출처: Ainslie et al. (2023), v3, p. 4, Figure 5.

이 그림은 다룬 과제에서 변환 직후와 추가 학습 뒤의 성능을 비교합니다. 관찰된 차이를 모든 모델과 task에서 같은 폭으로 회복된다는 보장으로 읽어서는 안 됩니다. 이 절차의 핵심은 기존 checkpoint를 활용하면서 구조 변경 뒤의 성능을 추가 학습으로 보완한다는 데 있습니다.

4. 줄어드는 것은 주로 KV 저장량입니다

같은 모델·문맥 길이·자료형에서 MHA 대신 GQA-8을 쓰면 K/V head 수가 줄어 KV 캐시에 저장할 값도 함께 줄어듭니다. 앞의 설명용 32 query·8 K/V 예에서는 K/V head 축의 저장량이 MHA의 4분의 1입니다. 이 비율은 KV head 수만 놓고 비교한 것이지, 실제 장치 사용량을 측정한 결과는 아닙니다.

또한 GQA는 캐시에 담긴 값을 저비트로 인코딩하거나 일부 토큰을 버리는 압축이 아닙니다. K/V projection의 구조를 바꾸는 방식입니다. Query head와 그 계산은 남아 있고, 출력 projection이나 다른 모델 연산도 같은 비율로 작아지지 않습니다. 그러므로 캐시가 4분의 1이 된다는 예에서 전체 모델 계산이나 응답 시간이 4분의 1이 된다고 결론 내릴 수는 없습니다.

5. T5의 평균 47.2와 47.1은 정답률 비교가 아닙니다

논문은 T5.1.1 Large와 XXL encoder–decoder 모델에서 요약, 번역, 질문응답을 평가했습니다. 표의 평균 점수는 MHA-XXL 47.2, GQA-8-XXL 47.1입니다. 숫자가 가깝다는 것은 이 평가에서 보고된 여러 점수의 집계가 비슷하다는 뜻입니다. 정답률이나 보편적인 품질 보존율을 뜻하지 않습니다.

평가에는 요약용 ROUGE-1, 번역용 BLEU, 질문응답용 F1이 섞여 있습니다. 각 과제에서 점수의 단위와 의미가 다르므로 평균 하나를 보편적인 품질 척도로 보기는 어렵습니다. GQA가 모든 과제에서 MHA보다 높거나 같은 것도 아닙니다. 여기서 안전한 결론은 해당 T5 실험의 집계 점수가 가깝다는 정도입니다.

Fine-tuning 뒤 dev 성능이 가장 높은 checkpoint를 골랐고, Table 1도 dev 결과이므로 독립적인 blind test 결과는 아닙니다.

실험 범위도 중요합니다. 저자들은 decoder의 self-attention과 encoder 출력을 읽는 cross-attention에 GQA를 적용했지만 encoder의 self-attention은 유지했습니다. 따라서 이 논문만으로 decoder-only 모델이나 다른 구조의 성능까지 확인했다고 말할 수는 없습니다.

6. 5.39배는 TPU 실험의 조건부 비율입니다

논문 Table 1은 TPUv4 칩당 샘플 처리 시간을 보고합니다. 표기된 단위를 따르면 MHA-XXL은 1.51초, GQA-8-XXL은 0.28초입니다. 이는 전체 작업에 걸린 wall-clock 시간이 아닙니다. 1.51을 0.28로 나누면 약 5.39가 됩니다. 저자 측정에서 보고한 시간 지표의 비율이지, 모든 사용자가 GQA를 적용하면 언제나 5.39배 빠르게 응답을 받는다는 뜻은 아닙니다.

시간은 8개 TPUv4를 사용한 실험에서 측정했으며, 모델마다 들어갈 수 있는 batch 크기와 병렬 배치를 따로 조정했습니다. TPU당 batch를 최대 32까지 허용한 조건입니다. 따라서 같은 batch의 단일 요청 지연시간이나 일반 서비스의 토큰 생성 속도를 직접 나타내지는 않습니다. 같은 표에서 GQA-8은 MQA보다도 0.28초 대 0.24초로 조금 더 오래 걸렸습니다.

단위 표기에도 주의가 필요합니다. Table 1은 초라고 쓰지만 Figure 3은 비슷한 크기의 시간을 밀리초라고 표시합니다. 원문은 이 차이를 해소하지 않으므로 한쪽을 임의로 정답으로 선택하지 않습니다. 위 비율은 Table 1의 표시값끼리 계산한 것입니다.

7. 8개 그룹은 실험의 절충이지 만능 설정은 아닙니다

Figure 6. 입력 길이 2,048과 출력 길이 512 조건에서 그룹 수에 따른 시간을 비교합니다. 출처: Ainslie et al. (2023), v3, p. 4, Figure 6.

그룹 수를 늘리면 공유하는 K/V가 줄어 MHA 쪽 구조에 가까워지고, 시간과 저장량의 절충도 달라집니다. 논문은 T5 실험에서 GQA-8을 유용한 중간 선택으로 제시합니다. 그러나 Figure 6의 시간 곡선만으로 다른 모델에서도 8개가 최선이라고 정할 수는 없습니다. 목표 모델과 가속기, 문맥 길이, batch 조건에 맞춰 다시 측정해야 합니다.

GQA가 보여 주는 요점은 query의 다양성을 유지하면서 K/V head를 그룹별로 공유할 수 있다는 것입니다. T5 결과는 MHA와 비슷한 집계 점수와 더 작은 보고 시간 지표를 함께 제시하지만, 그 두 수치는 실험 범위와 측정 조건 안에서 읽어야 합니다. 구조를 바꾸고 추가 학습한 모델의 결과이며, 모든 모델에 바로 적용되는 무손실 캐시 압축이나 보편적 지연시간 보장은 아닙니다.

References

Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training generalized multi-query Transformer models from multi-head checkpoints. In H. Bouamor, J. Pino, & K. Bali (Eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (pp. 4895–4901). Association for Computational Linguistics. https://doi.org/10.18653/v1/2023.emnlp-main.298