기술 요약: PowerStep: ℓp-노름 최강 강하를 통한 메모리 효율적 적응 최적화

문제 제기

적응형 최적화 알고리즘, 특히 Adam 과 AdamW 는 누적된 1 차 및 2 차 모멘트 통계를 사용하여 그래디언트 업데이트를 전처리함으로써 대규모 신경망 (예: Transformer) 학습의 표준이 되었습니다. 이러한 방법들은 효과적이지만, 모델 파라미터와 동일한 차원을 가진 2 차 모멘트 누적기 ($v_t$) 를 유지하며 일반적으로 안정성을 위해 전체 32 비트 부동소수점 (fp32) 정밀도를 요구하므로 상당한 메모리 오버헤드를 초래합니다. 이러한 "메모리 벽"은 모멘텀이 있는 확률적 경사 하강법 (SGD) 에 비해 최적화기 상태의 발자국을 두 배로 늘려, 수천억 개의 파라미터를 가진 모델의 확장성을 제한합니다.

이 문제를 완화하기 위한 기존 시도들로는 2 차 모멘트 행렬의 인수분해, 블록 단위 공유, 또는 저정밀도 양자화 적용 등이 있습니다. 그러나 2 차 모멘트 추정기 (특히 역제곱근 연산) 가 양자화 오차에 매우 민감하기 때문에, 양자화는 종종 적응형 방법에서 학습 불안정성이나 발산을 초래합니다. 대안으로 일부 방법은 2 차 모멘트를 완전히 제거하기도 하지만, 대규모 작업에서 Adam 과 같은 수렴 속도를 달성하지 못하는 경우가 많습니다.

방법론

저자들은 2 차 모멘트 통계를 저장하지 않고 좌표별 적응성을 달성하는 메모리 효율적 최적화기인 PowerStep을 소개합니다. 이 방법은 ℓp-노름 최강 강하 (steepest descent) 원리에서 도출되었습니다.

이론적 기반

일반 노름 공간에서 ℓp-노름 제약 하에 목적 함수 $f(\theta)$를 최대한 감소시키는 최강 강하 방향 $v^*$는 다음과 같이 주어집니다:
$$v^*_i \propto -\text{sign}(g_i) \cdot |g_i|^{\frac{1}{p}-1}$$
여기서 $g_i$는 그래디언트 성분입니다. 이는 $\beta = 1/(p-1) \in [0, 1]$인 부호付き 거듭제곱 변환 (signed power transform) $\Phi_\beta(x) = \text{sign}(x) \odot |x|^\beta$로 이어집니다.

알고리즘 설계

원시 그래디언트에 이 변환을 적용하는 Powerball 방법과 달리, PowerStep 은 무거운 공 모멘텀 버퍼 (heavy-ball momentum buffer)($m_t$) 에 이 변환을 적용합니다. 업데이트 규칙은 다음과 같습니다:

1. 모멘텀 업데이트: $m_t = \gamma m_{t-1} + g_t$
2. 변환된 업데이트: $\theta_t = \theta_{t-1} - \eta_t \Phi_\beta(m_t)$

이 설계는 대각 전처리 행렬 $D_t = \text{diag}(|m_t|^{\beta-1})$을 생성합니다. $\beta \in (0, 1)$인 경우, 지수 $\beta-1$은 음수이므로 스텝 크기는 누적 모멘텀의 크기에 반비례합니다. 작은 모멘텀 크기 (평탄한 방향) 는 증폭되고, 큰 크기 (가파른 방향) 는 감쇠됩니다. 이는 Adam 과 유사한 좌표별 적응성을 제공하지만 2 차 모멘트 추정기 없이 모멘텀 버퍼에서 즉시 계산되므로 2 차 모멘트 추정기가 필요 없습니다.

저정밀도 구현

PowerStep 의 단일 버퍼 설계는 공격적인 양자화에 자연스럽게 적합합니다. 저자들은 부호付き 거듭제곱 변환과 파라미터 업데이트는 전체 정밀도로 수행하면서 int8정밀도로 모멘텀 버퍼 $m_t$만 저장할 것을 제안합니다. 이는 Adam 과 대조적인데, Adam 의 경우 제곱근과 나눗셈 연산의 비선형성으로 인해 2 차 모멘트 버퍼 $v_t$를 양자화하면 치명적인 오차 누적이 발생합니다.

주요 기여

