Online importance sampling for stochastic gradient optimization

## Abstract

Abstract Machine learning problems rely heavily on stochastic gradient descent (SGD) for optimization. The effectiveness of SGD is contingent upon accurately estimating gradients from a mini-batch of data samples. Instead of the commonly used uniform sampling, adaptive or importance sampling reduces noise in gradient estimation by forming mini-batches that prioritize crucial data points. Previous research has suggested that data points should be selected with probabilities proportional to their gradient norm. Nevertheless, existing algorithms have struggled to efficiently integrate importance sampling into machine learning frameworks.
In this work, we make two contributions. First, we present an algorithm that can incorporate existing importance functions into our framework. Second, we propose a simplified importance function that relies solely on the loss gradient of the output layer. By leveraging our proposed gradient estimation techniques, we observe improved convergence in classification and regression tasks with minimal computational overhead. We validate the effectiveness of our adaptive and importance-sampling approach on image and point-cloud datasets.

## 1 Introduction

Stochastic gradient descent (SGD) combined with back-propagation and efficient gradient techniques—such as Adam []—has unlocked a realm of possibilities.
Its importance lies in its ability to optimize complex models by iteratively updating parameters based on the gradient of the loss function.
However, despite its significance, SGD has notable limitations.
Convergence rates depend on various factors, notably the noise in gradient estimation, which significantly impacts both robustness and speed. Addressing this noise, and effectively reducing it, is an active area of research  [; ; ; ; ].

Several approaches have been proposed to mitigate the noise in gradient estimation, including data diversification [; ], adaptive batch sizes, weighted sampling [] or importance sampling []. These methods aim to improve the quality of gradient approximations and accelerate convergence in noisy optimization landscapes.

Our proposed idea centers around the concept of importance sampling, which involves constructing mini-batches using a non-uniform data-point selection, i.e., selecting certain data points with higher likelihood. The objective is to strategically allocate computational resources to data points that exert the most significant influence on the optimization task.

In this paper, we introduce a novel approach that leverages information extracted from the network’s output to quantify the extent of model modification required for each data sample. This information guides our importance sampling strategy, resulting in substantial improvements in convergence across various tasks. Our approach has a lower computational overhead compared to state-of-the-art methods [; ].

In summary, our contributions can be distilled into the following key points:

- •
We propose an efficient and robust strategy for adaptive and importance sampling.
- •
We introduce an algorithm with minimal overhead for implementing adaptive and importance sampling.
- •
We demonstrate the effectiveness of our approach across classification and regression problems, encompassing both importance sampling and adaptive sampling techniques.

## 2 Related work

Gradient estimation serves as a cornerstone in the realm of machine learning, underpinning the optimization of models. In practical scenarios, computing the exact gradient is often unfeasible due to the sheer volume of data, leading to the reliance on mini-batch approximations. Improving these approximations to obtain more accurate and lower-variance estimates is an enduring challenge in the field. The ultimate goal is to expedite gradient descent optimization by enhancing the quality of gradient approximations, representing a core objective in the domain of machine learning.

#### Importance sampling.

Sampling data points proportionally to the norm of the gradient is regarded as the optimal choice in optimization.
Bordes et al. [] developed an online algorithm (LASVM) that uses importance sampling to train kernelized support vector machines.
Zhao & Zhang []; Needell et al. []; Wang et al. []; Alain et al. [] proved that importance sampling proportional to the gradient norm is the optimal sampling strategy. However, even with this optimal strategy, the resulting estimation is not error-free as multiple parameters derivatives are estimated simultaneously. Nevertheless, this approach represents a practical trade-off that provides a good compromise in terms of optimization quality across all dimensions simultaneously.

In practice, estimating the gradient for each data point can be computationally intensive. This has prompted the search for more efficient sampling strategies that can approximate the gradient norm without incurring excessive computational costs. Ideally, such alternative strategies should mimic the gradient norm’s behavior while remaining computationally lightweight, thus contributing to the scalability and practicality of machine learning algorithms Katharopoulos & Fleuret [].
Loshchilov & Hutter [] proposed to use a ranking data based on their respective loss. This rank is then use to create an importance sampling strategy giving more importance to data with high rank (high loss).
Katharopoulos & Fleuret [] proposed to importance sample the loss function. Dong et al. [] proposed a re-sampling based algorithm to reduce the number of backpropagation computation. A sub-selection of data is performed based on the loss. Similarly, Zhang et al. [] proposed to use a re-sampling based on multiple heuristics to reduce the number of backward propagations and focus on more contributing data. Katharopoulos & Fleuret [] proposed an upper bound to the gradient norm that can be used as an importance function. They proposed to re-sample data based on importance computed on the last layer. All the re-sampling method reduce the number of unnecessary backward propagation but still suffer from forward computation.

We propose an efficient algorithm and an importance function which when used for importance or adaptive sampling, shows significant improvements.

#### Adaptive data weighting/sampling.

