SGD at the Edge of Stability: Stochastic Stabilization with Large Learning Rates

기술 요약: 안정성의 경계에 있는 SGD: 큰 학습률을 통한 확률적 안정화

문제 정의

현대 딥러닝은 고전적인 매끄러운 볼록 최적화 이론이 정당화하는 수준보다 훨씬 더 큰 학습률을 일상적으로 사용하며, 이는 **안정성의 경계(Edge of Stability, EoS)**라고 알려진 현상이다. 결정론적 경사 하강법(Gradient Descent, GD)의 EoS 동작은 곡률에 의한 진동 단계와 이후의 자기 안정화(self-stabilizing) 단계로 이어지는 과정으로 분석되어 왔으나, 확률적 설정에서의 연구는 여전히 미개척 상태로 남아 있다.

확률적 경사 하강법(Stochastic Gradient Descent, SGD)에서는 미니 배치 노이즈로 인한 무해한 확률적 변동과 곡률 대비 큰 스텝 사이즈로 인한 실제 불안정성 사이의 구분이 모호하다. 손실 함수의 비단조성(non-monotonicity)은 이 두 가지 원인 모두에서 발생할 수 있으며, 이로 인해 고전적인 결정론적 안정성 임계값(Hessiel의 최대 고유값에 기반한)은 정의하기 어렵거나 불충분하게 된다. 결과적으로, 다중 클래스 교차 엔트로피 손실(multi-class cross-entropy loss)에 대해 임의로 큰 학습률 하에서 SGD가 어떻게 행동하는지에 대한 원칙적인 이해가 부족하다. 핵심 질문은 다음과 같다: 교차 엔트로피 손실을 사용하는 SGD에 적합한 안정성의 개념은 무엇이며, 고전적인 작은 학습률 범위를 넘어선 역학을 어떻게 규명할 수 있는가?

방법론

저자들은 임의의 큰 스텝 사이즈($\eta > 0$) 하에서 선형 분류기와 2층 신경망에 적용되는 다중 클래스 교차 엔트로피 손실에 대한 SGD를 분석하기 위한 이론적 프레임워크를 개발한다.

1. 확률적 리아푸노프 안정성 (Stochastic Lyapunov Stability):
   실현된 손실의 점별 단조성(pointwise monotonicity)에 의존하는 대신(노이즈가 존재하는 상황에서는 이를 보장하는 것이 불가능함), 본 논문은 확률적 리아푸노프 안정성 관점을 채택한다. 손실 함수 $L(W)$를 리아푸노프 함수로 취급한다. 안정성은 $L(W_{t+1}) \leq L(W_t)$가 아니라, 필트레이션 $\mathcal{F}_t$가 주어졌을 때 기대 손실이 감소한다는 조건으로 정의된다:
   $$\mathbb{E}[L(W_{t+1}) - L(W_t) \mid \mathcal{F}_t] \leq 0$$
   이 기준은 곡률과 노이즈 사이의 균형을 포착한다. 이는 풀 배치(full-batch) 극한에서는 결정론적 안정성 조건으로 환원되지만, 확률적 설정에서는 미니 배치 변동을 고려한다.

2. 2단계 역학 분석 (Two-Phase Dynamics Analysis):
   분석은 훈련 역학을 두 가지 영역으로 분리한다:

   • 안정성의 경계(EoS) 영역: 반복 계산이 곡률에 의한 진동을 경험할 수 있고 손실이 단조롭게 감소하지 않는 초기 단계이다.
   • 안정 영역 (Stable Regime): 기대 손실이 1단계 하강(one-step descent)을 허용하는 영역이다. 저자들은 스텝 사이즈, 배치 크기 및 데이터 마진에 따라 결정되는 $\tilde{L}$에 대해 $L(W_t) \leq \tilde{L}$을 만족하는 안정 집합 $\mathcal{S}$를 정의한다.

3. 확률적 안정화 메커니즘 (Stochastic Stabilization Mechanism):
   본 논문의 핵심적인 이론적 기여는 SGD가 **자기 안정화 속성(self-stabilizing property)**을 가진다는 것을 증명하는 것이다. 미니 배치 노이즈가 일시적으로 반복 계산을 안정 집합 $\mathcal{S}$ 밖으로 밀어내더라도, 역학은 높은 확률로 고정된 유한한 반복 횟수 내에 안정 영역으로 돌아온다. 이러한 복귀 메커니즘은 곡률과 노이즈 크기 모두가 손실이 0에 가까워짐에 따라 감소하는 교차 엔트로피 손실의 기하학적 구조에 의해 구동된다.

