On the Learning Dynamics of Two-layer Linear Networks with Label Noise SGD

Tongcheng Zhang; Zhanpeng Zhou; Mingze Wang; Andi Han; Wei Huang; Taiji Suzuki; Junchi Yan

cs.LGMar 11, 2026

On the Learning Dynamics of Two-layer Linear Networks with Label Noise SGD

Authors: Tongcheng Zhang, Zhanpeng Zhou, Mingze Wang, Andi Han, Wei Huang, Taiji Suzuki, Junchi Yan

Organizations: Sch. of Computer Science and Zhiyuan College, Shanghai Jiao Tong University · Peking University · University of Sydney · RIKEN Center for Advanced Intelligence Project ·

Abstract

One crucial factor behind the success of deep learning lies in the implicit bias induced by noise inherent in gradient-based training algorithms. Motivated by empirical observations that training with noisy labels improves model generalization, we delve into the underlying mechanisms behind stochastic gradient descent (SGD) with label noise. Focusing on a two-layer over-parameterized linear network, we analyze the learning dynamics of label noise SGD, unveiling a two-phase learning behavior. In Phase I, the magnitudes of model weights progressively diminish, and the model escapes the lazy regime; enters the rich regime. In Phase II, the alignment between model weights and the ground-truth interpolator increases, and the model eventually converges. Our analysis highlights the critical role of label noise in driving the transition from the lazy to the rich regime and minimally explains its empirical success. Furthermore, we extend these insights to Sharpness-Aware Minimization (SAM), showing that the principles governing label noise SGD also apply to broader optimization algorithms. Extensive experiments, conducted under both synthetic and real-world setups, strongly support our theory. Our code is released at https://github.com/a-usually/Label-Noise-SGD.

Figures & tables

Figure 1: (Left).Label noise SGD ( Algorithm 1 ) leads to better generalization. Test loss $\mathcal{L}(\bm{\theta}(t))$ and accuracy ${\operatorname{Acc}}(\bm{\theta}(t))$ vs. training epochs $t$ . (Right). Label noise SGD leads to sparser solutions. Testing accuracy of pruned model ${\operatorname{Acc}}(\bm{\theta}(T;\alpha))$ vs. the percentage of remaining parameters $\alpha$ . Here, $\bm{\theta}(T;\alpha)$ represents the pruned model derived from the pretrained model $\bm{\theta}(T)$ , with $\alpha$ % of parameters remaining. We use both vanilla SGD and label noise SGD to train the models, with no weight decay or momentum. The learning rate is set to $0.1$ , and the total number of epochs is $160$ . Exponential moving average is employed to smooth the test accuracy curves. Results are presented for ResNet-18 [ 22 ] trained on CIFAR-10 [ 28 ] , across different label noise probabilities, $(\tau\in\{0.05,0.1,0.2\})$ . As shown in figure (a) and figure (b), label noise SGD consistently outperforms vanilla SGD in both test loss and accuracy across different values of the label flipping probability $\tau$ , providing an around $1.5$ % improvement in test accuracy. As shown in Figure 1 (c) , models trained with label noise SGD maintain higher performance at the same sparsity level compared to those trained with vanilla SGD.

I	The magnitudes of model weights progressively diminish;	Lemma 4.3
I	The model escapes the lazy regime; enters the rich regime.	Theorem 4.2
II	The alignment between model weights and the ground-truth	Lemma 4.5
II	interpolator increases; the model eventually converges.	Lemma 4.6

Table 1: Overview of the two-phase picture and corresponding theoretical results.