Adaptive weighting/sampling operates by dynamically adjusting the contribution of individual data samples during optimization. This dynamic alteration of weights aims to prioritize specific data points that exert more influence on the gradient descent process. This alteration can be performed by either increasing or decreasing their weight contributions to the estimator.
Adaptive sampling is different from importance sampling, primarily in the way the gradient estimator is defined for each respective strategy.
While adaptive weighting/sampling significantly accelerates convergence, it does so by introducing a bias into the gradient estimation.

To compute adaptive weights within a given mini-batch, Santiago et al. [] proposed a method that maximizes the effective gradient of the mini-batch. This approach strategically allocates weights to data points, aligning their contributions with the optimization objective. While this may introduce bias, it allows for a faster and more efficient convergence towards an optimal solution, making adaptive weighting a valuable strategy in machine learning optimization.
Our proposed algorithm and the corresponding adaptive weighting approach naturally focuses on data points with high sampling probabilities.

## 3 Gradient estimation

In machine learning, optimization is key to refining the models. The goal is to find optimal parameters $\theta$ for a model function $m(x,\theta)$, with $x$ a data sample, that minimize a loss function ${\mathcal{L}}$ over a dataset ${\Omega}$. The optimization is typically expressed as

$$ $\displaystyle\theta^{*}=\underset{\theta}{\mathrm{argmin}}\,L_{\theta},\quad \text{where}\quad L_{\theta}=\frac{1}{{|\Omega|}}\int_{{\Omega}}{\mathcal{L}}( m(x,\theta),y)\,\mathrm{d}x.$ (1) $$

The loss function $\mathcal{L}$ quantifies the difference between model predictions $m(x,\theta)$ and actual data $y$. The factor in front of the integral normalizes the total loss $L_{\theta}$ by the dataset size.

In practice, the minimization of the total loss is tackled via iterative gradient descent. At each step $t$, the gradient $\nabla L_{\theta_{t}}$ of that loss with respect to the current model parameters $\theta_{t}$ is computed, and the parameters are updated as

$$ $\theta_{t+1}=\theta_{t}-\lambda\nabla L_{\theta_{t}},$ (2) $$

where $\lambda>0$ is the learning rate. This iterative procedure can be repeated until convergence.

### 3.1 Monte Carlo gradient estimator

