기술적 요약: 낙관적 이중 평균화 (Optimistic Dual Averaging) 가 현대 최적화기를 통합한다

문제 제기

딥러닝 최적화는 두 가지 주요 축을 따라 진화해 왔습니다: 기하학과 적응 (geometry and adaptation) (모델 구조를 반영하는 암시적 노름 또는 전처리기를 가진 업데이트를 설계하는 것, 예: Adam, Muon) 과 성분들의 조합 (composition of ingredients) (안정성을 위한 그래디언트 피드백, 모멘텀, 스케줄을 결합하는 것, 예: Lion, schedule-free 방법). 많은 최첨단 최적화기 (Muon, Lion, AdEMAMix, NAdam 등) 는 서로 다르게 보이지만, 그들의 관계를 설명하고 하이퍼파라미터 튜닝을 위한 원칙적인 지침을 제공하는 통합된 이론적 프레임워크는 부재합니다.

구체적이고 지속적인 과제는 **가중치 감쇠 (weight decay)**의 튜닝입니다. 여러 에포크 훈련에서 가중치 감쇠는 반복점 (iterate) 노름을 제약하는 정규화제로 작용합니다. 그러나 과적합이 덜 우려되는 단일 에포크 훈련에서도 가중치 감쇠는 종종 유익하게 작용합니다. 최근의 스케일링 규칙은 모델 차원 $d$에 비례하는 가중치 감쇠를 설정할 것을 제안하지만, 모델 차원 $d$에 대한 의존성과 훈련 기간 $n$에 대한 의존성을 분리하는 것은 어렵습니다. 또한, 최적의 가중치 감쇠 값은 최적화기에 따라 크게 달라집니다 (예: Lion 은 AdamW 보다 훨씬 큰 감쇠가 필요함). 이로 인해 각 최적화기마다 비용이 많이 드는 튜닝이 필요합니다.

방법론: SODA

저자들은 SODA(Stochastic Optimistic Dual Averaging) 를 소개합니다. 이는 그래디언트의 이중 처리와 반복점의 원시 (primal) 처리를 통합하는 낙관적 이중 평균화 (ODA) 의 일반화입니다.

이론적 프레임워크

SODA 는 표준 이중 평균화 업데이트를 일반화합니다:
$$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$$
이를 위해 원시 외삽 시퀀스 (primal extrapolation sequence) $(y_k)$와 낙관적 이중 변수 (optimistic dual variable) $(\bar{m}_{k+1})$를 도입합니다. 알고리즘은 다음을 유지합니다:

1. 이중 처리: 모멘텀 기반 그래디언트 평균 $m_{k+1}$과 현재 그래디언트 $\nabla f(y_k, \xi_k)$를 포함하는 낙관적 평균 $\bar{m}_{k+1}$.
2. 원시 처리: 매개변수 $\lambda_k$와 $\bar{\lambda}_k$에 의해 제어되는 반복점 업데이트 $x_{k+1}$과 외삽된 점 $y_{k+1}$.

업데이트 규칙은 다음과 같습니다:
$$\begin{aligned} m_{k+1} &= (1 - \alpha_k)m_k + \alpha_k \nabla f(y_k, \xi_k) \\ \bar{m}_{k+1} &= (1 - \bar{\alpha}_k)m_{k+1} + \bar{\alpha}_k \nabla f(y_k, \xi_k) \quad (\text{낙관성}) \\ z_{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}_{k+1}) \\ x_{k+1} &= (1 - \lambda_k)x_k + \lambda_k z_{k+1} \\ y_{k+1} &= (1 - \bar{\lambda}_k)x_{k+1} + \bar{\lambda}_k z_{k+1} \quad (\text{원시 외삽}) \end{aligned}$$
여기서 $h_k$는 기하학을 정의하는 정규화제 (미러 맵) 입니다 (예: Lion 의 경우 $\ell_\infty$, Muon 의 경우 스펙트럼).

기존 최적화기의 통합

이 논문은 특정 매개변수와 기하학을 선택함으로써 SODA 가 여러 현대 최적화기를 특수한 경우로 복원함을 보여줍니다:

• NAdam / Lion / Muon: 이들은 "현대화 (Modernized)" 영역에서 복원되는데, 이때 $\bar{\lambda}_k = 0$(그래디언트 평가점에 대한 원시 외삽 없음) 이고 $\bar{\alpha}_k > 0$(낙관성) 입니다.
  • Lion: $\ell_\infty$ 기하학 ($h_k = K^*$) 과 특정 모멘텀 매개변수를 사용합니다.
  • Muon: 스펙트럼 기하학과 Nesterov 모멘텀 ($\bar{\alpha}_k = \alpha_k$) 을 사용합니다.
  • NAdam: 적응형 스케일링을 사용하는 부드러운 $\ell_\infty$ 기하학을 사용합니다.
• 논문의 표 1은 이러한 방법들을 그래디언트 평균화 ($\alpha_k$), 낙관성 ($\bar{\alpha}_k$), 기하학 ($h_k$) 에 대한 해당 선택과 명시적으로 매핑합니다.

SODA Wrapper

주요 실용적 기여는 새로운 하이퍼파라미터를 도입하지 않고 어떤 기본 최적화기 (예: Adam, Lion, Muon) 를 감싸는 SODA Wrapper(알고리즘 1) 입니다.

