Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Implicit Regularization of SGD with Preconditioning for Least Square Problems

Junwei Su, Difan Zou, Chuan Wu

TL;DR

This work investigates why SGD can exhibit weaker generalization than ridge in least-squares problems due to unequal optimization across covariance directions. It proposes a simple preconditioning design $\mathbf{G}= (\beta\mathbf{H}+\mathbf{I})^{-1}$ to rebalance updates and derives non-asymptotic excess-risk bounds for both preconditioned SGD and ridge, including cases where the covariance $\mathbf{H}$ must be estimated from unlabeled data. The authors prove that there exist choices of $\beta$ and learning rate $\eta$ making preconditioned SGD comparable to ridge (and preconditioned ridge), and they extend these results to settings where $\mathbf{H}$ is unknown but estimated with high probability bounds. Empirical results on Gaussian least-squares problems corroborate the theory, showing robust improvements and demonstrating practical viability of learning the preconditioner from unlabeled data.

Abstract

Stochastic gradient descent (SGD) exhibits strong algorithmic regularization effects in practice and plays an important role in the generalization of modern machine learning. However, prior research has revealed instances where the generalization performance of SGD is worse than ridge regression due to uneven optimization along different dimensions. Preconditioning offers a natural solution to this issue by rebalancing optimization across different directions. Yet, the extent to which preconditioning can enhance the generalization performance of SGD and whether it can bridge the existing gap with ridge regression remains uncertain. In this paper, we study the generalization performance of SGD with preconditioning for the least squared problem. We make a comprehensive comparison between preconditioned SGD and (standard \& preconditioned) ridge regression. Our study makes several key contributions toward understanding and improving SGD with preconditioning. First, we establish excess risk bounds (generalization performance) for preconditioned SGD and ridge regression under an arbitrary preconditions matrix. Second, leveraging the excessive risk characterization of preconditioned SGD and ridge regression, we show that (through construction) there exists a simple preconditioned matrix that can make SGD comparable to (standard \& preconditioned) ridge regression. Finally, we show that our proposed preconditioning matrix is straightforward enough to allow robust estimation from finite samples while maintaining a theoretical improvement. Our empirical results align with our theoretical findings, collectively showcasing the enhanced regularization effect of preconditioned SGD.

Improving Implicit Regularization of SGD with Preconditioning for Least Square Problems

TL;DR

This work investigates why SGD can exhibit weaker generalization than ridge in least-squares problems due to unequal optimization across covariance directions. It proposes a simple preconditioning design to rebalance updates and derives non-asymptotic excess-risk bounds for both preconditioned SGD and ridge, including cases where the covariance must be estimated from unlabeled data. The authors prove that there exist choices of and learning rate making preconditioned SGD comparable to ridge (and preconditioned ridge), and they extend these results to settings where is unknown but estimated with high probability bounds. Empirical results on Gaussian least-squares problems corroborate the theory, showing robust improvements and demonstrating practical viability of learning the preconditioner from unlabeled data.

Abstract

Stochastic gradient descent (SGD) exhibits strong algorithmic regularization effects in practice and plays an important role in the generalization of modern machine learning. However, prior research has revealed instances where the generalization performance of SGD is worse than ridge regression due to uneven optimization along different dimensions. Preconditioning offers a natural solution to this issue by rebalancing optimization across different directions. Yet, the extent to which preconditioning can enhance the generalization performance of SGD and whether it can bridge the existing gap with ridge regression remains uncertain. In this paper, we study the generalization performance of SGD with preconditioning for the least squared problem. We make a comprehensive comparison between preconditioned SGD and (standard \& preconditioned) ridge regression. Our study makes several key contributions toward understanding and improving SGD with preconditioning. First, we establish excess risk bounds (generalization performance) for preconditioned SGD and ridge regression under an arbitrary preconditions matrix. Second, leveraging the excessive risk characterization of preconditioned SGD and ridge regression, we show that (through construction) there exists a simple preconditioned matrix that can make SGD comparable to (standard \& preconditioned) ridge regression. Finally, we show that our proposed preconditioning matrix is straightforward enough to allow robust estimation from finite samples while maintaining a theoretical improvement. Our empirical results align with our theoretical findings, collectively showcasing the enhanced regularization effect of preconditioned SGD.
Paper Structure (27 sections, 25 theorems, 205 equations, 1 figure)

This paper contains 27 sections, 25 theorems, 205 equations, 1 figure.

Table of Contents

  1. Introduction
  2. Our contributions.
  3. Related Work
  4. Problem Setup and Preliminaries
  5. Main Results
  6. Excessive Risk of Preconditioned SGD and Ridge Regression
  7. Design of The Precondition Matrix
  8. Effectiveness of Preconditioned SGD in Theoretical Setting
  9. Effectiveness of Preconditioned-SGD with Estimation
  10. Empirical Study
  11. Discussion
  12. Excessive Risk Analysis of Preconditioned SGD and Ridge
  13. Ecessive Risk of SGD with precondition
  14. Ecessive Risk of ridge regression with precondition
  15. Analysis of precondition SGD with standard ridge regression
  16. ...and 12 more sections

Key Result

Theorem 4.1

Consider ridge regression with parameter $\lambda > 0$ and precondition matrix $\mathbf{M}$. Suppose Assumptions assump:data_distribution, assump:fourth_moment and assump:model_noise hold. Let $\widehat{\mathbf{H}} = \mathbf{H}^{1/2}\mathbf{M}^{-1}\mathbf{H}^{1/2}, \quad \text{and} \quad \widehat{ where $\widehat{\lambda}_1,\dots, \widehat{\lambda}_d$ are the sorted eigenvalues for $\widehat{\ma

Figures (1)

  • Figure 1: Generalization performance comparison among SGD, Ridge regression, PreSGD and PreSGD-Est, where the stepsize $\eta$, controlling factor $\beta$ and regularization parameter $\lambda$ are fine-tuned to achieve the best performance. The problem dimension is $d=200$ and the variance of model noise is $\sigma^2=1$. We consider $6$ combinations of $2$ different covariance matrices and $3$ different ground truth model vectors. The generalization performance (excessive risk) is measured using mean squared error (MSE) on an independent test set of size $1000$ that is generated from the same problem instance. A set of unlabelled data of the same size as the training set is drawn from the same problem instance to estimate $\widetilde{\mathbf{G}}$ in (PreSGD-Est). The plots are averaged over $10$ independent runs.

Theorems & Definitions (25)

  • Theorem 4.1: Excessive risk lower bound for ridge with preconditioning
  • Theorem 4.2: Excessive risk upper bound of preconditioned SGD
  • Theorem 4.3: Preconditioned-SGD comparable to standard ridge regression
  • Theorem 4.4: Preconditioned-SGD comparable to preconditioned ridge regression
  • Theorem 4.5: Preconditioned-SGD with estimated $\widetilde{\mathbf{G}}$ comparable to standard ridge regression
  • Theorem 4.6: Preconditioned-SGD with estimated $\widetilde{\mathbf{G}}$ comparable to preconditioned ridge regression
  • Theorem A.1: Extension of Theorem 5.1 in zou2021benign
  • Theorem A.2: Theorem B.2 in zou2021benefits
  • Lemma B.1
  • Lemma B.2
  • ...and 15 more