Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD

Konstantinos Christopher Tsiolis, Alireza Mousavi-Hosseini, Murat A. Erdogdu

TL;DR

This paper characterize the relationship between learning rate(s) and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates and demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.

Abstract

To understand feature learning dynamics in neural networks, recent theoretical works have focused on gradient-based learning of Gaussian single-index models, where the label is a nonlinear function of a latent one-dimensional projection of the input. While the sample complexity of online SGD is determined by the information exponent of the link function, recent works improved this by performing multiple gradient steps on the same sample with different learning rates -- yielding a non-correlational update rule -- and instead are limited by the (potentially much smaller) generative exponent. However, this picture is only valid when these learning rates are sufficiently large. In this paper, we characterize the relationship between learning rate(s) and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates. We demonstrate that, in certain cases, there is a phase transition from an "information exponent regime" with small learning rate to a "generative exponent regime" with large learning rate. Our framework covers prior analyses of one-pass SGD and SGD with batch reuse, while also introducing a new layer-wise training algorithm that leverages a two-timescales approach (via different learning rates for each layer) to go beyond correlational queries without reusing samples or modifying the loss from squared error. Our theoretical study demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.

From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD

TL;DR

This paper characterize the relationship between learning rate(s) and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates and demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.

Abstract

To understand feature learning dynamics in neural networks, recent theoretical works have focused on gradient-based learning of Gaussian single-index models, where the label is a nonlinear function of a latent one-dimensional projection of the input. While the sample complexity of online SGD is determined by the information exponent of the link function, recent works improved this by performing multiple gradient steps on the same sample with different learning rates -- yielding a non-correlational update rule -- and instead are limited by the (potentially much smaller) generative exponent. However, this picture is only valid when these learning rates are sufficiently large. In this paper, we characterize the relationship between learning rate(s) and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates. We demonstrate that, in certain cases, there is a phase transition from an "information exponent regime" with small learning rate to a "generative exponent regime" with large learning rate. Our framework covers prior analyses of one-pass SGD and SGD with batch reuse, while also introducing a new layer-wise training algorithm that leverages a two-timescales approach (via different learning rates for each layer) to go beyond correlational queries without reusing samples or modifying the loss from squared error. Our theoretical study demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.
Paper Structure (31 sections, 16 theorems, 111 equations, 4 figures, 2 algorithms)

This paper contains 31 sections, 16 theorems, 111 equations, 4 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Our Contributions
  3. Notation.
  4. Related Work
  5. Feature Learning and Single-Index Models.
  6. Learning Rate and Generalization.
  7. Problem Setup
  8. Sample Complexity of a Generic Online Algorithm
  9. Examples
  10. Online SGD
  11. Batch Reuse SGD
  12. Alternating SGD
  13. Proof Sketch
  14. Conclusion
  15. Technical Background
  16. ...and 16 more sections

Key Result

Lemma 2.3

Suppose there exists an orthonormal polynomial basis of the space $L^2((\sigma_*)_{\#} \mathcal{N}(0,1))$. Then there exists $I \in \mathbb{N}$ such that $\mathrm{IE}(\sigma_*^I) = p_*$. Moreover, if $\sigma_*$ is polynomial of degree at most $q$, then $I \leq C_q$ for some constant $C_q$ depending

Figures (4)

  • Figure 1: Combinations of learning rate $\eta$ and sample size $n$ which achieve alignment $\langle\boldsymbol{w},\boldsymbol{\theta}_*\rangle \geq 0.5$ for a network \ref{['eq:two_layer_nn']} with $N = 1$ trained by alternating SGD (Algorithm \ref{['alg:alternating_sgd']}) in the setting $\sigma_* = \sigma = \mathsf{He}_3$ and $d=50$.
  • Figure 2: Alignments $\langle\boldsymbol{w},\boldsymbol{\theta}_*\rangle$ greater than 0.5 for alternating SGD with different choices of $\eta$ and $n$. The hyperparameter $\gamma$ is chosen according to Corollary \ref{['cor:alternating_sgd']}. Results are averaged over 10 runs.
  • Figure 3: Alignments $\langle\boldsymbol{w},\boldsymbol{\theta}_*\rangle$ greater than 0.5 for batch reuse SGD with different choices of $\eta$ and $n$. The hyperparameter $\gamma$ is chosen according to Corollary \ref{['cor:batch_reuse_sgd']}. Results are averaged over 10 runs.
  • Figure 4: Alignments $\langle\boldsymbol{w},\boldsymbol{\theta}_*\rangle$ greater than 0.5 for online SGD with different choices of $\gamma$ and $n$. Note that $\eta = 0$ in this case. Results are averaged over 10 runs.

Theorems & Definitions (30)

  • Definition 2.1: Information Exponent, benarous2021online
  • Definition 2.2: Generative Exponent, damian2024computational
  • Lemma 2.3: lee2024neural Proposition 6, Lemma 8
  • Theorem 3.3
  • Remark 3.4: On the Optimal Choices of $\gamma$ and $\eta$
  • Corollary 4.1
  • Corollary 4.2
  • Corollary 4.3
  • Definition A.1: Probabilist's Hermite Polynomials o2021analysis
  • Lemma A.2: o2021analysis
  • ...and 20 more