기술 요약: Pro-KLShampoo

문제 정의

현대 대규모 언어 모델 (LLM) 사전 학습은 기울기의 행렬 구조를 활용하는 최적화기에 크게 의존합니다. 두 가지 뚜렷한 최전선이 등장했습니다:

1. 명시적 크로네커 인수분해 전처리 (Explicit Kronecker-factored preconditioning): KL-Shampoo가 대표적입니다. 이는 실제 기울기 공분산과 크로네커 구조 근사치 간의 크룰백 - 라이블러 (KL) 발산을 최소화하여 전처리기를 추정합니다.
2. 기울기 모멘텀의 직교화 (Orthogonalization of gradient momentum): Muon이 대표적입니다. 이는 스펙트럼 노름 하에서 가장 가파른 하강 (steepest descent) 을 실현하기 위해 모멘텀을 직교화합니다.

이러한 접근법들은 일반적으로 고립되어 개발됩니다. KL-Shampoo 는 KL 최소화를 통해 강력한 이론적 기반을 제공하지만, 완전한 크로네커 인수를 유지함으로써 높은 계산 및 메모리 비용을 초래합니다. 반면, Muon 과 같은 직교화 방법은 효율적이지만 KL 추정 관점 없이 다른 기하학적 프레임워크 (스펙트럼 노름 기반 가장 가파른 하강) 에서 작동합니다. 본 논문은 KL-Shampoo 의 통계적 효율성과 직교화 방법의 계산적 효율성을 모두 달성하기 위해 이러한 접근법들을 통합할 필요성을 다룹니다.

방법론: Pro-KLShampoo

저자들은 Pro-KLShampoo(직교화를 통해 백색화를 회복한 투영 KL-Shampoo) 를 제안합니다. 핵심 혁신은 KL-Shampoo 전처리기의 구조적 관찰에서 비롯됩니다: 그 고유값 스펙트럼은 "스파이크 - 플랫 (spike-and-flat)" 형태를 보입니다. 이는 소수의 지배적인 고유값 ("스파이크") 을 followed by 대략적으로 균일한 꼬리 ("플랫") 로 구성됩니다. 이 구조는 계층과 학습 단계 전반에 걸쳐 경험적으로 관찰되며, 랭크-$\rho$ 신호 + 노이즈 기울기 모델 하에서 이론적으로 정확합니다.

1. 구조적 제한

Pro-KLShampoo 는 KL-Shampoo 의 크로네커 인수 중 하나 (특히 $m \le n$인 직사각형 계층의 경우 오른쪽 인수 $R$) 를 스파이크 - 플랫 구조와 정렬된 매개변수 계열로 제한합니다:

• 투영 부분공간 ($r$차원): 전처리기는 $U \in \text{St}(n, r)$로 생성된 추적된 부분공간에서 완전한 스펙트럼 구조를 유지합니다.
• 보완 부분공간 ($n-r$차원): 나머지 방향은 단일 공유 고유값 $\mu_\perp$로 축소됩니다.

이 제한은 오른쪽 인수 저장 공간을 $O(n^2)$에서 $O(nr)$로 줄이고, QR 분해 비용을 $O(n^3)$에서 $O(r^3)$로 감소시킵니다.

2. 직교화를 통한 회복

스칼라 꼬리에 대한 단순한 제한은 꼬리가 완벽하게 균일하지 않을 경우 정보를 잃게 됩니다. 보완 부분공간에서 완전한 KL-Shampoo 전처리기의 대수적 형태를 회복하기 위해, Pro-KLShampoo 는 보완 기울기 성분에 직교화(극 분해) 를 적용합니다.

업데이트 규칙은 다음과 같습니다:
$$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$$
여기서:

• 첫 번째 항은 추적된 부분공간 내의 완전한 KL-Shampoo 업데이트를 나타냅니다.
• 두 번째 항은 보완 부분에 직교화를 적용합니다.
• 저자들은 식 (12) 의 항등식을 증명하여, 이 직교화가 보완 부분에서 완전한 KL-Shampoo 의 방향별 백색화 (whitening) 의 대수적 형태를 회복함을 보여줍니다. 이는 Muon 스타일의 직교화와 KL 추정 프레임워크를 효과적으로 연결합니다.

3. 부분공간 추적

최적 부분공간 $U$는 행렬 $\Phi_L = E[G^\top L^{-1} G]$의 고유공간으로 식별됩니다. 스파이크 - 플랫 가설 하에서, $\Phi_L$의 상위-$r$고유공간이 최적의 선택입니다. 알고리즘은 $\Phi_L$의 실행 추정치에 대한 단일 단계의 거듭제곱 반복 (QR 분해를 통해) 을 사용하여 이 부분공간을 추적합니다.

