기술 요약: GQA-µP: 그룹화 쿼리 어텐션을 위한 최대 파라미터화 업데이트

문제 제기

모델 아키텍처 간 하이퍼파라미터 전이, 특히 최대 업데이트 파라미터화 (µP) 는 작은 프록시 모델을 스윕하여 대규모 모델의 최적 하이퍼파라미터를 결정할 수 있게 합니다. µP 는 모델 너비 간 학습률 전이에 성공적으로 적용되어 왔으나, 새로운 아키텍처와 가중치 감쇠 (weight decay) 및 깊이 (depth) 와 같은 다른 하이퍼파라미터에 적용하는 것은 여전히 도전적입니다.

**그룹화 쿼리 어텐션 (Grouped-Query Attention, GQA)**에 대한 특정 격차가 존재합니다. GQA 는 쿼리 헤드 그룹 간 키 (key) 와 값 (value) 헤드를 공유하여 메모리 사용량을 줄이기 위해 현대의 대규모 언어 모델 (LLM) 에서 널리 사용되는 메커니즘입니다. 표준 µP 를 GQA 에 적용하려는 이전 시도들은 중요한 실패를 드러냈습니다. 구현이 표준 "좌표 점검 (coordinate checks)" (활성화 노름 유효성 검증) 을 통과했음에도 불구하고, 서로 다른 KV 헤드 수 간 학습률을 효과적으로 전이하지 못했습니다. 또한, 기존 µP 프레임워크는 "레이지 러닝 (lazy-learning)" 가정에 의존하지 않고 가중치 감쇠와 깊이에 대한 원칙 있는 스케일링을 도출하는 데 어려움을 겪었습니다. 핵심 이론적 어려움은 GQA 가중치 행렬이 종종 랭크 퇴화 (rank-degenerate, 저랭크) 되어 있어, 표준 스펙트럼 노름이 훈련 중 입력의 기대 변형을 정확하게 포착하지 못하기 때문에 발생합니다.

방법론

저자들은 Yang 등 (2023a) 의 스펙트럼 특징 학습 관점을 바탕으로 새로운 스케일링 법칙을 도출합니다. 방법론은 세 가지 주요 이론적 및 분석적 진전을 포함합니다:

1. 스펙트럼 노름을 통한 특징 학습의 재정의:
   이 논문은 가중치 행렬에 대한 스펙트럼 노름 조건을 휴리스틱에서 특징 학습의 공식적 정의로 격상시킵니다. 저자들은 가중치 업데이트의 스펙트럼 노름이 올바르게 스케일링되도록 보장하는 것이 활성화 노름만 분석하는 것보다 특징 학습에 대해 더 강력하고 신뢰할 수 있는 조건이라고 주장합니다. 이러한 전환은 레이지 러닝 가정에 의존하지 않고 완전-P (Complete-P) 깊이 및 가중치 감쇠 스케일링을 도출할 수 있게 합니다.

2. 기대 연산자 노름의 도입:
   GQA 행렬의 랭크 퇴화 문제를 해결하기 위해 저자들은 수정된 노름인 기대 연산자 노름 (expected operator norm) ($||A||_E$) 을 도입합니다.

   • 표준 스펙트럼 노름 ($||A||$) 은 입력 벡터의 최대 변형을 측정합니다. 풀랭크 행렬의 경우 이는 점근적으로 달성됩니다.
   • 랭크 퇴화 행렬 (GQA 와 같은) 의 경우, 최대 변형을 유발하는 벡터들은 훈련 중 (입력이 i.i.d. 이므로) 확률 0 으로 발생합니다. 결과적으로 스펙트럼 노름은 실제 변형을 과대평가합니다.
   • 기대 연산자 노름은 입력 분포에 대한 출력 노름과 입력 노름의 비율의 기대값으로 정의됩니다: $||A||_{E,\Omega,p} := \mathbb{E}_{x \sim \Omega} [\frac{||Ax||_p}{||x||_p}]$.
   • 이 노름은 훈련 중 실제로 마주치는 변형을 정확하게 반영하여 GQA 에 대한 스케일링 법칙의 유효성을 복원합니다.

3. GQA-µP 스케일링 도출:
   기대 연산자 노름을 사용하여 저자들은 GQA 레이어에 대한 구체적인 스케일링 인자를 도출합니다.

   • $r$개의 키 - 값 헤드가 반복되는 GQA 레이어의 경우, 학습률 ($\eta$) 과 가중치 감쇠 ($\lambda$) 는 표준 어텐션과 다르게 스케일링되어야 합니다.
   • 도출된 학습률 스케일링은 $\eta = \Theta(\frac{1+\sqrt{r}}{n})$이며, 여기서 $n$은 임베딩 차원입니다.
   • 가중치 감쇠 스케일링은 정규화와 그래디언트 업데이트 간의 균형을 유지하기 위해 $\lambda = \Theta(\frac{2n}{1+\sqrt{r}})$로 조정됩니다.
   • 이러한 도출은 논문의 표 1 에 요약되어 있으며, GQA 를 가진 트랜스포머에 대한 통합된 파라미터화를 제공합니다.

4. 깊이 스케일링 (Complete-P):
   이 프레임워크는 잔차 네트워크에 적용되어 깊이 스케일링을 도출합니다. 잔차 블록의 합성에 스펙트럼 조건을 적용함으로써, 저자들은 신호 폭발이나 소실을 방지하기 위해 잔차 분기의 스케일링 인자 $\beta$가 $\Theta(L^{-1})$ (여기서 $L$은 깊이) 여야 함을 보여줍니다. 이는 "레이지 러닝 없음" 가정을 요구하지 않고 Dey 등 (2025) 의 동일한 결과를 복원합니다.