The parameter update step in [Eq. 2](#S3.E2) involves evaluating the total-loss gradient $\nabla L_{\theta_{t}}$. This requires processing the entire dataset ${\Omega}$ at each of potentially many (thousands of) steps, rendering the optimization computationally infeasible. In practice one has to resort to mini-batch gradient descent which estimates the gradient from a small set $\{x_{i}\}_{i=1}^{B}\subset{\Omega}$ of randomly chosen data points in a Monte Carlo fashion:

$$ $\nabla L_{\theta}\approx\frac{1}{|\Omega|\cdot{B}}\sum_{i=1}^{B}w(x_{i})\nabla {\mathcal{L}}(m(x_{i},\theta),y_{i})=\langle\nabla L_{\theta}\rangle,\quad \text{with}\quad x_{i}\propto p(x_{i}).$ (3) $$

Here, $\nabla{\mathcal{L}}(m(x_{i},\theta),y_{i})$ is the gradient (w.r.t. $\theta$) of the loss function for sample $x_{i}$ selected following a probability density function (pdf) $p$ (or probability mass function in case of a discrete dataset). Setting the weighting function to $w(x_{i})=1/p(x_{i})$ makes $\langle\nabla L_{\theta}\rangle$ an unbiased estimator for the total loss, i.e., $\text{E}[\langle\nabla L_{\theta}\rangle]=\nabla L_{\theta}$. Mini-batch gradient descent uses $\langle\nabla L_{\theta}\rangle$ in place of the true gradient $\nabla L_{\theta}$ in [Eq. 2](#S3.E2) to update the model parameters at every iteration. The batch size ${B}$ is typically much smaller than the dataset, enabling practical optimization. To preserve energy, we impose the condition $E[w(x)]=|\Omega|$, signifying that the average weight value should equal the cardinality of $\Omega$, which represents the number of elements in the dataset. It ensure the weight $w(x)$ compensate for the factor $1/|\Omega|$ in front of the summation.

### 3.2 Sampling and weighting strategies

Mini-batch gradient notoriously suffers from Monte Carlo noise in the gradient estimate ([3](#S3.E3)), which can make the parameter-optimization trajectory erratic and convergence slow. That noise comes from the often vastly different contributions of different samples $x_{i}$ to that estimate. Oblivious to this, classical mini-batch gradient descent assigns equal importance to all samples, selecting and weighting them uniformly. Several strategies can be employed to efficiently and effectively reduce the estimation noise, based on judicious, non-uniform sample selection/weighting.

#### Importance sampling

Typically, the selection of samples that go into a mini-batch is done with constant probability $p(x_{i})=1/{|\Omega|}$. Importance sampling is a technique for using a non-uniform pdf to strategically pick samples proportionally on their contribution to the gradient, to reduce estimator variance. Setting $w(x_{i})=1/p(x_{i})$ maintains unbiasedness for any valid importance distribution $p$. A distribution $p$ is valid if it assigns non-zero probability to samples with non-zero contribution.

#### Adaptive sampling.

Adaptive sampling represents an alternative approach to data selection, where samples are chosen based on non-uniform sampling distributions without applying normalization to their contributions. In this method, as outlined in [Eq. 3](#S3.E3), each data point’s weight, denoted as $w(x_{i})$, remains constant at ${|\Omega|}$. Consequently, adaptive sampling naturally focuses on data points with high sampling probabilities, effectively modifying the problem statement [Eq. 1](#S3.E1) by defining a revised domain $\Omega^{\prime}$. This domain is denser around data points with elevated sampling probabilities, making it adaptable to evolving probabilities during the optimization process.

Adaptive weighting employs a uniform sampling distribution, where $p(x_{i})=1/{|\Omega|}$, but introduces an adaptive weight normalization factor, $w(x_{i})$. In contrast to adaptive sampling, which focuses on non-uniform data selection, adaptive weighting prioritizes data samples by assigning them varying levels of importance through the weight normalization. Much like its counterpart, this method results in a modification of the problem statement, allowing for the emphasis of specific data samples during the optimization process. If carefully chosen, this emphasis can significantly accelerate optimization convergence.

### 3.3 Practical importance/adaptive sampling algorithm

We propose an algorithm to efficiently perform importance and adaptive sampling for mini-batch gradient descent, outlined in [Algorithm 1](#alg1).
Similarly to Loshchilov & Hutter [] and Schaul et al. [], it is designed to use an importance function that relies on readily available quantities for each data point, introducing only negligible memory and computational overhead over classical uniform mini-batching.

Figure: Algorithm 1 Gradient descent using importance or adaptive sampling

We store a set of persistent *un-normalized importance* scalars $q=\{q_{i}\}_{i=1}^{|\Omega|}$ that are updated continuously during the optimization.
We begin with a first epoch to processes all data points once and determine their initial importance (line 6).
After that, at each mini-batch optimization step $t$, we normalize the importance values to obtain the probability density function (PDF) $p$ (line 10).
We then extract ${B}$ data samples (with replacement) using this PDF $p$ (line 11).
The loss ${\mathcal{L}}$ is evaluated for each selected data sample (line 12), and backpropagated to compute the corresponding loss gradient (line 13).
Depending on whether we want to perform importance sampling or adaptive sampling, the per-sample weight is selected (line 14) and used in the gradient estimation ([3](#S3.E3)) (line 15).
For importance sampling $w(x)=1/p(x)$ and for adaptive sampling $w(x)=|\Omega|$.
Finally, the network parameters are updated using the estimated gradient (line 16).
In line 17, the subroutine ComputeSampleImportance$(x,y)$ returns the sample importance for each data sample from the mini-batch.
Any importance heuristic such as the gradient norm [; ; ; ] or the loss [; ; ], or more advanced importance [] can be implemented in this subroutine.
For efficiency, our algorithm reuses the forward pass computations made during line 12 to execute ComputeSampleImportance$(x,y)$ subroutine, thereby updating $q$ only for the current mini-batch samples.
The weighting parameter $\alpha$ ensures weight stability as discussed in [Eq. 7](#S5.E7).
At the end of each epoch (line 18), we add a small value to the un-normalized weights of all data to ensure that every data point will be eventually evaluated, even if its importance is deemed low by the importance metric.

It is importance to note that the first epoch is done without importance sampling to initialize each sample importance. This does not create overhead as it is equivalent to a classical epoch running over all data samples.
While similar schemes have been proposed in the past, they often rely on a multitude of hyperparameters, making their practical implementation challenging. This has led to the development of alternative methods like re-sampling [; ; ].
Our proposed sampling strategy has only a few hyperparameters. Tracking importance across batches and epochs minimizes the computational overhead, further enhancing the efficiency and practicality of the approach.

Figure: Figure 1: Visualization of the importance sampling at 3 different epoch and the underlying classification task. For each presented epoch, 800 data-point are presented with a transparency proportional to their weight according to our method.
Refer to caption: x1.png

## 4 Loss-gradient-based importance function

Figure: Algorithm 2 subroutine for cross entropy loss importance metric

Classification assigns discrete labels to data points, relying on identifying complex decision boundaries in the feature space. Accurate boundary identification is crucial, as minor parameter changes can significantly impact results. Importance sampling in classification emphasizes gradients along these boundaries, where parameter modifications have the greatest impact. [Figure 1](#S3.F1) illustrates this concept, showing iterative refinement of the sampling distribution to focus on boundary decisions in comparison to data within classes. The rightmost column illustrates the sampling distribution of the DLIS method of Katharopoulos & Fleuret [] at epoch 100. Both methods iteratively increase the importance of the sampling around the boundary decision compare to data inside the classes.

Our approach differs from that of in that we compute the gradient norm with respect to the network’s output logits. This approach often allows gradient computation without requiring back-propagation or graph computations, streamlining optimization.

#### Cross-entropy loss gradient.

Cross entropy is a widely used loss function in classification tasks. It quantifies the dissimilarity between predicted probability distributions and actual class labels. Specifically, for a binary classification task, cross entropy is defined as:

$$ $\mathcal{L}(m(x_{i},\theta))=-\sum_{j=1}^{J}y_{j}\log({s}_{j})\text{\; where\; }{s}_{j}=\frac{\exp(m(x_{i},\theta)_{j})}{\sum_{k=0}^{J}\exp(m(x_{i},\theta)_{ k})}$ (4) $$

where $m(x_{i},\theta)$ is an output layer, $x_{i}$ is the input data and $J$ means the number of classes.
The derivative of the loss $\mathcal{L}$ of a data point $x_{i}$ with respect to the network output layer $m(x_{i},\theta)_{j}$ reads

$$ $\frac{\partial\mathcal{L}}{\partial m(x_{i},\theta)_{j}}={s}_{j}-y_{j}$ (5) $$

This equation can be directly computed from the network output without any graph back-propagation. [Algorithm 2](#alg2) show a short pseudo-code of the importance computation for [Algorithm 1](#alg1). Proof of the derivation can be found in the [Appendix A](#A1).

Similarly to the approach of Katharopoulos & Fleuret [], the norm of the gradient across the network can be bounded as follows:

$$ ${\color[rgb]{0.0,0.0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0.0,0.0,0} \pgfsys@color@rgb@stroke{0.0}{0.0}{0}\pgfsys@color@rgb@fill{0.0}{0.0}{0}\left \|\frac{\partial\mathcal{L}(x_{i})}{\partial\theta}\right\|=\left\|\frac{ \partial\mathcal{L}(x_{i})}{\partial m(x_{i},\theta)}\cdot\frac{\partial m(x_{ i},\theta)}{\partial\theta}\right\|\leq\left\|\frac{\partial\mathcal{L}(x_{i}) }{\partial m(x_{i},\theta)}\right\|\cdot\left\|\frac{\partial m(x_{i},\theta)} {\partial\theta}\right\|}$ (6) $$

While this bound may not offer conclusive proof of optimality or a direct correlation with the gradient norm, it establishes a discernible link between the two. The key lies in the initial gradient’s selection at the start of the chain rule. This gradient essentially sets the tone for all subsequent gradients in the network. Consequently, when variations in this initial gradient of the loss are pronounced, energy propagates through the network, elevating the gradient norm. Essentially, the intricacies of gradient variation at the outset have a rippling effect, significantly influencing the overall gradient norm, despite the bound not providing a definitive measure of optimality.

#### Generalized loss gradient.

We can extend our framework to other problems, e.g., regression, using the automatic gradient computation operator. Our framework can compute the loss gradient w.r.t. the output layer using such autograd operators.
To demonstrate the generalization capacity of our framework, we also perform experiments on the regression problem.

## 5 Experiments

In this section, we delve into the experimental outcomes of our proposed algorithm and sampling strategy. Our evaluations encompass diverse classification tasks, spanning both importance sampling and adaptive sampling. We benchmarked our approach against those of Katharopoulos & Fleuret [] and Santiago et al. [], considering various variations in comparison. Distinctions in our comparisons lie in assessing performance at equal steps/epochs and equal time intervals. The results presented here demonstrate the loss and classification error, computed on test data that remained unseen during the training process.

Figure: Figure 2: We compare loss and classification error metrics for the MNIST dataset between the resampling algorithm by Katharopoulos & Fleuret [] (DLIS) and our algorithm. At equal epochs, the resampling algorithm with importance sampling works better than uniform sampling for DLIS. However, at equal time, the resampling cost is too high, making DLIS even slower than standard uniform sampling. Our algorithm outperforms all existing methods.
Refer to caption: x2.png

Figure: Figure 3: When comparing on CIFAR-10 and Point cloud ModelNet40 [] classification datasets, DLIS performs poorly at equal time due to the resampling overhead. Unlike DLIS, we use standard uniform sampling which is faster. We also compare against another adaptive scheme by Santiago et al. [] (LOW). Our adaptive (Ours AS) and importance sampling (Ours IS) shows improvements on the ModelNet40 dataset against other methods. Overall, our adaptive variant achieves lower classification errors with minimal overhead compared to others.
Refer to caption: x3.png

Figure: Figure 4: On Oxford flower 102 classification dataset [], where the number of classes are high, our approach shows significant improvement compared to other methods. DLIS performs worse due to the sparsity in the data, which hampers their resampling strategy.
Refer to caption: x4.png

Figure: Figure 5: On CIFAR-100 classification dataset, instead of comparing the DLIS resampling algorithm, we use DLIS weights in our algorithm. We display zoom-in of the end of the curves to highlight the differences. At equal epochs (left), our methods (Ours IS & AS) show improvements compared to LOW [] and DLIS weights. At equal time (right), LOW and the DLIS weights takes longer to converge. Overall our approach shows faster convergence with lower importance computation.
Refer to caption: x5.png

### 5.1 Implementation details

For fair comparisons, we implement our method and all baselines in a single PyTorch framework.
Experiments run on a workstation with an NVIDIA GeForce RTX 2080 graphics card and an Intel(R) Core(TM) i7-9700 CPU @ 3.00GHz.
The baselines include uniform sampling, DLIS [] and LOW [].
Uniform means we sample every data point from a uniform distribution.
DLIS importance samples the data mainly depending on the norm of the gradient on the last output layer.
We use functorch [] to accelerate this gradient computation.
LOW is based on adaptive weighting that maximizes the effective gradient of the mini-batch using the solver in Vandenberghe [].

#### Weight stability.

Updating the persistent per-sample importance $q$ directly sometime leads to a sudden decrease of accuracy during training.
To make the training process more stable, we update $q$ by linearly interpolating the importance at the previous and current steps:

$$ $q(x)=\alpha\cdot q_{prev}(x)+(1-\alpha)\cdot q(x)$ (7) $$

where $\alpha$ is a constant for all data samples. In practice, we use $\alpha\in\{0.0,0.1,0.2,0.3\}$ as it gives the best trade-off between importance update and stability. This can be seen as a momentum evolution of the per-sample importance to avoid high variation. Utilizing an exponential moving average to update the importance metric prevents the incorporation of outlier values. This is particularly beneficial in noisy setups, like situations with a high number of class or a low total number of data.

#### MNIST dataset.

The MNIST database contains 60,000 training images and 10,000 testing images.
We train a 3-layer fully-connected network (MLP) for image classification over 50 epochs with an Adam optimizer [].

#### CIFAR dataset.

CIFAR-10 Krizhevsky et al. [] contains 60,000 32x32 color images from 10 different object classes, with 6,000 images per class.
CIFAR-100 Krizhevsky et al. [] has 100 classes containing 600 images each, with 500 training images and 100 testing images per class. For both datasets, we train the same ResNet-18 network [].
We use the SGD optimizer with momentum 0.9, initial leaning rate 0.01, and batch size 64.
We divide the initial learning rate by 10 after 70 epochs for CIFAR-10 and train the network for a total of 100 epochs. Additionally, we trained a Vision Transformer (ViT) [] on CIFAR-10 using the Adam optimizer with an initial learning rate 0.0001 and a cosine annealing scheduler [].
For CIFAR-100, we divide the learning rate by 10 after 100, 200 epochs and train for a total of 300 epochs.
For both datasets, we use random horizontal flip and random crops to augment the data on the fly.

#### Point-cloud classification dataset.

We train a PointNet Qi et al. [] with 3 shared-MLP layers and one fully-connected layer, on the ModelNet40 dataset Wu et al. [].
The dataset contains point clouds from 40 categories.
The data are split into 9,843 for training and 2,468 for testing.
Each point cloud has 1,024 points.
We use the Adam optimizer Kingma & Ba [], with batch size 64, weight decay 0.001, initial learning rate 0.00002 divided by 10 after 100, 200 epochs.
We train for 300 epochs in total.

#### Oxford 102 flower dataset.

The Oxford 102 flower dataset Nilsback & Zisserman [] contains flower images from 102 categories.
We follow the same experiment setting of  Zhang et al. [; ]. We use the original test set for training (6,149 images) and the original training set for testing (1,020 images).
In terms of network architecture, we use the pre-trained VGG-16 network Simonyan & Zisserman [] for feature extraction and only train a two-layer fully-connected network from scratch for classification.
We use the Adam optimizer Kingma & Ba [] with a learning rate 0.001 and train the two-layer fully-connected network for 100 epochs.

### 5.2 Results

In [Fig. 2](#S5.F2), we compare our algorithm against DLIS []. DLIS applies resampling to both uniform and their importance sampling. This increases the overall computation overhead for their approach.
Standard uniform sampling is much faster than the resampling approach.
We assess cross-entropy loss and classification error in terms of epoch count and equal time.
The DLIS method shows similar performance to ours at equal epochs but incurs high computational costs due to the need for a large dataset forward pass during re-sampling.

Figure: Figure 6: Image regression using a 5-layer SIREN network Sitzmann et al. [], which is trained for predicting RGB values from 2D pixel coordinates. Our importance and adaptive sampling strategies based on autograd gradients show clear improvements compared to uniform sampling at equal-epoch loss curves. We further show the error map and zoom-in results using uniform sampling and our importance sampling. Equal-time comparisons show similar improvements (see appendix).
Refer to caption: x6.png

Figure: Figure 7: Comparisons on CIFAR-10 using Vision Transformer (ViT) []. The results show consistent improvement of Ours IS/AS over LOW [] and DLIS [].
Refer to caption: x7.png

In [Fig. 3](#S5.F3), we compare uniform sampling, DLIS, LOW, and our importance and adaptive sampling strategies for CIFAR-10 and point cloud classification tasks.
Our adaptive sampling excels in classification error across epochs and optimization time, especially when compared at equal time due to our low overhead compared to DLIS and LOW. DLIS performs worse than uniform sampling even with an equal number of epochs for CIFAR-10, highlighting the challenges in effective re-sampling strategies, even with 5-fold increase in data sampling.
Our algorithm is versatile, can work on different data-formats (images, point clouds). In [Fig. 7](#S5.F7), we also perform a comparison with a Vision Transformer architecture [] on CIFAR-10 dataset, where our approach is showing consistent improvements.
DLIS convergence is hampered due to resampling which discards many *less important* data samples, causing training to focus on a subset of the dataset.

We conduct a similar experiment on the Oxford flower classification task in  [Fig. 4](#S5.F4)).
This dataset is known for its complexity due to a large number of classes with few images per class.
Our method employing adaptive sampling again emerged as the top performer.
Notably, DLIS exhibits under-performance, likely due to the challenges of re-sampling in a dataset with a high number of classes with only (10 data samples per class). With this distribution of data, re-sampling does not achieve a good estimation of the gradient. This causes visible over-fitting in the convergence curve. This highlights the robustness of our sampling metric as well as the use of memory based algorithm.

Our algorithm allows using weights from any existing method as an importance function.
To demonstrate this feature, in [Fig. 5](#S5.F5) we use DLIS weights in our algorithm.
Here we present a comparison of classification error for CIFAR-100 dataset.
Although all methods perform comparably at equal epochs, Ours AS stands out with the best results.
At equal time, both DLIS and LOW are hampered by their respective weight computation overhead, leading to slightly inferior performance.
It should be noted that DLIS can be implemented through our algorithm and can benefit from the memory even if there is still an overhead in the weight computation.
Transitioning from re-sampling to a memory-based algorithm proves advantageous in such scenarios.
Additional comparisons with the same setup can also be found for MNIST in [Appendix C](#A3).

[Uncaptioned image]: x8.png

Zhao & Zhang [] have shown that importance weights wrt the gradient norm gives the optimal sampling distribution. On the right inline figure, we show the difference between various weighting strategies and the gradient norm wrt all parameters.
In this experiment, all sampling weights are computed using the same network on an MNIST optimization task. Our proposed sampling strategies, based on the loss gradient are the closest approximation to the gradient norm.

All the results shown have used our analytic importance function from [Eq. 5](#S4.E5).
To show that our importance function can generalize to other problems,
we use a simple autograd operation to compute the loss gradient at the output layer.
We show it’s applicability on the image regression problem in [Fig. 6](#S5.F6).
In this task, we trained a 5-layer SIREN network Sitzmann et al. [] to predict RGB values from given 2D pixel coordinates.
The left side of the figure illustrates the loss evolution wrt the training epochs.
On the right, we showcase the reference image we train the network on, the error image (difference from the reference), and a crop for both Uniform sampling and our importance sampling method.
To use our method in this example with an $L_{2}$ norm loss function, we employed autograd computation with respect to the output of the network.
[Appendix C](#A3) provides a more detailed comparison between analytic formulation and the autograd for the cross-entropy loss.
Our method exhibits better loss convergence despite a minor computational overhead. The error image shows more uniform and reduced amplitude errors compared to uniform sampling. This difference is evident in the crop, where our method provides sharper fur details, contrasting with the blurred result of uniform sampling. This experiment highlights the versatility of non-uniform sampling, expanding its use beyond classification tasks.

#### Discussions

Both ours and DLIS importance metrics are highly correlated but ours is simpler and efficient to evaluate. The efficiency comes from using the memory vector of the importance metric which avoids the need of resampling. The resulting algorithm gives higher performance at equal time and has more stable convergence.
Our adaptive sampling achieve better convergence properties than LOW. Adaptive sampling generally outperforms adaptive weighting. Using adaptive sampling mini-batch are composed with multiple high importance data instead of a weighting based on the respective contribution. Thus the average in [Eq. 3](#S3.E3) is done with more contributing data. Additionally, our loss derivative-based weighting is closer to the gradient norm than loss-based importance.

#### Limitations

As the algorithm rely on past information to drive a non-uniform sampling of data, it requires seeing the same data multiple times. This creates a bottleneck for architectures that rely on progressive data streaming. More research is needed to design importance sampling algorithms for data streaming architectures, which is a promising future direction.
Non-uniform data sampling can also create slower runtime execution. The samples selected in a mini-batch are not laid out contiguously in memory leading to a slower loading. We believe a careful implementation can mitigate this issue.

## 6 Conclusion

In conclusion, our work introduces an efficient sampling strategy for machine learning optimization, including both importance and adaptive sampling. This strategy, which relies on the gradient of the loss and has minimal computational overhead, was tested across various classification as well as regression tasks with promising results.
Our work demonstrates that by paying more attention to samples with critical training information, we can speed up convergence without adding complexity. We hope our findings will encourage further research into simpler and more effective importance/adaptive sampling strategies for machine learning.

## Appendix A Derivative of cross-entropy loss

[Uncaptioned image]: x9.png

Machine learning frameworks take data $x$ as input, performs matrix multiplication with weights and biases added.
The output layer is then fed to the softmax function to obtain values ${s}$ that are fed to the loss function.
$y$ represents the target values.
We focus on the categorical cross-entropy loss function for the classification problem (with $J$ categories) given by:

$$ ${\mathcal{L}}$ \xspace _ cross-ent $\xspace=-\sum_{i}y_{i}\log{s}_{i}\text{\; where\;}{s}_{i}=\frac{\exp(m(x_{i}, \theta)_{l})}{\sum_{l}^{J}\exp(m(x_{i},\theta)_{l})}$ (8) $$

For backpropagation, we need to calculate the derivative of the $\log{s}$ term wrt the weighted input $z$ of the output layer. We can easily derive the derivative of the loss from first principles as shown below:

$$ $\displaystyle\frac{\partial{\mathcal{L}}_{\text{cross-ent}}}{\partial m(x_{i}, \theta)_{j}}$ $\displaystyle=-\frac{\partial}{\partial m(x_{i},\theta)_{j}}\left(\sum_{i}^{J} y_{i}\log{s}_{i}\right)=-\sum_{i}^{J}y_{i}\frac{\partial}{\partial m(x_{i}, \theta)_{j}}\log{s}_{i}=-\sum_{i}^{J}\frac{y_{i}}{{s}_{i}}\frac{\partial{s}_{i }}{\partial m(x_{i},\theta)_{j}}$ (9) $\displaystyle=-\sum_{i}^{J}\frac{y_{i}}{{s}_{i}}{s}_{i}\cdot(\mathbf{1}\{i==j \}-{s}_{j})\text{, can be easily derived from first principles,}$ (10) $\displaystyle=\sum_{i}^{J}{y_{i}}\cdot{s}_{j}-\sum_{i}^{J}y_{i}\cdot(\mathbf{1 }\{i==j\})={s}_{j}\sum_{i}^{J}{y_{i}}-y_{j}={s}_{j}-y_{j}$ (11) $$

The partial derivative of the cross-entropy loss function wrt output layer parameters has the form:

$$ $\displaystyle\frac{\partial{\mathcal{L}}_{\text{cross-ent}}}{\partial m(x_{i}, \theta)_{j}}$ $\displaystyle={s}_{j}-y_{j}$ (12) $$

For classification tasks, we directly use this analytic form of the derivative and compute it’s norm as weights for adaptive and importance sampling.

## Appendix B Algorithm details

Figure: Algorithm 3 Subroutine for initialization for [Algorithm 1](#alg1)

## Appendix C Additional results

Figure: Figure 8: Comparison of the loss evolution for the image regression problem ([Fig. 6](#S5.F6)) at both equal epoch and equal time. While our method have a small overhead cause by the use of autograd, we achieve lower loss for both equal epoch and equal time.
Refer to caption: x10.png

Figure: Figure 9: MNIST comparison of adaptive sampling and importance sampling using our method. We compare to DLIS weights using our algorithm and LOW. Loss and Classification error results are presented for equal epoch and equal time. While Ours IS, DLIS, and LOW perform similarly at equal epoch, their computational overhead causes them to perform less at equal time.
Refer to caption: x12.png

Figure: Figure 10: Comparisons between our analytic gradient and autograd numerical gradient on Point cloud classification. Using analytic and autograd gradient results in similar convergence rate and the analytic one is faster without the need of computing per sample gradient using autograd. Meanwhile, it demonstrates that our method is not limited to classification tasks. We show an example in [Fig. 6](#S5.F6) using autograd gradient for importance and adaptive sampling on a regression problem.
Refer to caption: x14.png

## Peer Reviews

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
questions: - Why (6) is only for binary classification tasks (as indicated in Line 205)? Could $J$ be larger than 2? 
- In [1], the authors mention that ``The individual parameter derivatives vary uniquely across the data points, and estimation using a single distribution inevitably requires making a trade-off" and advocate for the multiple importance sampling (MIS) approach. Could you please comment on this? Could you experimentally compare their MIS-based algorithm with your algorithm?

[1] Salaün, Corentin, Xingchang Huang, Iliyan Georgiev, Niloy J. Mitra, and Gurprit Singh. "Multiple importance sampling for stochastic gradient estimation." arXiv preprint arXiv:2407.15525 (2024).
soundness: 1
strengths: - The idea of using importance sampling weights for data pruning is interesting.
- The plots in Figure 1 numerically verify that the learned importance sampling weights are somewhat meaningful and provide intuitions of why this method could improve upon uniform sampling.
summary: This paper proposes a new adaptive method of importance sampling for stochastic gradient estimation in multi-class classification. The sampling weights of this method do not depend on the costly full backpropagation on each data point. The importance sampling weights can also be used for data pruning, which can be further combined with the importance sampling for gradient estimation. The authors conducted experiments on classification tasks to verify the effectiveness of their algorithm compared to SGD with uniform sampling and previous importance sampling methods.
weaknesses: - A major concern with the importance sampling method proposed in this paper is that it remains "loss-based". To be specific, the weight of each data $x$ is proportional to $\\|\frac{\partial \mathcal{L}(x)}{\partial m(x,\theta)}\\|\_2 = \sqrt{\sum_{j=1}^J(s\_j(x) - y\_j(x))^2}$, where $s\_j(x)$ is the predicted probability of data $x$ belongs to class $j$ while $y\_j(x)$ is the groundtruth. Thus, the importance sampling weight of data $x$ can be viewed as its $\ell_2$ loss on label prediction. However, it is unclear how this approach relates to the theoretically optimal importance sampling weight based on gradient norms. If the gradient w.r.t. the output does not take the specific form as in the logistic loss, does it still make sense to sample based on the norm of the gradient w.r.t the output?
- There is no formal convergence analysis of the proposed algorithm. So the algorithm remains heuristic. 
- The experiments in this paper were not repeated with different random seeds, resulting in a lack of error bars on reported values and curves. This makes their experimental results less reliable.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: 1. Can you clarify the novelty of your proposed metric and why should it be better than those used in the literature?
2. Can you add more comparisons in the experiments with more recent works ([2] and a lot more) and also the most simple random reshuffling with a reduced number of epochs?
3. Can you add more details regarding sampling with or without replacement? It is said within a batch, it is sampled without replacement, so the samples will not repeat. But what about for different batches from the same epoch. To me, it seems like samples can repeat within an epoch, and this can lead to inferior performance compared to those will not repeat within an epoch such as [2].
soundness: 2
strengths: 1. Overall, the writing is clear and easy to follow.
summary: The authors propose to use the derivative wrt to the network output for importance sampling. They also propose that their method can be used for online data pruning. They demonstrate their performance is stronger than some of the baselines.
weaknesses: There are many papers in the literature regarding importance sampling and the authors seem not to know many of them. It is true that the methods proposed in the earlier years have the problem of computation efficiency, but nowadays there are many methods with little overhead.


1. The proposed method has limited novelty. While the authors claim that the derivative wrt to the network output is different from what is used in (Katharopoulos&Fleuret,2018), from my understanding, it is pretty similar if not identical, and there are many other works that utilize something similar like the EL2N score in [1]. Even if there is some difference, the novelty seems limited and the changes are not argued or justified. (Like the difference may just be taking the derivative to the pre- or post- activation output)
2. The experiment comparison is weak. They do not compare with some more recent work like [2]. Additionally, they do not compare with some simple and already used baseline in practice. For example, random shuffling with a reduced number of epochs is often stronger than many of these importance sampling methods. (This is sometimes much stronger than uniform sampling as the samples won’t repeat within an epoch. Also, it is important that the learning rate schedule changes so that the learning rate decays)



[1] Paul, M., Ganguli, S., & Dziugaite, G. K. (2021). Deep learning on a data diet: Finding important examples early in training. Advances in neural information processing systems, 34, 20596-20607.

[2] Qin, Z., Wang, K., Zheng, Z., Gu, J., Peng, X., Zhou, D., ... & You, Y. InfoBatch: Lossless Training Speed Up by Unbiased Dynamic Data Pruning. In The Twelfth International Conference on Learning Representations.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See above
soundness: 3
strengths: The authors evaluated their proposed framework on multiple datasets and different tasks. The results look promising and show some gains compared to the competitors that seem consistent across tasks.
summary: This work studies stochastic gradient descent with online importance sampling and data pruning. The authors propose a practical metric that requires little computation and use it both for the importance weights and pruning scores. This metric is updated on the fly during training. The authors then evaluate their framework on multiple tasks such as classification and regression and popular benchmark datasets such as Cifar10/100 and MNIST.
weaknesses: 1) I found the writing in Sections 3 and 4 to be unclear and somewhat lacking in precision; they would need rewriting and clarifications in my opinion.
For example, equation (1) is confusing, as it seems different from the typical loss minimized, which is generally expressed as $\mathbb{E}(\mathcal{L}(m(x, \theta), y))$, where $x,y$ follow a given data generating process. I fail to understand this renormalization and what $p(x,y)$ refers to here: is it the data-generating process or importance sampling weights? Could you explain and define $p(x,y)$.
Besides, there are some assumptions that are not clearly stated and appear here and there in the text, for example, that the model is Lipschitz (Section 4.1).  

2) The proposed metric seems to be the same as the EL2N proposed in Deep Learning on a Data Diet, Paul et al. Could the authors explain the difference?

3) Although the proposed method seems to outperform the other methods consistently, the differences are sometimes very small, and it is difficult to know if we can not attribute it to statistical noise. The authors indicate that they average over 3 runs; it would be interesting to quantify the variability of the results.