극한 효율을 위한 행렬‑그라디언트 최적화: PolarGrad와 통합 전처리 관점

초록

본 논문은 딥러닝 파라미터를 벡터와 행렬로 구분해 전처리(preconditioning) 전략을 재정립한다. 기존 Adam 계열은 헤시안(곡률) 조건수를 낮추는 ‘벡터 전처리’에 초점을 맞추는 반면, Shampoo·Muon과 같은 구조‑인식 옵티마이저는 그라디언트 자체의 조건수를 개선하는 ‘행렬 전처리’를 수행한다. 이러한 시각을 바탕으로, 그라디언트의 극분해(polar decomposition)와 핵노름 스케일링을 이용한 새로운 클래스인 PolarGrad를 제안한다. PolarGrad는 Muon을 특수 경우로 포함하며, QDWH·ZOLO‑PD와 같은 고효율 극분해 알고리즘을 활용한다. 다양한 행렬 최적화 문제와 대규모 언어 모델 사전학습 실험에서 PolarGrad가 Adam·Muon보다 빠른 수렴과 안정성을 보인다.

상세 분석

논문은 먼저 딥러닝 모델 파라미터가 스칼라·벡터·행렬·텐서 등 서로 다른 대수적 구조를 가진다는 사실을 강조한다. 기존의 Adam·AdamW와 같은 적응형 옵티마이저는 모든 파라미터를 하나의 긴 벡터로 평탄화하고, 대각선 형태의 전처리 행렬을 통해 헤시안의 역제곱근을 근사한다. 이는 ‘곡률 전처리(curvature preconditioning)’라 부르며, 주로 조건수가 큰 헤시안을 완화해 수렴 속도를 높이는 데 목적이 있다. 그러나 행렬 파라미터에 대해선 그라디언트 자체가 높은 조건수를 가질 수 있다. 예를 들어, Transformer의 attention weight는 크기가 수천인 행렬이며, 그라디언트의 특이값 분포가 매우 비대칭적이다. 이 경우 헤시안 전처리만으로는 충분히 조건을 개선하지 못한다.

이에 저자들은 ‘그라디언트 전처리(gradient preconditioning)’라는 개념을 도입한다. 행렬 업데이트 방향 X에 대해 κ₂(X)=σ_max/σ_min가 클 경우, X는 ill‑conditioned라 하여 수렴이 느려진다. 행렬을 반직교(orthogonal) 혹은 반직교 투영(semi‑orthogonal projection)으로 변환하면 κ₂가 1이 되어 최적의 조건을 달성한다. Muon은 이러한 아이디어를 실제 구현한 사례로, 그라디언트를 극분해 A=UP, 여기서 U는 반직교, P는 양의 대칭 행렬이며, 업데이트에 U만을 사용한다. 논문은 Muon이 사실상 ‘그라디언트 전처리’를 수행함을 수학적으로 증명하고, Adam이 행렬 파라미터에 적용될 때 발생하는 학습 불안정성(특히 초기 단계에서의 발산)과 학습률 워밍업 필요성을 이 관점에서 설명한다. Adam은 헤시안 전처리와 달리 그라디언트 조건을 개선하지 못해, 초기 그라디언트가 급격히 변동할 때 큰 스텝을 취하게 되고, 이는 손실 발산으로 이어진다.

PolarGrad는 이러한 통합 관점을 바탕으로 설계된 새로운 옵티마이저 군이다. 핵심 아이디어는 그라디언트 G∈ℝ^{m×n}에 대해 극분해 G=U·P를 수행하고, 업데이트 방향을
ΔW = -η· (‖P‖* )^{-1}·U·M
여기서 ‖P‖는 핵노름(특이값의 합)이며, M은 모멘텀(예: Adam‑style 1차 모멘트)이다. 핵노름 스케일링은 전통적인 Adam의 2차 모멘트(분산) 추정과 유사한 역할을 하면서, 행렬의 전체 스펙트럼 정보를 반영한다. 특수 경우 η·‖P‖_^{-1}=1이면 PolarGrad는 Muon과 동일해진다.

알고리즘 구현 측면에서 저자들은 기존 Muon에서 사용된 뉴턴‑슐즈 반복이 수렴 속도와 수치 안정성에서 취약함을 지적하고, QDWH(Quasi‑Diagonal‑Weighted‑Halley)와 ZOLO‑PD(Zero‑Order‑Lipschitz‑Optimized Polar Decomposition)와 같은 최신 극분해 방법을 도입한다. 이들 알고리즘은 반복 횟수가 적고, 파라미터 차원에 거의 선형적으로 스케일되며, GPU에서 효율적인 배치 연산이 가능하도록 설계되었다. 논문은 각 알고리즘의 수렴 이론을 정리하고, 실제 학습에서의 연산 비용을 Adam·Muon 대비 1020% 정도 증가시키면서도 전체 학습 단계 수를 3040% 감소시키는 결과를 보고한다.

실험에서는 (1) 저차원 랜덤 행렬 회귀, (2) 저‑랭크 행렬 인수분해, (3) 대규모 언어 모델(예: 1.3B 파라미터 GPT‑like) 사전학습, (4) MoE 기반 15B 모델 훈련을 포함한 네 가지 벤치마크를 수행했다. 모든 실험에서 PolarGrad는 최종 퍼플렉시티·손실에서 Adam보다 평균 23% 개선, Muon보다 0.51% 개선을 보였으며, 특히 초기 10% 학습 단계에서 손실 감소 속도가 크게 앞섰다. 또한 학습률 워밍업 없이도 안정적으로 수렴했으며, 워밍업을 적용하면 추가적인 미세 개선이 관찰되었다. Ablation study에서는 핵노름 스케일링을 제거하거나, 극분해 대신 SVD 기반 직교화만 사용할 경우 성능이 현저히 떨어짐을 확인했다.

결론적으로, 논문은 딥러닝 최적화에서 “벡터‑행렬 전처리의 차별화”라는 새로운 패러다임을 제시하고, PolarGrad라는 실용적이며 이론적으로 정당화된 옵티마이저를 통해 기존 Adam·Muon의 한계를 극복한다는 점을 입증한다. 향후 연구는 Tensor‑level 전처리, 자동 전처리 선택 메커니즘, 그리고 더 큰 규모의 멀티모달 모델에 대한 적용 가능성을 탐색할 여지를 남긴다.