주요 기여

1. GQA 에 대한 최초의 µP 도출: 이 논문은 저자들의 지식 범위 내에서 그룹화 쿼리 어텐션에 대한 µP 스케일링의 원칙 있는 첫 번째 도출을 제공하며, 표준 µP 구현이 학습률 전이에 실패했던 불일치를 해결합니다.
2. 스펙트럼 µP 의 이론적 확장: 이 작업은 Yang 등 (2023a) 의 스펙트럼 µP 이론을 확장하여 스펙트럼 노름 조건을 특징 학습의 정의로 격상시키고, 랭크 퇴화 가중치를 처리하기 위해 기대 연산자 노름을 도입합니다. 이를 통해 가중치 감쇠, 재귀 블록 및 GQA 에 대한 스케일링 도출이 가능해집니다.
3. 전이의 실증적 검증: 저자들은 도출된 스케일링이 서로 다른 GQA 반복 횟수 ($r$) 간 학습률과 모델 크기 간 가중치 감쇠 전이를 가능하게 함을 입증합니다. 그들은 바닐라 Adam-µP 구현이 GQA 에 대해 스펙트럼 좌표 점검을 통과하지 못하는 반면, 제안된 스케일링은 이를 통과함을 보여줍니다.

실험 결과

저자들은 광범위한 실험을 통해 이론적 도출을 검증합니다:

• 학습률 전이: KV 헤드 수 (1 에서 12 까지) 를 변화시킨 실험은 표준 파라미터화 및 바닐라 Adam-µP 가 일관된 최적 학습률로 수렴하지 못함을 보여줍니다. 반면, 제안된 GQA-µP 스케일링은 최적 학습률 영역을 크게 좁혀 효과적인 전이를 가능하게 합니다.
• 좌표 점검: 활성화 업데이트 노름 ($||\Delta h||$) 분석은 바닐라 구현이 정확하다고 시사했으나, 가중치 행렬 스펙트럼 노름 ($||\Delta W||$) 분석은 KV 헤드 수에 대한 비선형 의존성을 드러내 전이 실패를 설명했습니다. 제안된 GQA-µP 스케일링은 이 의존성을 제거합니다.
• 가중치 감쇠 및 $\tau_{\text{epoch}}$ 전이: 이 논문은 올바른 스케일링을 사용하면 가중치 감쇠와 Wang & Aitchison (2024) 이 도입한 훈련 시간 상수 $\tau_{\text{epoch}}$가 모델 크기 간 전이 가능함을 보여줍니다. 제안된 구현의 최적 하이퍼파라미터 분산은 베이스라인에 비해 현저히 낮습니다.
• GQA 의 노이즈: 저자들은 학습률 전이가 가능하지만, GQA 훈련은 KV 헤드 수가 감소함에 따라 더 두드러지는 고유한 노이즈를 도입한다고 관찰합니다. 그들은 트레이드오프를 지적합니다: 도출된 스케일링을 통한 특징 학습을 달성하면 $r$이 감소함에 따라 더 노이즈가 많은 역학이 발생하지만, 훈련을 안정화하기 위해 분산을 제약하면 최적 학습률이 이동합니다.

중요성 및 주장

이 논문은 현대의 생산급 LLM 아키텍처에 µP 를 적용할 수 있는 중요한 격차를 해소했다고 주장합니다. 저랭크 행렬에 대한 표준 스펙트럼 노름의 한계를 식별하고 기대 연산자 노름을 제안함으로써, 저자들은 GQA 와 같은 복잡한 아키텍처에 대한 µP 스케일링 도출을 위한 견고한 이론적 기반을 제공합니다.

이 작업의 중요성은 다음과 같습니다:

• 실용적 유용성: 이는 실무자가 작은 프록시 모델을 사용하여 대규모 GQA 기반 모델 (예: LLaMA 3, Mistral 7B) 의 최적 하이퍼파라미터를 결정할 수 있게 하여, 하이퍼파라미터 튜닝의 계산 비용을 줄입니다.
• 이론적 엄밀성: 이는 활성화 기반 점검을 넘어 가중치 행렬 기반 스펙트럼 조건으로 특징 학습의 정의를 정제하여, 구현 유효성 검증에 더 효과적임을 입증합니다.
• 광범위한 적용 가능성: 이 프레임워크는 가중치 감쇠 및 깊이 (Complete-P) 에 대한 기존 스케일링을 성공적으로 복원하면서 이를 GQA 로 확장하여, 향후 아키텍처 혁신을 위한 통합된 접근법을 시사합니다.

저자들은 점근적 분석이 타당하지만, 반복 매개변수 $r$의 실용적 범위는 유한하며 일반적으로 작음 (예: $r \le 16$) 을 겸손하게 지적합니다. $r$에 대한 스케일링의 주된 목표는 이론적 무한 극한을 모델링하는 것이 아니라, 반복 횟수를 변경할 때 학습률 드리프트를 방지하는 실증적 목적입니다. 또한 그들은 GQA 훈련이 본질적으로 노이즈가 많음을 강조하며, 매우 다른 GQA 구성 간 학습률 전이 시 주의가 필요함을 시사합니다.