Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Full-Batch Gradient Descent Outperforms One-Pass SGD: Sample Complexity Separation in Single-Index Learning

Filip Kovačević, Hong Chang Ji, Denny Wu, Mahdi Soltanolkotabi, Marco Mondelli

TL;DR

This work considers learning a d-dimensional single-index model with a quadratic activation, for which it is known that one-pass SGD requires n samples to achieve weak recovery, and shows that full-batch GD outperforms one-pass SGD (with the same activation) in statistical efficiency.

Abstract

It is folklore that reusing training data more than once can improve the statistical efficiency of gradient-based learning. However, beyond linear regression, the theoretical advantage of full-batch gradient descent (GD, which always reuses all the data) over one-pass stochastic gradient descent (online SGD, which uses each data point only once) remains unclear. In this work, we consider learning a $d$-dimensional single-index model with a quadratic activation, for which it is known that one-pass SGD requires $n\gtrsim d\log d$ samples to achieve weak recovery. We first show that this $\log d$ factor in the sample complexity persists for full-batch spherical GD on the correlation loss; however, by simply truncating the activation, full-batch GD exhibits a favorable optimization landscape at $n \simeq d$ samples, thereby outperforming one-pass SGD (with the same activation) in statistical efficiency. We complement this result with a trajectory analysis of full-batch GD on the squared loss from small initialization, showing that $n \gtrsim d$ samples and $T \gtrsim\log d$ gradient steps suffice to achieve strong (exact) recovery.

Full-Batch Gradient Descent Outperforms One-Pass SGD: Sample Complexity Separation in Single-Index Learning

TL;DR

This work considers learning a d-dimensional single-index model with a quadratic activation, for which it is known that one-pass SGD requires n samples to achieve weak recovery, and shows that full-batch GD outperforms one-pass SGD (with the same activation) in statistical efficiency.

Abstract

It is folklore that reusing training data more than once can improve the statistical efficiency of gradient-based learning. However, beyond linear regression, the theoretical advantage of full-batch gradient descent (GD, which always reuses all the data) over one-pass stochastic gradient descent (online SGD, which uses each data point only once) remains unclear. In this work, we consider learning a -dimensional single-index model with a quadratic activation, for which it is known that one-pass SGD requires samples to achieve weak recovery. We first show that this factor in the sample complexity persists for full-batch spherical GD on the correlation loss; however, by simply truncating the activation, full-batch GD exhibits a favorable optimization landscape at samples, thereby outperforming one-pass SGD (with the same activation) in statistical efficiency. We complement this result with a trajectory analysis of full-batch GD on the squared loss from small initialization, showing that samples and gradient steps suffice to achieve strong (exact) recovery.
Paper Structure (53 sections, 39 theorems, 525 equations, 2 figures)

This paper contains 53 sections, 39 theorems, 525 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Our contributions
  3. Related work
  4. Learning single and multi-index models.
  5. Benefit of data reuse.
  6. Problem setup
  7. Notation.
  8. Weak and strong recovery.
  9. Optimization algorithms.
  10. Sample complexity separation in spherical gradient flow
  11. Quadratic activation requires $\Omega(d\log d)$ samples
  12. $\Theta(d)$ samples suffice for truncated activation
  13. Strong recovery and iteration complexity
  14. Concluding remarks
  15. Proof of Theorem \ref{['thm:imp']}
  16. ...and 38 more sections

Key Result

Theorem 3.1

Consider the quadratic activation $\sigma(z)=z^2$. Let $\theta(t)$ be the solution at time $t$ of the spherical gradient flow ODE in eq:flowdef with initialization sampled uniformly from the sphere, i.e., $\theta(0)\sim {\rm Unif}({\mathcal{S}}^{d-1})$. Assume that $n=o(d\log d)$, i.e., $\lim_{n\to\

Figures (2)

  • Figure 1: Overlap achieved by minimizing the empirical correlation loss on the sphere as a function of $\delta=n/d$. We run spherical gradient descent with learning rate $\eta=0.1$ for $T=1000\log^2 d$ steps; experiments are averaged across 128 random seeds. Left: for the unbounded quadratic activation, increasing $d$ yields a larger threshold $\delta$ for weak recovery; we include a spline fit (solid lines) to smooth out the fluctuations. Middle: for the truncated activation ($M=8$) the overlap at fixed $\delta$ is almost $d$-independent. Right: threshold $\delta=n/d$ required to achieve target squared overlap values $\{0.1, 0.2, 0.3, 0.4, 0.5\}$ for $\sigma(z) = z^2$ (extracted from Figure \ref{['fig:unbounded']}), where we observe a clear $\delta\simeq\log d$ fit.
  • Figure 2: Overlap and parameter norm vs. number of GD steps. We use the truncated quadratic activation \ref{['eq:sigmadef-noc']} with $M=8$, run Euclidean gradient descent with learning rate $\eta = 0.1 / M^2$ and initialization scale $1/d^2$, and fix $\delta=n/d=10$; experiments are averaged across 1024 random seeds. Left, Middle: Observe that the time required for non-trivial overlap and norm growth increases with $d$. Right: number of GD steps required to achieve target squared overlap values $\{0.1, 0.2, 0.3, 0.4, 0.5\}$ (extracted from Figure \ref{['fig:overlap']}), where we observe a clear $T\simeq\log d$ fit.

Theorems & Definitions (68)

  • Theorem 3.1
  • Theorem 3.2
  • Theorem 4.1
  • Theorem A.1
  • proof
  • proof : Proof of Theorem \ref{['thm:imp']}
  • Proposition B.1
  • proof
  • Lemma B.2
  • proof
  • ...and 58 more