[논문 리뷰] A mean-field theory of lazy training in two-layer neural nets: entropic regularization and controlled McKean-Vlasov dynamics.

Belinda Tzen, Maxim Raginsky|arXiv (Cornell University)|2020. 02. 05.

Stochastic Gradient Optimization Techniques참고 문헌 27인용 수 11

한 줄 요약

이 논문은 두 층 신경망에서 레이지 트레이닝을 위한 평균장 이론을 개발하며, 가중치의 확률측도 위에서 자유 에너지 최소화 문제로 학습 문제를 공식화함으로써 근사 오차와 가우시안 사전분포로부터의 KL 발산을 균형 잡는다. 이는 엔트로피 최적 운반(스체로딩어 다리 문제)을 해결하는 제어된 McKean–Vlasov 동역학을 도입하고, 레이지 영역에서의 확률적 경사하강법이 정량적인 $L^2$ 오차 한계를 갖는 채로 이 동역학을 탐욕스럽게 근사함을 보여준다.

ABSTRACT

We consider the problem of universal approximation of functions by two-layer neural nets with random weights that are nearly in the sense of Kullback-Leibler divergence. This problem is motivated by recent works on lazy training, where the weight updates generated by stochastic gradient descent do not move appreciably from the i.i.d. Gaussian initialization. We first consider the mean-field limit, where the finite population of neurons in the hidden layer is replaced by a continual ensemble, and show that our problem can be phrased as global minimization of a free-energy functional on the space of probability measures over the weights. This functional trades off the $L^2$ approximation risk against the KL divergence with respect to a centered Gaussian prior. We characterize the unique global minimizer and then construct a controlled nonlinear dynamics in the space of probability measures over weights that solves a McKean--Vlasov optimal control problem. This control problem is closely related to the Schrodinger bridge (or entropic optimal transport) problem, and its value is proportional to the minimum of the free energy. Finally, we show that SGD in the lazy training regime (which can be ensured by jointly tuning the variance of the Gaussian prior and the entropic regularization parameter) serves as a greedy approximation to the optimal McKean--Vlasov distributional dynamics and provide quantitative guarantees on the $L^2$ approximation error.

연구 동기 및 목표

• 두 층 신경망에서의 레이지 트레이닝을 평균장 근사로 형식화하여 이산적인 뉴런 집단을 가중치 위의 연속적 확률측도로 대체한다.
• 자기 중심 가우시안 사전분포로부터의 $L^2$ 근사 위험과 KL 발산을 조합한 자유 에너지 기능의 전역 최소화자로 최적의 가중치 분포를 특성화한다.
• 가중치 공간에서 최적의 분포 진동을 실현하는 제어된 McKean–Vlasov 동역학을 구성한다.
• 레이지 영역에서의 SGD와 이 최적 동역학 사이의 이론적 연결을 확립하고, 정량적인 $L^2$ 근사 오차 보장을 제공한다.

제안 방법

• 가중치 위의 확률측도 공간에서 자유 에너지 기능의 전역 최소화 문제로 평균장 근사에서의 학습 문제를 공식화한다.
• 자유 에너지 기능을 $L^2$ 근사 위험과 중심 가우시안 사전분포에 대한 Kullback-Leibler 발산 항의 합으로 정의한다.
• 변분법과 깁스 측도의 성질을 이용하여 자유 에너지의 유일한 전역 최소화자를 특성화한다.
• 최소화자 쪽으로 가중치 분포를 진동시키는 비선형 확산 과정(맥클린–브라운 동역학)을 제어 가능한 비선형 확산 과정으로 구성한다.
• 제어 문제를 슈뢰딩어 다리 문제와 연결하여, 그 값이 최소 자유 에너지와 같음을 보인다.
• 레이지 영역에서의 SGD가 이 제어된 동역학을 탐욕스럽게 근사함을 보이며, 가우시안 사전분포의 분산과 엔트로피 정규화 파라미터에 따라 의존하는 정량적 $L^2$ 오차 한계를 제공한다.

실험 결과

연구 질문

• RQ1 두 층 신경망에서의 레이지 트레이닝을 위한 평균장 근사에서 최적의 가중치 분포는 무엇인가?
• RQ2 가중치 분포의 진동 동역학은 확률측도 공간에서 제어된 McKean–Vlasov 과정으로 어떻게 모델링할 수 있는가?
• RQ3 최적 제어 문제와 엔트로피 최적 운반(슈뢰딩어 다리) 사이의 연결 고리는 무엇인가?
• RQ4 레이지 트레이닝 영역에서의 확률적 경사하강법은 최적의 McKean–Vlasov 동역학을 얼마나 잘 근사하는가?
• RQ5 제어 가능한 하이퍼파rameter 조정 하에 SGD에 대한 정량적 $L^2$ 근사 오차 한계는 무엇인가?

주요 결과

• 자유 에너지 기능의 유일한 전역 최소화자는 근사 정확도와 사전 충실도를 균형 잡는 깁스 측도이다.
• 최적의 가중치 분포는 슈뢰딩어 다리 문제를 해결하는 제어된 McKean–Vlasov 동역학을 통해 달성된다.
• 최적 제어 문제의 값은 최소 자유 에너지와 같으며, 이는 최적 해의 변분적 특성화를 제공한다.
• 레이지 영역에서의 SGD는 최적의 McKean–Vlasov 동역학을 탐욕스럽게 근사하는 데 기여한다.
• SGD의 $L^2$ 근사 오차는 정량적으로 한정되어 있으며, 가우시안 사전분포의 분산과 엔트로피 정규화 파라미터를 조정함으로써 제어할 수 있다.