OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

OptMuon: 확률적 최적화를 위한 폐쇄 루프 직교화 모멘텀 방법론의 기술 요약 (Zero-Noise Optimality 포함)

1. 문제 정의

본 논문은 행렬 값 비볼록 목적 함수(matrix-valued nonconvex objectives)에 대한 확률적 1차 최적화 문제를 다룹니다:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
목표는 $T$번의 확률적 그래디언트 호출 내에서 기대 프로베니우스 노름(expected Frobenius norm)의 그래디언트 $\mathbb{E}[\|\nabla f(X_{out})\|_F]$를 최소화하는 무작위 반복값(random iterate) $X_{out}$을 찾는 것입니다.

Muon 스타일의 옵티마이저들은 모멘텀 행렬을 근사적으로 직교하는 방향으로 변환(극 분해를 통해)함으로써 경험적인 안정성을 입증해 왔으나, 기존 방식들은 일반적으로 이러한 방향을 개방형 루프(open-loop) 또는 고정된 크기 규칙과 결합합니다. 이러한 규칙은 대개 규정된 매끄러움 상수($L$)나 분산 수준($\sigma^2$)에 의존하므로, 하이퍼파라미터 튜닝에 민감하며 관찰된 최적화 궤적에 기반하여 업데이트 크기를 명시적으로 교정하지 못합니다. 반면, STORM이나 AdaGrad-Norm과 같은 적응형 방법들은 "폐쇄 루프(closed-loop)" 스칼라 적응 및 노이즈 적응성을 제공하지만, 일반적으로 직교화된 모멘텀 방향이 아닌 가공되지 않은(raw) 그래디언트 상에서 작동합니다.

핵심 과제는 폐쇄 루프 스칼라 적응이 기하학적 안정성을 유지하면서도 노이즈 적응성 및 제로 노이즈 최적성(zero-noise optimality)을 달성하기 위해 Muon 스타일의 직교화된 모멘텀과 성공적으로 결합될 수 있는지 여부입니다.

2. 방법론: OptMuon

저자들은 업데이트 방향과 업데이트 크기를 분리하는 프레임워크인 OptMuon을 제안합니다.

• 방향 (직교화): 업데이트 방향은 모멘텀 행렬 $M_t$의 극 인자(polar factor)인 $\text{Orth}(M_t) = U W^\top$입니다 (여기서 $M_t = U \Sigma W^\top$는 SVD임). 이는 방향이 기하학적으로 정규화됨을 보장하며, 방향(orientation)과 특이값 스케일링(singular-value scaling)을 분리합니다.
• 크기 (폐쇄 루프 적응): 스텝 사이즈는 고정된 문제 상수가 아닌, 관찰된 궤적(그래디언트 및 모멘텀 이력)으로부터 유도된 스칼라 계수 $\gamma_t$에 의해 결정됩니다.
  • 계수 스케줄: 크기 제어의 핵심은 실행 최대값 보정 AdaGrad-Norm(running-maximum corrected AdaGrad-Norm) 스케줄입니다. 이는 노이즈가 심한 구간에서 스텝을 억제하기 위해 분모에 확률적 그래디언트 에너지를 누적시키지만, 분자에 $\max g_i^2$ 항을 포함하여 고립된 그래디언트 스파이크로 인해 계수가 과도하게 붕괴되는 것을 방지합니다.
  • 리프시츠-프리(Lipschitz-Free): 이 알고리즘은 파라미터 선택을 위한 매끄러움 상수 $L$, 분산 $\sigma^2$, 또는 유계 그래디언트 상수 $\hat{G}$에 대한 지식이 필요하지 않습니다 (비록 이론적 상수에는 이들이 등장할 수 있으나).

두 가지 변형:

1. OptMuon-A (평균 매끄러움): 평균 매끄러움 가정(average smoothness assumption)을 위해 설계되었습니다. 매 반복마다 하나의 확률적 그래디언트를 사용하며 표준 모멘텀 누적($M_t = G_t + (1-\alpha_t)M_{t-1}$)을 수행합니다. 스케줄의 지수는 $q=1/2$입니다.
2. OptMuon-I (개별 매로움): 개별 매끄러움 가정(individual smoothness assumption)을 위해 설계되었습니다. 분산을 줄이기 위해 동일한 샘플에 대해 두 번의 그래디언트 평가를 사용하는 STORM 유형의 재귀적 모멘텀 추정기($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)) + G_t$)를 채택합니다. 지수는 $q=2/3$입니다.

