Research map / Optimization / Optimizers
Online Learning and Information Exponents: The Importance of Batch size & Time/Complexity Tradeoffs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies how batch size affects the iteration time of one-pass SGD for learning single/multi-index targets with two-layer networks, and proposes Correlation loss SGD to escape the self-interaction limit for large batches, with an ODE description of the dynamics.
Problem
Time/sample complexity tradeoff of one-pass SGD with large batch sizes when learning multi-index targets with two-layer networks on isotropic Gaussian data; standard SGD stops benefiting from larger batches beyond about d^{ℓ/2}.
Approach
Characterizes the optimal batch size and learning-rate scaling as a function of the information exponent, shows vanilla SGD's self-interaction term blocks speedups for large batches, and proposes correlation-loss SGD (vanishing second-layer scale) that suppresses auto-correlation terms and reaches weak recovery in polylog(d) steps for batch sizes up to d^{ℓ-1}.
Key results
- Vanilla SGD with batch size n_b ≲ d^{ℓ/2} and optimal learning rate γ=O(n_b d^{-ℓ/2}) achieves T=O(I(ℓ)/n_b) iterations without changing total sample complexity.
- Correlation loss SGD reaches weak recovery in T=polylog(d) steps for batch sizes n_b = O(d^{ℓ-1}) (per the stated theory), beyond vanilla SGD's d^{ℓ/2} limit.
- Numerical experiments with ℓ=3 confirm Correlation loss SGD achieves faster weak recovery than vanilla SGD across batch sizes.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase batch size up to n_b ≲ d^(ℓ/2) with optimally scaled learning rate | One-pass SGD · Single-index weak recovery | iterations T: T(n_b) = \sfrac{I(\ell)}{n_b} | Helped |
| increase batch size up to n_b ≲ d^(ℓ/2) | One-pass SGD · Single-index weak recovery | samples N: N = O(d^{\ell-1}) | No clear effect |
| increase batch size beyond d^(ℓ/2) for ℓ>2 using standard SGD | One-pass SGD · Single-index weak recovery, ℓ>2 | samples N: N = O(n_bd^{\sfrac{\ell}{2}-1}) | Hurt |
| replace square loss with correlation loss | One-pass SGD · Single-index weak recovery, ℓ>2 | iterations T: T = O(\sfrac{d^{\ell-1}}{n_b}) | Helped |
| replace square loss with correlation loss | One-pass SGD · Single-index weak recovery, ℓ≤2 | iterations T | No clear effect |
| use correlation loss and increase batch size to n_b = O(d^max(ℓ-1,1)) | One-pass SGD · Single-index weak recovery | iterations T: T=\mathrm{polylog}(d) | Helped |
| use one gradient step with batch size n_b=O(d^ℓ) | Two-layer neural network · Single-index weak recovery | iterations T: T=1 | Helped |
| increase learning rate above the optimal rate for the batch size | One-pass SGD · Single-index weak recovery | target correlation | Hurt |
| use a lower learning rate than the optimal rate | One-pass SGD · Single-index weak recovery | iterations T | Hurt |
| switch from correlation loss to square loss after escaping the initialization plateau | Adaptive SGD · Single-index target learning | full target learning | Mixed |
Surprising findings
- Larger batch sizes beyond roughly d^{ℓ/2} do not reduce time for vanilla one-pass SGD, due to negative self-interaction terms.
- The ODE description for large batches coincides with the n_b=1 case up to rescaling, with intra-batch correlations contributing only vanishing corrections at finite d.
Limitations
- Correlation loss SGD does not learn the full target better than vanilla SGD; it only accelerates weak recovery and is limited by the loss function at the global minimum.
- Analysis restricted to p=k=1 for rigorous results; larger p and k only supported by numerical simulations.
- Benefits of correlation loss are limited to information exponent ℓ>2.
Experimental setup
Datasets: synthetic Gaussian covariates with multi-index teacher labels
Metrics: teacher-student weight correlation m_t vs iteration time number of iterations T to weak recovery sample complexity N
Baselines: vanilla projected SGD with square loss one-step GD (giant step) n_b=1 online SGD
Builds on: information exponent (Ben Arous et al.) Saad and Solla online learning ODEs one-step giant-step analysis (Damian et al.)
Closest papers by meaning
- The Benefits of Reusing Batches for Gradient Descent in Two-Layer Networks: Breaking the Curse of Information and Leap Exponents (ICML 2024 · similarity 0.90)
- Learning Orthogonal Multi-Index Models: A Fine-Grained Information Exponent Analysis (NeurIPS 2025 · similarity 0.90)
- From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD (NeurIPS 2025 · similarity 0.90)
- Emergence and scaling laws in SGD learning of shallow neural networks (NeurIPS 2025 · similarity 0.89)
- Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit (NeurIPS 2024 · similarity 0.88)
- Learning quadratic neural networks in high dimensions: SGD dynamics and scaling laws (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.