4. 신경망으로의 확장:
   이 프레임워크는 평균장 스케일링(mean-field scaling) 하의 2층 신경망으로 확장된다. 저자들은 활성화 함수의 미분 가능성, 매끄러움, 근사 동질성(near-homogeneity)에 대한 완만한 조건을 부과하고, 활성화 도함수의 하한을 고려하는 "유효 마진" $\tilde{\gamma}$를 도입한다. 이를 통해 비선형 설정에서도 선형 수렴 보장을 전이할 수 있다.

주요 기여

• 확률적 안정성 기준: 본 논문은 조건부 리아푸노프 드리프트(conditional Lyapunov drift)에 기반하여 SGD를 위한 새로운 안정성 정의를 제안한다. 이는 파라미터가 발산하면서도 손실은 0으로 수렴할 수 있는 분리 가능한 교차 엔트로피 최소화 문제에 적합하다.
• 큰 학습률에 대한 수렴 보장: 저자들은 임의의 큰 스텝 사이즈($\eta > 0$)를 가진 다중 클래스 교차 엔트로피 손실에 대한 SGD의 첫 번째 이론적 수렴 보장을 제공한다. 이들은 스텝 사이즈의 크기에 관계없이 최적의 반복 계산(best iterate)이 $\tilde{O}(1/t)$ 수렴율을 달성함을 입증한다.
• 안정 영역의 특성 규명: SGD가 고정된 반복 횟수 내에 안정 영역에 진입함을 증명한다. 일단 진입하면, 마지막 반복 계산(last iterate)은 아선형(sublinear) 속도로 수렴한다.
• 확률적 안정화 증명: SGD가 "자기 안정화"됨을 엄밀하게 증명한다. 노이즈로 인해 궤적이 안정 집합을 벗어나더라도, 정해진 단계 내에 안정성으로 복귀한다. 이는 결정론적 GD에서 일단 안정 단계에 진입하면 구조적으로 그곳에 머무르는 것과 대조적이다.
• 신경망 분석: 모델의 비선형성에도 불구하고, 활성화 함수가 특정 "리키(leaky)" 조건을 만족하는 경우 큰 스텝 사이즈의 수렴 보장이 유지됨을 보여줌으로써 결과를 2층 신경망으로 확장한다.

결과

• 이론적 경계: 정리 2.2는 선형 분류기에 대한 최적 반복 계산의 기대 손실 상한을 설정하며, $\tilde{O}(1/t)$ 속도를 보여준다. 정리 4.1과 정리 5.4는 역학이 안정 영역에 진입한 후의 마지막 반복 계산에 대해 유사한 보장을 제공한다.
• 안정성 복귀: 정리 4.2(선형 모델용)와 정리 5.5(신경망용)는 안정 집합에 머물 확률과 탈출이 발생했을 때 복귀하는 데 필요한 최대 단계를 정량화한다. 이 경계값들은 배치 크기, 스텝 사이즈 및 현재 손실 값에 의존한다.
• 실험적 검증: 합성 데이터, MNIST(2층 NN), CIFAR-10(8층 NN)에 대한 실험을 통해 이론적 발견을 검증한다. 실험 결과는 다음을 보여준다:
  • EoS 단계에서의 비단조적 손실 감소 후 안정 영역에서의 수렴.
  • 큰 스텝 사이즈를 사용할 때 더 빠른 손실 수렴.
  • 작은 스텝 사이즈와 비교하여 큰 스텝 사이즈에서 대등하거나 향상된 테스트 정확도.
  • 더 큰 배치 사이즈를 사용할 때 감소하는 노이즈와 진동.

의의 및 주장

본 논문은 현대 딥러닝 최적화에 대한 이론적 이해의 중요한 공백을 메운다고 주장한다. 이전 연구들이 결정론적 GD에 대한 EoS 현상을 설명했다면, 본 연구는 실질적인 딥러닝의 표준인 확률적 설정으로 이를 확장한다.

저자들은 자신들의 결과가 확률성이 최적화의 2단계 구조를 어떻게 변화시키는지에 대한 원칙적인 이해를 제공한다고 강조한다. 저자들은 "자기 안정화" 메커니즘이 교차 엔트로피 손실을 사용하는 SGD의 고유한 특성이며, 손실이 시간이 지남에 따라 감소하는 한 큰 학습률에서도 알고리즘이 발산하지 않도록 보장한다고 주장한다. 이 연구는 딥러닝의 모든 최적화 문제를 해결하려는 것이 아니라, 큰 스텝 사이즈 영역에서 SGD를 분석하기 위한 엄밀한 프레임워크를 제공하며, 큰 학습률이 실제로 효과적이고 안정적이라는 경험적 관찰을 뒷받침한다. 이론적 보장은 분리 가능한 데이터와 교차 엔트로피 손실에 특화되어 있으며, 신경망에 대한 분석은 평균장 스케일링 하의 2층 구조로 제한된다.