# The Optimization Landscape of SGD Across the Feature Learning Strength

arXiv: https://arxiv.org/abs/2410.04642 · HTML: https://mlsanity.com/paper/2410.04642/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 54 of its 67 references that are in the map.
- First posted on arXiv 6 Oct 2024.

## Builds on (most cited first, 30 of 54)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel](https://mlsanity.com/paper/2010.15110/) (2020)
- [Finite Versus Infinite Neural Networks: an Empirical Study](https://mlsanity.com/paper/2007.15801/) (2020)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [A Dynamical Model of Neural Scaling Laws](https://mlsanity.com/paper/2402.01092/) (2024)
- [Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity](https://mlsanity.com/paper/2106.09524/) (2021)
- [A Spectral Condition for Feature Learning](https://mlsanity.com/paper/2310.17813/) (2023)
- [Grokking as the Transition from Lazy to Rich Training Dynamics](https://mlsanity.com/paper/2310.06110/) (2023)
- [Neural Networks as Kernel Learners: The Silent Alignment Effect](https://mlsanity.com/paper/2111.00034/) (2021)
- [Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks](https://mlsanity.com/paper/2205.09653/) (2022)
- [Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks](https://mlsanity.com/paper/2310.02244/) (2023)
- [Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity](https://mlsanity.com/paper/2106.15933/) (2021)

## Built on by (most cited first, 13 of 13)

- [Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful](https://mlsanity.com/paper/2507.07101/) (2025)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks](https://mlsanity.com/paper/2506.06489/) (2025)
- [Position: Solve Layerwise Linear Models First to Understand Neural Dynamical Phenomena (Neural Collapse, Emergence, Lazy/Rich Regime, and Grokking)](https://mlsanity.com/paper/2502.21009/) (2025)
- [Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models](https://mlsanity.com/paper/2502.09863/) (2025)
- [The Importance of Being Lazy: Scaling Limits of Continual Learning](https://mlsanity.com/paper/2506.16884/) (2025)
- [Saddle-To-Saddle Dynamics in Deep ReLU Networks: Low-Rank Bias in the First Saddle Escape](https://mlsanity.com/paper/2505.21722/) (2025)
- [Learning richness modulates equality reasoning in neural networks](https://mlsanity.com/paper/2503.09781/) (2025)
- [Learnable Multipliers: Freeing the Scale of Language Model Matrix Layers](https://mlsanity.com/paper/2601.04890/) (2026)
- [Feature learning is decoupled from generalization in high capacity neural networks](https://mlsanity.com/paper/2507.19680/) (2025)
- [Over-Alignment vs Over-Fitting: The Role of Feature Learning Strength in Generalization](https://mlsanity.com/paper/2602.00827/) (2026)
- [A Theory of Saddle Escape in Deep Nonlinear Networks](https://mlsanity.com/paper/2605.01288/) (2026)
- [Explaining Near-Zero Hessian Eigenvalues Through Approximate Symmetries in Neural Networks](https://mlsanity.com/paper/2607.07845/) (2026)

## Abstract (arXiv)

We consider neural networks (NNs) where the final layer is down-scaled by a fixed hyperparameter γ. Recent work has identified γ as controlling the strength of feature learning. As γ increases, network evolution changes from "lazy" kernel dynamics to "rich" feature-learning dynamics, with a host of associated benefits including improved performance on common tasks. In this work, we conduct a thorough empirical investigation of the effect of scaling γ across a variety of models and datasets in the online training setting. We first examine the interaction of γ with the learning rate η, identifying several scaling regimes in the γ-η plane which we explain theoretically using a simple model. We find that the optimal learning rate η^* scales non-trivially with γ. In particular, η^* propto γ2 when γ ll 1 and η^* propto γ2/L when γ gg 1 for a feed-forward network of depth L. Using this optimal learning rate scaling, we proceed with an empirical study of the under-explored "ultra-rich" γ gg 1 regime. We find that networks in this regime display characteristic loss curves, starting with a long plateau followed by a drop-off, sometimes followed by one or more additional staircase steps. We find networks of different large γ values optimize along similar trajectories up to a reparameterization of time. We further find that optimal online performance is often found at large γ and could be missed if this hyperparameter is not tuned. Our findings indicate that analytical study of the large-γ limit may yield useful insights into the dynamics of representation learning in performant models.