Shallow mlps follow a simple feature ode when learning basic functions
A "feature ODE" describing the learning behavior of shallow MLPs on simple functions
Machine LearningArtificial Intelligence
Summary
Training two-layer neural networks to learn certain simple functions is usually very complex to understand mathematically. The authors propose a simpler way to describe how the network learns, using a special equation they call a "feature ODE." Surprisingly, even though the equation is a big simplification, the network's learning patterns closely match what this equation predicts. The authors talk about why this might happen and how it could be useful for understanding neural network training.
What this means in practice
- •For machine learning engineers: Model training can be analyzed more simply using the feature ODE to predict learning progress on simple function tasks.
- •For data scientists: The feature ODE provides an intuitive tool to understand neural network behavior on structured Gaussian data inputs.
Authors
Joseph Turnbull, Berkan Ottlik, Jamie Simon
Abstract
We study the gradient flow training dynamics of two-layer multi-layer perceptrons (MLPs) on anisotropic Gaussian data. Even in this simple setting, solving the dynamics in full is notoriously difficult, so we try something different. We propose a "feature ODE" that might be intuitively expected to describe the learning dynamics of a shallow MLP. We find that MLPs trained to learn simple target functions (low-order Hermite polynomials and staircase functions) actually roughly follow the loss trajectories from this feature ODE. This agreement is a surprise because of the size of the simplification and the fact that it does not come from a rigorous chain of approximations. We discuss this surprise and speculate on its reasons and uses.