기술 요약: "딥러닝에서 언제 어떤 Schatten-$p$ 노름을 사용해야 하는가?"

문제 정의

Muon과 같은 Schatten-$\infty$ 기반 옵티마이저는 특정 딥러닝 설정, 특히 짧은 훈련 과정이나 큰 배치 크기를 포함하는 환경에서 유망한 경험적 성능을 입증해 왔습니다. 그러나 이들의 일반적인 유용성에 대해서는 상충하는 관찰 결과들이 존재합니다. 이들은 특정 레짐(예: CIFAR-10 스피드런)에서는 탁월한 성능을 보이지만, 방대한 토큰 예산을 가진 대규모 훈련 시나리오에서는 Adam이나 SGD와 같은 표준적인 방법들을 일관되게 능가하지 못합니다. 나아가, 유한한 Schatten-$p$ 노름을 활용하는 최근의 변형들(예: HTMuon, Soft-Muon)은 경쟁력 있거나 더 우수한 결과를 보여주며, 반드시 Schatten-$\infty$ 노름만을 사용해야 할 필요성에 의문을 제기하고 있습니다.

본 논문이 다루는 핵심 질문은 다음과 같습니다: 어떤 Schatten-$p$ 노름을 옵티마이저가 사용해야 하며, 어떤 훈련 레짐에서 그것이 최적인가? 기존의 이론적 프레임워크, 구체적으로 SODA(Schedule-free Optimistic Dual Averaging) 프레임워크는 확률적 노이즈 하에서 $p > 2$인 $p$-노름 기하학을 분석하는 능력이 부족했습니다. 이는 Muon의 성공과 그 한계에 대한 메커니즘을 이론적으로 설명하지 못하게 만들었습니다.

방법론

본 논문은 현대적인 옵티마이저 구성 요소인 낙관적 이중 평균(optimistic dual averaging), 네스테로프 모멘텀(Nesterov momentum), 가중치 감쇠(weight decay)를 통합하는 SODA 프레임워크(Pethick et al., 2026)를 확장합니다. 저자들은 SODA의 분석을 $p \ge 2$인 $p$-균등 볼록스(p-uniformly convex) 참조 함수로 일반화하여, 벡터 $\ell_p$ 및 행렬 Schatten-$p$ 기하학을 모두 다룹니다.

주요 방법론적 구성 요소는 다음과 같습니다:

1. 알고리즘 템플릿: 분석은 참조 함수 $h(x) = \frac{1}{p}\|x - z_0\|_p^p$의 펜첼 공액(Fenchel conjugate)의 그래디언트에 의해 정의되는 쌍대 맵(dual map)을 사용하는 프라이멀-듀얼(primal-dual) 업데이트 규칙을 사용합니다.
   • $p=2$인 경우, 유클리드 기하학을 회복합니다.
   • $p \to \infty$인 경우, Muon에서 사용되는 선형 최소화 오라클(LMO, $UV^\top$)을 회복합니다.
   • 유한한 $p > 2$인 경우, HTMuon 및 Soft-Muon과 같은 방법론에 대응합니다.
2. 노이즈 가정: 분석에는 헤비 테일(heavy-tailed) 노이즈 가정을 포함합니다. 구체적으로 그래디언트가 $r$-모멘트 조건($r \in (1, 2]$)을 만족한다고 가정합니다. 이를 통해 쌍대 지수(dual exponent) $q = p/(p-1)$에 의존하는 수렴 속도를 도출할 수 있습니다.
3. 레짐 분석: 저자들은 두 가지 뚜렷한 레짐 하에서 오라클 복잡도 바운드를 도출합니다:
   • 결정론적(Deterministic): 노이즈가 무시할 수 있는 수준인 경우.
   • 확률적(Stochastic): 노이즈가 존재하며, 분산을 관리하기 위해 배치 크기가 조절되는 경우.
4. 스케일링 법칙: 이론적 바운드는 모델 행렬의 유효 차원 $d$와 총 토큰 예산 $N$을 연관시키는 현대적 딥러닝 스케일링 법칙(예: Chinchilla)에 적용됩니다.

주요 기여

1. 임의의 $p$에 대한 노이즈 강건한 가속화

저자들은 SODA가 확률적 노이즈에 강건하면서도 $p$-노м 기하학 상의 가속된 수렴 속도를 달 achieves함을 증명합니다. 이들은 다음과 같은 차원 독립적(dimension-free) 속도를 확립합니다:
$$O\left( \frac{L_p R_p^2}{n^{1+2/p}} + \frac{\sigma_q R_p}{n^{1/p}} \right)$$
이 결과는 선택한 기하학 $p$와 노이즈 테일의 헤비함(쌍대 지수 $q$를 통해) 사이의 근본적인 연결을 밝혀냅니다. 이는 임의의 $p$-노름 기하학에 대한 최초의 노이즈 강건 가속화 결과입니다.

