Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shota Imai, Sota Nishiyama, Masaaki Imaizumi

TL;DR

This work analyzes gradient-based neural network training under SGD from a high-dimensional, infinite-width perspective. By deriving a two-dimensional ODE for macroscopic variables and revealing a fast-slow decomposition, the authors show that feature unlearning emerges from slow dynamics along a critical manifold, with a precise scaling law for the decay of alignment and growth of second-layer weights. They ground the analysis with Tensor Programs and singular perturbation theory, validate it numerically, and corroborate it with experiments on real networks. The findings quantify when unlearning occurs, relate it to the data-generating nonlinearity and initialization, and illuminate how learning can transition into a lazy regime despite ongoing optimization. The results have implications for understanding long-term feature retention, SGD dynamics, and stability of learned representations in deep networks.

Abstract

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, a 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 updated 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 a 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 of the feature unlearning. Our results yield the following insights: (i) the strength of the primary nonlinear term in data induces the feature unlearning, and (ii) an initial scale of the second-layer weights mitigates the feature unlearning. Technically, our analysis utilizes Tensor Programs and the singular perturbation theory.

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

TL;DR

This work analyzes gradient-based neural network training under SGD from a high-dimensional, infinite-width perspective. By deriving a two-dimensional ODE for macroscopic variables and revealing a fast-slow decomposition, the authors show that feature unlearning emerges from slow dynamics along a critical manifold, with a precise scaling law for the decay of alignment and growth of second-layer weights. They ground the analysis with Tensor Programs and singular perturbation theory, validate it numerically, and corroborate it with experiments on real networks. The findings quantify when unlearning occurs, relate it to the data-generating nonlinearity and initialization, and illuminate how learning can transition into a lazy regime despite ongoing optimization. The results have implications for understanding long-term feature retention, SGD dynamics, and stability of learned representations in deep networks.

Abstract

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, a 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 updated 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 a 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 of the feature unlearning. Our results yield the following insights: (i) the strength of the primary nonlinear term in data induces the feature unlearning, and (ii) an initial scale of the second-layer weights mitigates the feature unlearning. Technically, our analysis utilizes Tensor Programs and the singular perturbation theory.
Paper Structure (42 sections, 12 theorems, 104 equations, 8 figures)

This paper contains 42 sections, 12 theorems, 104 equations, 8 figures.

Table of Contents

  1. Introduction
  2. Notation
  3. Setup
  4. Online supervised learning
  5. Conditions
  6. Teacher model
  7. Student model and training process
  8. Preparation: Reduction to fast-slow dynamics
  9. ODE of macroscopic variables
  10. Macroscopic variables of neural networks
  11. ODE with infinite-width limit
  12. Validation of ODE
  13. Fast-slow reformulation
  14. Empirical observation of multi time scale
  15. Fast-Slow Ansatz
  16. ...and 27 more sections

Key Result

Proposition 1

Let $R_{\tau,i}^m := R_i^m(\lfloor m \tau / \gamma \rfloor), a_{\tau,i}^m := a_i^m(\lfloor m \tau / \gamma\rfloor)$. Then, for any finite $\tau \ge 0$ and $i = 1,...,m$, $R_\tau, a_\tau$ satisfying the ODE eq:original satisfies the following asymptotic equalities

Figures (8)

  • Figure 1: Fast-slow dynamics of first-layer alignment ($R_\tau$) and second-layer weights ($a_\tau$) in time $\tau$, explaining the evolution of alignment and test loss. In the space of $R_\tau$ and $a_\tau$, the $R-a$ space in the left panel, we find a critical manifold (green curve). Each trajectory starts from its initial point $(0,\bar{a})$ and, after reaching the manifold, slowly evolves along it. On all panels, the red trajectory represents feature learning, while the blue trajectory represents feature unlearning. The solid line represents slow dynamics, and the dashed line represents fast dynamics. In the feature learning case, test loss decreases like a staircase when $|R_\tau|$ increases via the fast dynamic. In the feature unlearning case, $|R_\tau|$ initially increases, then converges to zero due to the slow dynamics on the manifold.
  • Figure 2: Multi time-scale appears in numerical simulations of \ref{['eq:original']}. In the early stage of the dynamics, $R_\tau$ quickly moves away from $0$, while $a_\tau$ stays around the initial value $\bar{a}$. We set $\bar{k}_\star = \bar{k} = 5$ and $c = (1, 1, 1, 1, 1)$, and also set $c_{\star} = ( 1, -1, 1, -1, 1 ), \bar{a} = 4$ (left), and $c_\star = ( 2,4,6,8,10 ), \bar{a} = 1$ (right).
  • Figure 3: Trajectories of the model \ref{['eq:original']} on the $R-a$ space. The dots on the $a$-axis are the initial values $\bar{a}$. The red, yellow, pink trajectories show feature learning, and the blue trajectories shows feature unlearning. We set $\bar{k}_\star = \bar{k} = 5$ and $c = (1, 1, 1, 1, 1)$, and also set $c_{\star} = ( 1, 1, 1, 1, 1 )$ (left) or $c_\star = ( 1, -1, 1, -1, 1 )$ (right).
  • Figure 4: Simulated trajectories, alignments, second-layer weights, and losses of the model \ref{['eq:original']}. We set $\bar{k}_\star = \bar{k} = 5$ and $c = (1, 1, 1, 1, 1), c_{\star} = ( 1, -1, 1, 1, 1 )$, and $\bar{a} \in \{5, 9, 10\}$. We can observe that the learning dynamics proceeds differently for each case.
  • Figure 5: Phase maps for the feature unlearning by \ref{['eq:original']}. We set $\bar{k}_\star = \bar{k} = 5$ and $c = (1, 1, 1, 1, 1)$, and also set $c_{\star, 1} = c_{\star, 4} = c_{\star, 5} = 1, c_{\star, 2} = -1$ (left), or $c_{\star, 1} = c_{\star, 2} = c_{\star, 5} = 1, c_{\star, 3} = -1$ (right).
  • ...and 3 more figures

Theorems & Definitions (22)

  • Proposition 1
  • Definition 1: Feature unlearning
  • Definition 2: Critical manifold
  • Theorem 2: Feature unlearning
  • Theorem 3: Scaling law
  • Proposition 4: Difference equation of macroscopic variable
  • Lemma 5: Symmetricity of macroscopic variables
  • Lemma 6
  • proof : Proof of Lemma \ref{['lem:euler']}
  • proof : Proof of Proposition \ref{['prop:ode_validation']}
  • ...and 12 more