[论文解读] Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks

Kaitong Hu, Zhenjie Ren+2 more|arXiv (Cornell University)|May 19, 2019

Markov Chains and Monte Carlo Methods参考文献 58被引用 30

一句话总结

本文引入了平均场朗之万动力学（MFLD），将其视为在概率测度的2-沃瑟斯坦空间中的连续时间梯度流，证明该过程的分布以指数速度收敛到唯一平稳分布，该分布使能量泛函最小化。收敛性通过拉萨勒不变性原理与HWI不等式的创新应用得以证明，无需对称或卷积型相互作用势，且建立了有限维与无限维优化问题之间的O(1/N)误差界。

ABSTRACT

Our work is motivated by a desire to study the theoretical underpinning for the convergence of stochastic gradient type algorithms widely used for non-convex learning tasks such as training of neural networks. The key insight, already observed in the works of Mei, Montanari and Nguyen (2018), Chizat and Bach (2018) as well as Rotskoff and Vanden-Eijnden (2018), is that a certain class of the finite-dimensional non-convex problems becomes convex when lifted to infinite-dimensional space of measures. We leverage this observation and show that the corresponding energy functional defined on the space of probability measures has a unique minimiser which can be characterised by a first-order condition using the notion of linear functional derivative. Next, we study the corresponding gradient flow structure in 2-Wasserstein metric, which we call Mean-Field Langevin Dynamics (MFLD), and show that the flow of marginal laws induced by the gradient flow converges to a stationary distribution, which is exactly the minimiser of the energy functional. We observe that this convergence is exponential under conditions that are satisfied for highly regularised learning tasks. Our proof of convergence to stationary probability measure is novel and it relies on a generalisation of LaSalle's invariance principle combined with HWI inequality. Importantly, we assume neither that interaction potential of MFLD is of convolution type nor that it has any particular symmetric structure. Furthermore, we allow for the general convex objective function, unlike, most papers in the literature that focus on quadratic loss. Finally, we show that the error between finite-dimensional optimisation problem and its infinite-dimensional limit is of order one over the number of parameters.

研究动机与目标

• 为非凸学习任务中随机梯度类算法的收敛性提供理论基础，特别是针对深度神经网络的训练。
• 通过将有限维非凸问题提升至概率测度的无限维空间，分析神经网络的能量景观。
• 利用线性泛函导数，建立能量泛函最小化解的存在性与唯一性。
• 在较弱正则性条件下，证明MFLD过程可指数收敛至对应全局最小化解的平稳分布。
• 量化有限维优化与其中 mean-field 极限之间的近似误差，其量级为 O(1/N)，其中 N 为参数数量。

提出的方法

• 将训练神经网络的有限维非凸优化问题，提升至概率测度空间上的无限维问题。
• 在概率测度空间上定义能量泛函，并通过线性泛函导数的首阶条件刻画其唯一最小化解。
• 将平均场朗之万动力学（MFLD）定义为2-沃瑟斯坦度量下的梯度流，以描述系统分布的演化。
• 应用广义拉萨勒不变性原理，证明其边缘分布收敛至平稳分布。
• 利用HWI不等式，在损失函数与势函数满足正则性假设下，建立指数收敛速率。
• 推导出有限维优化与其中 mean-field 极限之间的 O(1/N) 误差界，适用于一般凸目标函数，不仅限于二次型。

实验结果

研究问题

• RQ1 能否通过 mean-field 极限严格证明在过参数化神经网络中随机梯度下降的收敛性？
• RQ2 在概率测度空间中的能量泛函是否具有唯一最小化解，且能否通过泛函导数刻画？
• RQ3 在何种条件下，平均场朗之万动力学可指数快速收敛至全局最小化解？
• RQ4 有限维训练与 mean-field 极限之间的近似误差如何随参数数量变化？
• RQ5 收敛性证明是否可不假设相互作用势具有对称性或卷积结构？

主要发现

• 定义在概率测度空间上的能量泛函具有唯一最小化解，其由包含线性泛函导数的首阶条件刻画。
• 平均场朗之万动力学（MFLD）可指数快速收敛至对应能量泛函全局最小化解的平稳分布。
• 收敛性证明依赖于拉萨勒不变性原理与HWI不等式的创新结合，且在较弱正则性条件下成立。
• 有限维优化问题与其无限维 mean-field 极限之间的误差界为 O(1/N)，其中 N 为参数数量。
• 结果适用于一般凸目标函数，不限于二次损失，且不要求相互作用势为卷积型或对称。
• MFLD 的平稳分布恰好是能量泛函的最小化解，从而在动力学与解的最优性之间建立了直接联系。