[논문 리뷰] Function approximation by neural nets in the mean-field regime: Entropic regularization and controlled McKean-Vlasov dynamics

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

Stochastic Gradient Optimization Techniques인용 수 4

한 줄 요약

이 논문은 두 계층 신경망에서 러닝 트레이닝(lazy training)의 평균장 이론을 수립한다. 함수 근사 문제를 $L^2$ 근사 위험과 가우시안 사전분포로부터의 KL 발산을 균형 잡는 자유 에너지 최소화 문제로 공식화한다. 엔트로피 정규화를 통한 통제 가능한 McKean–Vlasov 동역학을 도입하여, SGD가 최적의 분포 동역학을 정량적 $L^2$ 오차 한계로 근사함을 보인다.

ABSTRACT

We consider the problem of function approximation by two-layer neural nets with random weights that are "nearly Gaussian" in the sense of Kullback-Leibler divergence. Our setting is the mean-field limit, where the finite population of neurons in the hidden layer is replaced by a continuous ensemble. We show that the problem can be phrased as global minimization of a free energy functional on the space of (finite-length) paths over probability measures on the weights. This functional trades off the $L^2$ approximation risk of the terminal measure against the KL divergence of the path with respect to an isotropic Brownian motion prior. We characterize the unique global minimizer and examine the dynamics in the space of probability measures over weights that can achieve it. In particular, we show that the optimal path-space measure corresponds to the Föllmer drift, the solution to a McKean-Vlasov optimal control problem closely related to the classic Schrödinger bridge problem. While the Föllmer drift cannot in general be obtained in closed form, thus limiting its potential algorithmic utility, we illustrate the viability of the mean-field Langevin diffusion as a finite-time approximation under various conditions on entropic regularization. Specifically, we show that it closely tracks the Föllmer drift when the regularization is such that the minimizing density is log-concave.

연구 동기 및 목표

• 두 계층 신경망에서 러닝 트레이닝 영역에서 함수 근사에 대한 평균장 이론을 개발하기 위해.
• 최적의 가중치 분포를 $L^2$ 위험과 가우시안 사전분포로부터의 KL 발산을 조합한 자유 에너지 기능의 최소화자로 특성화하기 위해.
• 엔트로피 정규화를 통해 최적의 분포 동역학을 실현하는 통제 가능한 McKean–Vlasov 확률미분방정식을 구성하기 위해.
• SGD가 최적의 평균장 동역학에 대한 탐욕적 근사로 작용할 때의 $L^2$ 근사 오차에 대한 정량적 보장을 확립하기 위해.
• 최적 제어 문제를 슈뢰딩거 다리 문제와 Föllmer의 드리프트에 연결하여 평균장 극한의 동적 해석을 가능하게 하기 위해.

제안 방법

• 자유 에너지 기능의 전역 최소화 문제로 평균장 근사 문제를 공식화: $\mathcal{F}_{\tau,\beta}(\mu) = \mathbb{E}_{w \sim \mu}[\|f(w,\cdot) - f^*\|_{L^2}^2] + \beta \cdot D_{\mathrm{KL}}(\mu \| \gamma_\tau)$, 여기서 $\gamma_\tau$ 는 중심이 0인 가우시안 사전분포이다.
• 정규성 조건 하에서 자유 에너지 기능의 전역 최소화자의 존재성과 유일성을 증명하며, $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ 와 $\mathcal{W}_2(\mu^*, \gamma_\tau)$ 에 대한 명시적 상한을 제공한다.
• 종단 비용이 $\beta$ 와 $\tau$ 를 모두 포함하는 스토크라스틱 최적 제어 문제를 해결하는 드리프트 $\varphi(w,t)$ 를 갖는 통제 가능한 McKean–Vlasov SDE를 구성한다.
• 최적 제어 법칙이 Föllmer 드리프트임을 식별하며, 이는 엔트로피 최적 운반(스튀뢰딩거 브리지) 문제의 해와 대응한다.
• 가중치 분포와 가치 함수의 진화를 기술하는 결합된 정방향(Fokker–Planck) 및 역방향(Hamilton–Jacobi–Bellman) PDE를 유도한다.
• 적절하게 조정된 학습률 $\eta$, 네트워크 너비 $N$, 정규화 $\beta$ 를 갖는 SGD가, 고확률적으로 $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$ 의 오차로 최적의 McKean–Vlasov 동역학을 근사함을 보인다.

실험 결과

연구 질문

• RQ1 두 계층 신경망에서의 러닝 트레이닝 영역은 어떻게 평균장 자유 에너지 최소화 프레임워크를 통해 형식화할 수 있는가?
• RQ2 근사 오차와 가우시안 사전분포로부터의 이탈 간의 트레이드오프를 최소화하는 최적의 가중치 분포의 구조는 어떠한가?
• RQ3 최적의 분포 동역학은 통제 가능한 McKean–Vlasov 과정으로 실현될 수 있으며, 엔트로피 최적 운반과는 어떤 관계가 있는가?
• RQ4 스토크라스틱 경사하강법(SGD)은 러닝 트레이닝 영역에서 최적의 평균장 동역학을 어느 정도 근사하는가?
• RQ5 최적의 분포 동역학을 근사할 때 SGD의 유한 샘플 및 고확률 오차 한계는 무엇인가?

주요 결과

• 자유 에너지 기능은 전역 최소화자 $\mu^*$ 를 유일하게 가지며, KL 발산 $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ 와 $L^2$ 워샤르슈타인 거리 $\mathcal{W}_2(\mu^*, \gamma_\tau)$ 는 각각 $O(\beta)$ 와 $O(\sqrt{\tau D})$ 로 유계이다.
• McKean–Vlasov 동역학의 최적 제어 법칙은 Föllmer 드리프트이며, 이는 슈뢰딩거 브리지 문제의 해이자 종단 분포가 최소화자 $\mu^*$ 와 일치하도록 보장한다.
• 통제 가능한 McKean–Vlasov 동역학은 Fokker–Planck 및 Hamilton–Jacobi–Bellman PDE의 결합된 시스템에 의해 지배되며, 가치 함수는 최소 자유 에너지를 코딩한다.
• SGD는 $N$ 개의 뉴런을 갖는 경우 시간 $t \in [0,1]$ 동안 고확률 $1-\delta$ 로 최적의 평균장 동역학을 근사하며, $L^2$ 오차 한계 $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$ 를 달성한다.
• 근사 오차는 $N$ 과 $\delta$ 에 대해 유리하게 스케일링되며, 가우시안 사전분포의 분산 $\tau$ 와 정규화 파라미터 $\beta$ 를 조정함으로써 작게 만들 수 있다.
• 이 분석은 러닝 트레이닝, 엔트로피 정규화, McKean–Vlasov 동역학 간의 엄밀한 연결을 제공하며, 신경미분핵 및 평균장 극한에 관한 이전 연구를 확장한다.