Stein-Rule Shrinkage for Stochastic Gradient Estimation in High Dimensions

TL;DR

This work reframes stochastic gradient estimation as a high-dimensional risk problem and introduces a Stein-rule shrinkage gradient estimator that contracts noisy mini-batch gradients toward a momentum-based restricted estimator. It proves risk dominance for $p \,\ge\, 3$ and minimax optimality under Gaussian noise, and integrates the estimator into Adam to form SR-Adam with negligible overhead. Empirical results on CIFAR-10/100 demonstrate consistent gains in large-batch and noisy settings, with ablations showing that selective shrinkage of high-dimensional convolutional layers is crucial. By linking high-dimensional decision theory to modern optimization, the paper provides a principled variance-reduction mechanism for deep learning.

Abstract

Stochastic gradient methods are central to large-scale learning, yet their analysis typically treats mini-batch gradients as unbiased estimators of the population gradient. In high-dimensional settings, however, classical results from statistical decision theory show that unbiased estimators are generally inadmissible under quadratic loss, suggesting that standard stochastic gradients may be suboptimal from a risk perspective. In this work, we formulate stochastic gradient computation as a high-dimensional estimation problem and introduce a decision-theoretic framework based on Stein-rule shrinkage. We construct a shrinkage gradient estimator that adaptively contracts noisy mini-batch gradients toward a stable restricted estimator derived from historical momentum. The shrinkage intensity is determined in a data-driven manner using an online estimate of gradient noise variance, leveraging second-moment statistics commonly maintained by adaptive optimization methods. Under a Gaussian noise model and for dimension p>=3, we show that the proposed estimator uniformly dominates the standard stochastic gradient under squared error loss and is minimax-optimal in the classical decision-theoretic sense. We further demonstrate how this estimator can be incorporated into the Adam optimizer, yielding a practical algorithm with negligible additional computational cost. Empirical evaluations on CIFAR10 and CIFAR100, across multiple levels of label noise, show consistent improvements over Adam in the large-batch regime. Ablation studies indicate that the gains arise primarily from selectively applying shrinkage to high-dimensional convolutional layers, while indiscriminate shrinkage across all parameters degrades performance. These results illustrate that classical shrinkage principles provide a principled and effective approach to improving stochastic gradient estimation in modern deep learning.

Figures (10)

• Figure 1: Bar chart comparison of best test accuracy on CIFAR10 under different label noise levels (0.0, 0.05, 0.1). Each subplot corresponds to one noise level. Bars show mean accuracy over five independent runs, and error bars indicate one standard deviation.
• Figure 2: Bar chart comparison of best test accuracy on CIFAR100 under different label noise levels (0.0, 0.05, 0.1). Despite the increased task difficulty, SR-Adam consistently outperforms baseline optimizers, particularly in noisy regimes.
• Figure 3: SimpleCNN on CIFAR10: test accuracy vs. epoch across noise levels. Mean $\pm$ std over runs.
• Figure 4: SimpleCNN on CIFAR10: test loss vs. epoch across noise levels. Mean $\pm$ std over runs.
• Figure 5: SimpleCNN on CIFAR100: test accuracy vs. epoch across noise levels. Mean $\pm$ std over runs.

Theorems & Definitions (5)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Lemma 5