1. 알고리즘 설계: PowerStep 은 ℓp-노름 최강 강하 원리에서 도출되었으며, 무거운 공 모멘텀 버퍼에 부호付き 거듭제곱 변환을 적용함으로써 2 차 모멘트 통계 없이 좌표별 적응성을 달성함을 보여줍니다.
2. 이론적 분석: 이 논문은 표준 규칙성 가정 (L-스무스, 유계 그래디언트 등) 하에서 PowerStep 이 비볼록 확률적 최적화에 대해 최적의 $O(1/\sqrt{T})$수렴 속도를 달성함을 증명합니다.
3. 실증적 검증: 밀집 및 전문가 혼합 (Mixture-of-Experts) 아키텍처를 포함한 1 억 2,400 만에서 2,350 억 개 파라미터에 이르는 Transformer 모델에 대한 광범위한 실험에서, PowerStep 이 AdamW 와 동일한 수렴 속도를 달성하면서 최적화기 상태 메모리는 절반만 사용함을 보여줍니다.
4. 양자화 강건성: 공격적인 int8 양자화와 결합될 때, PowerStep 은 수치적으로 안정적으로 유지되며 AdamW 대비 약 8 배의 최적화기 메모리를 줄입니다. 반면, 동일한 양자화 조건 하에서 AdamW 는 발산합니다.

실험 결과

• 수렴 속도: 1 억 2,400 만에서 2,350 억 개 파라미터에 이르는 모델 전반에 걸쳐 PowerStep 은 AdamW 의 학습 손실 궤적과 일치합니다. 관련 메모리 효율적 방법들 (pbSGDM, SignSGD, AdamS) 은 더 큰 모델에서 더 느린 수렴이나 불안정성을 보였습니다.
• 메모리 효율성: PowerStep 은 두 개의 버퍼 대신 하나의 버퍼만 저장하므로 AdamW (fp32) 대비 최적화기 상태 메모리 발자국을 50% 줄입니다.
• 양자화 성능: 최적화기 상태의 int8 양자화 하에서:
  • AdamW: 2 차 모멘트 추정기에서의 양자화 오차 누적으로 인해 학습이 즉시 붕괴됩니다.
  • PowerStep: 안정적으로 유지되며 전체 정밀도 버전과 동일한 수렴을 보입니다.
  • 메모리 감소: PowerStep (int8) 은 AdamW (fp32) 대비 최적화기 메모리를 약 8 배 감소시킵니다. 예를 들어, 2,350 억 개 파라미터 모델에서 NPU 당 메모리는 AdamW 의 6,768 MB 에서 PowerStep int8 의 872.4 MB 로 감소합니다.
• 하이퍼파라미터 민감도: 이 방법은 모멘텀 계수 $\gamma$가 $[0.85, 0.95]$범위일 때 강건합니다. 거듭제곱 지수 $\beta = 0.1$은 빠른 초기 진행과 장기적 안정성 사이의 최적 균형을 제공하며, $\beta=0$(SignSGD 와 동일) 은 붕괴를 초래하고 $\beta=0.2$는 수렴이 너무 느립니다.

의의 및 주장

이 논문은 PowerStep 이 대규모 학습을 위한 원칙적이고 확장 가능하며 자원 효율적인 대안을 제공한다고 주장합니다. 적응형 최적화와 ℓp-노름 기하학 사이의 직접적인 연결을 확립함으로써, PowerStep 은 2 차 모멘트 추정의 메모리 비용 없이 적응형 방법의 이론적 최적성을 달성합니다.

저자들은 2 차 모멘트 버퍼의 제거가 메모리 사용량을 절반으로 줄일 뿐만 아니라 저정밀도 양자화에 대한 강건성을 근본적으로 향상시킨다고 강조합니다. 이는 수렴 속도나 최종 모델 품질을 희생하지 않으면서 추가적인 차수 (약 8 배) 의 메모리 감소 (AdamW 대비 총 약 8 배) 를 가능하게 합니다. 이 연구는 비선형 변환과 결합된 무거운 공 모멘텀이 Adam 의 복잡한 2 차 모멘트 메커니즘을 대체할 수 있음을 시사하며, 향후 대규모 모델 학습을 위한 더 단순하고 확장 가능한 경로를 제시합니다.

저자들이 식별한 향후 연구 방향으로는 구조화된 행렬 값 업데이트 (예: Muon 최적화기) 와 PowerStep 의 통합, 그리고 더 공격적인 양자화 (예: int4) 탐색 등이 포함됩니다.