Paper:
Feature selection problems have been extensively studied in the setting of linear estimation, for instance LASSO, but less emphasis has been placed on feature selection for non-linear functions. In this study, we propose a method for feature selection in non-linear function estimation problems. The new procedure is based on directly penalizing the 0 norm of features, or the count of the number of selected features. Our 0 based regularization relies on a continuous relaxation of the Bernoulli distribution, which allows our model to learn the parameters of the approximate Bernoulli distributions via gradient descent. The proposed framework simultaneously learns a non-linear regression or classification function while selecting a small subset of features. We provide an information-theoretic justification for incorporating Bernoulli distribution into our approach. Furthermore, we evaluate our method using synthetic and real-life data and demonstrate that our approach outperforms other embedded methods in terms of predictive performance and feature selection.

Review:
The author rebuttal sufficiently addresses my concerns, so I am upgrading my score.
***
The paper considers the problem of embedded feature selection for supervised learning with nonlinear functions. A feature subset is evaluated via the loss function in a "soft" manner: a fraction of an individual feature can be "selected". Sparsity in the feature selection is enforced via a relaxation of l0 regularization. The resulting objective function is differentiable both in the feature selection and learned function making (simultaneous) gradient-based optimization possible. A variety of experiments in several supervised learning tasks demonstrates that the proposed method has superior performance to other embedded and wrapper methods.
My decision is to reject, but I'm on the fence regarding this paper. I'm not clearly seeing the motivation for an embedded feature selection method for neural network models: for the datasets considered in the paper, it would seem that training a nonlinear model that used all the features would result in performance at least as good as training the nonlinear model with a prepended STG layer. Perhaps there is evidence that filtering features, e.g., irrelevant features, results in higher accuracy and that the prepended STG layer achieves this accuracy, but that evidence is missing from the paper. Also, there could be downstream computational savings, e.g., at prediction time, if the dimension was very large, but this is not the setting tested in the experiments. I suppose interpretability could be considered motivation, but, even so, isn't there at least one simpler, deterministic approach (described below) that also "solves" the problem? Finally, it isn't clear how the method scales with increasing sample size and dimension as all the datasets tested are relatively small in these respects.
***
Questions and suggestions related to decision:
* The performance values using all features should be included in the experimental results so that the value added by STG can be assessed.
* Why not use the simpler deterministic and differentiable relaxation z = \sigma(\mu), where \sigma() is a "squashing" function from the real numbers to [0,1] applied element-by-element to the vector \mu? What specifically is/are the advantage(s) that the randomness in the definition of z at the bottom of pg. 3 provide over this deterministic alternative?
* Though well-described and methodologically rigorous, the experimental comparison is none-the-less a little disappointing: one dataset for classification and half the datasets for regression are synthetic and low-dimensional. The remaining regression datasets are real but also low-dimensional. The survival analysis dataset is also low-dimensional (as described in the supplementary material). This leaves one real classification dataset which was on the order of 20,000 examples and 2500 features. Why were larger sample-size and dimensionality datasets not tested? These should be readily available. For example, the gisette dataset from the NIPS 2003 feature selection challenge has 5000 features. See "MISSION: Ultra Large-Scale Feature Selection using Count-Sketches" by Aghazadeh & Spring et al. (2018) for other high-dimensional datasets. Even a single run for each large dataset would have provided some evidence of scalability.
***
Other minor comments not related to decision:
* "Concrete Autoencoders for Differentiable Feature Selection and Reconstruction" by Abid et al. (2019) targets unsupervised feature selection but has enough similarities in the approach that it should be considered related work.
* [Typo?] The unnumbered equation after (5) should not have a sum over d in the second term. Perhaps a sum over k was intended? Also, in this equation, the gradient of the loss wrt/ z samples, average of gradients over z samples times..., does not seem to match what the gradient would be given the algorithmic description in the supplementary material, a gradient of the (sample) average z times...
* The abstract states the paper is proposing a method for high-dimensional feature selection, but all of the experiments have datasets with max. dimensionality 2538.
* Some discussion of how the regularization parameter can be selected by a user of the proposed method would be good to include.