주요 기여

1. 구조적 관찰 및 매개변수 제한: 논문은 KL-Shampoo 전처리기의 "스파이크 - 플랫" 고유값 구조를 식별합니다 (저랭크 기울기 모델 하에서 경험적으로 강력하고 이론적으로 정확함) 그리고 이를 KL 목적 함수를 매개변수 계열로 제한함으로써 활용합니다.
2. 직교화를 통한 대수적 회복: 보완 부분공간에 직교화를 적용하면 완전한 KL-Shampoo 전처리기의 대수적 형태가 회복됨을 보여줍니다. 이는 Muon 스타일 직교화와 KL 추정 간의 이론적 다리를 제공합니다.
3. 수렴 보장: 저자들은 연산자 노름 매끄러움 하에서 이상화된 알고리즘에 대해 $O(1/\sqrt{T})$ 속도의 비볼록 수렴 보장을 제공합니다.
4. 경험적 성능: GPT-2(124M/350M) 및 LLaMA(134M/450M) 모델에 대한 광범위한 평가에서 Pro-KLShampoo 는 모든 테스트된 부분공간 랭크 ($r \in \{32, 64, 128\}$) 에서 검증 손실, GPU 당 피크 메모리, 그리고 일치된 손실 수준 도달에 필요한 월클락 시간 측면에서 KL-Shampoo 를 일관되게 능가했습니다.

실험 결과

저자들은 네 가지 사전 학습 구성에서 AdamW, Muon, KL-Shampoo 에 대해 Pro-KLShampoo 를 평가했습니다.

• 검증 손실: Pro-KLShampoo 는 모든 모델 규모에서 테스트된 모든 랭크에서 가장 낮은 검증 손실을 달성했습니다.
• 메모리 효율성: GPU 당 피크 메모리 사용량은 KL-Shampoo 보다 낮았으며 Muon 과 유사했습니다.
• 월클락 시간: Pro-KLShampoo 는 KL-Shampoo 보다 특정 검증 손실 수준에 더 빠르게 도달했습니다.
  • GPT-2에서는 주로 단계당 시간 감소 ($O(n^3)$ QR 을 $O(r^3)$로 대체) 에 의해 속도 향상이 주도되었습니다.
  • LLaMA에서는 단계당 효율성 외에도 더 빠른 수렴 (더 낮은 학습 손실 곡선) 에 의해 속도 향상이 주도되었습니다.
• 애블레이션 연구:
  • 직교화를 제거한 경우 (Smok-Hop 변형), 특히 대규모 (LLaMA 450M) 에서 성능 저하가 발생하여 방향별 백색화를 회복하기 위한 직교화의 필요성을 확인했습니다.
  • 부분공간 전용 또는 보완 전용 변형 중 어느 것도 전체 Pro-KLShampoo 의 성능을 달성하지 못했으며, 이는 두 구성 요소가 모두 필수적임을 나타냅니다.

중요성 및 주장

이 논문은 Pro-KLShampoo 가 이전에 고립되어 있던 두 가지 최적화 최전선을 성공적으로 통합했다고 주장합니다. 기울기 공분산의 내재적 "스파이크 - 플랫" 구조를 활용함으로써, KL-Shampoo 의 통계적 효율성을 달성하면서도 직교화를 활용하여 계산적 효율성을 유지합니다.

저자들은 스칼라 꼬리 업데이트를 직교화로 대체하는 것이 휴리스틱이 아니라 완전한 KL-Shampoo 형태의 대수적 회복이라고 강조합니다. 이를 통해 이 방법은 완전한 크로네커 인수 유지와 관련된 계산 오버헤드를 크게 줄이면서도 KL 기반 전처리의 이점 (예: 스텝 크기 견고성) 을 유지할 수 있습니다. 이 연구는 저차원 스펙트럼 구조가 심층 네트워크 최적화 전반에 걸쳐 만연하며, LLM 사전 학습을 위한 더 효율적인 최적화기 설계에 효과적으로 활용될 수 있음을 시사합니다.

논문은 제한 사항에 대해 겸손하게 서술하며, 랭크 $r$은 경험적으로 설정된다는 점과, 뉴턴 - 슈ultz 반복 및 네스테로프 모멘텀을 포함하는 실제 변형에 대한 혼합 가중치 $\alpha_{kl}$의 계층별 적응적 보정 및 수렴 분석에 대한 추가 연구가 필요하다고 지적합니다.