# Stein-Rule Shrinkage for Stochastic Gradient Estimation in High Dimensions

arXiv: https://arxiv.org/abs/2602.01777 · HTML: https://mlsanity.com/paper/2602.01777/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 2 of its 15 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Builds on (most cited first, 2 of 2)

- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)

## Abstract (arXiv)

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.