The Optimization Landscape of SGD Across the Feature Learning Strength (2410.04642v3)

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.

• The paper demonstrates scaling regimes in the gamma-eta plane, showing that the optimal learning rate scales as gamma^2 for small gamma and as gamma^(2/L) in deeper networks.
• The paper uncovers that in the ultra-rich feature learning regime, networks initially plateau in loss before a dramatic drop, resulting in improved performance.
• The paper validates empirical observations with analytical models, emphasizing the importance of tuning gamma to optimize online representation learning.

The Optimization Landscape of SGD Across the Feature Learning Strength

The paper under review explores the influence of the hyperparameter $\gamma$ on neural networks within the context of stochastic gradient descent (SGD), exploring how it modulates feature learning dynamics. By varying $\gamma$ across many orders of magnitude, the research investigates the transition from lazy kernel dynamics to rich feature-learning dynamics, with the latter being crucial for enhanced model performance across tasks, particularly in an online training setting.

Core Contributions and Findings

1. Interplay Between $\gamma$ and Learning Rate $\eta$: The authors systematically explore the $\gamma$-$\eta$ plane, revealing distinct scaling regimes. Notably, the optimal learning rate $\eta^*$ is shown to scale proportionally to $\gamma^2$ when $\gamma$ is much less than 1. In the ultra-rich feature-learning regime, where $\gamma$ vastly exceeds 1, $\eta^*$ scales with $\gamma^{2/L}$, with $L$ representing network depth. This finding contributes to our understanding of how $\gamma$ can be optimized to harness the full learning capability of deep networks.
2. Ultra-Rich Feature Learning Regime: In the $\gamma \gg 1$ domain, networks exhibit characteristic loss trajectories, often starting with a plateau before undergoing dramatic loss drop-offs. The paper finds that networks trained in this regime along similar trajectories offer improved performance, thereby highlighting the importance of tuning $\gamma$ to realize the full potential of modern deep learning architectures.
3. Theoretical and Empirical Analysis: An analytical model helps to verify the observed phenomena in linear networks, reinforcing the empirical findings. The paper contributes critical insights into the optimization dynamics by investigating the interaction between different learning rates and feature-learning strengths. It identifies a potential for progressive sharpening which explains the large $\gamma$ regime behavior.
4. Implications for Representation Learning: One key implication of this research is its emphasis on the optimal tuning of $\gamma$ to either match or outperform standard networks in an online setting. This underlines the potential for $\gamma$-tuned models to avoid missing optimal performance windows due to a poorly selected feature-learning strength.

Implications and Future Directions

The findings hold significant implications for both theoretical explorations and practical implementations of neural networks. A deeper understanding of $\gamma$'s impact across various architectures and datasets can influence how models are scaled and optimized, paving the way for more efficient large-scale models.

Future research might include extending these analyses to other optimizers beyond SGD or exploring hybrid parameterization techniques within real-world applications, such as natural language processing or computer vision tasks. Additionally, further investigation into the behavior of networks in the ultra-rich regime could unlock new strategies for training models more effectively across varied computational budgets and architectures.

In summary, this work provides a comprehensive exploration into the effect of feature-learning strength, delivering actionable insights for tuning deep learning models in complex, large-scale data environments. The exploration of $\gamma$ and its impact stands as a significant contribution to the landscape of neural network optimization.