SFESS: Score Function Estimators for $k$-Subset Sampling

Thank you for your valuable feedback. We address your comments and questions below.

W1: Where is $k$-subset sampling used?

The first paragraph of the introduction lists some examples of where differentiable $k$-subset sampling can be used. Generally, it is used to include $k$-hot vectors in differentiable optimization by sampling the vector and estimating its gradient. There are three examples in the experiment section: choosing multiple options at once (feature selection), learning a sparse code (VAE), and a top-$k$ operation (stochastic k-NN). The main limitations of existing methods are their biased gradients and reliance on the downstream function’s gradients (see the paragraph starting on line 136 and Table 1). As for applications, $k$-subset sampling has broad uses, e.g. MRI reconstruction [1], microarray data [2], and explainability [3].

Q1: Importance of properties in Table 1

Here we clarify why the properties in Table 1 are important and preferred:

• Exact samples. This is preferred for two reasons (also mentioned on line 146): (1) it makes the downstream model consistent at training and test time (facilitating generalization), and (2) it allows for training with downstream functions that only accept discrete inputs and cannot be relaxed. Consider, for instance, a reinforcement learning setting where the action space is a subset (e.g. press $k$ out of $n$ buttons). In fact, obtaining exact samples was a motivating idea behind SIMPLE.
• Unbiased. An unbiased estimator’s expected value is equal to the true parameter. This is generally a desirable property in statistical estimators and particularly important when analyzing convergence (see W1 in our response to Reviewer W34X).
• Non-differentiable. This indicates whether the gradient estimate can be computed for a non-differentiable downstream function $f$ (see the table caption). The column header is now changed to “Non-differentiable $f$” in the revised paper to improve clarity. See W2 in our response to Reviewer kBPo for particular use cases of this property.
• Tuning-free. This property is desirable because it (1) simplifies the model and (2) reduces the computational burden of hyperparameter tuning. In GS, STGS, and IMLE, tuning the temperature parameter for a given problem is crucial, often requiring repeated training runs.

We have revised the caption of Table 1 to clarify points that were not already in the paper.

Q2: Estimating gradients

We are unsure about the intended question, so, we considered two possible interpretations below and answered them accordingly. Let us know if they do not answer your question.

Interpretation 1: Estimating the gradient of the $k$-subset sampling method

To clarify, the gradient required for GS, STGS, and SIMPLE is the gradient of the downstream function $f$. We do not differentiate the methods themselves. The key difference between these methods and our proposed method is that it does not require the downstream function’s gradient and can thus be used in a broader setting.

Interpretation 2: Using estimated gradients of $f$

This is interesting. In principle, it is possible to combine the methods requiring gradients with a gradient estimator for the pathwise gradient if it is unavailable. However, we are not aware of any such attempts, and trivial solutions would be costly and likely prone to error propagation due to using an estimate to compute another.

Q3. Interpreting Figure 5

Figure 5 highlights the importance and effectiveness of variance reduction. In previous work score function estimators have been dismissed due to their high variance, and Figure 5 shows us why — variance reduction is crucial for constructing a score function estimator. Our proposed estimator SFESS+VR shows that variance reduction (along with our proposed method for efficiently computing the score function) makes score function estimators a viable option for gradient estimation in $k$-subset sampling. The viability of score function estimators with variance reduction is the main message of the paper.

References

[1] Van Gorp, Hans, Iris Huijben, Bastiaan S. Veeling, Nicola Pezzotti, and Ruud JG Van Sloun. Active deep probabilistic subsampling. In International Conference on Machine Learning, 2021.

[2] Zena M. Hira 1, Duncan F. Gillies. A Review of Feature Selection and Feature Extraction Methods Applied on Microarray Data. In Adv Bioinformatics, 2015.

[3] Jianbo Chen, Le Song, Martin Wainwright, and Michael Jordan. Learning to Explain: An Information-Theoretic Perspective on Model Interpretation. In International Conference on Machine Learning, 2018.

Thank you for your valuable feedback. We address your comments and questions below.

W1: Theoretical foundation

Algorithm 2 is a standard application of stochastic gradient descent using the proposed gradient estimates. The algorithmic motif follows that of training a variational autoencoder [1] where gradient estimates are used within SGD (this is sometimes called doubly stochastic optimization due to the two sources of variance: the subsampling of data and the gradient estimates [2]). This optimization approach has become widespread in deep learning in the last decade. For example, the gradient estimators GS, STGS, IMLE, and SIMPLE were all proposed in this setting without proof of convergence. In the same vein, SFESS is a gradient estimator that we evaluate as part of this common approach.

SFESS also inherits the theoretical foundation of score function estimators. Perhaps the most important property is its unbiased gradient, a proof of which is presented in Appendix A. Unbiased gradient estimates are a common assumption in the analysis of e.g. gradient descent algorithms which does not hold for GS, STGS, IMLE, or SIMPLE. As an example, Lemma 4.2 in Bottou et al. 2018 [3] shows the expected improvement between two iterates for continuously differentiable downstream functions $f$ with Lipschitz-continuous gradient $\nabla f$. For unbiased gradient estimators (like SFESS), the expected improvement is bounded depending on the learning rate, the Lipschitz constant of $\nabla f$, and the second moment of the gradient estimator — which motivates variance reduction to improve convergence.