2. 배치 크기 스케일링 규칙

유계 분산($r=2$) 하에서의 오라클 복잡도를 분석함으로써, 본 논문은 노름 의존적 배치 크기 스케일링 규칙을 도출합니다. 고정된 총 오라클 예산 $N$에 대해, 최적의 배치 크기 $B_p$는 다음과 같이 스케일링됩니다:
$$B_p \asymp N^{2(p-1)/(3p-2)}$$

• $p=2$ (유클리드)의 경우, $B_2 \asymp N^{1/2}$입니다.
• $p \to \infty$ (Schatten-$\infty$)의 경우, $B_\infty \asymp N^{2/3}$입니다.
  이는 왜 Schatten-$\infty$ 방법론들이 그 이점을 실현하기 위해 큰 배치 크기를 요구하는지를 설명합니다.

3. 레짐 의존적 최적성

본 논문은 $p$의 최적 선택이 유효 차원 $d$와 오라클 예산 $N$ 사이의 관계에 의해 결정되는 레임 의존적임을 보여줌으로써 Muon의 유용성에 대한 갈등을 해결합니다:

• 고차원 레짐 ($N \lesssim d$): 문제의 매끄러움(smoothness)을 기하학에 맞추는 것이 유리합니다. 만약 목적 함수가 Schatten-$\infty$ 노름에서 매끄럽다면, $p \to \infty$를 사용하는 것이 최적입니다.
• 저차원 레짐 ($N \gtrsim d$): 더 작은 $p$ 값이 선호됩니다.
  • 결정론적 설정에서는 $p=2$ (유클리드)가 더 강력한 가속을 제공합니다.
  • 헤비 테일 노이즈가 있는 확률적 설정에서는, 노이즈 가정과 더 잘 정렬되는 가장 작은 허용 가능한 기하학 $p = r/(r-1)$이 최적입니다.

4. 스케일링 법칙에 대한 함의

저자들은 이러한 이론적 레짐을 현대적 스케일링 법칙(예: Chinchilla)과 연결합니다. 모델 크기가 토큰 수($N$)에 따라 선형적으로 확장되는 일반적인 스케일링 레시피 하에서, 가중치 행렬의 유효 랭크는 $d \propto N^{1/3}$으로 성장합니다. 결과적으로 $N \gg d$가 되어, 표준적인 대규모 훈련을 저차원 레짐에 위치시킵니다.

• 결론: 이 레짐에서는 Schatten-$\infty$보다 작은 $p$ (유한한 Schatten-$p$)를 사용하는 것이 이론적으로 더 바람직합니다. 이는 왜 Muon의 이점이 규모가 커짐에 따라 감소하는지, 그리고 왜 유한한 $p$ 옵티마이저(HTMuon, Soft-Muon)들이 강력한 경쟁자로 부상하고 있는지에 대한 이론적 근거를 제공합니다.

결과 및 의의

본 논문은 Schatten-$\infty$ 옵티마이저의 경험적 성능에 나타난 외견상의 모순을 해결한다고 주장합니다. 이 연구의 의의는 다음과 같습니다:

• 이론적 통합: $p$-노름 기하학의 관점을 통해 Muon, Soft-Muon, 유클리드 옵티마이저의 동작을 설명하는 단일 프레임워크(SODA)를 제공합니다.
• 경험적 현상 설명:
  • 웜업(Warmup): 분석은 왜 Schatten-$\infty$ 방법론들이 스텝사이즈 웜업을 필요로 하지 않는지(유효 스텝사이즈가 $k^{-1}$로 감소함) 설명하는 반면, 더 작은 $p$ 기하학은 웜업과 유사한 단계를 보임을 설명합니다.
  • 배치 크기: Schatten-$\infty$ 방법론에 큰 배치가 필수적임을 정량화하고 임의의 $p$에 대한 스케일링 규칙을 제공합니다.
  • 스케일링: Schatten-$\infty$ 방법론이 기대만큼 잘 확장되지 않을 수 있다는 관찰에 대한 이론적 기초를 제공하며, 이는 Chinchilla와 같은 레시피 하에서 "저차원적" 특성이 유한한 $p$를 선호함을 시사합니다.
• 실무적 가이드라인: 결과는 훈련 중에 $p$를 조정하는 전략(예: $\infty$ 근처에서 시작하여 예산 $N$이 차원 $d$를 초과함에 따라 유한한 $p$로 전환)이 잠재력이 있음을 시사하며, 이는 Soft-Muon과 같은 최근의 경험적 스케줄과 일치합니다.

저자들은 자신의 분석이 최악의 경우를 가정한 노름 변환에 의존하고 있으며, 더 정교한 문제 의존적 관계가 경계선을 개선할 수 있음을 언급하며 겸허한 어조를 유지합니다. 또한 동적인 $p$-스위칭에 대한 실무적 탐구는 향후 과제로 남겨둡니다.