기술 요약: 딥러닝 최적화를 위한 모델 기반 프레임워크를 통한 적응형 메모리 모멘텀

1. 문제 제기

현대 딥러닝 모델은 주로 모멘텀 기반의 1 차 최적화 알고리즘 (예: Heavy Ball 가 포함된 SGD, AdamW, Lion) 에 의존합니다. 이러한 방법들은 최적화기의 "메모리"를 지배하는 모멘텀 계수 $\beta$를 사용하여, 과거 기울기가 현재 업데이트 방향에 얼마나 영향을 미치는지를 결정합니다. 실제 적용에서는 $\beta$가 거의 보편적으로 고정된 값 (일반적으로 $\beta = 0.9$) 으로 설정되며 훈련 내내 일정하게 유지됩니다.

저자들은 이러한 정적 접근 방식이 가진 두 가지 치명적인 한계를 지적합니다:

1. 비최적성: 경험적 증거에 따르면 단일 고정 $\beta$는 전체 훈련 궤적이나 서로 다른 데이터셋 및 모델 아키텍처에 걸쳐 거의 최적일 수 없습니다.
2. 불안정성: 최적 고정 $\beta$ 주변의 성능 지형은 매우 불안정합니다. 최적 값에서 약간만 벗어나도 수렴이 급격히 저하될 수 있습니다.
3. 이론적 격차: 모멘텀이 결정론적 볼록 설정에서 수렴을 가속화하지만, 확률적 비볼록 딥러닝 설정에 대한 이론적 보장은 종종 확률적 Heavy Ball 이 경험적 성공에도 불구하고 일반 SGD 보다 나은 수렴 속도를 달성하지 못한다고 보여줍니다.

제기된 핵심 질문은 다음과 같습니다: 모멘텀 계수가 최적화 지형에 더 잘 부합하도록 시간에 따라 적응할 수 있을까요?

2. 방법론

본 논문은 최적화 과정에서 온라인으로 계산되는 동적 시간 가변 계수 $\beta_t$로 상수 모멘텀 계수를 대체하는 Adaptive Memory (AM) 프레임워크를 제안합니다.

핵심 프레임워크: 모델 기반 근사

이 방법은 **근사 모델 기반 최적화 (proximal model-based optimization)**에서 파생되었습니다. 실제 손실 함수 $f(x)$를 직접 최소화하는 대신, 저자들은 두 개의 평면으로 구성된 대리 모델 $f^m_t(x)$를 사용하여 이를 근사합니다:

1. 기울기 평면 (Gradient Plane): 현재 기울기 $\nabla f(x_t)$에 기반합니다.
2. 모멘텀 평면 (Momentum Plane): 누적된 모멘텀을 인코딩하는 이전 하강 방향 $\frac{1}{\eta}(x_{t-1} - x_t)$에 기반합니다.

대리 모델은 다음과 같이 정의됩니다:
$$f^m_t(x) = \max \left\{ f(x_t) + g_t^\top(x - x_t), \quad \hat{f}(x_t) + \frac{1}{\eta}(x_{t-1} - x_t)^\top(x - x_t) \right\}$$
여기서 $g_t = \nabla f(x_t)$이며 $\hat{f}(x_t)$는 모멘텀 평면을 나타내는 편향 항입니다.

적응형 $\beta_t$ 유도

업데이트 단계는 현재와 이전 하강 방향 간의 정렬을 제어하기 위한 추가 정규화 항을 포함한 근사 최소화 문제로 공식화됩니다:
$$x_{t+1} \in \arg\min_x \left( f^m_t(x) + \frac{1}{2\eta}\|x - x_t\|^2 + \frac{\lambda}{\eta}\langle x_{t-1} - x_t, x - x_t \rangle \right)$$

이 최적화 문제를 풀면 적응형 모멘텀 계수 $\beta_t^*$에 대한 폐쇄형 표현이 도출됩니다. 이 해는 현재 기울기와 누적된 모멘텀 사이의 균형을 맞추는 2 차 계획법을 포함합니다. 결과적인 업데이트 규칙은 동적으로 계산된 $\beta_t$를 가진 Heavy Ball 방법과 유사합니다:
$$d_{t+1} = \frac{\beta_t^* + \lambda}{1 + \lambda} d_t + \frac{1 - \beta_t^*}{1 + \lambda} g_t$$
$$x_{t+1} = x_t - \eta d_{t+1}$$

실제 구현 세부 사항

확률적 설정에서 안정성과 효율성을 보장하기 위해 저자들은 몇 가지 실용적인 수정을 도입합니다:

• 클리핑 (Clipping): $\beta_t$는 $[0, \beta_{\max}]$ 구간으로 클리핑되며, 여기서 $\beta_{\max}$는 일반적으로 표준 고정 값 (예: 0.9) 으로 설정됩니다.
• 손실 근사: 동일한 배치에서 정확한 손실 차이 $\Delta f = f(x_{t-1}) - f(x_t)$를 계산하는 것은 비효율적이며, 다른 배치를 사용하면 노이즈가 발생합니다. 저자들은 이를 1 차 전개로 근사합니다: $\Delta f \approx \eta g_t^\top d_t$.
• 과대평가 전략: 수렴을 보장하기 위해 과소평가 평면을 사용하는 전통적인 번들 방법과 달리, 저자들은 과대평가 평면 ($\hat{f}(x_t) = f(x_{t-1})$) 을 사용하는 것이 실제에서 더 빠른 수렴으로 이어진다고 발견했습니다.
• AdamW 로의 확장: 이 프레임워크는 대각 사전 조건부 행렬 ($P_t$) 과 분리된 가중치 감쇠를 통합하여 AdamW 로 확장됩니다. 적응형 $\beta_t$는 대규모 언어 모델 (LLM) 의 계층 이질성을 고려하기 위해 계층별로 계산됩니다.