Figure 2: Two-phase dynamics of label noise SGD under synthetic setup. We replicate the synthetic problem setup from Section 4.5 . (a) Loss curves. Training $\mathcal{L}_{\mathcal{D}_{\rm train}}(\bm{\theta}(t))$ and test loss $\mathcal{L}_{\mathcal{D}_{\rm test}}(\bm{\theta}(t))$ vs. training iteration $t$ . (b) Learning dynamics on average. The mean neuron norm $\text{Avg}_{i\in[m]}(\|\bm{w}_{i}(t)\|_{2})$ and the mean neuron alignment $\text{Avg}_{i\in[m]}(\langle\bm{w}_{i}(t),\bm{\theta}^{\star}\rangle)$ vs. training iteration $t$ . (c) Learning dynamics of $i$ -th neuron. The alignment of the $i$ -th neuron $\langle\bm{w}_{i}(t),\bm{\theta}^{\star}\rangle$ vs. its weight norm $\|\bm{w}_{i}(t)\|_{2}$ , with darker points indicating larger $t$ . (Bottom) Complete view of dynamics of each neuron. This plot is similar to (c); yet instead of focusing on a single neuron, we plot the status over iterations.

Figure 3: Label noise SGD induces the rich regime. (a, b). Training $\mathcal{L}_{\mathcal{D}_{\rm train}}(\bm{\theta}(t))$ and test loss $\mathcal{L}_{\mathcal{D}_{\rm test}}(\bm{\theta}(t))$ vs. training epochs $t$ . Label noise SGD induces the progressively diminishing phenomenon. (c). The first-layer weight norm $\|\bm{W}(t)\|_{F}$ vs. training epochs $t$ . We use GD to train the models with NTK parameterization [ 25 ] , both with and without label noise. We also train a linearized model with GD as baseline. Results are presented for WideResNets trained on a random subset of 64 images from CIFAR-10 due to the $O(n^{2})$ computational complexity of NTK.

Figure 4: Two-phase dynamics of SAM under synthetic setup. We replicate the synthetic problem setup from Section 4.5 , replacing label noise SGD with SAM. (a) Loss curves. Training $\mathcal{L}_{\mathcal{D}_{\rm train}}(\bm{\theta}(t))$ and test loss $\mathcal{L}_{\mathcal{D}_{\rm test}}(\bm{\theta}(t))$ vs. training iteration $t$ . (b) Learning dynamics on average. The averaged neuron norm $\text{Avg}_{i\in[m]}(\|\bm{w}_{i}(t)\|_{2})$ and the averaged neuron alignment $\text{Avg}_{i\in[m]}(\langle\bm{w}_{i}(t),\bm{\theta}^{\star}\rangle)$ vs. training iteration $t$ . (c) Learning dynamics of $i$ -th neuron. The alignment of $i$ -th neuron $\langle\bm{w}_{i}(t),\bm{\theta}^{\star}\rangle$ vs. its weight norm $\|\bm{w}_{i}(t)\|_{2}$ , with darker points indicating larger iteration $t$ .

Figure 5: SAM induces the rich regime. Training $\mathcal{L}_{\mathcal{D}_{\rm train}}(\bm{\theta}(t))$ and test loss $\mathcal{L}_{\mathcal{D}_{\rm test}}(\bm{\theta}(t))$ vs. training epochs $t$ . We use both GD and full-batch SAM to train the models with NTK parameterization [ 25 ] . We also train a linearized model with GD as baseline. Results are presented for WideResNets trained on a random subset of 64 images from CIFAR-10 due to the $O(n^{2})$ computational complexity of NTK.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Neuron dynamics in alternating noisy/noiseless SGD training. We replicate the synthetic problem setup from Subsection "Setup and Overview: A Two-Layer Linear Network", changing the standard SGD label noise to training with alternating noisy and noiseless SGD every 5000 steps. (a) Learning dynamics on average. The averaged neuron norm $\text{Avg}_{i\in[m]}(\|\bm{w}_{i}(t)\|_{2})$ and the averaged neuron alignment $\text{Avg}_{i\in[m]}(\langle\bm{w}_{i}(t),\bm{\theta}^{\star}\rangle)$ vs. training iteration $t$ . (b) Learning dynamics of $i$ -th neuron. The $i$ -th neuron norm $\|\bm{w}_{i}(t)\|_{2}$ and the corresponding absolute value $|a_{i}(t)|$ in the second layer vs. training iteration $t$ .

Two-Layer Neural Networks

See all 14 papers

Apr 7, 2026

Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networks

