Stein-Rule Shrinkage for Stochastic Gradient Estimation in High Dimensions

TL;DR

This paper introduces a Stein-rule shrinkage framework for stochastic gradient estimation in high-dimensional deep learning, leading to an improved optimizer called SR-Adam that outperforms standard Adam in noisy, large-batch settings.

Contribution

It develops a novel high-dimensional shrinkage estimator for stochastic gradients, integrated into Adam, with theoretical optimality and practical improvements demonstrated on image classification tasks.

Findings

SR-Adam outperforms Adam in large-batch regimes.

Shrinkage applied to convolutional layers yields most gains.

The method is minimax-optimal under Gaussian noise assumptions.

Abstract

Stochastic gradient methods are central to large-scale learning, but they treat mini-batch gradients as unbiased estimators, which classical decision theory shows are inadmissible in high dimensions. We formulate gradient computation as a high-dimensional estimation problem and introduce a framework based on Stein-rule shrinkage. We construct a gradient estimator that adaptively contracts noisy mini-batch gradients toward a stable estimator derived from historical momentum. The shrinkage intensity is determined in a data-driven manner using an online estimate of gradient noise variance, leveraging statistics from adaptive optimizers. Under a Gaussian noise model, we show our estimator uniformly dominates the standard stochastic gradient under squared error loss and is minimax-optimal. We incorporate this into the Adam optimizer, yielding SR-Adam, a practical algorithm with negligible computational cost. Empirical evaluations on CIFAR10 and CIFAR100 across multiple levels of input noise show consistent improvements over Adam in the large-batch regime. Ablation studies indicate that 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 approach to improving stochastic gradient estimation in deep learning.