# Stochastic Gradient Optimization with Model-Assisted Sampling

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 34 references that are in the map.
- First posted on arXiv 25 Jun 2026.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Minimizing Finite Sums with the Stochastic Average Gradient](https://mlsanity.com/paper/1309.2388/) (2013)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm](https://mlsanity.com/paper/1310.5715/) (2013)
- [Revisiting Small Batch Training for Deep Neural Networks](https://mlsanity.com/paper/1804.07612/) (2018)
- [Stochastic Optimization with Importance Sampling](https://mlsanity.com/paper/1401.2753/) (2014)
- [On the Generalization Benefit of Noise in Stochastic Gradient Descent](https://mlsanity.com/paper/2006.15081/) (2020)
- [Variance Reduction in SGD by Distributed Importance Sampling](https://mlsanity.com/paper/1511.06481/) (2015)
- [The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects](https://mlsanity.com/paper/1803.00195/) (2018)

## Abstract (arXiv)

This work addresses the problem of variance in stochastic gradient estimation for machine learning optimization. Deep learning relies on mini-batch methods such as stochastic gradient descent, which approximate full gradients but introduce noise, creating trade-offs between convergence stability, speed, and generalization. Existing methods, including variance reduction techniques (e.g., SVRG and SAG) and adaptive optimizers, aim to mitigate gradient noise but may introduce additional computational overhead. We propose a model-assisted sampling framework that interprets mini-batch gradients through survey sampling theory, treating the dataset as a fixed finite population and gradients as sample-based estimates. Our aim is to bridge machine learning optimization and survey sampling theory by combining their perspectives on sample-based estimation and variance reduction. By incorporating auxiliary gradient-prediction models, we construct more efficient gradient estimators, with uniform sampling arising as a special case when no auxiliary information is used. Our approach integrates easily with existing optimizers, improving efficiency without altering their dynamics. Empirical results on synthetic and six benchmark datasets show performance gains in 71-86% of the experiments, particularly for medium-sized input spaces in our benchmarks. Notably, with momentum-based optimizers such as AdamW, the proposed estimator achieves clearly better generalization in roughly half the training epochs compared to baseline estimator.