기술 요약: 그룹 쿼리 전문가 (Grouped Query Experts, GQE)

문제 정의

셀프 어텐션(Self-attention)은 트랜스포머 아키텍처, 특히 긴 컨텍스트 시퀀스에서 쌍별 토큰 상호작용이 이차적으로(quadratically) 스케일링됨에 따라 계산상의 병목 현상이 됩니다. 표준 밀집 어텐션(dense attention) 메커니는 토큰의 정보 함량이나 난이도와 관계없이 모든 토큰에 동일한 어텐션 헤드를 적용합니다. 이러한 균일한 활성화는 낮은 정보량을 가진 토큰(예: 구두점, 불용어)이 정보가 풍부한 단어와 동일한 자원을 소비하게 만들어 계산 효율성을 저하시킵니다. 혼합 전문가(Mixture-of-Experts, MoE)가 트랜스포머 피드 포워드(MLP) 블록에 조건부 연산을 성공적으로 도입한 반면, 어텐션 메커니즘에 유사한 희소성(sparsity)을 적용하는 것은 여전히 해결되지 않은 과제로 남아 있습니다. 기존의 그룹 쿼리 어텐션(Grouped-Query Attention, GQA)은 여러 쿼리 헤드가 하나의 KV 헤드를 공유하도록 하여 메모리 대역폭과 KV 캐시 크기를 줄이지만, 여려 여전히 모든 토들은에 대해 모든 쿼리 헤드를 평가합니다.

방법론: 그룹 쿼리 전문가 (GQE)

저자들은 GQA 모델의 어텐션 블록에 MoE 원리를 통합하는 방법인 **그룹 쿼리 전문가(Grouped Query Experts, GQE)**를 제안합니다. 핵심 설계 철학은 GQA의 밀집 KV 캐시 이점을 유지하면서 쿼리 헤드 연산을 조건부로 만드는 것입니다.

아키텍처 및 라우팅

1. 그룹 전문가 풀 (Grouped Expert Pool): 각 고정된 GQA 그룹 내에서(여러 쿼리 헤드가 단일 KV 헤드를 공유하는 구조), 쿼리 헤드는 "전문가"의 풀로 취급됩니다.
2. 조건부 활성화 (Conditional Activation): 각 토큰에 대해, 라우터는 각 GQA 그룹으로부터 $k$개의 쿼리 헤드 전문가로 구성된 희소 서브셋을 선택합니다.
3. 밀집 KV 경로 (Dense KV Path): 결정적으로, 키-값(Key-Value) 투영과 KV 캐시는 밀집 상태를 유지하며 변경되지 않습니다. 희소성은 오직 쿼리 측 연산에만 적용됩니다.
4. 출력 구성: 안정적인 학습을 보장하기 위해 특정 메커니즘을 통해 출력을 구성합니다:
   • 하드 연결 (Hard Concatenation): 각 그룹에서 선택된 $k$개 전문가의 출력을 연결합니다.
   • 재정규화된 가중 합 (Renormalized Weighted Sum): 라우터 확률을 사용하여 선택된 전문가 출력의 가중 합을 계산합니다. 이는 이산적인 top-$k$ 선택의 미분 불가능성 문제를 해결하여, 언어 모델링 손실(loss)이 라우터에 도달할 수 있는 미분 가능한 경로를 제공합니다.
   • 항상 켜져 있는 공유 헤드 (Always-On Shared Head): 라우팅과 관계없이 모든 토큰에 대해 추가적인 어텐션 헤드가 계산됩니다. 이는 학습 중 안정적인 앵커 역할을 하여, 라우터가 단일 전문가로 붕괴되는 것을 방지하고 항상 기본적인 어텐션 패턴을 사용할 수 있도록 합니다.

학습 및 손실

표준 언어 모델링 손실 외에도, GQE는 라우터가 각 그룹 내의 전문가 풀에 토큰을 고르게 분배하도록 유도하여 "전문가 붕괴(expert collapse)"를 방지하는 로드 밸런싱 보조 손실(load-balancing auxiliary loss)을 사용합니다.

주요 기여

