Reproduce “Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent” on OpenResearch

About this paper

Behrad Moniri, Hamed Hassani

This work by researchers at the University of Pennsylvania analyzes feature learning in linear-width two-layer neural networks after two steps of gradient descent, demonstrating that it enables the learning of multiple feature directions. The study shows that reusing training data across steps is crucial for capturing genuinely non-linear features, with specific step-size scalings determining the number of features learned.

Key insights

• A second step of gradient descent can induce the learning of multiple feature directions, quantitatively determined by the interplay of the two gradient step-sizes.
• Reusing training data (reused batch setting) is critical for enabling the model to learn genuinely non-linear features from target functions that lack a linear component.
• The spectrum of the first-layer weight matrix evolves to display multiple outlying singular values, each corresponding to a distinct learned feature direction.

Summary by alphaXiv.