The Optimization Landscape of SGD Across the Feature Learning Strength
Keywords: Feature Learning, Maximal Update Parameterization, Optimization, Edge of Stability, Catapult Effect, Silent Alignment
TL;DR: We empirically sweep over the feature learning hyperparameter for models trained online, identifying a variety of novel scaling relationships with sharpness, optimal learning rate, and performance, and explaining these in a linear network model.
Abstract: We consider neural networks (NNs) where the final layer is down-scaled by a fixed hyperparameter $\gamma$.
Recent work has identified $\gamma$ as controlling the strength of feature learning.
As $\gamma$ 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 $\gamma$ across a variety of models and datasets in the online training setting.
We first examine the interaction of $\gamma$ with the learning rate $\eta$, identifying several scaling regimes in the $\gamma$-$\eta$ plane which we explain theoretically using a simple model.
We find that the optimal learning rate $\eta^*$ scales non-trivially with $\gamma$. In particular, $\eta^* \propto \gamma^2$ when $\gamma \ll 1$ and $\eta^* \propto \gamma^{2/L}$ when $\gamma \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'' $\gamma \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 $\gamma$ values optimize along similar trajectories up to a reparameterization of time.
We further find that optimal online performance is often found at large $\gamma$ and could be missed if this hyperparameter is not tuned.
Our findings indicate that analytical study of the large-$\gamma$ limit may yield useful insights into the dynamics of representation learning in performant models.
Supplementary Material: zip
Primary Area: optimization
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 2795
Loading