그룹화된 쿼리 어텐션(GQA)이란 무엇인가요?

Ainslie et al은 트랜스포머 신경망을 확립한 2017년 논문 "Attention is All You Need"에 소개된 혁신적인 자기 주의 알고리즘인 멀티헤드 어텐션(MHA)의 최적화로 그룹화된 쿼리 어텐션을 고안했습니다. 보다 구체적으로, MHA의 초기 최적화인 멀티 쿼리 어텐션(MQA)의 일반화 및 보다 제한적인 적용으로 제안되었습니다.

표준 멀티헤드 어텐션은 머신 러닝, 자연어 처리(NLP), 생성형 AI에서 진화적 도약을 촉진했지만, 컴퓨팅 리소스와 메모리 대역폭에 엄청난 부담을 주었습니다. LLM이 더 크고 정교해짐에 따라 이러한 메모리 사용 요구 사항은 진행에 병목 현상이 되었고, 특히 텍스트 생성, 요약 및 기타 생성형 AI 작업에 사용되는 자기 회귀 디코더 전용 LLM의 경우 이는 더욱 걸림돌이 되었습니다.

후속 연구는 멀티헤드 어텐션을 강화하거나 간소화하는 기술에 초점을 맞췄습니다. 플래시 어텐션 및 링 어텐션과 같은 일부 기능은 모델을 훈련시키고 실행하는 데 사용되는 GPU가 계산 및 메모리 스토리지를 처리하는 방식을 개선합니다. GQA 및 MQA와 같은 다른 기술들은 트랜스포머 아키텍처가 토큰을 처리하는 방식에 대한 변화를 모색했습니다.

그룹화된 쿼리 어텐션은 표준 멀티헤드 어텐션과 멀티 쿼리 어텐션 간의 균형을 맞추는 것을 목표로 합니다. 전자는 메모리 대역폭 오버헤드가 증가하고 속도가 저하되는 대신 정확도를 극대화합니다. 후자의 경우 정확도는 떨어지지만 속도와 효율성은 극대화됩니다.

그룹화된 쿼리 어텐션이 트랜스포머 모델을 최적화하는 방법을 이해하려면 먼저 멀티 헤드 어텐션이 일반적으로 어떻게 작동하는지 이해하는 것이 중요합니다. GQA와 MQA는 모두 MHA의 핵심 방법론을 대체하는 것이 아니라 개선할 뿐입니다.

LLM과 트랜스포머 아키텍처를 사용하는 다른 모델의 원동력은 시퀀스의 각 토큰 간의 관계를 이해하기 위한 수학적 프레임워크인 셀프 어텐션입니다. 셀프 어텐션을 통해 LLM은 정적 기준 정의뿐만 아니라 다른 단어와 구문이 제공하는 컨텍스트를 통해 텍스트 데이터를 해석할 수 있습니다.

텍스트 생성에 사용되는 자동 회귀 LLM에서 어텐션 메커니즘은 모델이 그 순간에 가장 "주의할" 가치가 있는 이전 토큰을 결정하여 시퀀스의 다음 토큰을 예측하는 데 도움이 됩니다. 가장 관련성이 높다고 간주되는 토큰의 정보에는 더 큰 어텐션 가중치가 부여되는 반면, 관련이 없는 것으로 간주되는 토큰의 정보에는 0에 가까운 어텐션 가중치가 부여됩니다.

트랜스포머 모델을 애니메이션화하는 멀티헤드 어텐션 메커니즘은 어텐션 계층을 여러 개의 어텐션 헤드로 분할하고 셀프 어텐션을 여러 번 병렬로 계산하여 풍부한 컨텍스트 정보를 생성합니다.

"Attention is All You Need"의 저자는 관계형 데이터베이스의 용어인 쿼리, 키, 값을 사용하여 어텐션 메커니즘을 설명했습니다. 관계형 데이터베이스는 관련 데이터의 저장 및 검색을 간소화하도록 설계되었습니다. 각 데이터에 고유 식별자("키")를 할당하고 각 키는 해당 값과 연결됩니다.관계형 데이터베이스의 목표는 각 쿼리를 적절한 키와 일치시키는 것입니다.

시퀀스의 각 토큰에 대해 멀티헤드 어텐션에는 3개의 벡터 생성이 필요합니다.