• 구성: $\bar{\lambda}_k = 0$( "현대화" 영역과 일치), 균일 반복점 평균화 $\lambda_k = 1/(k+2)$, 그리고 스텝사이즈 스케일링 $\gamma_k = (k+2)\eta_k$를 설정합니다.
• 가중치 감쇠 제거: 이 래퍼는 원점이 아닌 초기화 $z_0$에 업데이트를 고정합니다. 결과적으로 생성된 업데이트 규칙은 $\mu_k = 1/(k+2)$의 가중치 감쇠 스케줄을 효과적으로 구현합니다.
• 해석: 저자들은 이 프레임워크에서 가중치 감쇠가 단순히 정규화제가 아니라 **원시 평균화 (primal averaging)**의 한 형태라고 주장합니다. 이론적으로 근거된 $1/k$ 감쇠는 별도의 가중치 감쇠 하이퍼파라미터를 튜닝할 필요성을 제거합니다.

주요 기여

1. 낙관성과 이중화를 통한 통합: SODA 는 Muon, Lion, NAdam 등을 낙관적 이중 평균화 사례로 복원하는 단일 템플릿을 제공하여 구조적 관계를 명확히 합니다.
2. 이론적 보장: 이 프레임워크는 $L$-스무스 볼록 함수에 대한 수렴 분석에 이론적으로 근거한 하이퍼파라미터 선택 (특히 $1/k$ 감쇠 스케줄) 을 유도합니다. 가중치 감쇠를 원시 평균화를 위한 메커니즘으로 재해석합니다.
3. 파라미터 없는 래퍼: 제안된 SODA 래퍼는 성능을 유지하거나 향상시키면서 모든 기본 최적화기에 대한 가중치 감쇠 튜닝 필요성을 제거합니다.
4. 실증적 검증: 광범위한 실험은 모델 크기와 훈련 기간 전반에 걸쳐 일관된 개선을 보여줍니다.

실험 결과

실험은 FineWeb100데이터셋을 사용하여 nanoGPT에서 수행되었으며, SODA 로 감싸진 최적화기를 튜닝된 베이스라인과 비교했습니다.

• 기간 간 전이: Muon 의 최적 가중치 감쇠는 훈련 기간이 증가함에 따라 감소하는 것으로 나타났으며, 이는 정적 스케일링 규칙 (예: $1/d$) 과 모순됩니다. SODA 래퍼의 동적 $1/(k+2)$ 감쇠는 Muon 이 최적으로 튜닝된 가중치 감쇠를 부여받았을 때도 Muon 을 일관되게 능가했습니다 (그림 1).
• 일반적 적용성: 이 래퍼는 추가 튜닝 없이 다양한 훈련 기간 (1$\times$ 및 4$\times$ Chinchilla 스텝) 에서 Adam, Muon, Scion의 성능을 향상시켰습니다 (그림 2).
• 스케일링: 모델 크기를 64M 에서 1B 파라미터로 확장할 때 (Chinchilla 스케일링 준수), SODA(특히 낙관적 변형 SODA$^\dagger$) 는 기본 최적화기 (Scion/uScion) 를 일관되게 능가했으며, 모델 크기가 커질수록 성능 격차가 벌어졌습니다 (그림 4).
• 단순화: SODA$^\dagger$ 구성은 모든 레이어에서 Adam 을 Lion 으로 대체하고 단일 모멘텀 하이퍼파라미터를 사용하여 Muon 베이스라인을 단순화했음에도 불구하고 경쟁력 있거나 우수한 성능을 달성했습니다.

중요성과 주장

이 논문은 SODA 가 하이퍼파라미터 튜닝 없이 다양한 훈련 기간과 모델 규모 간에 최적화 전략을 전이할 수 있는 원칙적인 메커니즘을 제공한다고 주장합니다.

• 가중치 감쇠의 재해석: 저자들은 관점의 전환을 제기합니다: 가중치 감쇠는 과적합에 대한 정규화제로서만 이해되는 것이 아니라, 최적화 궤적을 초기화에 고정하는 원시 평균화의 한 형태로 이해되어야 합니다.
• 실질적 영향: 현대 최적화기를 SODA 프레임워크 하에 통합함으로써 저자들은 가중치 감쇠를 튜닝하는 번거롭고 비용이 많이 드는 과정을 제거하는 "플러그 앤 플레이" 래퍼를 제공합니다. 이는 모델 크기가 커져 수동 튜닝이 prohibitively expensive(부적절하게 비쌈) 해짐에 따라 특히 중요합니다.
• 겸손함: 저자들은 이론적 분석이 현재 볼록 (convex) 설정으로 제한되어 있음 (딥러닝에 대해 경험적으로 유익함) 과 단일 에포크 훈련에 초점을 맞추고 있음을 인정합니다. 그들은 과적합 우려로 인해 다중 에포크 훈련은 여전히 정규화 강도에 대한 하이퍼파라미터를 필요로 할 수 있다고 제안합니다.

요약하자면, SODA 는 낙관적 이중 평균화의 렌즈를 통해 현대 최적화기의 설계 공간을 통합하고, 규모 전반에 걸쳐 훈련 안정성과 성능을 개선하는 이론적으로 근거된 파라미터 없는 방법을 제공합니다.