Finally, to clarify the wording in Algorithm 2: training until convergence refers to its usual meaning in deep learning and not a formal notion of convergence, i.e. training until some metric no longer improves significantly.

Q1: How was $k$ chosen?

Note that, the subset size $k$ is not chosen. It is given as a part of the problem definition (gradient estimation of $k$-subset sampling). As such, it exists for all other methods as well. Importantly, SFESS works for any $k$, unlike GS and STGS which deteriorate with larger $k$ (see line 427). As for the experimental setup, we compare all methods in various problem settings ($n$ and $k$) as shown in Tables 2-4.

Q2: How were the random seeds chosen?

The random seeds were not chosen deliberately. The random seeds $\\{0, 1, 2, 3, 4\\}$ were used for repeated experiments. We have made the code available through an anonymous repository in a separate comment. The code will also be made publicly available upon publication to facilitate reproducibility.

Q3: Performance comparisons

We would like to emphasize that we do not claim that SFESS provides superior performance to other methods in all settings. The paper proposes that score function estimators should not be dismissed as they produce comparable performance to state-of-the-art methods (GS, STGS, SIMPLE) in the showcased experiments while providing qualitative advantages as mentioned in W1 of response to Reviewer kBPo.

To give some examples of the broader applicability of score function estimators, consider the following downstream functions that are not differentiable:

• A black-box function or e.g. a network with quantized weights.
• Directly optimizing a non-differentiable metric (e.g. accuracy).
• A reinforcement learning environment with a non-differentiable reward function.

References

[1] Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In International Conference on Learning Representations, 2014.

[2] Titsias, Michalis, and Miguel Lázaro-Gredilla. Doubly stochastic variational Bayes for non-conjugate inference. In International Conference on Machine Learning, 2014.

[3] Bottou, Léon, Frank E. Curtis, and Jorge Nocedal. Optimization Methods for Large-Scale Machine Learning. In SIAM Review, 60.2 223-311, 2018.

Thank you for your valuable feedback. We address your comments and questions below.

W1: Novelty of score function estimators

We agree with the reviewer that score function estimators are a well-established technique in general, which is why we would like to bridge this research gap for the specific task of $k$-subset sampling. Interestingly, prior works on $k$-subset sampling [1, 2, 3] mention score function estimators and their high variance or the cost of computing the score function. The following excerpts highlight this research gap:

Score-based methods such as REINFORCE [Williams, 1992] for estimating the gradient of such objectives typically have high variance. [1] (page 1)

SFE uses faithful samples and exact marginals (which is feasible only when m is very small) and converges much more slowly than the other methods [...] [2] (page 8)

[The score function estimator] is typically avoided due to its notoriously high variance, despite its apparent simplicity. [3] (page 3)

Besides identifying this research gap, we address the two limitations above by (1) proposing an efficient method for computing the score function, and (2) using control variates to reduce the estimator's variance. We show that this novel combination of techniques results in comparable results to existing state-of-the-art methods on various tasks in sections 4.2 to 4.4 while providing valuable qualitative advantages: unbiased gradient estimates, applicability to non-differentiable $f$, and not requiring hyperparameter tuning (see Table 1).

SFE in SIMPLE is an example of how score function estimators for $k$-subset sampling were not viable prior to our contributions. In SIMPLE. the score function was computed by enumerating all subsets in a toy problem (the problem adopted in section 4.1). This is not feasible for larger experiments, which is most likely why SFE was not included in SIMPLE's larger experiments (Tables 2 and 3 in [3]). In addition, its performance without variance reduction would have been poor (see SFESS vs. SFESS+VR in our experiments).

W2: Results compared to other methods

We would like to emphasize that we do not claim that SFESS provides superior performance to other methods in all settings. The paper proposes that score function estimators should not be dismissed as they produce comparable performance to state-of-the-art methods (GS, STGS, SIMPLE) in the showcased experiments while providing qualitative advantages as mentioned in response to W1 above.

To give some examples of the broader applicability of score function estimators, consider the following downstream functions that are not differentiable:

• A black-box function or e.g. a network with quantized weights.
• Directly optimizing a non-differentiable metric (e.g. accuracy).
• A reinforcement learning environment with a non-differentiable reward function.

References

[1] Sang Michael Xie and Stefano Ermon. Reparameterizable Subset Sampling via Continuous Relaxations. In International Joint Conference on Artificial Intelligence, 2019.

[2] Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions. In Advances in Neural Information Pro- cessing Systems, 2021.

[3] Kareem Ahmed, Zhe Zeng, Mathias Niepert, and Guy Van den Broeck. SIMPLE: A Gradient Estimator for k-Subset Sampling. In International Conference on Learning Representations, 2023.