1. GQA에서의 MoE 정식화: 본 논문은 모든 KV 헤드를 밀집 상태로 유지하면서 각 토큰당 $k$개의 전문가를 활성화하도록 GQA 그룹 내에 특화된 MoE를 정식화했습니다. 이는 근본적인 GQA 구성과 비교를 일치시키고 GQA의 메모리 프로필을 보존합니다.
2. 품질 저하 없는 계산 감소: 저자들은 GQE가 모든 전문가가 활성화된 GQA 베이스라인의 다운스트림 정확도와 일치하면서도, 희소한 쿼리 헤드 전문가 서브셋만을 활성화할 수 있음을 입증했습니다.
3. 안정화 메커니즘: 논문은 성공적인 어텐션 내 희소 라우팅을 위해서는 두 가지 특정 구성 요소, 즉 라우터에 적절한 학습 신호를 제공하기 위한 재정규화된 가중 합 슬롯과 레이어를 고정하기 위한 항상 켜져 있는 공유 헤드가 필요함을 식별했습니다.

실험 결과

실험은 FineWeb-Edu 데이터셋을 사용하여 250M 파라미터 규모에서 30B 토큰 예산을 기준으로 수행되었습니다.

• 정확도: 주요 설정(16개 쿼리 헤드, 8개 KV 그룹, $k=1$)에서 GQE는 8개의 라우팅된 전문가와 1개의 공유 헤드를 활성화하여, 총 16개의 쿼리-어텐션 연산 중 9개를 수행합니다. 이러한 감소에도 불구하고(전체 쿼리 연산의 약 56% 활성화), GQE는 다운스트림 정확도(HellaSwag, PIQA, ARC-Easy)에서 모든 전문가가 활성화된 GQA 베이스라인과 일치하는 성능을 보였으며, 베이스라인의 55.86에 비해 평균 56.04를 달성했습니다.
• 절제 연구 (Ablation Studies): 논문은 재정규화된 가중 합 슬롯이나 공유 헤드를 제거할 경우 성능이 저하됨(각각 약 55.18 및 55.43으로 하락)을 보여주며, 이는 특정 아키텍처 안정 장치 없이는 희소성만으로는 불충분하다는 것을 확인시켜 줍니다.
• 처리량 및 속도 향상: GQE는 긴 컨텍스트 환경에서 상당한 속도 향상을 보여줍니다.
  • 짧은 시퀀스(2k 토큰)에서는 라우팅 오버헤드로 인해 속도 향상이 완만합니다(~1.15×).
  • 긴 컨텍스트 환경(4k에서 1024k 토큰)에서 GQE는 GQA 베이스라인 대비 1.67배에서 1.80배의 속도 향상을 달enc합니다. 이는 불필요한 쿼리 계산을 건너뛰기 때문이며, 어텐션 비용이 이차적으로 증가함에 따라 시퀀스 길이가 길어질수록 절감 효과가 커집니다.

의의 및 주장

본 논문은 GQE가 모델 품질을 희생하지 않으면서 긴 컨텍스트 트랜스포머의 셀프 어텐션 계산 비용을 줄일 수 있는 실행 가능한 경로를 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 효율성: 테스트된 구성에서 쿼리 헤드 계산을 약 절반으로 줄이면서도 GQA의 성능과 동등한 수준을 유지합니다.
• 확장성: 속도 향상은 컨텍스트 길이가 증가함에 따라 더욱 가치 있어지며, 이는 어텐션의 이차적 스케일링 병목 현상을 해결합니다.
• 설계 통찰: 본 연구는 MoE를 어텐션으로 전이하기 위해서는 나이브한 희소 어텐션 시도에서 흔히 보이는 정확도 저하를 피하기 위해 세심한 아키텍처 설계(특히 공유 헤드와 재정규화된 라우팅 신호)가 필요함을 강조합니다.

저자들은 결과가 250M 규모에서의 결과임을 언급하며 더 큰 규모에서의 검증이 필요하다고 겸손하게 밝히며, 이를 명확한 개선이라기보다 베이스라인과의 "일치"로 프레임화하여, 동일한 성능을 위해 계산량을 줄이는 성과를 강조했습니다.