• 토큰이 "찾고 있는" 정보를 나타내는 쿼리 벡터 Q입니다. 예를 들어, 명사에 대한 쿼리 벡터는 명사를 설명하는 형용사에 대한 검색을 나타낼 수 있습니다.
• 토큰에 포함된 정보를 나타내는 키 벡터 K입니다. 각 토큰의 키 벡터와 다른 토큰의 쿼리 벡터의 관련성을 나타내는 정렬 점수는 어텐션 가중치를 계산하는 데 사용됩니다.
• 값 벡터 V는 다른 토큰의 키 벡터에서 어텐션 가중치 기여에 의해 업데이트될 컨텍스트 정보를 나타냅니다.

어텐션 메커니즘에 의해 매개되는 이 세 가지 벡터 간의 수학적 상호 작용은 모델이 각 토큰에 대한 맥락별 이해를 조정하는 방식입니다.
 

주어진 토큰에 대한 이 3개 벡터를 각각 생성하기 위해 모델은 해당 토큰의 원래 벡터 임베딩으로 시작합니다. 이는 벡터의 각 차원이 토큰의 의미론적 의미의 일부 추상적 요소에 해당하는 숫자 인코딩입니다. 이러한 벡터의 차원 수는 미리 정해진 하이퍼매개변수입니다.

각 토큰에 대한 Q, K 및 V 벡터는 첫 번째 어텐션 계층 이전의 선형 계층을 통해 원래 토큰 임베딩을 전달하여 생성됩니다. 이 선형 계층은 모델 가중치의 3가지 고유한 행렬(W_Q, W_K 및 W_V)로 분할되고, 그 안의 특정 가중치 값은 텍스트 예제의 방대한 데이터 세트에 대한 자기 지도 사전 훈련을 통해 학습됩니다.

토큰의 기존 벡터 임베딩을 W_Q, W_K 및 W_V로 곱하면 각각 해당 쿼리 벡터, 키 벡터 및 값 벡터가 생성됩니다. 각 벡터가 포함하는 차원 수 d는 각 가중치 행렬의 크기에 따라 결정됩니다. Q 및 K는 동일한 차원 수 d_k를 갖습니다.

그런 다음 이 3개의 벡터는 어텐션 계층으로 전달됩니다.

어텐션 계층에서 Q, K 및 V 벡터는 시퀀스의 각 위치에 있는 각 토큰 간의 정렬 점수를 계산하는 데 사용됩니다. 그런 다음 이러한 정렬 점수는 소프트맥스 함수를 사용하여 어텐션 가중치로 정규화됩니다.

시퀀스의 각 토큰 x에 대해 정렬 점수는 해당 토큰의 쿼리 벡터 Q_x와 다른 각 토큰의 키 벡터 K의 점곱 을 계산하여 즉, 두 벡터를 곱하여 계산합니다. 두 토큰 간의 의미 있는 관계가 각 벡터 사이의 유사성에 반영되어 있다면 두 토큰을 곱하면 큰 값을 얻을 수 있습니다. 두 벡터가 정렬되지 않은 경우, 함께 곱하면 작거나 음의 값이 산출됩니다. 대부분의 트랜스포머 모델은 확장된 점곱 어텐션라는 변형을 사용하는데, 여기서는 QK는 확장된 즉, 다음과 같이 곱해집니다. $\frac{1}{\sqrt{d_K}}$  학습 안정성을 개선합니다.

그런 다음 이러한 쿼리 키 정렬 점수가 소프트맥스 함수에 입력됩니다.소프트맥스는 모든 입력을 0과 1 사이의 값으로 정규화하여 모두 합이 1이 되도록 합니다. 소프트맥스 함수의 아웃풋은 어텐션 가중치이며, 각각은 해당되는 다른 토큰이 토큰 x로부터 받는 어텐션 점유율(1 이내)을 나타냅니다. 토큰의 어텐션 가중치가 0에 가까우면 무시됩니다. 어텐션 가중치 1은 토큰이 x의 전체 어텐션을 받고 다른 모든 것은 무시됨을 의미합니다.

마지막으로, 각 토큰의 값 벡터에 어텐션 가중치를 곱합니다. 이렇게 각 이전 토큰의 어텐션 가중치가 적용된 기여도를 모두 평균하여 토큰 x의 기존 임베딩 벡터에 추가합니다. 이제 토큰 x의 임베딩은 관련성이 있는 시퀀스의 다른 토큰이 제공한 컨텍스트를 반영하도록 업데이트됩니다.

업데이트된 벡터 임베딩은 그 자체의 가중치 행렬 W_Z를 갖는 다른 선형 계층으로 전송되고, 여기서 컨텍스트 업데이트된 벡터는 일관된 차원 수로 다시 정규화되고 다음 어텐션 계층으로 전송됩니다. 각 진행형 어텐션 계층은 더 큰 맥락적 뉘앙스를 포착합니다.
 

