The Dimension Strikes Back with Gradients: Generalization of Gradient Methods in Stochastic Convex Optimization
Matan Schliserman, Uri Sherman, Tomer Koren
TL;DR
We study how gradient methods generalize in stochastic convex optimization as a function of problem dimension. The authors construct a dimension-polynomial learning problem where full-batch GD trained with $n$ samples can converge to a bad ERM with constant population excess risk, implying a lower bound of $Ω(√d)$ samples for nontrivial generalization; a parallel, differentiable construction yields a similar bound for one-pass SGD on empirical risk. The core technique embeds exponentially many nearly orthogonal ERMs, augments the loss with non-smooth components, memorizes the training data inside iterates, and uses a round-robin, multi-subspace scheme to ensure progress toward a bad solution; a smoothing step renders the construction differentiable without altering the dynamics. The results are tight up to polylogarithmic factors relative to known upper bounds and resolve open questions about the dimension dependence of gradient-based methods in SCO, highlighting that dimension-driven generalization limits arise even for convex, differentiable losses.
Abstract
We study the generalization performance of gradient methods in the fundamental stochastic convex optimization setting, focusing on its dimension dependence. First, for full-batch gradient descent (GD) we give a construction of a learning problem in dimension $d=O(n^2)$, where the canonical version of GD (tuned for optimal performance of the empirical risk) trained with $n$ training examples converges, with constant probability, to an approximate empirical risk minimizer with $Ω(1)$ population excess risk. Our bound translates to a lower bound of $Ω(\sqrt{d})$ on the number of training examples required for standard GD to reach a non-trivial test error, answering an open question raised by Feldman (2016) and Amir, Koren, and Livni (2021b) and showing that a non-trivial dimension dependence is unavoidable. Furthermore, for standard one-pass stochastic gradient descent (SGD), we show that an application of the same construction technique provides a similar $Ω(\sqrt{d})$ lower bound for the sample complexity of SGD to reach a non-trivial empirical error, despite achieving optimal test performance. This again provides an exponential improvement in the dimension dependence compared to previous work (Koren, Livni, Mansour, and Sherman, 2022), resolving an open question left therein.