Rethinking Bregman Divergences in Kronecker-Factored Optimizers

기술 요약: 크로네커 인수 분해 최적화 도구에서의 브레그만 다이버전스 재고

문제 정의
Shampoo 스타일의 최적화 도구들은 계산 복잡도를 줄이기 위해 그래디언트 공분산 행렬 $C = E[\text{vec}(G)\text{vec}(G)^\top]$를 크로네커 인수 분해 구조인 $L \otimes R$을 사용하여 근사한다. 최근 연구(Lin et al., 2026)는 $C$와 $L \otimes R$ 사이의 브레그만 행렬 다이버전스(Bregman matrix divergences)를 최소화하는 문제로 프레이밍함으로써 다양한 Shampoo 변형들을 통합하였다. 그러나 일반적인 조밀한(dense) 공분산 행렬의 경우, 정확한 크로네커 인수 분해가 불가능하므로 근사 오차는 피할 수 없다. 서로 다른 브레그만 다이버전스가 이 오차를 공분산 스펙트럼 전체에 어떻게 배분하는지, 그리고 결과적으로 생성된 프리컨디셔너(preconditioner)가 스펙트럼의 모든 부분을 동등하게 신뢰하는지는 불분명하다. 또한, 그래디언트 공분산은 실제 헤시안(Hessian)과 크게 다를 수 있으므로, 모든 스펙트럼 성분이 프리컨디셔닝에 적합한 것은 아니다. 본 논문이 다루는 핵심 질문은 다음과 같다: 정확한 크로네커 매칭이 불가능할 때, 구조화된 최적화 도구는 공분산 스펙트럼의 어느 부분을 근사하고 신뢰해야 하는가?

방법론 및 분석
저자들은 두 가지 주요 분석적 관점, 즉 브레그만 근사의 스펙트럼 기하학과 공분산 고유 공간(eigenspaces)의 경험적 신뢰성을 통해 이 문제를 접근한다.

1. 브레그만 근사의 스펙트럼 기하학:
   본 논문은 Frobenius, von Neumann, LogDet 다이버전스가 근사 오차를 어떻게 다르게 처벌하는지 규명한다. 다이버전스를 주변부 스펙트럼 항과 교차 스펙트럼 결합 항 $g_F(\lambda_i, \omega_j)$로 분해함으로써, 저자들은 다음을 보여준다:

   • Frobenius 다이버전스: 결합 항 $-\lambda_i \omega_j$는 쌍선형 크기 패널티를 부과하여 상위(큰 고윳값) 스펙트럼을 강력하게 강조한다.
   • von Neumann 다이버전스: $-\lambda_i \log \omega_j$ 항은 상위 스펙트럼에 대한 선호도를 유지하면서도 이를 로그 단위로 완화한다.
   • LogDet 다이버전스: $\lambda_i / \omega_j$ 항은 상대적 비율 왜곡을 생성하여, 근사가 꼬리(tail) 스펙트럼(작은 고윳값)에 매우 민감하게 만든다.

   정리 2(Theorem 2)는 이러한 최소화 문제의 정체점(stationary points)이 가공되지 않은 잔차 $(S^* - C)$가 아니라, 다이버전스 가중 잔차 $\nabla^2 F(S^*)[S^* - C]$에 의해 결정됨을 확립한다. 결과적으로, 생성된 크로네커 인수는 선택된 다이버전스의 특정 스펙트럼 선호도(예: LogDet 인수는 작은 고윳값에 더 민감하고, Frobenius 인수는 지배적인 부분 공간에 집중함)와 일치하게 된다.

2. 공분산 프리컨디셔닝의 스펙트럼 신뢰성:
   저자들은 SST-2 데이터셋으로 학습된 트랜스포머 모델에서 그래디언트 공분산과 헤시안 고유 공간 사이의 정렬(alignment)을 경험적으로 조사한다. 결과(표 3)에 따르면, 상위 공분산 고유 공간은 하위 고윳값 성분보다 헤시안과 훨씬 더 잘 정렬되어 있다. 꼬리 스펙트럼은 프리컨디셔닝에 있어 노이즈가 많고 신뢰도가 낮은 것으로 관찰되었으며, 이는 기존 연구(Kunstner et al., 2019; Thomas et al., 2020)와 일치한다.

