MLPGradientFlow: going with the flow of multilayer perceptrons (and finding minima fast and accurately)
This work provides a tool for researchers and practitioners in machine learning to optimize neural networks more efficiently, though it is incremental as it builds on existing numerical methods and focuses on specific network types.
The authors tackled the problem of efficiently and accurately finding minima in multilayer perceptrons by developing MLPGradientFlow, a software package that numerically solves gradient flow differential equations, showing that adaptive integration methods outperform Adam in accuracy and speed, and that Newton's method is more effective for finding fixed points, with gradients computed up to 5x faster than PyTorch for small networks.
MLPGradientFlow is a software package to solve numerically the gradient flow differential equation $\dot θ= -\nabla \mathcal L(θ; \mathcal D)$, where $θ$ are the parameters of a multi-layer perceptron, $\mathcal D$ is some data set, and $\nabla \mathcal L$ is the gradient of a loss function. We show numerically that adaptive first- or higher-order integration methods based on Runge-Kutta schemes have better accuracy and convergence speed than gradient descent with the Adam optimizer. However, we find Newton's method and approximations like BFGS preferable to find fixed points (local and global minima of $\mathcal L$) efficiently and accurately. For small networks and data sets, gradients are usually computed faster than in pytorch and Hessian are computed at least $5\times$ faster. Additionally, the package features an integrator for a teacher-student setup with bias-free, two-layer networks trained with standard Gaussian input in the limit of infinite data. The code is accessible at https://github.com/jbrea/MLPGradientFlow.jl.