두 변형 모두 다음과 같이 파라미터를 업데이트합니다:
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. 주요 기여

본 논문은 다음과 같은 구체적인 기여를 주장합니다:

• 폐쇄 루프 모멘텀 직교화: 행렬 기하학(극 인자를 통해)과 크기 교정(자기 정규화 스칼라 스케줄을 통해)을 깔끔하게 분리하는 프레임워크를 구축했습니다.
• 실행 최대값 보정 스케줄: AdaGrad-Norm의 자기 정규화 억제 기능을 유지하면서도, 단일 그래디언트 스파이크로 인한 계수 붕괴 위험을 완화하는 새로운 계수 설계를 제시했습니다.
• 노이즈 적응형 수렴 속도: 분산 $\sigma$가 감소함에 따라 수렴 바운드의 확률적 항이 자동으로 개선됨을 보여주는 이론적 보증을 제공합니다.
• 제로 노이즈 최적성(Zero-Noise Optimality): $\sigma = 0$일 때, 바운드가 수동 재조정 없이도 거의 최적인 결정론적 속도 $\tilde{O}(T^{-1/2})$로 자동 축소됨을 입증했습니다.

4. 이론적 결과

하한 존재성, 유계 분산을 가진 비편향 확률적 그래디언트, 매끄러움(평균 또는 개별), 그리고 거의 확실하게 유계인 확률적 그래디언트 가정 하에, 본 논문은 다음을 증명합니다:

• OptMuon-A (평균 매끄러움): 다음의 수렴 속도를 달성합니다:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• Opt-Muon-I (개별 매끄러움): 더 정교한 확률적 속도를 달성합니다:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

결과의 의의:

• 제로 노이즈 회복: 노이즈가 없는 환경($\sigma = 0$)에서 두 바운드 모두 $\tilde{O}(T^{-1/2})$로 축소되며, 이는 최적의 결정론적 1차 속도와 일치합니다. 이는 고정된 스텝 사이즈를 요구하거나 $\tilde{O}(T^{-1/4})$를 보이는 기존 Muon 분석들과 대조됩니다.
• 노이즈 적응성: 확률적 항($\sigma^{1/2}T^{-1//4}$ 및 $\sigma^{1/3}T^{-1/3}$)은 적응적입니다. 즉, 노이즈 수준이 낮아짐에 따라 자동으로 타이트해지며, 이는 고정된 스텝 사이즈를 사용하는 방법들과 다릅니다.
• 비교: 본 논문은 OptMuon-A가 저노이즈 환경에서 최근의 분산 감소 Muon 변형들(예: Chang et al., 2025)이 확립한 $\tilde{O}(T^{-1/4})$ 속도를 개선하면서도 제로 노이즈 최적성을 유지한다고 언급합니다.

5. 의의 및 한계점

의의:
본 논문은 폐쇄 루프 스칼라 적응이 기하학적 안정성을 희생하지 않고도 Muon 스타일의 직교화된 모멘텀과 결합될 수 있음을 입증했습니다. 이는 이전에 스칼라 적응형 스텝 사이즈 규칙에서 발견되었던 "제로 노이즈 적응성" 개념을 극 인자 모멘텀 업데이트로 확장합니다. 이를 통해 수동 개입 없이 노이즈 수준에 자동으로 적응하고 최적의 결정론적 성능을 회복하는 강건한 최적화를 가능하게 합니다.

한계점 (저자들이 명시한 내용):

• 정확한 직교화: 이론은 정확한 극 인자 계산을 가정합니다. Newton-Schulz 반복법을 통해 이를 근사할 수 있지만, 현재 분석에는 부정확한 직교화로 인해 발생하는 오차에 대한 정량화가 포함되어 있지 않습니다.
• 유계 그래디언트: 분석은 거의 확실하게 유계인 확률적 그래디언드 조건(가정 4)에 의존합니다. 본 논문은 그래디언트 유계가 필요 없는 이론을 주장하지 않습니다.
• 경험적 검증: 본 논문은 이론적 수렴 보증에 집중하며, 대규모 경험적 평가나 복합/제약 조건 목적 함수로의 확장은 포함하지 않았으며, 이는 향후 연구 방향으로 언급되었습니다.