3. 제안된 솔루션: 부분 공간 인식 크로네커 최적화 도구 (BregTop):
   다이버전스에 의해 유도된 스펙트럼 선호도와 실제 스펙트럼의 신뢰성 사이의 불일치에서 영감을 얻어, 저자들은 BregTop을 제안한다. 이 최적화 도구는 두 가지 관찰의 강점을 결합한다:

   • 상위 민감형 다이버전스(예: von Neumann 또는 Frobenius)를 사용하여 크로네커 인수를 선택한다.
   • 고윳값 기반 프리컨디셔닝을 오직 상위 부분 공간(결합 크로네커 스펙트럼의 가장 큰 $\lceil \rho mn \rceil$ 개 항목) 내에서만 적용한다.
   • 나머지 하위 부분 공간은 잠재적으로 오도할 수 있는 곡률 정보를 사용하여 프리컨디셔닝을 시도하는 대신, 적응형 등방성 가속 상수($\chi_{\rho, q}$)를 적용하여 처리한다.

   이론적으로, 이는 프리컨디셔너를 곡률 스케일링된 상위 성분과 하위 공간을 위한 균일한 가속 성분으로 분해한다(정리 4).

주요 기여

• 스펙트럼 기하학적 특성화: 본 논문은 Frobenius, von Neumann, LogDet 다이버전스가 어떻게 서로 다른 고윳값 의존적 패널티를 유도하는지 공식적으로 규명하며, 정체된 크로네커 인수가 가공되지 않은 오차가 아닌 다이버전스 가중 잔차에 의해 지배됨을 보여준다.
• 경험적 스펙트럼 신뢰성: 공분산의 전체 스펙트럼은 프리컨디셔닝에 신뢰할 수 없지만, 상위 고유 공간은 헤시안과 강하게 정렬되어 있음을 입증하여 부분 공간 인식 접근 방식의 타당성을 제시한다.
• 부분 공간 인식 최적화 도구 설계: 저자들은 상위 스펙트럼에 대한 프리컨디셔닝과 노이즈가 많은 하위 공간에 대한 등방성 가속 사이를 적응적으로 전환하는 알고리즘인 BregTop을 도입한다.

실험 결과
저자들은 4개 층의 트랜스포머를 SST-2 전체 데이터셋으로 학습시켜 BregTop을 표준 Shampoo, KL-Shampoo, VN-Shampoo, F-Shampoo와 비교 평가하였다. 사용된 지표는 EMA 평활화된 훈련 손실(EMA-smoothed training loss)이 0.05 아래로 떨어지는 데 필요한 단계 수이다.

• BregTop-VN-v2는 목표 손실에 도달하는 데 1003 단계가 소요되어, 표준 Shampoo(1119 단계)와 베이스라인인 VN-Shampoo(1127 단계)보다 우수한 성능을 보였다.
• BregTop 변형들은 각각의 다이버전스 베이스라인(예: BregTop-VN은 VN-Shampoo보다, BregTop-F는 F-Shampoo보다)을 일관되게 개선하였으며, 이는 상위 부분 공간을 격리하여 프리컨디셔닝하는 것의 이점을 검증한다.

의의 및 주장
본 논문은 구조화된 최적화 도구의 설계가 크로네커 인수 분해에서의 불가피한 근사 오차와 공분산 스펙트럼의 가변적인 신뢰성을 고려해야 한다고 주장한다. 서로 다른 브레그만 다이버전스가 본질적으로 서로 다른 스펙트럼 영역을 우선시하며, 상위 스펙트럼만이 헤시안과 신뢰할 수 있게 정렬되어 있다는 점을 인식함으로써, 제안된 부분 공간 인식 접근 방식은 더 강력한 프리컨디셔닝 전략을 제공한다. 저자들은 이 설계가 하위 부분 공간에서 학습이 일어난다는 Song et al. (2025)의 학습 역학 관찰과 일치한다고 언급한다. 즉, 하위 공간에 등방성 가속을 적용함으로써 필수적인 그래디언트 흐름을 보존하는 동시에, 상위 부분 공간의 프리컨디셔닝을 통해 고곡률 방향에서의 노이즈 섞인 업데이트를 억제한다. 결과적으로 이 하이브리드 접근 방식은 완전한 2차 정보 없이도 단계 효율성을 향상시킨다.