Paper: Efficient Gradient Estimation via Adaptive and Importance Sampling (ICLR-2024)

Reviewer comment (Theory):
Sufficient details have not been provided to formally understand the guarantees of the subroutines in Algorithm 1, such as ComputeImportanceSampling.

Author response:
Please refer to the general response where we detail the algorithm. 

The subroutine $\text{ComputeImportanceSampling}$ is intentionally abstract, to represent any importance-sampling heuristic. The purpose of this subroutine is to compute the importance of each mini-batch sample and update the memory vector $q$. We will include pseudo-code implementation for our specific importance function.

Reviewer comment (Writing):
It is not clear to me what the general framework in Section 4 is proposing, perhaps additional pseudocode would be helpful.

Author response:
In Section 4 we first analytically derive the importance function for the categorical cross entropy loss. Then we generalize our idea to regression tasks. In particular, we only need to compute the loss gradient w.r.t. the network _output_ to apply our importance sampling strategy to other tasks.

Reviewer comment (Experiments):
There are a large number of popular methods for either acceleration or variance reduction for SGD that should be compared for a thorough empirical evaluation (for instance, any number of the references in the related work section).

Author response:
We have compared our method against state-of-the-art importance and adaptive sampling methods like DLIS, LOW and loss-based importance sampling. Our approach, in principle, can be combined with other variance reduction strategies. However, we believe that this would require additional exploration and goes beyond the scope for this work.

Reviewer comment (Reproducibility):
After the approximate probabilities for importance sampling are obtained, how does the algorithm efficiently use these values to sample a data point? That is, are the sampling probabilities organized into a data structure of some sort that maps each real number (possibly implicitly through a set of intervals) to a specific data point?

Author response:
We first normalize the persistent importance distribution $q$ to obtain a probability distribution $p$ (Algorithm 1, L6) that is then used to select data points (L7). Standard routines can be used to perform the importance sampling given $p$. Our implementation uses `numpy.random.choice`; another option is `torch.multinomial`.

Note that when sampling multiple data points simultaneously, it is important to use sampling with replacment, otherwise the prescribed probability distribution $p$ will not be obeyed.

Reviewer comment (Theory):
The only issue I see is with the specific motivation of the proposed importance function. I had a very tough time parsing the text below Eq 5 on page 5. For eg: The gradient norm in eq has two quantities $x$ and $x_i$. Also, the lipschitz constant is written as $l_{m(x, \theta)}$ indicating that the constant of the network is input dependent. Is this intended? Additionally, this bound on the gradients is not informative about the chosen importance function. Eqn 6 makes it seem that the proposed importance function is valid only for Lipschitz neural networks. Is my understanding correct? If so, the proposed method is limited in its applicability.

Author response:
# Correction of Eqn. 6


We regret the typo in this equation. We have corrected this equation and it has the following form:
$\left\Vert \frac{\partial \mathcal{L}(x_i)}{\partial \theta} \right\Vert = \left\Vert \frac{\partial \mathcal{L}(x_i)}{\partial m(x_i,\theta)} \cdot \frac{\partial m(x_i,\theta)}{\partial \theta} \right\Vert \leq \left\Vert \frac{\partial \mathcal{L}(x_i)}{\partial m(x_i,\theta)} \right\Vert \cdot \left\Vert \frac{\partial m(x_i,\theta)}{\partial \theta} \right\Vert$.

Since the gradient of the loss wrt the parameter $\theta$ is bounded by the gradient wrt the output layer, we chose to importance sample this term: our analytic importance function (shown in eqn. 5).
The purpose of this equation 6 is to provide a bound on the derivative with respect to the model parameters, not the input. We will add this clarification in the revised version.

Reviewer comment (Theory):
A discussion on the theoretical properties of the importance function is needed: why is this importance function better than the one in DLIS? In the absence of this, it is tough to attribute the successes of this method. For eg: the way the importances are tracked in Algorithm1 (Lines 13, 14) and the subsequent sampling in Line 7 (instead of the more involved update process in DLIS paper) may be more "implementation friendly" and thus the noticed performance improvements.

Author response:
# Imporvement over DLIS

The main distinctions and reasons for improvement lie in the algorithm and the computation of weights. Utilizing the derivative of the loss with respect to the model output is significantly faster than DLIS's importance metric for comparable distributions. 
Additionally, unlike DLIS, our algorithm does not rely on resampling which avoids the computational cost of forward passes on samples without subsequent backward loops. 
These two differences account for the superior performance of our approach at equal time comparison.

Figure 1 shows that both DLIS and our approach have similar importance distribution, but ours is computationally efficient. 

Figure 7 further compares the weights. Ours and DLIS weights are closer to the optimal weights.

Reviewer comment (Experiments):
The experiments are limited in the range of experiments covered. The paper covers experiments on MLPs and CNNS. It would interesting to see how the method performs on modern transformers.

Author response:
# Transformer

We have performed additional experiments on CIFAR-10 using vision transformers (ViT). The results show consistent improvement of ours IS/AS over LOW and DLIS. We will add the plot to the appendix.

