How Does Gradient Descent Learn Features -- A Local Analysis for Regularized Two-Layer Neural Networks
This provides theoretical insights into feature learning mechanisms in neural networks, which is an incremental contribution to the existing literature on training dynamics.
The paper tackles the problem of understanding how gradient descent learns features in two-layer neural networks, showing that feature learning can occur both at the initial steps and towards the end of training through a local convergence analysis with a regularized objective.
The ability of learning useful features is one of the major advantages of neural networks. Although recent works show that neural network can operate in a neural tangent kernel (NTK) regime that does not allow feature learning, many works also demonstrate the potential for neural networks to go beyond NTK regime and perform feature learning. Recently, a line of work highlighted the feature learning capabilities of the early stages of gradient-based training. In this paper we consider another mechanism for feature learning via gradient descent through a local convergence analysis. We show that once the loss is below a certain threshold, gradient descent with a carefully regularized objective will capture ground-truth directions. We further strengthen this local convergence analysis by incorporating early-stage feature learning analysis. Our results demonstrate that feature learning not only happens at the initial gradient steps, but can also occur towards the end of training.