어텐션 가중치가 적용된 컨텍스트의 각 부분을 개별적으로 고려하는 대신 다른 토큰의 어텐션 가중치 기여도의 평균을 사용하는 것은 수학적으로 효율적이지만 세부 정보가 손실됩니다.

이를 보상하기 위해 트랜스포머 네트워크는 기존 입력 토큰의 임베딩을 균일한 크기의 h 조각으로 분할합니다. 마찬가지로 W_Q, W_K 및 W_V를 각각 쿼리 헤드, 키 헤드 및 값 헤드라고 하는 h 하위 집합으로 분할합니다. 각 쿼리 헤드, 키 헤드 및 값 헤드는 기존 토큰 임베딩의 일부를 받습니다. 쿼리 헤드, 키 헤드 및 값 헤드의 이러한 병렬 삼중항 각각에서 생성된 벡터는 해당 어텐션 헤드로 공급됩니다.결국, 이러한 h 병렬 회로의 아웃풋은 전체 토큰 임베딩을 업데이트하기 위해 다시 연결됩니다.

학습 과정에서 각 회로는 의미론적 의미의 개별적인 측면을 포착하는 별개의 가중치를 학습합니다. 이를 통해 모델은 단어의 의미가 주변 다른 단어의 컨텍스트에 따라 영향을 받을 수 있는 다양한 방식을 처리할 수 있습니다.

표준 멀티헤드 어텐션의 단점은 결정적인 결함이 있다기보다는 최적화가 전혀 이루어지지 않았다는 점입니다. MHA는 최초의 알고리즘이었으며 어텐션 계산을 위한 일반적인 메커니즘 중 가장 복잡한 실행을 나타냅니다.

MHA의 비효율성은 대부분 계산과 모델 매개변수가 너무 많다는 데서 비롯됩니다. 표준 MHA에서는 각 어텐션 블록의 각 쿼리 헤드, 키 헤드, 값 헤드에 고유한 가중치 매트릭스가 있습니다. 예를 들어, 각 어텐션 계층에 8개의 어텐션 헤드가 있는 모델(대부분의 최신 LLM보다 훨씬 적은 수)의 경우 계층의 Q, K, V 헤드에만 24개의 고유 가중치 매트릭스가 필요합니다. 이는 각 계층에서 엄청난 수의 중간 계산을 수반합니다.

이 구성의 한 가지 결과는 계산 비용이 많이 든다는 것입니다. MHA에 대한 컴퓨팅 요구 사항은 시퀀스 길이와 관련하여 2차식으로 확장되므로 입력 시퀀스의 토큰 수를 2배로 늘리려면 복잡성을 4배로 늘려야 합니다. 이로 인해 컨텍스트 창의 크기에 대한 실질적인 제한이 생깁니다.

MHA는 또한 시스템 메모리에 큰 부담을 줍니다. GPU에는 각 후속 처리 단계에서 호출해야 하는 방대한 양의 중간 계산의 아웃풋을 저장할 수 있는 온보드 메모리가 많지 않습니다. 이러한 중간 결과는 대신 GPU 칩 자체에 위치하지 않는 고대역폭 메모리(HBM)에 저장됩니다. 따라서 메모리에서 키와 값을 읽어야 할 때마다 약간의 지연 시간이 발생합니다. 트랜스포머 모델이 수십억 개의 매개변수로 확장되기 시작하면서 추론을 학습하고 실행하는 데 필요한 시간과 컴퓨팅이 모델 성능에 병목 현상을 일으켰습니다.

더 많은 발전을 이루기 위해서는 난해하고 복잡한 언어 패턴을 학습하고 재현하는 트랜스포머의 용량을 줄이지 않으면서 계산 단계의 수를 줄일 수 있는 방법이 필요했고, 이러한 상황에서 MQA와 GQA가 도입되었습니다.

멀티 쿼리 어텐션(MQA)은 메모리 사용량과 중간 계산을 줄이기 위해 멀티헤드 어텐션을 단순화하는 보다 계산 효율적인 어텐션 메커니즘입니다. 각 어텐션 헤드에 대해 고유한 키 헤드와 값 헤드를 학습하는 대신 MQA는 각 계층에서 단일 키 헤드와 단일 값 헤드를 사용합니다. 따라서 키 벡터와 값 벡터는 한 번만 계산됩니다. 이 단일 키 및 값 벡터 세트는 모든 h 어텐션 헤드에서 공유됩니다.