Reviewer comment (Experiments):
What is value of $\alpha$ used in Eqn 7? How important is it?

Author response:
# Parameter $\alpha$ in Eq.7

We choose $\alpha$ via grid search in $\{0,0.1,0.2,0.3\}$ depending on the task.

Reviewer comment (Evaluation):
Why is the DLIS performance on Oxford flowers getting worse after a point? What is sparsity of dataset?

Author response:
# DLIS underperformance on Oxford flower 102

The Oxford flower 102 dataset comprises 102 classes across 1020 images, i.e. 10 images/class, which is sparse compared to other datasets. In this example, DLIS has a poor performce as the mini-batch size is small. This, in combination with a resampling-based algorithm with high class count, yields poor gradient estimation, making it hard for the method to converge.

Reviewer comment (Writing):
In page 2: paragraph on importance sampling: what does "inherently multivariate" mean?

Author response:
# Naming convention

The gradient of a network lives in parameter space. The gradient can be seen as a multivariate vector. We generally refer to the different parameters as dimensions in this parameter space. We will clarify.

Reviewer comment (Writing):
In page 2: paragraph on importance sampling: what does "good compromise between optimization quality across all dimensions simultaneously" mean? What dimensions?

Author response:
The error in the parameter-gradient estimation depends on the sampling strategy. A sampling distribution may yield an estimation that is accurate for parameters and poor for others. A good distribution achieves low error across all parameters at the same time. We refer to this as compromise of quality across dimensions.

Reviewer comment (Writing):
The title of the paper sounds problematic from a grammatical point of view as the word adaptive is an adjective unlike importance.

Author response:
# Title choice

We mean "adaptive sampling" and "importance sampling", both of which are standard terms in literature. We are happy fully spell these out in the title, unless reviewers think this would make it overly verbose.

Reviewer comment (Theory):
While the sampling scheme is well founded, it seems that the presented approach can suffer from overfitting issues. The authors explicitly propose to add an \epsilon to the value of the importance to avoid focusing on a small set of the data points.

Author response:
# Other concerns

From our experiments, the proposed method does not suffer from more overfitting than other methods. In all our presented experiments, we have reported the metric tested on data never seen during training.

Reviewer comment (Writing):
What is meant with the sentence: In this method, as outlined in eq. (3), each data point’s weight, denoted as w(xi), remains constant at N.?

Author response:
# What is meant with the sentence: In this method, as outlined in eq. (3), each data point’s weight, denoted as w(xi), remains constant at N. ?

In adaptive sampling, achieving a proper average estimator requires the weight $w$ to compensate for the factor $1/|\Omega|$ in front of the summation. Hence, we set the weight to the dataset size: $w(x_i) = |\Omega|$.

Reviewer comment (Evaluation):
Beyond the weight calculation overhead what are limitations and potential drawbacks of the proposed method?

Author response:
# Limitations/drawbacks beyond weight calculation overhead

Our proposed algorithm relies on a memory of the data importance, with size proportional to the dataset size. In cases where data is streamed and not stored, the algorithm cannot be applied. It can also create outdated importance for data if this memory is not updated often enough. This can happen with a high learning rate and result in a poor sampling distribution. In the worse case, the error in the gradient estimation can be higher than without importance sampling.

Another limitation is related to non-uniform data sampling: The samples that go into a mini-batch are not laid out contiguously in memory, which slows down their loading. We believe that a careful implementation can mitigate this issue.

We will add a paragraph at the end of the result section to discuss the limits and drawbacks of the algorithm and sampling strategy.

Reviewer comment (Novelty):
The paper lacks discussions of related paper. For example, https://arxiv.org/pdf/2104.13114.pdf also considers the importance sampling problem by sampling data points proportionally to the loss, instead of norm of gradient.

Author response:
We will discuss the recommended citations in the related work section.

Reviewer comment (Presentation):
It is hard to tell the algorithm performance differences for some figures. For example, in Figure 5 (left), the authors claim that at equal epochs (left), our methods (Ours IS & AS) show improvements compared to LOW Santiago et al. (2021) and DLIS weights. It is really invisible to see the difference. The authors may consider plotting the log scale results.

Author response:
We will add zoom-ins in Figure 5 to clearly show the differences in convergence behavior near the tails of the plots.

Reviewer comment (Theory):
I feel a bit confused in Ours IS vs Ours AS. Does Ours IS mean that you set $w(x) =1 / p(x)$  and Ours AS mean that you set$w(x) =1 / N$? If so, why do you claim your adaptive sampling approach subsumes the adaptive weighting of Santiago et al. in page 3? I think they are different because Ours AS just do the sampling with non-uniform distribution and there is no weighting.

Author response:
Adaptive sampling selects samples non-uniformly while weighting them equally. The same sample can be selected multiple times during training, giving it a higher "weight". Adaptive weighting directly assigns higher weight to more important samples, while selecting the samples uniformly. The two strategies are very similar, and our framework can be tailored to perform either. Applying weighting through sampling is known to achieve better performance than explicit weighting.