3. 주요 기여

1. 동기 부여: 이 논문은 경험적으로 고정된 모멘텀이 비최적이며 최적 정적 $\beta$를 찾는 것이 번거롭고 비강건함을 입증합니다.
2. 새로운 방법론: 손실 함수의 2 평면 근사를 사용하는 근사 프레임워크를 도입하여 폐쇄형 적응형 모멘텀 계수를 유도합니다. 이는 모멘텀 기반 최적화를 위해 이러한 근사 프레임워크를 적용한 최초의 사례입니다.
3. 알고리즘 변형: 저자들은 SGD (AM-MGD) 와 AdamW (AM-AdamW) 모두에 대해 Adaptive Memory 변형을 구현했습니다.
4. 이론적 보장: 표준 매끄러움과 유계 기울기 가정 하에 볼록 및 비볼록 설정 모두에 대해 $O(1/\sqrt{T})$의 수렴 속도를 확립하여, 이 방법을 확률적 1 차 방법의 표준 지형 내에 위치시킵니다.
5. 하이퍼파라미터 효율성: 이 방법은 표준 베이스라인 설정 이상으로 추가 하이퍼파라미터 튜닝이 필요하지 않습니다 (예: $\lambda$와 $\beta_{\max}$는 작업 간에 고정됨).

4. 실험 결과

저자들은 다양한 작업에 걸쳐 AM 을 평가했습니다:

• 볼록 문제: LIBSVM 데이터셋의 로지스틱 회귀 작업에서 AM-MGD 는 튜닝이 필요 없이 고정 모멘텀 베이스라인 (최적 그리드 검색된 $\beta^*$ 포함) 을 일관되게 능가했습니다.
• 이미지 분류: VGG, ResNet, Wide-ResNet 아키텍처를 사용한 CIFAR-10/100 및 ImageNet 에서 AM-MGD 는 더 빠른 최적화와 더 나은 일반화를 달성했습니다.
  • 동역학: 적응형 $\beta_t$는 자연스럽게 세 단계를 보였습니다: 노이즈가 많은 초기 훈련 동안의 낮은 모멘텀 (SGD 유사), 가속화 단계 동안의 높은 모멘텀, 그리고 기울기가 약해지는 보간점 근처에서의 빈번한 "소프트 리스타트"($\beta_t$ 감소).
• 대규모 언어 모델 사전 훈련: C4 에서 훈련된 LLaMA 모델 (20M 에서 1B 파라미터) 에서:
  • AM-AdamW 는 모든 규모에서 표준 AdamW 를 능가했습니다.
  • 워밍업 제거: 주목할 만한 점은 AM-AdamW 가 학습률 워밍업 스케줄 없이도 초기 훈련을 안정화시켰다는 것입니다. 이는 표준 AdamW 가 종종 실패하거나 신중한 튜닝이 필요한 영역입니다. 이는 AM 이 워밍업에 대한 하이퍼파라미터 없는 대안이 될 수 있음을 시사합니다.
• 어블레이션 연구:
  • 이 방법은 $\lambda$ 선택에 대해 강건합니다 ($[0.01, 1]$ 범위 내).
  • 성능 향상은 다양한 배치 크기와 학습률에서 지속되며, AM 은 정적 베이스라인이 실패하는 더 높은 학습률에서 수렴을 허용합니다.
  • 계산 오버헤드는 무시할 수준입니다 (LLM 의 경우 약 0.2% 시간 증가) 추가 메모리 사용량은 없습니다.

5. 중요성 및 주장

이 논문은 Adaptive Memory 가 정적 모멘텀에 대한 원칙적이고 튜닝이 필요 없는 대안을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 이론과 실전의 연결: 이는 손실 지형의 국소 기하학에 적응하는 메커니즘을 제공하여 이론적 수렴 속도와 경험적 성능 간의 괴리를 해소합니다.
• 훈련 파이프라인 단순화: 초기 훈련을 안정화하고 학습률 워밍업 스케줄의 필요성을 잠재적으로 제거함으로써 대규모 모델의 훈련을 단순화합니다.
• 강건성: 이 방법은 높은 학습률과 다양한 배치 크기에 대한 강건성을 보여주어, 딥러닝 훈련의 하이퍼파라미터 선택에 대한 민감도를 줄입니다.

저자들은 이 작업이 적응형 학습률을 넘어 적응형 모멘텀으로 나아가는 최적화에서 적응성을 유도하는 새로운 방식을 위한 문을 열었다고 결론지으며, 파라미터별 모멘텀 계수와 모멘텀 리스타트에 대한 추가 연구를 향후 과제로 제안합니다.