이러한 단순화는 모델이 계산하고 고대역폭 메모리에 저장해야 하는 선형 투영의 수를 크게 줄입니다. MQA를 소개한 2019년 논문에 따르면 MQA는 10~100배 더 작은 키-값 쌍 스토리지(또는 KV 캐시)와 12배 더 빠른 디코더 추론을 허용합니다. MQA의 메모리 사용량 감소는 또한 더 큰 배치 크기를 가능하게 하여 훈련 속도를 크게 높입니다.

이러한 장점에도 불구하고 MQA에는 몇 가지 불가피한 단점이 있습니다.

• 성능 저하: 당연히 학습 가능한 고유한 모델 매개변수의 수를 줄이면 모델의 지식 및 뉘앙스 용량이 감소합니다. MQA는 표준 MHA에 비해 정확도가 유의미하게 떨어지므로 특정 상황 및 사용 사례에는 적합하지 않습니다.
  
  
• 처음부터 훈련해야 함: 표준 MHA로 훈련된 모델은 단순히 MQA에 적용할 수 없으며 대신 처음부터 MQA로 훈련해야 합니다. 이는 MQA를 사용하여 기존 모델을 최적화할 수 없으며 새 모델에 대해 MQA를 실험할 때 상당한 기회 비용이 수반됨을 의미합니다.
  
  
• 텐서 병렬 처리의 중복성: GPU에서 트랜스포머 모델을 학습시킬 때의 주요 이점 중 하나는 여러 개의 복잡한 텐서 연산을 병렬로 수행할 수 있다는 점입니다. 이러한 연산을 수행하는 GPU 클러스터의 각 노드에 K 및 V 값이 존재해야 하므로 실제로는 각 노드에 대해 복제해야 합니다. 이는 표준 MHA보다 여전히 효율적이지만 컴퓨팅 리소스를 최적으로 사용하는 것은 아닙니다.

그룹화된 쿼리 어텐션은 모든 쿼리 헤드에서 하나의 키 및 값 집합을 공유하는 대신 쿼리 헤드를 각각 키와 값 집합을 공유하는 여러 그룹으로 분할하는 멀티쿼리 어텐션보다 더욱 일반적이고 유연한 형식입니다.

2023년 5월에 'GQA: 멀티 헤드 체크포인트에서 일반화된 멀티 쿼리 트랜스포머 모델 학습'이 발표된 후 많은 LLM이 빠르게 GQA를 채택했습니다. 예를 들어, Meta는 2023년 7월 Llama 2 모델에 GQA를 처음 채택했으며, 2024년에 출시된 Llama 3 모델에도 GQA를 적용했습니다. Mistral AI는 2023년 9월에 출시한 Mistral 7B 모델에 GQA를 사용했으며, IBM의 Granite 3.0 모델 또한 빠른 추론을 위해 GQA를 채택했습니다.

이론적으로 GQA는 표준 MHA와 전체 MQA 사이의 스펙트럼을 일반화한 것으로 생각할 수 있습니다. 어텐션 헤드와 키-값 헤드 그룹 수가 같은 GQA는 표준 MHA에 해당하며 헤드 그룹이 1개인 GQA는 MQA에 해당합니다.

실제로 GQA는 대개 그룹 수 자체가 중요한 하이퍼 파라미터인 중간 접근 방식을 의미합니다.

그룹화된 쿼리 어텐션은 다음과 같은 몇 가지 장점을 제공하여 주요 LLM에서 비교적 광범위하게 채택하고 있습니다.

• 효율적인 GPU 사용: GQA의 키-값 쌍 배포는 텐서 병렬 처리를 활용하여 중복 값을 복제함으로써 '낭비되는' 컴퓨팅의 양을 줄입니다.
  
• 효과적인 절충안: GQA는 MHA만큼 정확하면서도 MQA만큼 빠르다는 점에서 디코더 추론 속도와 성능 정확도 간의 이상적인 절충안을 제공합니다.
  
  
• 메모리 대역폭 오버헤드 감소: MQA와 마찬가지로 GQA는 추론 시간에 계산, 저장 및 검색해야 하는 중간 계산 수를 크게 줄입니다.
  
• 유연한 훈련: MQA와 달리 그룹 쿼리 어텐션은 해당 접근 방식을 사용하여 모델을 처음부터 훈련할 필요가 없습니다. 표준 MHA를 사용하여 사전 훈련된 모델은 "업트레이닝"이라는 미세 조정 프로세스를 통해 GQA를 사용하도록 조정할 수 있습니다.