Dichotomy of Feature Learning and Unlearning: Fast-Slow Analysis on Neural Networks with Stochastic Gradient Descent

The dynamics of gradient-based training in neural networks often exhibit nontrivial structures; hence, understanding them remains a central challenge in theoretical machine learning. In particular, the concept of feature unlearning, in which a neural network progressively loses previously learned features over long training, has gained attention. In this study, we consider the infinite-width limit of a two-layer neural network trained with a large-batch stochastic gradient, then derive differential equations with different time scales, revealing the mechanism and conditions for feature unlearning to occur. Specifically, we utilize the fast-slow dynamics: while an alignment of first-layer weights develops rapidly, the second-layer weights develop slowly. The direction of the flow on a critical manifold, determined by the slow dynamics, decides whether feature unlearning occurs. We give numerical validation of the result and derive theoretical grounding and scaling laws for the feature unlearning. Our results yield the following insights: (i) the strength of the primary nonlinear term in the data induces the feature unlearning, and (ii) an initial scale of the second-layer weights mitigates the feature unlearning. Our result should be understood as a population loss of alignment rather than finite-sample overfitting. Technically, our analysis utilizes Tensor Programs and singular perturbation theory.

Neural networks are often described as gradually learning better features from data. However, training can be more surprising: a network may first discover a useful feature and later lose it as training continues. This paper studies this phenomenon, called feature unlearning, in a simplified but mathematically analyzable setting. We consider a two-layer neural network trained with stochastic gradient descent on data generated from a hidden rule. Although the network has many parameters, in a large-system limit its training behavior can be summarized by two key quantities: how strongly the first layer is aligned with the true feature in the data, and the scale of the second-layer weights. Our analysis shows that training proceeds on two different time scales. In the early phase, the first layer rapidly becomes aligned with the true feature, while the second layer changes only slowly. After this fast phase, the system moves along a special curve in the two-dimensional description of training. The direction of this slow movement determines the final outcome. In some cases, the learned feature is preserved; in others, the alignment gradually decays, meaning that the network “unlearns” the feature it had previously acquired. The paper identifies mathematical conditions that predict when feature unlearning occurs and derives scaling laws describing how fast this loss of alignment happens. Numerical simulations of both the limiting equations and finite neural networks support the theory. Overall, the work shows that feature learning is not necessarily permanent: even in a clean population-level setting, the geometry of the training dynamics can cause a neural network to acquire a useful representation early and lose it later.