[论文解读] Function approximation by neural nets in the mean-field regime: Entropic regularization and controlled McKean-Vlasov dynamics

Belinda Tzen, Maxim Raginsky|arXiv (Cornell University)|Feb 5, 2020

Stochastic Gradient Optimization Techniques被引用 4

一句话总结

本文通过将函数逼近问题形式化为自由能最小化问题，建立了两层神经网络中懒惰训练的平均场理论，该问题在 $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$ 逼近误差保证。
• 将最优控制问题与 Schrödinger 桥问题及 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$ 为中心的高斯先验。
• 在正则性条件下证明自由能泛函的全局最小解的存在性与唯一性，并给出 $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ 与 $\mathcal{W}_2(\mu^*, \gamma_\tau)$ 的显式界。
• 构造一个具有漂移 $\varphi(w,t)$ 的受控 McKean–Vlasov SDE，其求解一个随机最优控制问题，其中终端代价编码了 $\beta$ 与 $\tau$。
• 识别最优控制律为 Föllmer 流，其对应于熵最优传输（Schrödinger 桥）问题的解。
• 推导出控制权重分布演化的前向（Fokker–Planck）与后向（Hamilton–Jacobi–Bellman）PDE 联立方程。
• 证明当学习率 $\eta$、网络宽度 $N$ 与正则化参数 $\beta$ 调节得当时，SGD 以高概率近似最优 McKean–Vlasov 动力学，误差为 $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)$。

实验结果

研究问题

• RQ1 如何通过平均场自由能最小化框架形式化两层神经网络中的懒惰训练范式？
• RQ2 在逼近误差与偏离高斯先验之间权衡下，最优权重分布的结构是什么？
• RQ3 最优分布动力学能否通过受控的 McKean–Vlasov 过程实现？其与熵最优传输有何关联？
• RQ4 在懒惰训练范式下，随机梯度下降（SGD）在多大程度上近似最优平均场动力学？
• RQ5 SGD 在近似最优分布动力学时，其有限样本与高概率误差界是什么？

主要发现

• 自由能泛函具有唯一的全局最小解 $\mu^*$，且 KL 散度 $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ 与 $L^2$ Wasserstein 距离 $\mathcal{W}_2(\mu^*, \gamma_\tau)$ 分别有 $O(\beta)$ 与 $O(\sqrt{\tau D})$ 的界。
• McKean–Vlasov 动力学的最优控制律为 Föllmer 流，其求解 Schrödinger 桥问题，并确保终端分布匹配最小化器 $\mu^*$。
• 受控 McKean–Vlasov 动力学由 Fokker–Planck 与 Hamilton–Jacobi–Bellman PDE 联立系统控制，其中值函数编码了最小自由能。
• SGD 在 $N$ 个神经元下，以高概率 $1-\delta$ 在时间 $t \in [0,1]$ 内近似最优平均场动力学，实现 $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 动力学之间建立了严格的联系，拓展了神经正切核与平均场极限的既有研究。