Deep linear networks (DLNs) are used as an analytically tractable model of the training dynamics of deep neural networks. While gradient descent in DLNs is known to exhibit saddle-to-saddle dynamics, the impact of stochastic gradient descent (SGD) noise on this regime remains poorly understood. We investigate the dynamics of SGD during training of DLNs in the saddle-to-saddle regime. We model the training dynamics as stochastic Langevin dynamics with anisotropic, state-dependent noise. Under the assumption of aligned and balanced weights, we derive an exact decomposition of the dynamics into a system of one-dimensional per-mode stochastic differential equations. This establishes that the maximal diffusion along a mode precedes the corresponding feature being completely learned. We also derive the stationary distribution of SGD for each mode: in the absence of label noise, its marginal distribution along specific features coincides with the stationary distribution of gradient flow, while in the presence of label noise it approximates a Boltzmann distribution. Finally, we confirm experimentally that the theoretical results hold qualitatively even without aligned or balanced weights. These results establish that SGD noise encodes information about the progression of feature learning but does not fundamentally alter the saddle-to-saddle dynamics.

Guillaume Corlouer, Avi Semler, Alexander Strang +1

Moirai · University of Oxford · University of California, Berkeley +1

Gradient Descent Dynamics Stochastic Gradient Descent Deep Linear Networks Representation Learning

Ask PDF

Oct 20, 2026

Convergence of stochastic gradient descent for learning linear neural networks

Stochastic gradient descent (SGD) is the cornerstone of modern deep learning, yet its theoretical convergence properties remain poorly understood in non-convex settings. To address this gap, we study SGD dynamics in deep linear neural networks - simplified models that retain key non-convex characteristics while remaining analytically tractable. Building on recent results for gradient descent, we extend convergence guarantees to SGD under regularized squared loss. We establish almost sure boundedness of SGD iterates. Using the framework of asymptotic pseudotrajectories from stochastic approximation theory, our analysis establishes then, for both deterministic and adaptive step-size choices, convergence of SGD to the set of critical points under mild conditions (without relying on Lipschitz properties or whitened data). To our knowledge, this is the first work providing almost sure convergence guarantees for SGD in deep linear networks, bridging the gap between theoretical optimization and practical stochastic training.

Gabin Maxime Nguegnang, Bubacarr Bah, Holger Rauhut +1

Stochastic Optimization Convergence Stochastic Gradient Descent Nonconvex Stochastic Optimization Deep Linear Networks

Ask PDF

Mar 9, 2026

Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization

We study the implicit bias of Sharpness-Aware Minimization (SAM) when training

L

-layer linear diagonal networks on linearly separable binary classification. For linear models (

L=1

), both

\ell_\infty

- and

\ell_2

-SAM recover the

\ell_2

max-margin classifier, matching gradient descent (GD). However, for depth

L = 2

, the behavior changes drastically -- even on a single-example dataset. For

\ell_\infty

-SAM, the limit direction depends critically on initialization and can converge to

\mathbf{0}

or to any standard basis vector, in stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate. For

\ell_2

-SAM, we show that although its limit direction matches the

\ell_1

max-margin solution as in the case of GD, its finite-time dynamics exhibit a phenomenon we call "sequential feature amplification", in which the predictor initially relies on minor coordinates and gradually shifts to larger ones as training proceeds or initialization increases. Our theoretical analysis attributes this phenomenon to

\ell_2

-SAM's gradient normalization factor applied in its perturbation, which amplifies minor coordinates early and allows major ones to dominate later, giving a concrete example where infinite-time implicit-bias analyses are insufficient. Synthetic and real-data experiments corroborate our findings.

Chaewon Moon, Dongkuk Si, Chulhee Yun

Kim Jaechul Graduate School of AI, KAIST · Mobilint, Inc.

Sharpness-Aware Minimization Implicit Bias Deep Linear Networks

Ask PDF

Abstract

Two-Layer Neural Networks

Explore similar work

Stochastic Gradient Descent in the Saddle-to-Saddle Regime of Deep Linear Networks

Convergence of stochastic gradient descent for learning linear neural networks

Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization