[논문 리뷰] Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent

Benjamin Gess, Sebastian Kassing외 1명|arXiv (Cornell University)|2023. 02. 14.

Markov Chains and Monte Carlo Methods인용 수 4

한 줄 요약

이 논문은 작은 학습률 영역에서 정규 분산 계수를 갖는 확률적 수정 유동(SMFs)이라는 새로운 SDE 클래스를 소개한다. 이는 이전의 확률적 수정 방정식(SMEs)에서 발생하는 비정규성 문제를 피하면서, 작은 학습률에서 확률적 경사 하강법(SGD)을 모델링한다. 또한, 무한한 넓이와 작은 학습률 스케일링 영역에서, 다중점 통계를 일치시키고 오버파rameterized 네트워크의 엄밀한 분석을 가능하게 하는, 분포에 의존하는 SMFs를 통해 SGD의 변동성이 있는 평균장 근사값을 확립한다.

ABSTRACT

We propose new limiting dynamics for stochastic gradient descent in the small learning rate regime called stochastic modified flows. These SDEs are driven by a cylindrical Brownian motion and improve the so-called stochastic modified equations by having regular diffusion coefficients and by matching the multi-point statistics. As a second contribution, we introduce distribution dependent stochastic modified flows which we prove to describe the fluctuating limiting dynamics of stochastic gradient descent in the small learning rate - infinite width scaling regime.

연구 동기 및 목표

• 오버파라미터화된 딥 러닝에서 분석을 방해하는, 기존의 확률적 수정 방정식(SMEs)의 비정규성 분산 계수 문제를 해결하기 위해.
• SGD의 다중점 통계를 일치시키는 제한 SDE 모델을 개발하여, 이전의 연속 근사값을 향상시키기 위해.
• 무한한 넓이와 작은 학습률 영역에서, 분포에 의존하는 확률적 수정 유동을 사용하여 SGD의 변동성이 있는 평균장 근사값을 수립하기 위해.
• 아키텍처와 데이터에 따라 변하는 데에도 강건한, SGD의 동역학을 분석하기 위한 통합 프레임워크를 제공하기 위해.

제안 방법

• 원소 손실 기울기에서 직접 유도된 분산 계수를 갖는 원형 브라운 운동에 의해 구동되는 SDE로, 확률적 수정 유동(SMFs)을 제안한다. 이는 비퇴화된 공분산 행렬 Σ의 역행렬이 아닌, 개별 손실 기울기에서 유도된다.
• SMF SDE는 dXₜ^η(x) = -∇(R(Xₜ^η(x)) + η/4 |∇R(Xₜ^η(x))|²) dt + √η ∫_Θ G(Xₜ^η(x),θ) W(dθ,dt) 로 정의되며, 여기서 G(x,θ) = ∇R̃(x,θ) - ∇R(x) 이다.
• 마틴갈 문제 설정을 통해 SMFs가 기존 SMEs와 동일한 통계적 행동을 재현하지만, 더 높은 정규성을 갖는다는 것을 보여준다.
• 유한한 넓이 극한에서 SGD의 경험 측도 동역학을 모델링하기 위해 분포에 의존하는 SMFs를 도입하며, 계수는 해의 법칙에 따라 달라진다.
• 와서슈타인-2 거리와 혼돈 전파 원리를 사용하여, SMF 해의 경험 측도와 유한한 SGD 동역학의 경험 측도를 비교한다.
• SDE 이론, 말리아빈 미적분학, 혼돈 전파의 조합을 활용하여 기댓값과 분포 기반 의미에서 수렴 속도를 증명한다.

실험 결과

연구 질문

• RQ1 기존의 확률적 수정 방정식에서 발생하는 분산 계수의 비정규성 문제를 피하면서, SGD에 대한 새로운 SDE 기반의 제한 동역학을 구성할 수 있는가?
• RQ2 제안된 확률적 수정 유동 모델이 SGD의 다중점 통계를 유지하는가? 이를 통해 더 정확한 연속 근사가 가능한가?
• RQ3 무한한 넓이와 작은 학습률 영역에서 SGD의 변동성이 있는 평균장 근사값은 무엇이며, 이를 분포에 의존하는 SDE로 어떻게 묘사할 수 있는가?
• RQ4 유한한 SGD 동역학이 제안된 제한 SDE로 수렴하는 것을 분포 기반 및 모멘트 기반 오차 측면에서 정량화할 수 있는가?

주요 결과

• 제안된 확률적 수정 유동(SMF) SDE는 개별 기울기 차이에서 직접 도출된 정규 분산 계수를 갖는다. 이는 공분산 행렬 Σ의 불안정한 제곱근 계산이 필요 없음을 의미한다.
• SMF는 SGD의 다중점 통계를 정확히 일치시키며, 고전적 SME와 동일한 마틴갈 문제를 만족하므로 통계적 일致성을 확보한다.
• 모든 f ∈ C⁴_b(ℝᵈ)에 대해, 유한한 시간 영역 동안 SMF와 SGD 동역학 하에서 f의 기댓값은 O(η²) 이내로 일치하며, 이는 이차 정확도를 보장한다.
• 무한한 넓이와 작은 학습률 영역에서, SGD의 경험 측도는 분포에 의존하는 SMF로 수렴하며, 이는 변동성이 있는 평균장 동역학을 포착한다.
• 경험 측도가 제한적 SMF로 수렴하는 속도는 기댓값 기반으로 O(η²)이며, 초기 측도의 모멘트와 표본 수 M에 따라 추가 오차 범위가 존재한다.
• 적절한 모멘트 및 정규성 조건 하에서, 초기 측도가 유한한 경험 측도로 근사된 경우에도 수렴 속도는 O(η²)로 유지되며, M과 η에 명시적인 의존성이 존재한다.