[논문 리뷰] Rigorous dynamical mean field theory for stochastic gradient descent methods

Cédric Gerbelot, Emanuele Troiani외 3명|arXiv (Cornell University)|2022. 10. 12.

Stochastic Gradient Optimization Techniques인용 수 6

한 줄 요약

이 논문은 고차원 설정에서 이산시간 스토하스틱 경사하강법(SGD) 및 관련 1차 최적화 방법에 대해 엄밀한 동적 평균장 이론(DMFT) 방정식을 수립한다. 일반적인 공분산 행렬을 가진 가우시안 데이터에서 SGD의 점근적 역학이 정확히 이산 DMFT 방정식으로 기술됨을 증명하며, 이는 자기일관성 있는 커널을 통해 기억 효과를 반영하고 자석화 및 노름 역학을 통한 일반화 성능의 정밀한 예측을 가능하게 한다.

ABSTRACT

We prove closed-form equations for the exact high-dimensional asymptotics of a family of first order gradient-based methods, learning an estimator (e.g. M-estimator, shallow neural network, ...) from observations on Gaussian data with empirical risk minimization. This includes widely used algorithms such as stochastic gradient descent (SGD) or Nesterov acceleration. The obtained equations match those resulting from the discretization of dynamical mean-field theory (DMFT) equations from statistical physics when applied to gradient flow. Our proof method allows us to give an explicit description of how memory kernels build up in the effective dynamics, and to include non-separable update functions, allowing datasets with non-identity covariance matrices. Finally, we provide numerical implementations of the equations for SGD with generic extensive batch-size and with constant learning rates.

연구 동기 및 목표

• 이산시간 스토하스틱 경사하강법(SGD) 및 관련 최적화 알고리즘에 대한 동적 평균장 이론(DMFT) 방정식의 엄밀한 유도를 제공한다.
• 이전의 i.i.d. 서브가우시안 데이터에 한정된 DMFT 결과를 일반적인 가우시안 데이터, 즉 임의의 공분산 행렬을 가진 데이터로 확장한다.
• 효과적인 확률적 적분미분방정식에 포함된 명시적 기억 커널을 유도하여 역학에서의 기억 효과를 고려한다.
• 고차원 극한에서 자석화 및 노름 역학을 통해 일반화 성능(예: 추정치와 진짜값 사이의余弦 유사도)의 정확한 예측을 가능하게 한다.
• 광범위한 배치 크기와 일정한 학습률을 가진 SGD에 대해 이산 DMFT 프레임워크를 수치적으로 검증한다.

제안 방법

• 반복적인 가우시안 조건부 기대와 섭동 이론을 사용하여 SGD에서의 가중치 갱신의 무한차원 극한을 분석함으로써 이산시간 DMFT 방정식을 도출한다.
• 추정치의 가중치 벡터와 잔차 오차의 효과적 역학을 기술하는 자기일관성 있는 확률적 적분미분방정식 시스템을 도입한다.
• 과거 기울기들에 대한 손실 함수의 함수적 도함수를 통해 기억 커널을 계산하여 역학에서의 비-마르코프 효과를 캡처한다.
• 기준 경로를 중심으로 선형 차수 섭동 전개를 사용하여 효과적 역학을 유도하며, 이는 고차원 극한에서 가우시안 과정로 수렴하는 항들을 포함한다.
• 자석화(진짜값과의 유사도)와 가중치 노름에 대한 폐쇄형 방정식을 유도하여 일반화 성능의 예측을 가능하게 한다.
• 일정한 학습률과 광범위한 배치 크기를 가진 SGD에 대해 유도된 DMFT 방정식을 수치적으로 구현하여 점근적 예측을 검증한다.

실험 결과

연구 질문

• RQ1 비.i.i.d. 데이터, 특히 항등행렬이 아닌 공분산 행렬을 가진 경우에도 이산시간 SGD에 대해 엄밀한 DMFT 방정식을 도출할 수 있는가?
• RQ2 기울기의 역사를 통해 유도되는 기억 효과가 고차원에서 SGD의 점근적 역학에 어떻게 나타나는가?
• RQ3 이산 DMFT 방정식이 고차원 극한에서 SGD의 일반화 성능을 얼마나 정확하게 기술하는가?
• RQ4 DMFT 프레임워크는 i.i.d. 서브가우시안 설계를 초월하여 비분리된 업데이트 함수와 비.i.i.d. 데이터를 포함하도록 확장될 수 있는가?
• RQ5 이산시간 SGD 역학과 연속시간 경사 흐름 간의 정확한 관계는 무엇인가?

주요 결과

• 논문은 이산시간 DMFT 방정식이 비.i.i.d. 가우시안 데이터를 포함한 고차원 극한에서 SGD 역학을 정확히 기술함을 증명한다.
• 기억 커널은 손실 함수의 기능적 도함수로부터 자연스럽게 유도되며, 명시적으로 계산되어 과거 기울기의 영향이 현재 갱신에 미치는 영향을 캡처한다.
• 일반화 성능은 자석화 $ m^t = \text{corr}(\bm{w}^*, \bm{w}^t) $ 로 완전히 특징지어지며, 이는 손실 도함수와 노이즈 공분산을 포함한 자기일관성 방정식에 따라 진화한다.
• 효과적 역학은 기억 커널 $ R_g(t,t') $, 공분산 $ C_g(t,t') $ 를 가진 노이즈 과정 $ u^t $, 그리고 자기정규화 항 $ \tilde{\nu}^t $ 를 포함하는 확률적 적분미분방정식 시스템으로 기술된다.
• 고차원 설정에서 추정치의 노름과 진짜값과의 여론 유사도의 진화를 수치적으로 구현한 DMFT 방정식이 정확하게 예측한다.
• 이 프레임워크는 경사하강법, 네스터로프 가속, 모멘텀 기반 알고리즘을 포함한 다양한 1차 방법에 대해 경험적 리스크 최소화 하에 적용 가능하다.