MLPGradientFlow: going with the flow of multilayer perceptrons (and finding minima fast and accurately)

TL;DR

MLPGradientFlow introduces a software package that numerically solves gradient flow equations for multilayer perceptrons, demonstrating superior accuracy and speed over traditional optimizers like Adam, especially with Newton's method and BFGS.

Contribution

The paper presents a new software package for solving gradient flow equations in neural networks, showing improved convergence and efficiency over standard optimization methods.

Findings

Adaptive Runge-Kutta methods outperform gradient descent in accuracy and speed.

Newton's method and BFGS are more effective for finding minima.

Gradients and Hessians are computed faster than in PyTorch for small networks.

Abstract

MLPGradientFlow is a software package to solve numerically the gradient flow differential equation $\dot \theta = -\nabla \mathcal L(\theta; \mathcal D)$, where $\theta$ 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.