[论文解读] 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)|Feb 5, 2020

Stochastic Gradient Optimization Techniques参考文献 27被引用 11

一句话总结

该论文通过将学习问题表述为在权重概率测度上的自由能最小化问题（平衡近似误差与相对于高斯先验的KL散度），为两层神经网络中的懒惰训练建立了一套平均场理论。该理论引入了一种受控的McKean–Vlasov动力学，其求解一个熵最优传输（Schrödinger桥）问题，并表明在懒惰区间内，随机梯度下降（SGD）以可量化的方式贪婪逼近该动力学，且具有$ 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 $近似风险与相对于中心高斯先验的KL散度项之和。
• 利用变分法和Gibbs测度的性质，表征自由能的唯一全局最小化器。
• 构建一种受控的非线性扩散过程（McKean–Vlasov动力学），使其在权重分布上向最小化器演化，该过程被构架为最优控制问题。
• 将控制问题与Schrödinger桥问题关联，证明其值等于最小自由能。
• 证明在懒惰区间内，SGD对受控动力学的逼近是贪婪的，且其$ L^2 $误差界可被证明，其依赖于先验的方差与熵正则化参数。

实验结果

研究问题

• RQ1 在两层神经网络的懒惰训练平均场极限下，最优权重分布是什么？
• RQ2 如何在概率测度空间中将权重演化动力学建模为受控的McKean–Vlasov过程？
• RQ3 最优控制问题与熵最优传输（Schrödinger桥）之间有何联系？
• RQ4 在懒惰训练区间内，随机梯度下降对最优McKean–Vlasov动力学的逼近程度如何？
• RQ5 在受控超参数调优下，能否为SGD推导出定量的$ L^2 $近似误差界？

主要发现

• 自由能泛函的唯一全局最小化器是一个Gibbs测度，其在近似精度与先验保真度之间实现平衡。
• 通过求解Schrödinger桥问题的受控McKean–Vlasov动力学，可实现最优权重分布。
• 最优控制问题的值等于最小自由能，从而为最优解提供了变分表征。
• 在懒惰区间内，SGD作为对最优McKean–Vlasov动力学的贪婪逼近。
• SGD的$ L^2 $近似误差可被定量界定，并可通过调节高斯先验的方差与熵正则化参数进行控制。