[논문 리뷰] High-dimensional limit theorems for SGD: Effective dynamics and critical scaling

Gérard Ben Arous, Reza Gheissari외 1명|arXiv (Cornell University)|2022. 06. 08.

Markov Chains and Monte Carlo Methods인용 수 7

한 줄 요약

이 논문은 일정 학습률을 가진 확률적 경사하강법(SGD)에 대해 고차원 극한 정리들을 수립하며, SGD의 요약 통계량이 초기화 및 단계 크기 스케일링에 따라 상미분방정식(ODE) 또는 확률미분방정식(SDE)으로 이끄는 효과적인 역학으로 수렴함을 보여준다. 기울기 흐름이 새로운 항에 의해 수정되는 중요한 스케일링 영역이 특정화되며, 이는 비자명한 단계도형과 양의 확률로 최적해가 아닌 해로 수렴하는 결과를 초래한다. 이는 과다매개변수화에 의해 완화된다.

ABSTRACT

We study the scaling limits of stochastic gradient descent (SGD) with constant step-size in the high-dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite-dimensional functions) of SGD as the dimension goes to infinity. Our approach allows one to choose the summary statistics that are tracked, the initialization, and the step-size. It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices. We show a critical scaling regime for the step-size, below which the effective ballistic dynamics matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate. We demonstrate our approach on popular examples including estimation for spiked matrix and tensor models and classification via two-layer networks for binary and XOR-type Gaussian mixture models. These examples exhibit surprising phenomena including multimodal timescales to convergence as well as convergence to sub-optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations. At the same time, we demonstrate the benefit of overparametrization by showing that the latter probability goes to zero as the second layer width grows.

연구 동기 및 목표

• 일정 학습률을 가진 고차원 환경에서 SGD의 스케일링 극한을 위한 통합 프레임워크를 개발하는 것.
• 초기화, 단계 크기 스케일링, 매개변수 영역 선택이 요약 통계량의 효과적 역학에 미치는 영향을 규명하는 것.
• 기울기 흐름이 새로운 항에 의해 수정되는 중요한 스케일링 영역을 특정화하는 것.
• 과다매개변수화가 최적해가 아닌 해로 수렴하는 것에 미치는 영향을 분석하는 것.
• 스피iked 행렬/텐서 모델과 가우시안 혼합 분류를 위한 두 층 신경망에서 이론을 적용하여 검증하는 것.

제안 방법

• 저자들은 유한차원 요약 통계량의 차원이 발산할 때 미약한 정규성 조건 하에서 극한 정리를 유도한다.
• 초기화 및 단계 크기 스케일링에 따라 궤적의 효과적 역학이 궤도상의 ODE 또는 SDE로 수렴함을 확립한다.
• 요약 통계량 궤적에 대한 기능적 중심극한정리에 의해 효과적 역학을 도출한다.
• 단계 크기가 $ O(1/ ext{dimension}) $로 스케일링되는 중요한 스케일링 영역이 특정화되며, 이 경우 ODE 극한에서 수정 항이 나타난다.
• 손실, 가중치 진폭, 진짜값과의 상관관계와 같은 특정 통계량의 추적을 가능하게 한다.
• 고차원 설정에서 가우시안 근사와 모멘트 경계를 활용하여, ReLU 및 시그모이드 활성화 함수에 대해서도 분석한다.

실험 결과

연구 질문

• RQ1 일정 단계 크기로 고차원 극한에서 SGD의 요약 통계량 궤적이 어떻게 수렴하는가?
• RQ2 효과적 역학이 궤도상(OED)인지 확산상(SDE)인지 결정하는 요소는 무엇인가?
• RQ3 단계 크기 스케일링이 효과적 역학과 수렴 행동을 결정하는 데 어떤 역할을 하는가?
• RQ4 비볼록 문제에서 무작위 초기화가 최적해가 아닌 해로 수렴하는 데 미치는 영향은 무엇인가?
• RQ5 과다매개변수화가 최적해가 아닌 해로 수렴할 확률을 제거할 수 있는가?

주요 결과

• 요약 통계량의 효과적 역학은 초기화 및 단계 크기 스케일링에 따라 ODE 또는 SDE로 수렴하며, 기울기 흐름이 새로운 수정 항이 포함된 중요한 스케일링 영역이 존재한다.
• 비판적 스케일링 영역에서 ODE 극한은 인구 손실에 대한 기울기 흐름과 일치하지만, 수정 항이 포함되어 단계도형을 변화시키며 최적해가 아닌 고정점으로 수렴할 수 있다.
• XOR 유형의 가우시안 혼합 모델에 대해 두 층 신경망에서 무작위 초기화를 사용한 SGD는 최적해가 아닌 해로 수렴할 확률이 0이 아닌 상한선을 가진다.
• 두 번째 층의 과다매개변수화는 최적해가 아닌 해로 수렴할 확률을 감소시키며, 두 번째 층의 너비가 증가함에 따라 이 확률은 0으로 수렴한다.
• 이 방법은 스피iked 행렬 및 텐서 모델에 적용 가능하며, 수렴에 대한 다중 모드 시간스케일과 복잡한 확산 극한을 드러낸다.
• 분석 결과, 확산 극한은 탈선할 수 있으며, 수렴 행동은 요약 통계량과 초기화의 선택에 의해 결정된다.