Stochastic Gradient Optimization with Model-Assisted Sampling Jonne Pohjankukka and Jukka Heikkonen University of Turku, Department of Computing
arXiv:2606.27171v1 [cs.LG] 25 Jun 2026
Abstract 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 minibatch 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.
Keywords: Stochastic gradient optimization, gradient estimation, variance reduction, model-assisted estimation, survey sampling.
1
Introduction
In practical deep learning (DL) scenarios, one typically operates with fixed datasets, assumed to be realizations from some unknown data-generating distribution to which we do not have direct access. Given the substantial computational burden associated with training modern models, often comprising millions or billions of parameters, it is standard practice to use only subsets of the data, i.e., minibatches, during training. This subsampling is not merely a convenience but a necessity: full-batch optimization would be computationally prohibitive and may even degrade generalization performance [16]. From a theoretical perspective, the dataset can be viewed as a finite population from which samples are repeatedly drawn. As a result, gradient estimates computed on mini-batches are inherently stochastic and subject to sampling variance, since they approximate the full population gradient using only partial information. This stochasticity can lead to oscillatory or unstable optimization trajectories and may slow convergence, particularly in smooth or ill-conditioned problems. On the other hand, in highly non-convex landscapes, gradient noise can provide beneficial exploratory behavior, helping the optimization process escape saddle points and sharp local minima. Consequently, there exists a fundamental trade-off between gradient variance, convergence stability, and generalization performance [31, 16, 35]. Within this stochastic setting, iterative optimization algorithms form the backbone of model training in machine learning (ML), particularly in DL. Due to computational constraints, gradient-based methods dominate, with much contemporary research focusing on extensions of classical gradient descent (GD), such as stochastic gradient descent (SGD), Adam [17, 22], and Nesterov’s accelerated momentum [25], among others [28]. These methods primarily aim to improve learning efficiency by adapting step sizes, incorporating momentum, or leveraging information from past updates. Despite their differences, they all rely on stochastic gradient estimates derived from mini-batches, and thus inherently contend with the variance introduced by subsampling. As a result, much of their effectiveness hinges on how well they balance fast progress toward minima with robustness to gradient noise, effectively determining the direction and magnitude of each update step to facilitate stable and efficient convergence [12]. Building on this, variance reduction of stochastic gradients has emerged in modern machine learning research. Standard stochastic optimizers exhibit a persistent variance floor due to noisy gradient 1
estimates, which can slow convergence and limit solution accuracy, especially near optima. To address this, a rich body of work has developed methods that explicitly reduce gradient variance while preserving computational efficiency. Techniques such as stochastic variance reduced gradient (SVRG) and stochastic average gradient (SAG) employ control variates constructed from past gradients or periodically computed full gradients, yielding unbiased estimators with significantly lower variance and provably faster convergence in certain settings [14, 30]. Extensions like SAGA further refine these ideas through memory-based corrections [9], while mini-batch strategies trade additional computation for variance reduction via averaging. In the context of DL, adaptive methods such as Adam can also be interpreted as implicitly mitigating variance through moment estimation [18]. More recent research explores hybrid and datadependent approaches, emphasizing that controlling both the magnitude and direction of gradient noise is crucial for achieving efficient, stable, and generalizable training [9, 3]. While variance reduction has been extensively studied within ML, the broader problem of reducing noise in sample-based estimators has a long history in classical statistics, particularly in the field of survey sampling. The general problem of reducing the variance of noisy sample-based population estimates has been extensively studied in the classical context of statistical survey sampling, which focuses on efficiently estimating population parameters from limited samples when access to the full population is infeasible [29]. The field encompasses a wide range of sampling-based methodologies, including design-based and model-based approaches, as well as hybrid model-assisted sampling methods that combine elements of both paradigms. In design-based sampling, inference is based on the randomization induced by the sampling design, and population parameters are treated as fixed but unknown quantities. These methods are generally robust and provide unbiased estimators under the sampling design, but may be statistically inefficient when strong auxiliary information is available but not fully utilized. In contrast, modelbased approaches assume an explicit statistical model for the population generation process, allowing potentially more efficient estimation and prediction by leveraging structural assumptions about the data [10]. However, their performance depends heavily on the validity of the assumed model and may become biased under model misspecification. Model-assisted approaches seek to combine the robustness of designbased inference with the efficiency gains of model-based methods. Common techniques in statistical survey sampling include simple random sampling, stratified sampling, systematic sampling, two-phase sampling or generalized regression estimator [29, 8]. In this work, we propose a novel model-assisted sampling framework to mitigate the variance-induced noise floor of stochastic gradient estimators. We show that the commonly used uniform sampling strategy employed in mini-batch SGD can be interpreted as a special case of our framework, corresponding to the use of a trivial prediction model for the unsampled loss gradients. Our approach formulates the empirical risk gradient field as an unknown and inaccessible population quantity that must be estimated from a stochastic sample drawn from a fixed population. This perspective enables the application of classical survey sampling methodology to stochastic optimization. Rather than proposing a new optimizer, our method improves existing optimization algorithms by replacing the standard mini-batch gradient estimator with a more statistically efficient estimator based on model-assisted sampling. To the best of our knowledge, the use of model-assisted survey sampling principles for mini-batch gradient estimation has received little attention in the ML optimization literature. The remainder of this paper is organized as follows. In Section 2, we briefly review the foundations of empirical risk minimization in ML, SGD, and classical survey sampling estimators, including the Horwitz–Thompson and difference estimators. We also link the gradient noise floor to the asymptotic convergence region of stochastic optimization, showing that lower-variance gradient estimators can reduce the residual error term and thereby improve convergence stability under standard smoothness and strongconvexity-type assumptions. In Section 3, we introduce the proposed model-assisted unbiased gradient estimator. Next in Section 4, we empirically validate our estimator against the baseline uniform-sampling gradient estimator on benchmark datasets. Finally, in Sections 5 and 6 we cover the discussion and conclusions.
2
Empirical risk minimization
The general goal in ML, given a hypothesis set H, a joint probability distribution P(x, y), and a loss function ℓ(h(x), y), is to find a hypothesis h ∈ H such that the risk function: Z R(h) = EP [ℓ(h(x), y)] = ℓ(h(x), y) dP(x, y), (1)
2
is minimized, where x, y are the corresponding explanatory and response variables. The tuple d = (x, y) is generally called a data point. Our task is thus to find a hypothesis h∗ which is defined as: h∗ := arg min R(h).
(2)
h∈H
Since P is practically always unknown to us, we need to rely on an estimator of R, i.e., R̂ which relies on a sample population S = {(xi , yi )} for i = 1, 2, ..., N , commonly defined as: N
R̂S (h) =
1 X ℓ(h(xi ), yi ). N i=1
(3)
Due to the law of large numbers R̂S converges to the true risk R as N → ∞. Thus, in practice, our target of optimization is R̂S rather than R, so the resulting optimization problem is to find a hypothesis h∗ such that: h∗ = arg min R̂(h). (4) h∈H
2.1
Gradient descent optimization
In many ML cases, analytical solutions of Eq. 4 are not possible, so we utilize iterative methods, most commonly using some variation of the gradient descent (GD) algorithm [4, 5]. Our hypothesis h is typically parameterized by a vector θ, i.e., h = h(θ). Thus, when searching for the optimal parameters θ ∗ , the estimated empirical risk can be written directly as a function of θ, i.e., R̂S (h) = R̂S (θ). In GD, we optimize the parameters with the general update rule as follows: θ t+1 = θ t − η∇θ R̂S (θ t )
(5)
where θ t represents the parameter vector at iteration t, η > 0 is the learning rate, ∇θ R̂S (θ t ) is the gradient of the empirical risk with respect to θ at step t. The goal is to iteratively update θ such that the value of R̂S (θ) decreases, ideally converging to a global minimum. The gradient of the empirical risk is clearly defined as: N N 1 X 1 X ∇θ ℓ(h(xi ), yi ) = g, (6) ∇θ R̂S (θ) = N i=1 N i=1 i where we have denoted gi = ∇θ ℓ(h(xi ), yi ). In practice, especially in the deep learning setting where θ may consist of a very large number of parameters, computing the gradient of the empirical risk using the full sample population S is typically infeasible due to e.g. memory constraints. This is because large amounts of intermediate activations and gradients must be stored in memory in order to carry out the parameter update in Eq. 5. For this reason, an unbiased mini-batch solution is typically used, the stochastic gradient descent (SGD) [3], in which a subsample of the data set is taken SI ⊂ S, and the corresponding gradient is calculated using this subsample: 1X ∇θ R̂SI (θ) = gi , (7) n i∈I
where I ⊂ [N ] = {1, 2, ..., N } is a random subset of data indexes from the full sample set S (called population) so that |I| = n for fixed batch size n. Thus, the actual update steps become: θ t+1 = θ t − η∇θ R̂SI (θ t ).
(8)
In line with recent work on gradient estimation and sampling strategies [24, 34, 1], our focus is on improving the gradient itself, which determines the direction of movement in the parameter space.
2.2
Horwitz-Thompson and difference estimators
In the field of survey sampling [8, 21], a common method for making an unbiased estimator, with data points having inclusion probabilities πi > 0, i.e., probability of being included in a sample, is by using the Horwitz-Thompson (also called the π-estimator) style estimators [29]: N
QHT =
1 X qi 1 X qi Ii = , N i=1 πi N πi i∈I
3
(9)
where Ii is a random variable taking either value 0 or 1 if sample point is included in sample or not, I ⊂ [N ], and qi is some fixed non-random data point from the full sample S, also called the population. The π-estimator is an unbiased estimator when averaged over the distribution D, i.e. sampling design, of possible sample selections: ED [QHT ] =
N N N q 1 X qi 1 X 1 X ED [Ii ] i = πi = q, N i=1 πi N i=1 πi N i=1 i
(10)
and it has the covariance matrix [29]: V(QHT ) =
N X N X πij − πi πj i=1 j=1
πi πj
qi q⊤ j ,
(11)
where πij = ED [Ii Ij ], i.e., the probability that both points i and j are included in the sample SI . There exists a potentially more efficient version of the π-estimator called difference estimator (πd estimator), which is a model-assisted estimator designed to incorporate model-learned estimation capability into correcting the error in the π-estimator. The difference estimator is defined as: ! ! N N N X q − q̂ X 1 X 1 X qi − q̂i i i Qdif = q̂ + = q̂ + , (12) Ii N i=1 i i=1 πi N i=1 i πi i∈I
where q̂i is a model-based estimate of qi . For example, if xi is a vector of explanatory variables with corresponding response vector qi ∀i, then q̂i = q(xi ; SI ) for some function q. The covariance matrix of the πd -estimator is: V(Qdif ) =
X πij − πi πj i,j
=
X
πi πj
(qi − q̂i )(qj − q̂j )⊤ (13)
pij ei e⊤j ,
i,j π −π π where we have denoted pij := ijπi πji j
and ei := qi − q̂i ∀ i, j. Notice that in a simple random sampling without replacement (SRSWOR) case, i.e., πi = n/N where |I| = n and πij = n(n − 1)/N (N − 1) with gradient model q := 0, the πd -estimator reduces to the sample average, i.e., corresponding to the uniform mini-batch gradient in SGD. The trace of the variance is: X X Tr{V(Qdif )} = pij Tr ei e⊤j = pij ⟨ei , ej ⟩ . (14) i,j
i,j
It is clear that Tr {V(Qdif )} → 0 as ∥ei ∥ → 0 ∀i. In other words, the better the model q, the more efficient the πd -estimator is. This estimator induces the main motivation for our work; by producing a good model for the (population) gradient, we obtain a more efficient estimator for it.
2.3
Theoretical motivation for model-assisted gradient estimation
The convergence behavior of stochastic optimization methods is closely tied to the variance of the stochastic gradient estimator [13, 11]. In particular, for strongly convex and smooth objective functions f (in our case, the loss function ℓ in Equation 1), the asymptotic convergence region of SGD-type parameter updates in Equation 8 is determined by the gradient noise variance. This provides theoretical motivation for variance-reduction strategies such as the proposed model-assisted gradient estimator. We present our motivating theorem by following the work of Gower et al. [13] beginning from few key assumptions and definitions. Considering the general optimization problem (as in Equation 3) # " N 1 X ∗ fi (x) , (15) x = arg min f (x) = N i=1 x∈Rd where each fi : Rd → R is smooth, but not necessarily convex. Also, it is assumed that f has a unique global minimizer x∗ and is µ-strongly (µ > 0) quasi-convex [15], i.e.: µ f (x∗ ) ≥ f (x) + ⟨∇f (x), x∗ − x⟩ + ∥x∗ − x∥2 (16) 2 for all x ∈ Rd . 4
Definition 1 (Expected smoothness). We say that f is L-smooth in expectation with respect to a sampling design D if there exists L = L(f, D) > 0 such that h i 2 ED ∥∇fv (x) − ∇fv (x∗ )∥ ≤ 2L f (x) − f (x∗ ) , (17) for all x ∈ Rd , where fv is an unbiased stochastic (mini-batch) estimator of f . When the expected smoothness conditions hold for f , we write (f, D) ∼ ES(L). Assumption 1 (Finite gradient noise). The gradient noise is assumed to be finite at global minimum, i.e.: h i 2 σ 2 := ED ∥∇fv (x∗ )∥ < ∞. (18) Next, we present the theorem by Gower et al. using our analogous notation from section 2.1: Theorem 1. Assume ℓ is µ-quasi-strongly convex and that (ℓ, D) ∼ ES(L). Choose ηt = η ∈ 1 0, 2L for all t. Then the iterates of SGD (Equation 8) satisfy 2ησ 2 ED ∥θ t − θ ∗ ∥2 ≤ (1 − ηµ)t ∥θ 0 − θ ∗ ∥2 + , µ h i where now σ 2 = ED ∇θ R̂SI (θ ∗ ) .
(19)
The first term in the bound decreases geometrically with the iteration index t, while the second term defines the asymptotic noise floor induced by stochastic gradient variance σ 2 . Consequently, even when the optimization process converges, the iterates remain confined to a neighborhood around the optimum whose size is directly proportional to the variance of the gradient estimator. This observation motivates the proposed model-assisted gradient estimator. Since the asymptotic error floor scales linearly with σ 2 , reducing the variance of the stochastic gradient estimator directly tightens the convergence bound and decreases the limiting optimization error. In this sense, among unbiased gradient estimators, the estimator achieving the smallest variance is theoretically optimal with respect to this convergence guarantee in SGD optimization. For the full-batch gradient, it is clear that σ 2 = 0, which sets the lower bound for stochastic gradient estimators. Using Theorem 19 and the model-assisted estimator of Equation 14 it is obvious that the model-assisted gradient estimator q ∗ ∈ Q which minimizes the prediction residuals ei in 14, thus also minimizing gradient noise, is explicitly defined as: N
1 X 2 ∥gi − q(xi )∥ , q∈Q N i=1
q ∗ = arg min σq2 = arg min q∈Q
(20)
where Q is the set of gradient models and σq2 is the variance of the model-assisted gradient estimator (using gradient model q) at θ ∗ .
3
Model-assisted gradient estimator
The artificially generated risk surface in Figure 1 gives visual motivation for our approach. Given a dataset defining a population risk function through a finite-sum objective [11], each individual sample contributes to the geometry of the overall loss surface in parameter space. As the number of sampled data points increases, the empirical risk surface converges toward the population risk surface. In the small-sample regime, however, the empirical risk may differ substantially from the true population risk, leading to inaccurate gradient estimates. Consequently, the optimization trajectory may exhibit noisy oscillations or move in suboptimal directions. More accurate risk estimates yield descent directions that better approximate the true population gradients, resulting in more stable optimization. The model-assisted gradient estimation is illustrated in Figure 2. As the (full-batch) population gradient is generally inaccessible in practical optimization settings, the true population risk surface and its exact gradients cannot be easily evaluated. Therefore, we rely on mini-batch gradient estimators constructed from small random subsets of the data. The proposed gradient estimator combines observed sample gradients with predictions from a gradient model estimating the unsampled components. The most similar study to ours was presented by Wang et al. [32] who proposed a control-variate method to reduce the variance of stochastic gradients. The key idea is to reduce stochastic-gradient 5
2% of data
10% of data
20% of data
45% of data
60% of data
75% of data
90% of data
100% of data
Figure 1: Artificially generated subsampled risk surfaces in a two-dimensional parameter space. Normalized steepest descent negative gradient directions are visualized as green arrows. The figure illustrates how the estimated artificial risk surface develops as increasingly larger fractions of the data are used to compute the loss and its descent directions, ranging from 2% of the data to the full dataset. Gradients at each sample data point
Noisy random sample (mini-batch) gradient
Estimated true population gradient
Direction of the true population (full-batch) gradient
(a)
Estimated gradients at unsampled data points
(b)
(c)
Figure 2: (a) The true population loss gradient calculated using all samples (full-batch). (b) Estimated gradient using a sample (i.e. mini-batch) of the population data. (c) Model-assisted gradient estimation. A mini-batch of data is sampled and used to estimate unsampled loss gradients. variance by adding a control-variate correction based on data statistics, such as low-order moments of the entire data. Their estimator replaces the noisy stochastic gradient with an unbiased corrected gradient by subtracting a control-variate term, which requires computing the control variate expectation hd (w) [32] and estimating the scaling matrix or coefficient A [32] so that the correction reduces variance. In our model-assisted approach from survey sampling theory, the estimators reduce the gradient variance by explicitly modeling the target gradient. Let S denote the data population as before. We define the corresponding gradient population as G = {gi = (xi , gi ) | di ∈ S, i ∈ [N ]} ,
(21)
where gi = ∇θ ℓ(h(xi ), yi ), xi ∈ Rm and gi ∈ Rd ∀i where m, d ∈ N is the number of input data dimensions and trainable parameters of hypothesis h respectively. In other words, G contains tuples pairing each input feature vector with its corresponding risk/loss gradient. Given these, we present the model-based gradient estimator in Algorithm 1. First, an initial subset of indices I 1 ⊂ [N ] is drawn uniformly at random. In the second stage, inclusion probabilities are assigned: elements in I 1 are always selected πi = 1, while the remaining elements receive probabilities 0 < πi < 1. A random sample GI = GI 1 ∪ GI 2 ⊂ G is then drawn according to sampling design D induced by these probabilities, corresponding to index sets I 1 and I 2 ⊂ [N ] \ I 1 , with |I 1 | + |I 2 | = n1 + n2 = n. Next, we train the gradient model q : Rm → Rd , x 7→ g using GI 1 . Thus, the gradient model with closed-form solution is 6
Algorithm 1 Model-assisted gradient estimator Require: Fixed gradient population G. Ensure: Unbiased gradient estimate gdif 1: Draw uniformly a random index set I 1 ⊂ [N ] 2: for each i ∈ I 1 do 3: Set inclusion probability πi = 1 4: end for 5: for each i ∈ / I 1 do 6: Choose 0 < πi < 1 7: end for 8: Define a sampling design D induced by the inclusion probabilities πi , i ∈ [N ] 9: Given g ∼ D, obtain a random sample GI = GI 1 ∪ GI 2 ⊂ G, where index sets I 1 and I 2 correspond to πi = 1 and 0 < πi < 1 cases respectively, with |I 1 | + |I 2 | = n1 + n2 = n 10: Fit gradient model q : Rm → Rd , x 7→ g using GI 1 11: for i = 1, . . . , N do 12: ĝi ← q(xi ) 13: end for 14: Compute difference estimator N X g − ĝ X 1 X k k (gj − ĝj ) + ĝ + gdif = N i=1 i πk 2 1 j∈I
k∈I
15: return gdif
deterministic w.r.t D. Finally, an unbiased difference estimator gdif of the population gradient is formed using Equation 12. The performance of the proposed estimator is fundamentally linked to the predictive accuracy of the gradient model q. Consequently, its effectiveness may deteriorate in high-dimensional settings due to curse-of-dimensionality (COD) effects, highly complex gradient landscapes, or insufficient sample sizes for reliable model construction. Nevertheless, the conventional uniform mini-batch gradient estimator n arises as a special case (q := 0, πi = N ∀i) of the proposed formulation, indicating that the proposed approach generalizes standard stochastic gradient estimation.
4
Empirical analysis
We empirically evaluate the behavior of our model-assisted gradient estimator against a baseline uniformly sampled gradient (i.e. mini-batch) over four different optimizers: SGD, SGD with momentum (SGD-M, momentum coefficient β = 0.9), Adam and AdamW. The experiments are tested over seven datasets: synthetic (a random 1D sinusoidal function with a downward-opening parabolic trend), Airfoil self-noise [6], Appliances energy [7], MNIST [20], Fashion-MNIST [33], CIFAR-10 and CIFAR-100 [19]. Our primary goal is to understand how the proposed variance-reduced gradient estimator compares to standard minibatch estimator and full-batch optimization in terms of convergence speed, stability, and final performance across different multilayer perceptron (MLP) and convolutional neural networks (CNN) with varying input space and model parameter sizes. The tested datasets and corresponding models are listed in Table 1. As the proposed model-assisted estimator in Algorithm 1 is based on a prediction model q, we will utilize kernel ridge regression (KRR, regularized kernel least squares) [2, 26] to model the gradient as we can solve the optimal solution analytically and cheaply (with small batch size). For the KRR-model, computations were performed using Scikit-learn library [27] with Gaussian kernel and default parameters (α = 0.1, γ = 1). We focus on three complementary metrics: (i) test loss, capturing generalization quality, (ii) loss standard deviation across runs, reflecting stability, and (iii) the epoch at which the loss reaches its minimum, indicating efficiency. These metrics allow us to jointly assess not only how well an optimizer performs but also how reliably and how quickly it reaches a good solution. The learning rate is fixed to η = 5×10−3 and run for 100 epochs in all experiments with batch sizes of 10 (n1 = 8, n2 = 2), 50 (n1 = 30, n2 = 20) and 100 (n1 = 80, n2 = 20). Reported results correspond to averaged performance trajectories over 400 distinct runs with random model initialization and data seeds (visualizations are shown up to 50 7
Table 1: Model architectures used for each dataset in the experiments. CIFAR-10 and CIFAR-100 use the same CNN architecture, differing only in the output layer dimensionality. Dataset(s)
Model
Input size
# Params
Architecture
Synthetic
MLP
1
321
Airfoil self-noise
MLP
5
385
Appliances energy
MLP
27
737
MNIST, FashionM- CNN NIST
784
13 978
CIFAR-10
CNN
3072
17 962
CIFAR-100
CNN
3072
19 492
Linear(1 → 16) + ReLU, Linear(16 → 16) + ReLU, Linear(16 → 1). Regression output without final activation. Linear(5 → 16) + ReLU, Linear(16 → 16) + ReLU, Linear(16 → 1). Regression output without final activation. Linear(27 → 16) + ReLU, Linear(16 → 16) + ReLU, Linear(16 → 1). Regression output without final activation. Conv2d(1 → 8, k = 3, p = 1) + ReLU + MaxPool(2), Conv2d(8 → 16, k = 3, p = 1) + ReLU + MaxPool(2), Flatten(16×7×7), Linear(784 → 16) + ReLU, Linear(16 → 10). Conv2d(3 → 8, k = 3, p = 1) + ReLU + MaxPool(2), Conv2d(8 → 16, k = 3, p = 1) + ReLU + MaxPool(2), Flatten(16 × 8 × 8), Linear(1024 → 16) + ReLU, Linear(16 → 10). Same architecture as CIFAR-10, except the final classification layer is Linear(16 → 100) instead of Linear(16 → 10).
epochs when curve trajectories have mostly stabilized). To make repeated experiments computationally feasible while still allowing variability across runs to be estimated reliably, each experiment was performed on a randomly sampled subset of 1,000 examples from the corresponding dataset. Of these, 800 examples were used for training and 200 were held out to evaluate generalization performance using the test loss. The results presented are intended to reveal consistent patterns across datasets rather than optimize performance for any single task. In particular, we are interested in whether model-assisted gradient estimators provide systematic improvements over standard baseline approach and how these improvements interact with optimizer choice and dataset characteristics. The Figures 3 and 4 represent the mean test loss curves for the SGD and AdamW optimizers, corresponding to the worst and best cases for our estimator respectively. For regression datasets, the loss is measured in terms of mean squared error and for image classification datasets, the loss is measured in terms of cross-entropy loss. The corresponding SGD-M and Adam figures are presented in the Appendix. Similarly, we illustrate the corresponding evolution of L2 distance to the full-batch gradient in Figures 5 and 6, which highlight how closely the two competing estimators follow the true population gradient. With vanilla SGD, our estimator does not show a clear overall advantage over the baseline, except for the synthetic sinusoid dataset. In this dataset case, we can see that our estimator achieves best test loss performance, as well as L2-distance to full-batch gradient, especially with higher batch sizes. The worst performance for our estimator with SGD occurs with higher dimensional input data cases (CIFAR). With momentum and past gradient utilizing optimizers (SGD-M, Adam, AdamW) we can notice a drastic change in the performance for our estimator. Although the model-assisted estimator shows a slightly larger average deviation from the full-batch gradient, it generally outperforms the baseline in terms of generalization performance, with the clearest and most statistically significant improvements observed on the MNIST datasets. In Tables 2, 3 and 4 we have listed detailed best results over all batch size cases respectively, which support the visual inspection of Figures 3, 4, 8 and 9. The tables’ cell colors are based on an equally weighted normalized score computed within each dataset from test loss, standard deviation, and epoch. Thus, the coloring highlights overall estimator performance in terms of accuracy, stability, and efficiency. With vanilla SGD, our estimator outperforms the baseline in approximately 14–29% of dataset cases. In contrast, when using momentum-based optimizers, the win rate increases substantially to 57–86%. The best observed performance is achieved with AdamW at batch size 100, where our estimator surpasses the baseline in 86% of cases. Across datasets, performance varies: CIFAR-100 represents the most challenging setting (win rate 25–50%), whereas MNIST yields the strongest results (75%), along with consistently lower and more stable test loss, achieved in nearly half the number of epochs. When results are aggregated over optimizers and datasets, the proposed estimator outperforms the baseline in 54% of the comparisons for batch sizes 10 and 50, and in 61% of the comparisons for batch size 100. When the optimizer is treated as a tunable hyperparameter on a per-dataset basis, the win rate further improves to 71–86%.
8
Batch size 10
20
30 epoch
40
50
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
0.08
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
0.10
30 epoch
40
50
Full-batch Our estimator Baseline
0.08
0.02 0 2.5
10
20
30 epoch
2.0
40
50
Full-batch Our estimator Baseline
20
30 epoch
40
50
Full-batch Our estimator Baseline
30 epoch
40
50
Full-batch Our estimator Baseline
20
30 epoch
40
50
Full-batch Our estimator Baseline
20
30 epoch
40
50
Full-batch Our estimator Baseline
2.00 5.0
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
30 epoch
40
50
4.60
0.02 0 2.5
10
20
30 epoch
40
50
30 epoch
40
50
40
50
2.0
0.50
Full-batch Our estimator Baseline
10
20
2.0
0.50 2.6
Full-batch Our estimator Baseline
10
20
30 epoch
Full-batch Our estimator Baseline
2.4
2.00 5.0
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
4.9 loss
loss
loss
20
Full-batch Our estimator Baseline
4.8
4.7 10
50
2.2
4.8
4.7
40
loss 10
4.9
4.8
30 epoch
0.08
1.0
2.2
10
20
1.5
2.4
2.2
10
loss 10
loss
2.4
50
Full-batch Our estimator Baseline
0.10
1.0
2.0
0.50 2.6
40
loss
20
30 epoch
1.5
1.0 10
20
loss 10
loss
1.0
10
0.04
1.5
loss
1.5
0.5 0 0.30 0.25 0.20 0.15 0.10 0.05 0.000
0.06
2.0
0.50
Full-batch Our estimator Baseline
1.5 loss
Full-batch Our estimator Baseline
1.0
loss
Fashion-MNIST CIFAR-10
50
1.5
1.0
CIFAR-100
40
loss
1.5
4.60
30 epoch
0.04
2.0
4.9
20
loss
0.04
2.00 5.0
10
0.06
loss
0.06
0.50 2.6
0.5 0 0.30 0.25 0.20 0.15 0.10 0.05 0.000
Batch size 100
1.0
loss
loss
Full-batch Our estimator Baseline
0.10
0.50
1.5 1.0
10
2.0
Full-batch Our estimator Baseline
loss
1.0
0.02 0 2.5
Batch size 50
loss
loss
1.5
0.5 0 0.30 0.25 0.20 0.15 0.10 0.05 0.000
2.0
Full-batch Our estimator Baseline
loss
MNIST
Appliances energy
Airfoil self-noise
Synthetic
2.0
4.7 10
20
30 epoch
40
50
4.60
10
20
30 epoch
40
50
Figure 3: The generalization performance for SGD optimizer with baseline (uniform mini-batch), modelassisted (our estimator) and full-batch gradients. Full-batch gradient case is listed for comparison purposes. The darker curve represents the average test loss of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
9
Batch size 10
1.2
Full-batch Our estimator Baseline
1.0
1.0
40
50
0 0.020
Full-batch Our estimator Baseline
loss 30 epoch
40
50
0 0.020
Full-batch Our estimator Baseline
40
50
0.0000 0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
50
2.5 2.0 1.5 1.0 0.5 0.00
50
3.0 2.5 2.0 1.5 1.0 0.50
50
20 30 epoch
40
20 30 epoch
40
50
Full-batch Our estimator Baseline
20
30 epoch
40
10
20 30 epoch
40
50
50
2.5 2.0 1.5 1.0 0.5 0.00
50
3.0 2.5 2.0 1.5 1.0 0.50
Full-batch Our estimator Baseline
10
20
30 epoch
40
Full-batch Our estimator Baseline
loss
Full-batch Our estimator Baseline
10
20
30 epoch
40
6
Full-batch Our estimator Baseline
10
20
30 epoch
40
6
Full-batch Our estimator Baseline
4
4
loss
5
loss
5
0.0000 0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
3
3
3
10
20
30 epoch
40
50
5.5
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
4.5 10
20
30 epoch
40
50
4.5 4.00
50
Full-batch Our estimator Baseline
10
20 30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
30 epoch
40
50
40
50
Full-batch Our estimator Baseline
20 6.0 5.5
10
20
Full-batch Our estimator Baseline
5.0
loss
loss
5.0
Full-batch Our estimator Baseline
40
4
5.5
5.0
4.00
20 6.0
20 30 epoch
loss
20 6.0
10
loss
10
Full-batch Our estimator Baseline
loss 10
loss
loss
Full-batch Our estimator Baseline
50
loss
10
40
0.005
loss
loss
Full-batch Our estimator Baseline
30 epoch
loss
20 30 epoch
20
0.010
loss 10
10
0.015
loss
5
20
0.005
0.0000 0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
3.0 2.5 2.0 1.5 1.0 0.50
0.4 10
0.015
0.005
2.5 2.0 1.5 1.0 0.5 0.00
0.6
0.010
loss
0.010
loss
Airfoil self-noise Appliances energy MNIST Fashion-MNIST CIFAR-10
30 epoch
0.015
6
CIFAR-100
20
Full-batch Our estimator Baseline
1.0
0.4 10
Batch size 100
0.8
0.6
0.4
1.2
Full-batch Our estimator Baseline
0.8
0.6
0 0.020
Batch size 50
loss
0.8
loss
Synthetic
1.2
4.5 10
20
30 epoch
40
50
4.00
10
20
30 epoch
Figure 4: The generalization performance for AdamW optimizer with baseline (uniform mini-batch), model-assisted (our estimator) and full-batch gradients. Full-batch gradient case is listed for comparison purposes. The darker curve represents the average test loss of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
10
50
20
30 epoch
40
50
50
10 8 6 4 2 00
50
10 8 6 4 2 00
40
Our estimator L2 Baseline L2
10
20
30 epoch
40
Our estimator L2 Baseline L2
10
20
30 epoch
40
L2 distance 10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
50
Our estimator L2 Baseline L2
L2 distance
L2 distance
40
2
1 20
30 epoch
40
50
00
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10 8 6 4 2 00
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10 8 6 4 2 00
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
8
30 epoch
40
50
Our estimator L2 Baseline L2
6
0 4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
3 2
1
10
0.00 1.50 1.25 1.00 0.75 0.50 0.25 0.000
2
3
2
Our estimator L2 Baseline L2
4
2 0 4
50
L2 distance
30 epoch
0.6
40
L2 distance
4
20
30 epoch
0.2
6
10
20
0.4
8
Our estimator L2 Baseline L2
10
L2 distance
30 epoch
Our estimator L2 Baseline L2
L2 distance
20
2
00
Our estimator L2 Baseline L2
L2 distance
10
4
3
50
L2 distance
Our estimator L2 Baseline L2
6
0 4
40
L2 distance
10 8 6 4 2 00
0.6
0.00 1.50 1.25 1.00 0.75 0.50 0.25 0.000
50
L2 distance
L2 distance
10 8 6 4 2 00
30 epoch
0.2
10
Batch size 100
L2 distance
0.2
0.00 1.50 1.25 1.00 0.75 0.50 0.25 0.000
20
L2 distance
L2 distance
Our estimator L2 Baseline L2
10
0.4
L2 distance
CIFAR-10
40
0.4
8
CIFAR-100
30 epoch
10 8 6 4 2 00
Our estimator L2 Baseline L2
L2 distance
20
Batch size 50
L2 distance
Our estimator L2 Baseline L2
10
10 8 6 4 2 00
L2 distance
L2 distance
Batch size 10
0.6
L2 distance
Synthetic Airfoil self-noise Appliances energy MNIST Fashion-MNIST
10 8 6 4 2 00
1
10
20
30 epoch
40
50
00
10
20
30 epoch
40
50
Figure 5: L2 distance to full-batch gradient with SGD optimizer. The darker curve represents the average of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
11
10
20
50
30 epoch
40
50
Our estimator L2 Baseline L2
0.3
20
30 epoch
40
50
20
30 epoch
40
50
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
4
30 epoch
3
40
50
Our estimator L2 Baseline L2
2
20
30 epoch
50
0 0
4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6 4
2 20
30 epoch
40
50
Our estimator L2 Baseline L2
4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6 4
2 30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
0 0 10 8 6 4 2 0 0
40
50
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
4
30 epoch
3
40
50
Our estimator L2 Baseline L2
0 0
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6
0 0 8
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6 2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
L2 distance
L2 distance
20
10
4
2
10
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
2
L2 distance
L2 distance
6
0 0 8
Our estimator L2 Baseline L2
0.3
4
2
10
50
1
L2 distance
Our estimator L2 Baseline L2
L2 distance
6
40
40
2
1
10
30 epoch
L2 distance
1
20
0.1
L2 distance
Our estimator L2 Baseline L2
10
L2 distance
10
00 0.4
0.2
L2 distance
Our estimator L2 Baseline L2
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
5
L2 distance
10
2
0 0 10 8 6 4 2 0 0
20
0.1
3
0 0 8
10
0.2
0.1
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
00 0.4 L2 distance
L2 distance
0.2
L2 distance
MNIST Fashion-MNIST
40
Our estimator L2 Baseline L2
0.3
0 0
CIFAR-10
30 epoch
5
L2 distance
Our estimator L2 Baseline L2
10
L2 distance
00 0.4
10
Our estimator L2 Baseline L2
15
0 0 10 8 6 4 2 0 0
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
L2 distance
5
Batch size 100
Our estimator L2 Baseline L2
L2 distance
L2 distance
L2 distance
10
4
CIFAR-100
Batch size 50 15
L2 distance
Appliances energy
Airfoil self-noise
Synthetic
Batch size 10 15
10
20
30 epoch
40
50
10
20
30 epoch
40
50
Figure 6: L2 distance to full-batch gradient with AdamW optimizer. The darker curve represents the average of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
12
Table 2: Results for batch size 10. Values are reported as mean minimum test loss ± standard deviation (epoch) over 400 distinct runs. Cell color is based on an equally weighted normalized combined score within each dataset, computed from loss, standard deviation, and epoch. Greener indicates better combined performance, while redder indicates worse combined performance. Full-batch rows are shown in gray as references. Bold indicates the lowest (best) score between our estimator and baseline within each dataset. The superscript ∗ indicates better score between our estimator and the baseline within the same dataset and optimizer. The final column reports dataset-wise win rates, and the bottom row reports optimizer-wise win rates for our estimator. The bottom-right gray cell reports the overall win rate of our estimator across all dataset–optimizer comparisons in the table. Dataset
Case
Synthetic
Our estimator Baseline Full-batch
Airfoil self-noise
Our estimator 0.0087 ± 0.0012 (100) 0.0050 ± 0.0008 (100) 0.0023 ± 0.0005 (100)∗ 0.0024 ± 0.0006 (99) Baseline 0.0082 ± 0.0010 (100)∗ 0.0050 ± 0.0008 (100) 0.0025 ± 0.0006 (100) 0.0026 ± 0.0006 (98)∗ Full-batch 0.0082 ± 0.0010 (100) 0.0049 ± 0.0008 (100) 0.0018 ± 0.0004 (98) 0.0017 ± 0.0004 (97)
Our estimator Appliances energy Baseline Full-batch
SGD
SGD-M
Adam
AdamW
0.62 ± 0.29 (100) 0.46 ± 0.11 (98)∗ 0.39 ± 0.058 (100)
3.06 ± 1.49 (1) 1.91 ± 1.38 (4)∗ 0.44 ± 0.59 (97)
0.22 ± 0.067 (98)∗ 0.24 ± 0.067 (100) 0.26 ± 0.075 (100)
0.24 ± 0.062 (100)∗ 0.26 ± 0.063 (100) 0.26 ± 0.078 (100)
Row win rate 50.0%
25.0%
0.015 ± 0.0030 (96) 0.014 ± 0.0025 (100)∗ 0.013 ± 0.0025 (100)
0.012 ± 0.0022 (92)∗ 0.012 ± 0.0021 (99) 0.012 ± 0.0021 (100)
0.010 ± 0.0022 (91)∗ 0.010 ± 0.0021 (97) 0.012 ± 0.0023 (10)
0.010 ± 0.0021 (92)∗ 0.010 ± 0.0022 (100) 0.012 ± 0.0024 (15)
75.0%
MNIST
Our estimator Baseline Full-batch
0.76 ± 0.27 (67) 0.59 ± 0.11 (51)∗ 0.61 ± 0.12 (43)
0.58 ± 0.18 (11)∗ 0.60 ± 0.18 (15) 0.66 ± 0.17 (7)
0.46 ± 0.13 (9)∗ 0.58 ± 0.34 (16) 0.69 ± 0.33 (2)
0.45 ± 0.14 (9)∗ 0.63 ± 0.46 (21) 0.73 ± 0.39 (2)
75.0%
Fashion-MNIST
Our estimator Baseline Full-batch
0.92 ± 0.24 (49) 0.75 ± 0.091 (50)∗ 0.76 ± 0.089 (43)
0.78 ± 0.11 (10)∗ 0.79 ± 0.12 (14) 0.78 ± 0.13 (6)
0.74 ± 0.14 (7)∗ 0.87 ± 0.34 (10) 0.91 ± 0.25 (2)
0.71 ± 0.11 (8)∗ 0.78 ± 0.22 (9) 0.88 ± 0.22 (2)
75.0%
CIFAR-10
Our estimator Baseline Full-batch
2.25 ± 0.12 (26) 2.07 ± 0.096 (60)∗ 2.07 ± 0.096 (47)
2.14 ± 0.11 (14)∗ 2.18 ± 0.12 (19) 2.10 ± 0.11 (10)
2.22 ± 0.13 (6) 2.22 ± 0.13 (9)∗ 2.24 ± 0.11 (1)
2.20 ± 0.14 (6) 2.21 ± 0.13 (10)∗ 2.22 ± 0.11 (1)
25.0%
CIFAR-100
Our estimator Baseline Full-batch
4.66 ± 0.031 (7)∗ 4.66 ± 0.027 (10) 4.67 ± 0.029 (14)
4.64 ± 0.022 (5)∗ 4.63 ± 0.021 (11) 4.64 ± 0.027 (7)
4.66 ± 0.034 (1) 4.63 ± 0.018 (4)∗ 4.70 ± 0.087 (1)
4.66 ± 0.044 (1) 4.64 ± 0.020 (4)∗ 4.71 ± 0.11 (1)
50.0%
14.3%
71.4%
71.4%
57.1%
53.6%
Column win rate
Table 3: Results for batch size 50. Notation, coloring, bolding, and superscript markers follow Table 2. Dataset
Case
SGD
SGD-M
Adam
AdamW
Synthetic
Our estimator Baseline Full-batch
0.57 ± 0.097 (95)∗ 0.59 ± 0.17 (100) 0.56 ± 0.11 (99)
1.15 ± 1.28 (100) 0.60 ± 0.63 (93)∗ 0.50 ± 0.65 (99)
0.33 ± 0.048 (98)∗ 0.35 ± 0.048 (98) 0.31 ± 0.056 (100)
0.33 ± 0.049 (100) 0.35 ± 0.052 (99)∗ 0.31 ± 0.051 (100)
Airfoil self-noise
Our estimator Baseline Full-batch
0.014 ± 0.0024 (100) 0.014 ± 0.0024 (100) 0.014 ± 0.0024 (100)
Our estimator 0.018 ± 0.0033 (100) Appliances energy Baseline 0.017 ± 0.0032 (100)∗ Full-batch 0.018 ± 0.0032 (100)
Row win rate 50.0%
0.0070 ± 0.00091 (100) 0.0030 ± 0.00071 (100) 0.0030 ± 0.00073 (100)∗ 0.0071 ± 0.00090 (100) 0.0033 ± 0.00083 (99)∗ 0.0034 ± 0.00082 (100) 0.0071 ± 0.00091 (100) 0.0026 ± 0.00074 (100) 0.0025 ± 0.00076 (100)
25.0%
0.013 ± 0.0022 (100)∗ 0.013 ± 0.0023 (100) 0.013 ± 0.0023 (100)
0.011 ± 0.0021 (96)∗ 0.011 ± 0.0020 (99) 0.012 ± 0.0021 (52)
0.011 ± 0.0021 (100) 0.011 ± 0.0021 (97)∗ 0.011 ± 0.0022 (65)
50.0%
MNIST
Our estimator Baseline Full-batch
0.75 ± 0.13 (100) 0.71 ± 0.13 (100)∗ 0.71 ± 0.12 (100)
0.59 ± 0.13 (30)∗ 0.62 ± 0.17 (39) 0.64 ± 0.17 (34)
0.50 ± 0.16 (17)∗ 0.66 ± 0.35 (26) 0.73 ± 0.45 (9)
0.49 ± 0.13 (17)∗ 0.64 ± 0.35 (30) 0.70 ± 0.34 (10)
75.0%
Fashion-MNIST
Our estimator Baseline Full-batch
0.84 ± 0.095 (100) 0.81 ± 0.084 (100)∗ 0.81 ± 0.085 (100)
0.76 ± 0.11 (28)∗ 0.76 ± 0.11 (43) 0.77 ± 0.12 (31)
0.72 ± 0.11 (13)∗ 0.77 ± 0.16 (20) 0.81 ± 0.23 (7)
0.72 ± 0.14 (13)∗ 0.77 ± 0.20 (18) 0.82 ± 0.28 (7)
75.0%
CIFAR-10
Our estimator Baseline Full-batch
2.17 ± 0.083 (100) 2.15 ± 0.085 (100)∗ 2.15 ± 0.086 (100)
2.10 ± 0.10 (39)∗ 2.12 ± 0.11 (57) 2.11 ± 0.11 (49)
2.20 ± 0.13 (10)∗ 2.21 ± 0.13 (17) 2.22 ± 0.11 (5)
2.19 ± 0.13 (10)∗ 2.18 ± 0.13 (17) 2.22 ± 0.12 (5)
75.0%
CIFAR-100
Our estimator Baseline Full-batch
4.67 ± 0.027 (43)∗ 4.67 ± 0.028 (62) 4.67 ± 0.030 (72)
4.65 ± 0.026 (18) 4.64 ± 0.021 (30)∗ 4.64 ± 0.023 (33)
4.65 ± 0.026 (2) 4.64 ± 0.024 (7)∗ 4.65 ± 0.027 (2)
4.65 ± 0.027 (2) 4.64 ± 0.027 (8)∗ 4.65 ± 0.027 (2)
25.0%
28.6%
57.1%
71.4%
57.1%
53.6%
Column win rate
13
Table 4: Results for batch size 100. Notation, coloring, bolding, and superscript markers follow Table 2. Dataset
Case
SGD
SGD-M
Adam
AdamW
Synthetic
Our estimator Baseline Full-batch
Airfoil self-noise
0.63 ± 0.085 (99)∗ 0.64 ± 0.13 (100) 0.62 ± 0.076 (99)
1.02 ± 1.15 (98) 0.55 ± 0.47 (91)∗ 0.54 ± 0.59 (100)
0.36 ± 0.070 (100)∗ 0.39 ± 0.068 (100) 0.34 ± 0.067 (100)
0.36 ± 0.065 (100)∗ 0.40 ± 0.065 (100) 0.34 ± 0.065 (100)
75.0%
Our estimator 0.018 ± 0.0033 (100) 0.0081 ± 0.0011 (100)∗ 0.0041 ± 0.0009 (100)∗ 0.0040 ± 0.0009 (100)∗ Baseline 0.018 ± 0.0032 (100)∗ 0.0082 ± 0.0012 (100) 0.0042 ± 0.0010 (100) 0.0042 ± 0.0010 (100) Full-batch 0.018 ± 0.0032 (100) 0.0082 ± 0.0012 (100) 0.0036 ± 0.0011 (100) 0.0037 ± 0.0011 (100)
75.0%
Our estimator 0.022 ± 0.0046 (100) Appliances energy Baseline 0.022 ± 0.0043 (100)∗ Full-batch 0.022 ± 0.0042 (100)
Row win rate
0.013 ± 0.0024 (100)∗ 0.013 ± 0.0025 (100) 0.013 ± 0.0025 (100)
0.012 ± 0.0022 (100) 0.012 ± 0.0021 (100)∗ 0.012 ± 0.0022 (99)
0.011 ± 0.0021 (95)∗ 0.011 ± 0.0020 (100) 0.012 ± 0.0024 (99)
50.0%
MNIST
Our estimator Baseline Full-batch
1.11 ± 0.22 (100) 1.07 ± 0.23 (100)∗ 1.07 ± 0.22 (100)
0.61 ± 0.13 (56)∗ 0.65 ± 0.16 (77) 0.66 ± 0.17 (68)
0.51 ± 0.14 (30)∗ 0.68 ± 0.36 (39) 0.75 ± 0.38 (20)
0.50 ± 0.13 (30)∗ 0.68 ± 0.34 (43) 0.72 ± 0.33 (21)
75.0%
Fashion-MNIST
Our estimator Baseline Full-batch
1.02 ± 0.13 (100) 0.97 ± 0.13 (100)∗ 0.97 ± 0.12 (100)
0.76 ± 0.094 (61)∗ 0.76 ± 0.098 (75) 0.77 ± 0.10 (64)
0.74 ± 0.15 (24)∗ 0.80 ± 0.22 (29) 0.87 ± 0.32 (15)
0.72 ± 0.11 (21)∗ 0.78 ± 0.21 (29) 0.82 ± 0.24 (14)
75.0%
CIFAR-10
Our estimator Baseline Full-batch
2.24 ± 0.084 (100) 2.22 ± 0.074 (100)∗ 2.22 ± 0.075 (100)
2.10 ± 0.099 (74)∗ 2.12 ± 0.10 (99) 2.10 ± 0.10 (100)
2.21 ± 0.13 (18) 2.19 ± 0.12 (23)∗ 2.24 ± 0.12 (11)
2.21 ± 0.11 (17)∗ 2.19 ± 0.13 (24) 2.23 ± 0.11 (10)
50.0%
CIFAR-100
Our estimator Baseline Full-batch
4.67 ± 0.031 (78) 4.67 ± 0.027 (99)∗ 4.67 ± 0.028 (100)
4.65 ± 0.024 (35)∗ 4.64 ± 0.025 (72) 4.64 ± 0.023 (66)
4.65 ± 0.027 (3) 4.64 ± 0.026 (9)∗ 4.65 ± 0.026 (3)
4.65 ± 0.025 (3) 4.64 ± 0.023 (8)∗ 4.65 ± 0.026 (3)
25.0%
14.3%
85.7%
57.1%
85.7%
60.7%
Column win rate
5
Discussion
The empirical results reveal dataset-dependent behavior for the proposed estimator, with performance patterns varying across datasets and optimizers. The model-assisted estimator generally improves generalization performance and reach their minimum loss in fewer epochs than standard mini-batch estimator. The proposed estimator also tend to reduce variance across runs, although this improvement is not observed in all settings. In controlled synthetic settings, both generalization performance and the L2-distance to the full-batch gradient behaved as intuitively expected, providing a validation of the model-assisted model optimization under well-specified conditions. Notably, optimizer improvements do not always seem correlate with smaller L2-distance to the fullbatch gradient. While L2-distance measures error magnitude to the population (full-batch) gradient, it does not seem to capture overall directional quality in the longer run. Estimators with larger norm errors relative to the full-batch gradient can still produce more reliable descent directions. This effect is especially relevant for adaptive methods such as Adam and AdamW, where gradients are rescaled using moment estimates. In these cases, the temporal consistency and directional structure of the gradient estimates may matter more than the magnitude of the gradient estimate at any single iteration. On smaller and larger datasets (e.g., Airfoil Self-Noise, CIFAR-10/100), the benefits are more limited, though still generally positive. In the CIFAR settings in particular, model-assisted estimators tend to more closely track the performance of full-batch optimization, including its eventual degradation as fullbatch optimization does not uniformly outperform stochastic methods, highlighting the beneficial role of noise for exploration and generalization. Finally, the improvements using model-assisted estimator comes with additional computational overhead from constructing model-assisted estimators, introducing a trade-off between efficiency and accuracy. While the auxiliary model is typically cheap to evaluate, its quality directly affects the accuracy of the resulting gradient estimates. Given that a single gradient computation can be expensive, especially in large models, the effectiveness of the approach depends on whether the low-cost model can provide sufficiently accurate guidance to justify its use in practice.
6
Conclusions
In this work, we examined gradient variance reduction using a model-assisted estimator inspired by survey sampling theory. Across multiple datasets and optimizers, overall, model-assisted gradient estimation improved training stability, convergence speed, and generalization, with the strongest gains observed for
14
momentum-based methods such as Adam, while vanilla SGD benefited less. At the same time, our results show that variance reduction alone is not sufficient to fully explain generalization, highlighting the need to balance variance control with beneficial stochasticity. We hope this work opens new directions for applying model-assisted estimator frameworks in ML optimization.
References [1]
Guillaume Alain, Yoshua Bengio, et al. “Variance Reduction in SGD by Distributed Importance Sampling”. In: arXiv preprint arXiv:1511.06481 (2015).
[2]
Mauricio A. Álvarez, Lorenzo Rosasco, and Neil D. Lawrence. Kernels for Vector-Valued Functions. Hanover, MA, USA: Now Publishers Inc., 2012. isbn: 1601985584.
[3]
Léon Bottou. “Large-Scale Machine Learning with Stochastic Gradient Descent”. In: Proceedings of COMPSTAT’2010. Ed. by Yves Lechevallier and Gilbert Saporta. Heidelberg: Physica-Verlag HD, 2010, pp. 177–186. isbn: 978-3-7908-2604-3.
[4]
Léon Bottou. “On-line learning and stochastic approximations”. In: On-Line Learning in Neural Networks. USA: Cambridge University Press, 1999, pp. 9–42. isbn: 0521652634.
[5]
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004. isbn: 0521833787.
[6]
Thomas Brooks, D. Pope, and Michael Marcolini. Airfoil Self-Noise. UCI Machine Learning Repository. DOI: https://doi.org/10.24432/C5VW2C. 1989.
[7]
Luis Candanedo. Appliances Energy Prediction. UCI Machine Learning Repository. 2017. doi: 10 .24432/C5VC8G.
[8]
William G. Cochran. Sampling Techniques. 3rd. John Wiley & Sons, 1977.
[9]
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. “SAGA: a fast incremental gradient method with support for non-strongly convex composite objectives”. In: Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 1. NIPS’14. MIT Press, 2014, pp. 1646–1654.
[10]
Michael Dumelle et al. “A comparison of design-based and model-based approaches for finite population spatial sampling and inference”. In: Methods in Ecology and Evolution 13.9 (2022), pp. 2018– 2029. doi: https://doi.org/10.1111/2041-210X.13919.
[11]
Guillaume Garrigos and Robert Michael Gower. “Handbook of convergence theorems for (stochastic) gradient methods”. In: arXiv preprint arXiv:2301.11235 (2023).
[12]
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. Book in preparation for MIT Press. MIT Press, 2016. url: http://www.deeplearningbook.org.
[13]
Robert Mansel Gower et al. “SGD: General Analysis and Improved Rates”. In: Proceedings of the 36th International Conference on Machine Learning. Ed. by Kamalika Chaudhuri and Ruslan Salakhutdinov. Vol. 97. Proceedings of Machine Learning Research. PMLR, 2019, pp. 5200–5209.
[14]
Rie Johnson and Tong Zhang. “Accelerating stochastic gradient descent using predictive variance reduction”. In: NIPS’13. Red Hook, NY, USA: Curran Associates Inc., 2013, pp. 315–323.
[15]
Hamed Karimi, Julie Nutini, and Mark Schmidt. “Linear Convergence of Gradient and ProximalGradient Methods Under the Polyak-Lojasiewicz Condition”. In: Machine Learning and Knowledge Discovery in Databases. Ed. by Paolo Frasconi et al. Cham: Springer International Publishing, 2016, pp. 795–811.
[16]
Nitish Shirish Keskar et al. “On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima.” In: CoRR abs/1609.04836 (2016). url: http://arxiv.org/abs/1609.04836.
[17]
Diederik P. Kingma and Jimmy Ba. “Adam: A Method for Stochastic Optimization”. In: Proceedings of the 3rd International Conference on Learning Representations (ICLR). 2015. url: https://ar xiv.org/abs/1412.6980.
[18]
Diederik P. Kingma and Jimmy Ba. “Adam: A Method for Stochastic Optimization.” In: ICLR (Poster). Ed. by Yoshua Bengio and Yann LeCun. 2015. url: http://dblp.uni-trier.de/db/co nf/iclr/iclr2015.html#KingmaB14.
15
[19]
Alex Krizhevsky. Learning multiple layers of features from tiny images. Tech. rep. University of Toronto, 2009.
[20]
Yann LeCun and Corinna Cortes. “MNIST handwritten digit database”. In: (2010). url: http: //yann.lecun.com/exdb/mnist/.
[21]
Sharon L. Lohr. Sampling: Design and Analysis. 2nd. Brooks/Cole, 2009.
[22]
Ilya Loshchilov and Frank Hutter. “Decoupled Weight Decay Regularization”. In: 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. url: https://openreview.net/forum?id=Bkg6RiCqY7.
[23]
Dominic Masters and Carlo Luschi. “Revisiting Small Batch Training for Deep Neural Networks”. In: ArXiv abs/1804.07612 (2018). url: https://api.semanticscholar.org/CorpusID:5032969.
[24]
Deanna Needell, Christopher De Sa, and Joel Tropp. “Stochastic gradient descent, weighted sampling, and the randomized Kaczmarz algorithm”. In: Mathematical Programming 155 (2016), pp. 549– 573.
[25]
Yu. E. Nesterov. “A method of solving a convex programming problem with convergence rate O k12 ”. In: Dokl. Akad. Nauk SSSR 269.3 (1983), pp. 543–547. url: http://mi.mathnet.ru/dan 46009.
[26]
Tapio Pahikkala. “New Kernel Functions and Learning Methods for Text and Data Mining”. PhD thesis. Turku, Finland: Turku Centre for Computer Science (TUCS), 2008.
[27]
Fabian Pedregosa et al. “Scikit-learn: Machine Learning in Python”. In: J. Mach. Learn. Res. 12.null (Nov. 2011), pp. 2825–2830. issn: 1532-4435.
[28]
Sebastian Ruder. An overview of gradient descent optimization algorithms. 2016. url: http://ar xiv.org/abs/1609.04747.
[29]
Carl-Erik Särndal, Bengt Swensson, and Jan Wretman. “Model Assisted Survey Sampling (Springer Series in Statistics)”. In: (2003). url: http://www.amazon.com/Assisted-Survey-Sampling-Sp ringer-Statistics/dp/0387406204/sr=8-1/qid=1172587067/ref=pd_bbs_sr_1/103-2111122 -6886251?ie=UTF8&s=books.
[30]
Mark Schmidt, Nicolas Le Roux, and Francis Bach. “Minimizing finite sums with the stochastic average gradient”. In: Math. Program. 162.1–2 (Mar. 2017), pp. 83–112. issn: 0025-5610. doi: 10 .1007/s10107-016-1030-6. url: https://doi.org/10.1007/s10107-016-1030-6.
[31]
Samuel L. Smith, Erich Elsen, and Soham De. “On the generalization benefit of noise in stochastic gradient descent”. In: Proceedings of the 37th International Conference on Machine Learning. ICML’20. JMLR.org, 2020.
[32]
Chong Wang et al. “Variance Reduction for Stochastic Gradient Optimization”. In: Advances in Neural Information Processing Systems. Ed. by C.J. Burges et al. Vol. 26. Curran Associates, Inc., 2013. url: https://proceedings.neurips.cc/paper_files/paper/2013/file/9766527f2b5d3 e95d4a733fcfb77bd7e-Paper.pdf.
[33]
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms. cite arxiv:1708.07747Comment: Dataset is freely available at https://github.com/zalandoresearch/fashion-mnist Benchmark is available at http://fashionmnist.s3-website.eu-central-1.amazonaws.com/. 2017. url: http://arxiv.org/abs/1708.0774 7.
[34]
Peilin Zhao and Tong Zhang. “Stochastic Optimization with Importance Sampling for Regularized Loss Minimization”. In: Proceedings of the 32nd International Conference on Machine Learning (ICML). 2015, pp. 1–9.
[35]
Zhanxing Zhu et al. The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Minima and Regularization Effects. 2019. url: https://openreview.net/forum?id=H1M7so ActX.
16
Appendix Kernel ridge regression in gradient modeling Given the set GI 1 := {gi : i ∈ I 1 } ⊂ G, the regularized least squares problem in a reproducing kernel Hilbert space (RKHS) is defined as: d X 1 X
m1 j=1
qj (xi ) − gji
2
+ λ∥q∥2K ,
(22)
i∈I 1
where m1 = |I 1 |, gji is the jth component of the ith gradient, and qj (xi ) is the corresponding model estimate evaluated at xi , and λ > 0 is the regularization coefficient. By the representer theorem, the solution admits the form X K(xi , x) ci , (23) q(x) = i∈I 1
where K : X × X → Rd×d is a matrix-valued kernel and ci ∈ Rd are the coefficient vectors associated with sample i. The function K corresponds to a positive semi-definite matrix with each component defined as (K(x, x′ ))d,d′ = k(x, x′ ) kT (d, d′ ),
(24)
where k(x, x′ ) is a scalar kernel acting on the input space X , and kT (d, d′ ) is a scalar kernel acting on the task (or output) indices that captures correlations between output dimensions. For tractability, we make the simplifying assumption that the output dimensions are uncorrelated, which corresponds to kT (d, d′ ) = δd,d′ , where δd,d′ denotes the Kronecker delta. Under this assumption, each block of the operator-valued kernel reduces to a scaled identity, K(x, x′ ) = k(x, x′ ) Id , and the full Gram matrix over m1 samples can be expressed compactly as K(X, X) = Id ⊗ KX , where (KX )ij = k(xi , xj ), i, j ∈ I 1 , KX ∈ Rm1 ×m1 , and consequently K(X, X) ∈ R(m1 d)×(m1 d) . It is clear that with the assumption of uncorrelated outputs, K(X, X) is block diagonal. The solution that minimizes 22 is then given by −1 c = (K(X, X) + λm1 I) y, (25) where the stacked vectors c and y are obtained by vectorizing the coefficient and output matrices C = d×1 ⊤ denoting the coefficient and output vectors [ ci ]⊤ i∈I 1 and Y = [ yi ]i∈I 1 , respectively, with ci , yi ∈ R 1 for sample i ∈ I . In other words, c = vec(C) and y = vec(Y), where vec(·) stacks the columns of a matrix into a single column vector, so that all elements corresponding to the first output dimension come first, followed by the second, and so on. Thanks to the Kronecker structure, the inverse in Equation 25 can be computed efficiently without explicitly forming the (m1 d) × (m1 d) matrix. Using the property (Id ⊗ (KX + λm1 Im1 ))−1 = Id ⊗ (KX + λm1 Im1 )−1 , we obtain c = Id ⊗ (KX + λm1 Im1 )−1 y, Thus, only the m1 × m1 matrix (KX + λm1 Im1 ) needs to be inverted, which greatly reduces the computational cost. Furthermore, the output dimensions coefficients can be solved in parallel due to our assumption of uncorrelated outputs. Let us denote A = (KX + λm1 Im1 )−1 and then the model coefficients are solved by C = AY. The prediction model, evaluated at new point x∗ , can then be written compactly as X q(x∗ ) = C⊤ kx∗ = Y⊤ A kx∗ = k(xi , x∗ ) ci , (26) i∈I 1 ⊤
where kx∗ = [k(xi , x∗ )] , i ∈ I 1 .
17
Input space effect on model-assisted gradient estimation
(b) m = 5, ρ = 0.5
(c) m = 5, ρ = 0.3
100000 80000 60000 40000 20000 0 0.52 0.99 5 10 50 100
160000 140000 120000 100000 80000 60000 40000 20000 0 0.52 0.99 5 10 50 100
250000 200000 150000 100000 50000 0 0.52 0.99 5 10 50 100
MSE
(a) m = 5, ρ = 0.7
MLP parameter count (thousands)
MLP parameter count (thousands)
(d) m = 50, ρ = 0.7
MLP parameter count (thousands)
(e) m = 50, ρ = 0.5
(f) m = 50, ρ = 0.3
1e6
5 10 50.2 99.9 MLP parameter count (thousands) (g) m = 500, ρ = 0.7
1e6
1.2 1.0 0.8 0.6 0.4 0.2 0.0 0.50 1
MSE
1.2 1.0 0.8 0.6 0.4 0.2 0.0 0.50 1
MSE
1e6
MSE
2.00 1e6 1.75 1.50 1.25 1.00 0.75 0.50 0.25 0.00 0.50 1
5 10 50 100 MLP parameter count (thousands)
5 10 50 100 MLP parameter count (thousands)
MSE
MSE
5 10 50 100 MLP parameter count (thousands)
800000 600000 400000 200000 0 0.50 1
MSE
300000 250000 200000 150000 100000 50000 0 0.50 1
MSE
MSE
In Figure 7 we compare the average performance of the proposed model-assisted gradient estimator against the standard uniform mini-batch gradient estimator in estimating the true population gradient over 100 randomly generated synthetic regression datasets. The plots show the mean squared error (MSE) between the estimated gradients and the true population gradient for the model-assisted estimator (blue) and the uniform mini-batch estimator (red). The experiments are performed using multilayer perceptron models with varying parameter counts. Furthermore, the input dimensionality is varied across 5, 50, and 500 features, and the fraction of sampled data (of population) used for gradient estimation is set to 0.7, 0.5, and 0.3. As the gradient prediction model q, we utilize kernel ridge regression (KRR) with an RBF kernel. To mitigate unstable extrapolation effects, the model prediction is constrained to decay toward zero whenever the queried parameter point lies sufficiently far from the gradient model training samples in the input space. The plots illustrate a clear benefit of using model-assisted gradient estimation especially when the input space size is smaller, but decreases clearly as the input space size decreases. The number of MLP seems to have relatively small effect in the difference in performance between model-assisted and minibatch gradients. The synthetic results suggest, as expected by the COD effect, that the input space has the clearest effect on the performance of the model-assisted estimator.
5 10 50.2 99.9 MLP parameter count (thousands) (h) m = 500, ρ = 0.5
2.0 1.5 1.0 0.5 0.0 0.50 1
5 10 50.2 99.9 MLP parameter count (thousands) (i) m = 500, ρ = 0.3
Figure 7: Comparison of the proposed model-assisted gradient estimator and the standard uniform minibatch gradient estimator in estimating the true population gradient over 100 randomly generated synthetic regression datasets. The plots report the mean squared error (MSE) between the estimated gradients and the true population gradient for the model-assisted estimator (blue) and the uniform mini-batch estimator (red). Experiments are conducted using multilayer perceptron models with varying parameter counts. The input dimensionality is varied as m ∈ {5, 50, 500}, while the sampling fraction used for gradient estimation is varied as ρ ∈ {0.7, 0.5, 0.3}, where ρ denotes the fraction of the population sample used for gradient estimation.
18
Computational complexity analysis with SGD and KRR model In order for the optimization of ML models to be efficient using model-assisted approach, the added computational complexity needs to be outweighted in terms of generalization error and/or required epochs. The computational overhead of the proposed model-assisted gradient estimator arises mainly from the gradient model training used to predict unsampled gradient contributions. In a kernel ridge regression (KRR, e.g. [26]) case, which we apply in this study, requires solving a dense linear system involving the kernel matrix, resulting in O(n3 ) for training the model, memory complexity of O(n2 ) and predicting the population gradients O(nN ). However, in our setting the number of sampled gradients used for constructing the predictor is usually much smaller than the whole population size, i.e. n ≪ N , making the computational overhead manageable. Common used batch sizes are in powers of two such as n ∈ {16, 32, 64, 128} which are significantly lower than usual full population (i.e., dataset) sizes. Smaller batch sizes are also recommended for better generalization performance [23]. Let ESGD denote the total number of epochs required by standard SGD to reach a given optimization criterion for the loss function ℓ. Let EMA denote the corresponding number of epochs required by the model-assisted optimizer. Furthermore, let N be the population dataset size, n the minibatch size, and d the computational cost of processing one sample gradient. The total computational complexity of standard SGD is CSGD = O(ESGD N d) ,
(27)
since each epoch processes all N samples and each individual data sample incurs computational cost O(d). For the proposed model-assisted gradient estimator, additional computational overhead is introduced through KRR training and prediction. The added worst-case (n1 = n) per-iteration computational complexity is approximately O n3 + nN , (28) where O(n3 ) corresponds to inversion/training of the kernel system and O(nN ) corresponds to predicting gradients for the remaining dataset points. Thus, the total computational complexity of the KRR-assisted optimizer is N 3 n + nN . (29) CMA = O EMA N d + n Equivalently, this can be simplified as CMA = O EMA N d + n2 + N .
(30)
The KRR-assisted optimizer is computationally advantageous when CMA < CSGD which yields straightforwardly a condition: EMA d . (31) < ESGD d + n2 + N The left-hand side represents the fraction of SGD epochs required by the model-assisted method, while the right-hand side gives the maximum allowable epoch ratio for the KRR-assisted optimizer to remain computationally cheaper than SGD. When d is large relative to n2 + N , the cost of ordinary gradient computation dominates the comparison. In this case, the additional KRR overhead is relatively small compared with the baseline cost of SGD, so a moderate reduction in epochs is enough to provide computational advantage. In contrast, when either n or N is large relative to d, the overhead caused by KRR becomes more significant: the n2 term reflects the cost associated with training the KRR model, while the N term reflects the cost of predicting gradients over the population. Consequently, the KRRassisted optimizer must achieve a larger reduction in epochs to offset these added training and prediction costs.
Result figures for SGD-M and Adam
19
Batch size 10
4
loss
2 1
1
1 10
20
30 epoch
40
00 0.04
50
Full-batch Our estimator Baseline
0.03
30 epoch
40
50
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
Full-batch Our estimator Baseline
0.000 0.040 0.035 0.030 0.025 0.020 0.015 0.0100
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
2.0
30 epoch
40
50
Full-batch Our estimator Baseline
20
30 epoch
40
50
Full-batch Our estimator Baseline
40
50
10
20
30 epoch
40
30 epoch
40
50
Full-batch Our estimator Baseline
10
20 30 epoch
40
50
4.70 4.68 4.66 4.64 4.62 4.600
50
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
2.0
0.5 0 2.5
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
2.0
50
50
4.70 4.68 4.66 4.64 4.62 4.600
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
loss
20
40
Full-batch Our estimator Baseline
2.5
0.50 3.0 2.8 2.6 2.4 2.2 2.00
50
loss 10
30 epoch
1.0
loss
loss
Full-batch Our estimator Baseline
0.50 3.0 2.8 2.6 2.4 2.2 2.00
20
loss
30 epoch
10
loss
loss
loss
20
0.000 0.040 0.035 0.030 0.025 0.020 0.015 0.0100
1.5
1.0 10
Full-batch Our estimator Baseline
0.03
loss 10
1.5
1.0
50
1.0
2.0
1.5
40
1.5
loss
loss
20
30 epoch
loss 10
1.0
10
20
0.01
2.5
0.5 0 2.5
10
0.02
1.5
1.0
4.70 4.68 4.66 4.64 4.62 4.600
00 0.04
50
loss
Full-batch Our estimator Baseline
1.5
0.50 3.0 2.8 2.6 2.4 2.2 2.00
40
loss
20
2.0
2.0
30 epoch
loss 10
2.5
0.5 0 2.5
20
0.01
loss
0.000 0.040 0.035 0.030 0.025 0.020 0.015 0.0100
10
0.03
0.01
Full-batch Our estimator Baseline
2
0.02
loss
0.02
Batch size 100
3 loss
Full-batch Our estimator Baseline
4
Full-batch Our estimator Baseline
3
2
00 0.04
Batch size 50
loss
3
loss
CIFAR-100
CIFAR-10
Fashion-MNIST
MNIST
Appliances energy
Airfoil self-noise
Synthetic
4
10
20 30 epoch
40
10
20 30 epoch
40
50
Figure 8: The generalization performance for SGD-M optimizer with baseline (uniform mini-batch), model-assisted (our estimator) and full-batch gradients. Full-batch gradient case is listed for comparison purposes. The darker curve represents the average test loss of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
20
Batch size 10
1.2
Full-batch Our estimator Baseline
1.0
1.0
50
loss 20
30 epoch
40
50
40
50
50
0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
50
2.5 2.0 1.5 1.0 0.5 0.00
50
3.5 3.0 2.5 2.0 1.5 1.0 0.50
Full-batch Our estimator Baseline
10
20 30 epoch
40
50
50
0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
50
2.5 2.0 1.5 1.0 0.5 0.00
50
3.5 3.0 2.5 2.0 1.5 1.0 0.50
Full-batch Our estimator Baseline
loss
Full-batch Our estimator Baseline
20 30 epoch
40
10
20 30 epoch
40
Full-batch Our estimator Baseline
loss
Full-batch Our estimator Baseline
20
30 epoch
40
10
20
30 epoch
40
Full-batch Our estimator Baseline
loss
Full-batch Our estimator Baseline
10
20
30 epoch
40
6
Full-batch Our estimator Baseline
10
20
30 epoch
40
6
Full-batch Our estimator Baseline
4
4
loss
5
loss
5
3
3
3
10
20
30 epoch
40
50
5.5
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
4.5 10
20
30 epoch
40
50
4.5 4.00
50
Full-batch Our estimator Baseline
10
20 30 epoch
40
50
Full-batch Our estimator Baseline
10
20 30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
Full-batch Our estimator Baseline
10
20
30 epoch
40
50
30 epoch
40
50
40
50
Full-batch Our estimator Baseline
20 6.0 5.5
10
20
Full-batch Our estimator Baseline
5.0
loss
loss
5.0
Full-batch Our estimator Baseline
40
4
5.5
5.0
4.00
20 6.0
30 epoch
loss
20 6.0
20
loss
10
10
loss
10
0 0.025 0.020 0.015 0.010 0.005 0.0000
loss
5
0.4 10
loss
20 30 epoch
loss
3.5 3.0 2.5 2.0 1.5 1.0 0.50
0 0.025 0.020 0.015 0.010 0.005 0.0000
loss
loss
10
loss
2.5 2.0 1.5 1.0 0.5 0.00
0.6
loss
40
loss
Airfoil self-noise Appliances energy MNIST Fashion-MNIST CIFAR-10
30 epoch
Full-batch Our estimator Baseline
0.0250 0.0225 0.0200 0.0175 0.0150 0.0125 0.01000
6
CIFAR-100
20
Full-batch Our estimator Baseline
1.0
0.4 10
Batch size 100
0.8
0.6
0.4
1.2
Full-batch Our estimator Baseline
0.8
0.6
0 0.025 0.020 0.015 0.010 0.005 0.0000
Batch size 50
loss
0.8
loss
Synthetic
1.2
4.5 10
20
30 epoch
40
50
4.00
10
20
30 epoch
Figure 9: The generalization performance for Adam optimizer with baseline (uniform mini-batch), modelassisted (our estimator) and full-batch gradients. Full-batch gradient case is listed for comparison purposes. The darker curve represents the average test loss of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
21
4
Our estimator L2 Baseline L2
30 epoch
40
50
00 0.5 0.4 0.3 0.2 0.1 0.00
50
0.6 0.5 0.4 0.3 0.2 0.1 0.00
50
20
30 epoch
40
L2 distance 40
50
Our estimator L2 Baseline L2
20
30 epoch
40
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
20
30 epoch
40
50
Our estimator L2 Baseline L2
6
L2 distance
L2 distance
10
8 4
20
30 epoch
40
50
Our estimator L2 Baseline L2
4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6 4
30 epoch
40
50
00
6
6 L2 distance
20
4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
20
30 epoch
40
50
Our estimator L2 Baseline L2
1.0
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
20
30 epoch
40
50
0.00
40
50
10
20
8
30 epoch
40
50
Our estimator L2 Baseline L2
6
0 0 8
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6
00
10
20
30 epoch
6
40
50
Our estimator L2 Baseline L2
00 2.0
10
20
30 epoch
40
50
40
50
Our estimator L2 Baseline L2
1.0
0.5
10
30 epoch
Our estimator L2 Baseline L2
1.5
1.0
0.5
20
2
L2 distance
1.5
L2 distance
1.5
00 2.0
10
L2 distance
10
0.6 0.5 0.4 0.3 0.2 0.1 0.00
4
2
Our estimator L2 Baseline L2
Our estimator L2 Baseline L2
2
4
2
50
4
2
10
40
L2 distance
2
30 epoch
2
L2 distance
L2 distance
6
0 0 8
20
4
2
10
10
L2 distance
2
00 0.5 0.4 0.3 0.2 0.1 0.00
L2 distance
10
4
0.00
30 epoch
L2 distance
L2 distance
Our estimator L2 Baseline L2
6
00
20
L2 distance
10
8
0 0 8
10
L2 distance
L2 distance
Our estimator L2 Baseline L2
1
L2 distance
20
Our estimator L2 Baseline L2
2
1
10
Batch size 100
3
2
1
0.6 0.5 0.4 0.3 0.2 0.1 0.00
4
Our estimator L2 Baseline L2
L2 distance
2
00 0.5 0.4 0.3 0.2 0.1 0.00
Batch size 50
3
L2 distance
3
00 2.0 CIFAR-100
Batch size 10
L2 distance
CIFAR-10
Fashion-MNIST
MNIST
Appliances energy
Airfoil self-noise
Synthetic
4
0.5
10
20
30 epoch
40
50
0.00
10
20
30 epoch
Figure 10: L2 distance to full-batch gradient SGD-M. The darker curve represents the average of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
22
Batch size 50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
30 epoch
40
50
Our estimator L2 Baseline L2
20
30 epoch
40
50
20
30 epoch
40
50
20
30 epoch
40
50
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
6
40
50
Our estimator L2 Baseline L2
L2 distance
L2 distance
Our estimator L2 Baseline L2
4
20
30 epoch
40
50
Our estimator L2 Baseline L2
4
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6 4
2 20
30 epoch
40
50
10
20
30 epoch
40
50
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
5 4 3 2 1 0 0
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
6
40
50
Our estimator L2 Baseline L2
0 0 8
10
20
30 epoch
40
50
Our estimator L2 Baseline L2
6
2
10
20
30 epoch
L2 distance
Our estimator L2 Baseline L2
00 5 4 3 2 1 00
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
4
2
10
Our estimator L2 Baseline L2
0.3
2
L2 distance
L2 distance
6
0 0 8
50
4
2
10
40
L2 distance
2
30 epoch
L2 distance
10
20
0.1
L2 distance
L2 distance
Our estimator L2 Baseline L2
5 4 3 2 1 0 0
10
L2 distance
10
00 0.4
0.2
L2 distance
L2 distance
Our estimator L2 Baseline L2
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
5
L2 distance
10
4
00 5 4 3 2 1 00
20
0.1
6
0 0 8
10
0.2
0.1
5 4 3 2 1 0 0
00 0.4 L2 distance
0.2
0.00 0.6 0.5 0.4 0.3 0.2 0.1 0.00
5
0.3
L2 distance
0.3
10
L2 distance
5
00 0.4
L2 distance
10
Our estimator L2 Baseline L2
15
40
50
Our estimator L2 Baseline L2
10
20
30 epoch
40
50
00 5 4 3 2 1 00
10
20
30 epoch
40
50
40
50
Our estimator L2 Baseline L2
L2 distance
10
Batch size 100
Our estimator L2 Baseline L2
15 L2 distance
L2 distance
15
L2 distance
CIFAR-100
CIFAR-10
Fashion-MNIST
MNIST
Appliances energy
Airfoil self-noise
Synthetic
Batch size 10
10
20
30 epoch
Figure 11: L2 distance to full-batch gradient Adam. The darker curve represents the average of 400 runs, the dark shading is the 95% confidence interval and lighter shading shows the standard deviation. Columns correspond to batch sizes and rows to datasets.
23
Difference estimator covariance V(Qdif ) = ED Qdif −
N X
! Qdif −
qk
= ED
N X
q̂i +
N X
i=1
Ii
i=1
qi − q̂i − πi
" = ED
X
q̂i q̂Tj +
X
i,j
T qj
j=1
k=1
N X
i,j
qi q̂Tj − q̂i q̂Tj
Ij
N X k=1
T ! N N N X X X q − q̂ i − qk qj q̂i + Ii i π i j=1 i=1 i=1
X q̂i qTj − q̂i q̂Tj − q̂i qTj πj i,j
(qi − q̂i )(qj − q̂j )T X qi qTj − q̂i qTj − Ii πi πi πj πi i,j i,j i,j # X X qi qTj − qi q̂Tj X T − qi q̂j − Ij + qi qTj π j i,j i,j i,j X X X X T T = q̂i q̂j + q̂i qTj − q̂i q̂j − q̂i qTj +
X
Ii
i,j
i,j
+
X
Ii Ij
i,j
i,j
X (qi − q̂i )(qj − q̂j )T X − qi qTj + q̂i qTj qi q̂Tj − πij + q̂i q̂Tj + π π i j i,j i,j i,j i,j i,j X X X X T T T T qi qj qi qj + qi q̂j + − qi q̂j − =
X
X
X
i,j
i,j
i,j
X πij − πi πj i,j
πi πj
i,j
(qi − q̂i ) qj − q̂j
T
24
.
(32)