Title: FEATURE SELECTION USING STOCHASTIC GATES

Abstract: 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.

1. INTRODUCTION:

Feature selection is a fundamental task in machine learning and statistics. Feature selection leads to a number of potential benefits: reducing experimental costs Min et al. (2014) , enhancing interpretability (Ribeiro et al., 2016) , computational speed up and even improving model generalization on unseen data (Chandrashekar & Sahin, 2014) . In biomedicine, scientists collect multitude datasets comprising of many biomarkers (e.g., genes or proteins) that require development of effective diagnostics or prognostics models. For instance, in Genome wide association studies (GWAS), feature selection can help identify such models and lead to improved risk assessment and reduced cost.

Feature selection methods may be classified into three major categories: filter methods, wrapper methods, and embedded methods. Filter methods attempt to remove irrelevant features prior to learning a model. These methods filter features using a per-feature relevance score that is created based on some statistical measure (Battiti, 1994; Peng et al., 2005; Estévez et al., 2009; Song et al., 2007; 2012; Chen et al., 2017) . Wrapper methods (Kohavi & John, 1997b; Stein et al., 2005; Zhu et al., 2007; Reunanen, 2003; Allen, 2013) use the outcome of a classifier to determine the relevance of each feature, which requires recomputing the classifier for each subset of features. This becomes computationally expensive for neural network based wrapper methods (Verikas & Bacauskiene, 2002; Kabir et al., 2010; Roy et al., 2015) .

Embedded methods aim to remove this burden by learning the model while simultaneously selecting the subset of relevant features. The Least Absolute Shrinkage and Selection Operator (LASSO) (Tibshirani, 1996) is a well-known embedded method, whose objective is to minimize the loss while enforcing an 1 constraint on the weights of the features. Although LASSO is scalable and widely used (Hans, 2009; Li et al., 2011; 2006) , it is restricted to the domain of linear functions. To allow the model to capture nonlinear interaction, it is appealing to consider the non-convex extension of the LASSO formulation using neural networks.

We develop a fully embedded feature selection method for nonlinear models. To the best of our knowledge it is the first 0 embedded feature selection method. Our method improves upon the LASSO formulation in two aspects: a) it captures nonlinear interactions between the features via neural network modeling and b) it employs an 0 -like regularization using gates whose weights are parametrized by a smooth variant of a Bernoulli distribution. Altogether these twofold improvements are formulated as a fully differentiable neural network. Specifically, our contributions are as follows:

1. By utilizing a recent development of continuous and differentiable approximation to discrete distributions (Maddison et al., 2016) , (Jang et al., 2017) , (Louizos et al., 2017) , we introduce a solution to the long standing problem of feature selection with an 0 regularization.

2. We present a novel simple relaxation of Bernoulli distribution to sparsify the input layer (the feature space) which we call stochastic gate (STG) and show its advantage over the distribution presented by (Louizos et al., 2017) both in performance and convergence time.

3. By applying these two relaxations to an input layer of a neural network, we perform embedded feature selection in classification, regression or survival analysis tasks and demonstrate its capabilities on artificial and real data sets.

4. We justify our probabilistic approach by analyzing the constrained Mutual Information maximization objective of feature selection. We demonstrate the applicability of our method using numerous examples (see Section 6 and Appendix).

Notation: We refer to vectors as bold lowercase x and random vectors as bold uppercase letters X. Scalars are non-bold case y, while random variables are capital case Y . A set is represented by script fonts X , Y, S. For example the n th vector-valued observation is denoted as x n whereas X d represents the d th feature of the vector-valued random variable X. Let [n] = 1, 2, . . . , n. For a set S ⊂ [D] let the vector s ∈ {0, 1} D be the characteristic function for the set. That is s i = 1 if i ∈ S and 0 otherwise. For two vectors x and z we denote x z to be the element-wise product between x and z. Thus, if we let s ∈ {0, 1} D be the characteristic vector of S, then we may define x S = x s.

The 1 norm of a vector is denoted as

x 1 = D i=1 |x i |.

Finally, the 0 norm of a vector is denoted as x 0 and counts the total number of non-zero entries in the vector x.

2. PROBLEM SETUP AND BACKGROUND:

Let X ∈ R D be the input domain with corresponding response domain Y. Given realizations from some unknown data distribution P X,Y , the goal of embedded feature selection methods is to simultaneously find a subset of indices S ⊂ {1, ...D} and construct a model that predicts Y based on the selected features X S .

2.1. RISK MINIMIZATION OBJECTIVE:

We assume that we are given a family of functions F such that any function f θ ∈ F is indexed by a set of parameters θ. Given some loss L, a selection of features S ⊂ [D], and a choice of parameters θ, we denote the risk of our model as

EQUATION

where we recall that s = {0, 1} D is a vector of indicator variables for the set S and denotes the point-wise product. Thus, the goal of the feature selection problem is to find the parameters θ and s that minimize R(θ, s) such that s 0 is small compared to D.

2.2. FEATURE SELECTION FOR LINEAR MODELS:

Before proceeding with our proposed method, we review the feature selection problem in the linear regression setting with the least squares loss. Thus, we restrict F to be the space of linear functions and the loss function to be the quadratic loss.

Given observations {x n , y n } N n=1 we may consider the constrained empirical risk minimization problem

min θ 1 N N n=1 L(θ T x n , y n ) s.t. θ 0 ≤ k.

(2)

Since the above problem is intractable, a number of authors replace the 0 constraint with a surrogate function Ω(θ) : R D → R + designed to penalize the number of selected features in θ. A natural choice for Ω is the 1 norm, which yields a convex problem and more precisely the LASSO optimization problem. The 1 is known to be the closest convex relaxation to the 0 . In fact, in certain settings the LASSO optimization and the 0 based objective have the same solution. While the original LASSO problem focuses on the constrained optimization problem, the regularized least squares problem, which is often used in practice, yields the following minimization objective:

EQUATION

The hyperparameter λ trades off the amount of regularization versus the fit of the objective1 . The 1 regularized method is very effective for feature selection and prediction; however, it achieves this through shrinkage of the coefficients. As a result, (Fan & Li, 2001) have considered non-convex choices for Ω that perform well both theoretically and empirically for prediction and feature selection.

Our goal is to apply such regularization techniques to perform feature selection while learning a non-linear function. Kernel methods have been considered (Yamada et al., 2014) , but scale quadratically in the number of observations. To alleviate this burden, (Gregorová et al., 2018) use random Fourier features to approximate the kernel. (Li et al., 2016) and (Scardapane et al., 2017) take an alternative approach by modeling f θ using a neural network with 1 regularization on the input weights. However, in practice, introducing an 1 penalty into gradient descent does not provide sufficient sparsification. Below, we discuss our method that works to directly use an 0 penalty.

3. PROPOSED METHOD:

We take a probabilistic approach to approximate the 0 norm, which can extend to non-linear models while remaining computationally efficient. To motivate such probabilistic approach, we provide theoretical support (see Section 4) based on a Mutual Information perspective of the feature selection problem.

To view the 0 regularized version of the risk (Eq. 1) from a probabilistic perspective, one can introduce a Bernoulli random vector S whose entries are independent and the d th entry satisfies

π d = P( Sd = 1) for d ∈ [D].

If we denote the empirical expectation over our observations as ÊX,Y , then, the empirical regularized risk (Eq. 1) becomes

EQUATION

where we have E S || S|| 0 = D d=1 π d and we constrain π d ∈ {0, 1}. Clearly, this formulation is equivalent to Eq. 1, with a regularized penalty on cardinality rather than an explicit constraint. We may then relax the discrete constraint on π d to be π d ∈ [0, 1]. Now, our goal is to find the model parameters θ * and Bernoulli parameters π * that minimize the empirical risk R(θ, π) via gradient descent. However, an optimization of a loss function which includes discrete random variables suffers from high variance (see F in the Appendix for more details). Therefore, inspired by a recently developed continuous approximation for discrete random variables, suggested by (Jang et al., 2017; Maddison et al., 2016) , we develop and use a novel and simple continuous distribution that is fully differentiable and suited for the task of feature selection.

3.1. CONTINUOUS RELAXATION:

Our continuous relaxation for the Bernoulli variables Sd for d ∈ [D] is termed stochastic gate (STG). The STG relies on the reparametrization trick, which is widely used for reducing the variance of gradient estimators (Miller et al., 2017; Figurnov et al., 2018) . To construct a continuous approximation to Bernoulli random variable via the reparametrization trick, we define z d = g(µ d + d ) = max(0, min(1, d + µ d + 0.5)) where d is drawn from a Gaussian distribution N (0, σ d ), where σ d is fixed throughout training. This approximation can be viewed as a clipped, mean-shifted, Gaussian random vector. Furthermore, the gradient of the objective with respect to µ d can be computed via the chain rule.

We can now rewrite the objective in Eq. 4 as

EQUATION

where Z is a random vector with D independent variables z d for [D]. To optimize the empirical surrogate of the objective (Eq. 5), we first differentiate it with respect to µ. Then, Monte Carlo sampling leads us to the following gradient estimator

∂ ∂µ d R(θ, µ) = 1 K K k=1 L (z k ) ∂z k d ∂µ d + λ K ∂ ∂µ d K k=1 Pr{z k d > 0},

where K is the number of Monte Carlo samples. Thus, we can update the parameters µ d for [D] via gradient descent. We note that if we replace ∂z k d ∂µ d with 1, the above gradient estimator for L is reduced to the Straight-Through estimator (Bengio et al., 2013) .

Under the continuous relaxation, the expected regularization term in the objective R(θ, µ) is simply the sum of the probability that the gates {z d } D d=1 are active, which is equal to

D d=1 Φ µ d + 1 2 σ d

, where Φ is the standard Gaussian CDF. To conclude, we can now optimize the objective in Eq. 5 using gradient descent over the model parameters θ and the parameters µ representing the Gaussian's mean (instead of the Bernoulli parameters π).

After training, to remove the stochasticity from the learned gates, we set ẑd = max(0, min(1, µ d + 0.5)), which informs what features are selected. Note that when |µ d | is less than 1 2 , ẑd returns the value between (0, 1). In such a case, we can treat the value of ẑd as feature importance or employ an additional thresholding (i.e. 1 if ẑd > 0.5 and 0 otherwise) depending on application-specific needs. In the Appendix, we provide the pseudo-code of our algorithm as well as the discussion of the choice of σ d .

4. CONNECTION TO MUTUAL INFORMATION:

In this section we show an equivalence between the Bernoulli formulation of the feature selection problem and the 0 regularized approach.

4.1. MUTUAL INFORMATION BASED OBJECTIVE:

From an information theoretic perspective, the goal of feature selection is to find a subset of features S that has the highest Mutual Information (MI) with the target variable Y . Recall that the MI between two random variables can be defined as I(X; Y ) = H(Y ) -H(Y |X) where H(Y ), H(Y |X) are the entropy of p Y (Y ) and the conditional entropy of p Y |X (Y |X), respectively (Cover & Thomas, 2006) . Then we can formulate the task as selecting S such that the mutual information between X S and Y are maximized:

EQUATION

)

where k is the hypothesised number of relevant features.

4.2. INTRODUCING RANDOMNESS:

We first demonstrate that under mild assumptions we can replace the deterministic search over the set S (or corresponding indicator vector s), by a search over the parameters of the distributions that model s. Our proposition is based on the following two assumptions:

Assumption 1: There exists a subset of indices S * with cardinality equal to k such that for any i ∈ S * we have

I(X i ; Y |X \{i} ) > 0. Assumption 2: I(X S * c ; Y |X S * ) = 0.

Discussion of assumptions: Assumption 1 that including an element from S * improves prediction accuracy. This assumption is equivalent to stating that feature i is strongly relevant (Kohavi & John, 1997a; Brown et al., 2012) . Assumption 2 simply states that S * is a superset of the Markov Blanket of the variable Y (Brown et al., 2012) . The assumptions are quite benign. For instance they are satisfied if X is drawn from a Gaussian with a non-degenerate covariance matrix and Y = f (X S * ) + w, where w is noise independent of X and f is not degenerate. With these assumptions in hand, we may present our result.

Proposition 1. Suppose that the above assumptions hold for the model. Then, solving the optimization in Eq. 6 is equivalent to solving the optimization

EQUATION

where the coordinates Si are drawn independently at random according to a Bernoulli distribution with parameter π i .

Due to length constraints, we leave the proof of this proposition and how it bridges the MI maximization (Eq. 6) and risk minimization (Eq. 2) to the Appendix (see Section C and D).

5. RELATED WORK:

The two most related works to this study are (Louizos et al., 2017) and (Chen et al., 2018) . In (Louizos et al., 2017) , they introduce the Hard-Concrete (HC) distribution as a continuous surrogate for Bernoulli distributions in the context of model compression. The HC distribution is induced by applying a hard sigmoid z = min(1, max(0, s)) to s, where we construct s by applying Sigmoid function to the logistic distribution as follows:

u ∼ U (0, 1), L = log(u) -log(1 -u), s = 1 1 + exp( -(log α+L) β ) , s = s(ζ -τ ) + τ

The interval (τ, ζ), with τ < 0 and ζ > 1, allows the distribution to have more probability mass on the edge of the support. The authors demonstrate how incorporating the HC in deep neural networks leads to fast convergence and improved generalization due to the sparsification effect. In this study, we introduce this type of sparsification for the long standing problem of feature selection with 0 regularization.We demonstrate that a simple relaxation of Bernoulli distributions is sufficient and works better than the HC distribution for feature selection. We find empirically that the HC yields high variance gradient estimates compared to the STG. This high-variance does not affect model sparsification as the sparsity pattern does not matter. Contrarily, for feature selection it is crucial as practitioners require a stable procedure. Furthermore, higher variance gradients will result in a slower convergence of the model (this has been demonstrated empirically in the supp. material). In (Chen et al., 2018) , the Gumbel-softmax trick is used to develop a framework for interpreting pre-trained models. Their method is focused on finding a subset of features given a particular instance, and therefore is not appropriate for general feature selection.

Some studies tackle embedded feature selection problems by extending LASSO and group LASSO to neural network models. The authors (Li et al., 2016; Scardapane et al., 2017) and (Feng & Simon, 2017 ) have a similar goal as ours in performing feature selection, but instead rely on the 1 relaxation to the 0 . Our approach, which utilizes stochastic gates coupled with the 0 norm achieves substantially better empirical performance compared against other 1 based baselines. This point is demonstrated in the next section.

6. EXPERIMENTS:

Here we provide empirical evaluation of our method in a wide range of settings. We begin with simple linear regression evaluating the capabilities of our approach to select a sparse set of relevant features (section 6.1). We then evaluate our method in nonlinear regression tasks (section 6.2). Next, we demonstrate the applicability of the method in a highly nonlinear classification task (section 6.3.

Finally, we present two potential applications (section 7).

We implement our method using both the STG and HC distributions and compare their applicability for feature selection. The hyperparamters of all the methods are optimized using validation sets via Optuna (Takuya Akiba & Koyama, 2019). In the Appendix, we provide all details of the Optuna based parameter tuning procedure. Additional experiments also appear in the Appendix, including high dimensional datasets and an extensive comparisons between the STG and HC (Louizos et al., 2017) distributions.

6.1. SUPPORT RECOVERY IN LINEAR REGRESSION:

In the setting of noisy linear regression, (Wainwright, 2009) have analyzed the probability of LASSO to correctly identify a sparse subset of active variables. The problem is know as support recovery and is formulated as follows; let β * ∈ R D be a fixed sparse vector, such that β * i ∈ {-0.5, 0.5} (with equal probability) if i ∈ S, and β * i = 0 otherwise. Suppose the cardinality of the support |S| = k is known. Given a matrix of measurements X ∈ R N ×D with values drawn independently from N (0, 1), the response y is defined as

EQUATION

where the values of the noise w i , i = 1, ..., N are drawn independently from N (0, 0.5).

Here, we reproduce this setting to evaluate the probability of the proposed approach to perfectly recover the support of β * . As suggested by (Wainwright, 2009) , we use a sparsity that scales with D such that k = 0.4D 0.75 . For each number of samples N in the range [10, 500], we run 200 simulations and count the portion of correctly recovered supports. We repeat this process for 2 different values of D and compare our performance to LASSO. For LASSO, the regularization parameter was set to its optimal value α N = 2σ2 log(D-k) log(k) N (Wainwright, 2009) . For STG and HC we set λ N = Cα N , such that C is a constant, which is selected using a grid search in the range [0.1,10]. As evident from Fig. 1 , even when restricting to linear functions our method has a clear advantage over LASSO. This implies that the 0 based penalty, although is non convex in nature, allows perfect recovery of the support using less samples. Furthermore, the proposed STG requires even less samples than the HC distribution and suffers from a smaller variance. 

6.2. NONLINEAR REGRESSION USING SYNTHETIC AND REAL DATASETS:

In this section, we evaluate our method for regression tasks against two other embedded feature selection methods: LASSO and Sparse Random Fourier Features (Gregorová et al., 2018) . Following the same format as (Gregorová et al., 2018) , the following functions are used to generate synthetic data: (SE1: 100/5) y = sin (x 1 + x 3 ) 2 sin (x 7 x 8 x 9 ) + N (0, 0.1). (SE2: 18/5) y = log((

EQUATION

3 ) + N (0, 0.01), where each x i is drawn from the standard Gaussian N (0, 1). For (SE3), the five consecutive coordinates are generated by x 5(j-1)+i = z j + N (0, 0.1), where z j ∼ N (0, 1) for i = 1, ..., 5 and j = 1, ..., 200. The numbers next to the experiment code indicate total dimensions/relevant dimensions in the feature space. We also evaluate our method using three real datasets: (RCP: 21/-), of computer systems activity, (REL 17/-) of F16 elevators and (RAI 39/-) of F16 ailernos all taken from the LIACC repository 2 . For each dataset, we generate 30 different replications and randomly split the data into train, validation, and test set (see Appendix for more details). The root mean squared error on the test set averaged over 30 random replicated datasets are reported in Table 1 . Our method outperforms all alternative methods for most cases. Note that (SE1) is generated using a sine function, which is in favor of the random Fourier feature based method (SRFF).

Table 1 : Regression performance comparison in terms of root mean squared error. The mean and standard deviation are computed across 30 resamples. N is 1000 for SE1,2,3 and 6000 for RCP. The values for LASSO and SRFF are borrowed from (Gregorová et al., 2018) 

6.3. NOISY BINARY XOR CLASSIFICATION:

In the following evaluation, we consider the problem of learning a binary XOR function for classification task. The first two coordinates x 1 , x 2 are drawn from a binary "fair" Bernoulli distribution. The response variable is set as an XOR of the first coordinates, such that y = x 1 ⊕ x 2 . The coordinates x i , i = 3, ..., D are nuisance features, also drawn from a binary "fair" Bernoulli distribution. The number of points we generate is N = 1, 500, of which 70 % are reserved for test and 10% of the remaining training set was reserved for validation. We compare the proposed method to four embedded feature selection methods (LASSO (Tibshirani, 1996) , C-support vectors (SVC) (Chang & Lin, 2011) , deep feature selection (DFS) (Li et al., 2016) , sparse group regularized NN (SG-L1-NN) (Scardapane et al., 2017) ). To provide more benchmarks, we also compare our embedded method against three wrapper methods (Extremely Randomized Trees (Tree) (Rastogi & Shim, 2000) , Random Forests (RF) (Strobl et al., 2008) ) and Extreme Gradient Boosting (XGBOOST) (Chen & Guestrin, 2016) .

To evaluate the feature selection performance, we calculate the Informative Features Weight Ratio (IFWR). IFWR is defined as the sum of weights W d over the informative features divided by the sum over all weights. In the case of binary weights the IFWR is in fact a recall measure for the relevant features (See the Appendix for more details.)

The experiment is repeated 20 times for different values of D, and the average test classification accuracy and standard deviation are presented in Fig. 2 (a), followed by the IFWR in Fig. 2(b ). The number of selected features affects the accuracy. Therefore, to treat all the methods in a fair manner, we tune the hyperparameter that controls the sparsity level using Optuna (Takuya Akiba & Koyama, 2019) which optimizes the overall accuracy across different Ds. For instance, the wrapper methods (Tree, RF and XGBOOST) has a threshold value to retain features. We retrain them using only such features whose weight is higher than the threshold. In terms of feature ranking (see Fig. 2 (c), only the tree based methods and the proposed (STG and HC based) provide the optimal median rank (which is 1.5) for the two relevant features. Nonetheless, the ranking provided by STG is the most stable comparing to all the alternative methods. 

7. APPLICATION:

7.1 PURIFIED POPULATIONS OF PERIPHERAL BLOOD MONOCYTES (PBMCS)

In (Zheng et al., 2017) , the authors have subjected more than 90, 000 purified cell populations of peripheral blood monocytes (PBMCs) to Single-cell RNA sequencing (scRNA-seq) analysis. Here we use this data and focus on classifying two subpopulations of T-cells, namely Naive and regulatory T-cells. We use the proposed method to select a subset of genes for which the network discriminates between Naive and regulatory T-cells. We first filter out the genes that are lowly expressed in the cells, which leaves us with D = 2538 genes (features). The total number of cells in these two classes is N = 20742, of which we only use 10% of the data for training. We apply the proposed method for different values of λ and report the number of selected features and classification accuracy on the test set. Here we compare our performance (STG and HC) to RF and LASSO. A least squares polynomial fit plot of the accuracy vs. number of selected features is presented in Fig. 3 . The accuracy obtained by a NN without feature selection is 91.09%, which is comparable to what we achieve with a small fraction of the features. We have also evaluated the performance of the Hard-Concrete applied to all layers (HC-Full), following the procedure in (Louizos et al., 2017) . Empirically we observed that using this type of regularization across all layers provides inferior capabilities in terms of feature selection. Moreover, when the HC-Full converges to a larger subset of features (d > 50), it does not generalize at all and the test accuracy is around 0.5. We incorporate our STG into DeepSurv (Katzman et al., 2018) to examine how this procedure improves survival analysis. The input to this analysis are gene expression profiles (Curtis et al., 2012) along with additional commonly used clinical variables, the goal is to predict survival time. In this setting, identifying a small subset of predictive covariates is important as it may lead to cheaper and more stable medical assays. See the Appendix for additional description of the data and the experimental setup.

We compared our method (Cox-STG and COX-HC) against three other methods: Cox model with 1 regularization (Cox-LASSO), Random Survival Forest (RSF) (Ishwaran et al., 2008) and the original DeepSurv. We evaluate the predictive ability of the learned models based on the concordance index (CI), a standard performance metric for model assessment in survival analysis, which measures the agreement between the rankings of the predicted and observed survival times. The performance in terms of the CI is reported in Table 2 . We see that Cox-HC and Cox-STG outperform the alternative methods, indicating that our method shrinks the feature size while learning a predictive model.

8. CONCLUSION:

In this paper, we propose a novel embedded feature selection method for neural networks and linear models based on stochastic gates. It has an advantage over previous 1 based regularization methods in terms of achieving a high level of sparsity while learning effective non linear models. We justify our probabilistic feature selection framework from the information theoretic perspective. In experiments, we demonstrate that our method consistently outperforms existing embedded feature selection methods in both synthetic and real datasets.

A APPENDIX B ALGORITHM:

Here we present the pseudo code of our method, presented in Algorithm 1. The loss L is typically negative-log likelihood for classification and squared loss for regression. N is the sample size, D is the number of features, K is the number of Monte Carlo samples.

Input: X ∈ R N ×D , target variables y ∈ R N , regularization parameter λ, number of epochs M , learning rate γ, Output: Trained model f θ and parameter µ ∈ R D .

1: Initialize the model parameter θ. Set µ = 0.

2: for i = 1, ..., M do 3:

for n = 1, ..., N do 

C PROOF OF PROPOSITION 1:

We now provide a proof for Proposition 1, showing the equivalence between the stochastic optimization (Eq.7) and the deterministic one (Eq. 6). Let S be a subset such that S * \ S = ∅. That is there exists some element in S * that is not in S. For any such set S we have that I(X S ; Y ) < I(X; Y ). Indeed, if we let i ∈ S * ∩ Sc then we have

I(X S ; Y ) ≤ I(X \{i} ; Y ) = I(X; Y ) -I(X i ; Y |X \{i} ) < I(X; Y ),

where the final inequality follows by Assumption 1. Assumption 2 also yields that for any set S such that S * ⊂ S, we have I(X S ; Y ) = I(X; Y ). Now, when we consider the Bernoulli optimization problem we have

max π I(X S; Y ) s.t. l π l ≤ k and 0 ≤ π l ≤ 1.

The mutual information can be expanded as

I(X S; Y ) = s I(X s; Y )p π (S = s),

where we have used the fact that S is independent of everything else. Recall that in optimization (Eq. 7) the coordinates of S are sampled at random. Therefore, the distribution that is being optimized over p π is a product distribution. Our goal is to understand the form of this distribution. To that end, we will consider a problem dropping the independence constraint. If we can show that the distribution found by solving this new optimization problem with less constraints is still a product distribution, then we obtain a solution to the original optimization (Eq. 6). Now, from above we know that the optimal value of the optimization is I(X S; Y ) for any set S * ⊂ S. Hence, any unconstrained distribution should place all of its mass on such subsets in order to maximize the mutual information. As a result l∈S * p(S l = 1) = k. However, there is an optimization constraint that E[ l S l ] ≤ k. Therefore, E[S l ] = 0 for any l / ∈ S * . Hence, the optimal solution is to select the distribution so that all of the mass is placed on the subset S * and no mass elsewhere. As this is also a product distribution, this complete the proof of the claim.

D BRIDGING THE TWO PERSPECTIVES:

To motivate the introduction of randomness into the risk, we have looked at the feature selection problem from a MI perspective. Based on the MI objective, we have observed that introducing randomness into the constrained maximization procedure, does not change the objective (Proposition 1). Here we provide a relation between the MI objective (Eq. 6) to the empirical risk (Eq. 1), which supports our proposed procedure.

We first note that the MI maximization over the set S can be reformulated as the minimization of the conditional entropy H(Y |X S ) since H(Y ) does not depend on S:

max S I(X S ; Y ) = max S H(Y ) -H(Y |X S ) ⇐⇒ min S H(Y |X S ).

Recall that X S = X S. By Proposition 1, we can rewrite the deterministic search over the set S by a search over the Bernoulli parameters π:

min π H(Y |X S) = min π E X,Y, S -log P θ * (Y |X S) = min θ min π E X,Y, S -log P θ (Y |X S),

where the expectation is over X, Y ∼ P θ * , which is the true data distribution, and S ∼ Bern( S|π). Put our model distribution as P θ , then we can rewrite the right hand side as:

E X,Y, S log P θ * (Y |X S) = E X,Y, S log P θ * (Y |X S) P θ (Y |X S) P θ (Y |X S) = E X,Y, S log P θ * (Y |X S) P θ (Y |X S) + E X,Y, S log P θ (Y |X S). Since KL(P θ * (Y |X S)||P θ (Y |X S)) is non-negative, E S KL(P θ * (Y |X S)||P θ (Y |X S))

is also non-negative because it is a weighted sum of non-negative terms. Noting that

E S KL(P θ * (Y |X S)||P θ (Y |X S)) = E X,Y, S log P θ * (Y |X S) P θ (Y |X S)

we can conclude that

E X,Y, S -log P θ * (Y |X S) ≤ E X,Y, S -log P θ (Y |X S).

If we consider the negative log likelihood of the target given the observations (i.e. -log P θ (Y |X S)) as a loss function L (which encodes the classification or regression function f θ ), then we see that minimizing the risk approximately maximizes the MI objective in Eq. 6.

E DETAILS OF REGULARIZATION TERM:

Here we provide the detail description of the regularization term. For the vector of stochastic gates z ∈ R D , the regularization term is expressed as follows:

E Z Z 0 = D d=1 P[z d > 0] = D d=1 P[µ d + σ d d > - 1 2 ] = D d=1 {1 -P[µ d + σ d d ≤ - 1 2 ]} = D d=1 {1 -Φ( -1 2 -µ d σ d )} = D d=1 Φ( µ d + 1 2 σ d ).

The derivative of the regularization term with respect to the distribution parameter µ d is simply the Gaussian PDF: d . In order to remove irrelevant features, this term has to be greater than the derivative of the loss with respect to µ d because otherwise µ d is updated in the incorrect direction. To encourage such behavior, we set σ = 0.5, which is around the maximum of the gradient during the initial phase as shown in Fig. 4 . Although the point that attains the maximum moves as µ changes, we empirically observe that setting σ = 0.5 performs well in our experiments when the regularization parameter λ is appropriately tuned (see Section I). 

EQUATION

z:{0,1} D N n=1 [L(f θ (z x n ), y n ] D d=1 π z d d (1 -π d ) 1-z d .

In practice, as the outer sum involves enumerating 2 D possibilities of the indicator variables, one can replace the outer sum with Monte Carlo samples from the product of Bernoulli distributions B(z|π). However, a Monte Carlo estimate of ∂ ∂π d R(θ, π) suffers from high variance. To see this, consider the following exact gradient of the empirical risk with respect to π d , which is

z:{0,1} D ,z d =1 L(z)p z i =d - z:{0,1} D ,z d =0 L(z)p z i =d ,

where p(

z i =d ) = D i =d π zi i (1 -π i ) 1-zi

, by absorbing the model f θ (•) and the data into L(•). Due to the discrete nature of z, we see that even the sign of the gradient estimate becomes inaccurate if we can only access a small number of Monte Carlo samples. While a score-function estimator such as REINFORCE (Williams, 1992) can be used, it is known that the reparametrization trick reduces the variance more in practice.

G COMPARISON TO HARD-CONCRETE DISTRIBUTION:

To evaluate the strength of the proposed continuous relaxation of the Bernoulli distribution, described in Subsection 3.1, we compare it with the Hard-Concrete distribution (Louizos et al., 2017) , another continuous surrogate for Bernoulli distributions, which was originally developed for neural network model compression. The details of the Hard-Concrete distribution is described below.

G.1 HARD-CONCRETE DISTRIBUTION:

The authors in (Louizos et al., 2017) introduce a modification of Binary Concrete, whose sampling procedure is as follows:

u ∼ U (0, 1), L = log(U ) -log(1 -U ) s = 1 1 + exp( -(log α+L) β ) s = s(ζ -τ ) + τ z = min(1, max(0, s))

where (τ, ζ) is an interval, with τ < 0 and ζ > 1. This induces a new distribution, whose support is [0, 1] instead of (0, 1). With 0 < β < 1, the probability density concentrates its mass near the end points, since values larger than 1-τ ζ-τ are rounded to one, whereas values smaller than -τ ζ-τ are rounded to zero. The CDF of s is

EQUATION

and so the CDF of s is

EQUATION

where φ = (β, log α, ζ, τ )

Now, the probability of being the gate z i being active is 1 -Q s(0|φ) and can be written as

EQUATION

G.2 THE EFFECT OF THE HEAVY TAIL DISTRIBUTION:

The main difference between our proposed distribution (STG) and the Hard-Concrete (Louizos et al., 2017) distribution is that the latter is based on the logistic distribution, which has a heavier tail than the Gaussian distribution we have employed. As shown in Fig 5 , the heavy-tailness results in instability during training. Furthermore, our method converges much faster and more reliably than the feature selection method using the Hard-Concrete distribution on a the two-moons, XOR and MADELON datasets (see Subsection J.1 and J.2). A similar phenomenon is also demonstrated in the MNIST experiment (see Subsection J.3). 

H NON-CONVEX DETERMINISTIC REGULARIZATION:

Our proposed objective for feature selection (summarized in Eq. 5) consists of three major differences compared to the LASSO formulation.

• Non-linearity, which is obtained by learning parameters of a multi-layer network with non linear activation's. • Feature gates, these are a specific form on non convex regularization.

• Stochasticity, which is achieved by injecting noise via the reparametrization trick.

In section 6.1, we have demonstrated that even in the linear regime the proposed formulation outperforms the LASSO in the task of identifying active features (see Fig. 1(a) ).

The stochastic property of the gates proposed in this paper is justified using a MI prospective (see Section 4). Nonetheless, in this section we evaluate whether the practical performance gain of our method stems from the non-convex regularization or from it's combination with the injected noise.

We define a deterministic gate using zd = g(µ d ) = max(0, min(1, µ d + 0.5)), where µ d is a parameter learned for each feature d = 1, ..., D. In the linear regression setting we can define the non-convex deterministic (NCD) as

EQUATION

)

where Φ is the standard Gaussian CDF. Combined with zd this non-convex regularized objective is deterministic and differentiable, and its solution can be approximated via gradient decent. We evaluate the applicability of such alternative objective for the task of feature selection in linear regression. The experimental setting follows the description in Section 6.1 using D = 64.

In Fig. 6 (b), we compare the values of the deteministic gates (DNC) and stochastic gates (STG) using N = 60. As depicted from this figure, the STG correctly identifies the active features (first 10) and sparsifies the non-active features. The deterministic version DNC does induce partial sparsity, but fails to remove some of the irrelevant features. Furthemore, it does not push the gates values of active features up to 1. Clearly, in the deterministic setting, shrinking the value of µ d while compensating with a large coefficient in θ is preferable for the objective in Eq. 12.

To evaluate the probability of support recovery using DNC, we apply additional thresholding to the gates. Specifically, we use a threshold of 0.5 to define if a feature is active or not. A comparison with STG, HC and LASSO in this setting appears in Fig. 6 (a). The DNC is superior to LASSO, but inferior to the HC and STG. Which means that there is a true advantage of stochacity in our formulation.

One aspect that is crucially difference between the STG and DNC is what we call "second chance".

For the deterministic formulation, if a feature is zeroed out in an early training phase, the gradient of the feature vanishes for further training and it will never be pulled back. For the STG, even if the feature has been pulled to zero, the gate may still be pushed back up to one at a later training phase. This is due to the injected noise which allows to reevaluate the gradient of each gate. In Fig. 6 (c), we demonstrate this "second chance" effect using N = 60 samples and presenting the gate's values (throughout training) for an active feature.

To conclude, this form of deterministic gates (Eq. 12) is not suitable for an embedded feature selection method. Using thresholding seems somewhat effective; however such thresholding was not required for the STG in any of our experiments. While the shrinkage phenomenon might be avoided using further regularization on θ this is out of the scope of the current study. 

I TUNING THE REGULARIZATION PARAMETER λ:

There have been several studies (Lederer & Müller, 2015; Coen et al., 2006) on the problem of tuning regularization parameter in regression and classification. Perhaps the most related regularized objective is the LASSO formulation. Studies such as (Chand, 2012; Fan & Tang, 2013) develop methods for tuning the regularization parameter in LASSO. In Eq. 5 as in the LASSO formulation, increasing the value of λ effectively forces a sparser solution with less features. For simplicity, we focus on one popular way for tuning λ in the LASSO. The idea is to minimize the generalization error using a cross validation procedure over a range of λ values.

In this section, we demonstrate that this procedure is also applicable for the proposed approach. We apply the STG to the RAI dataset and evaluate the MSE using a 5-folds cross validation procedure. We compute the MSE on the validation data using 100 values of λ logarithmically scaled in [10 -3 , 10 1 ].

As demonstrated in Fig. 7 , using a cross validation procedure for tuning λ provides a clear optimal value for this example.

Figure 7 : Demonstrating a cross validation procedure on the RAI dataset (see section I for details). We perform 5-folds cross validation and evaluate the MSE on the test test. The optimal λ = 0.135 which seems stable in this example based on the low standard deviation.

J ADDITIONAL EXPERIMENTS J.1 MADELON DATASET:

The MADELON dataset, first suggested for the NIPS 2003 feature selection problem, is a multivariate highly nonlinear binary classification problem. The MADELON dataset is generated using 32 groups of data points placed on a 5 dimensional hyper-cube and randomly labeling them by one of the two class labels. The first 5 informative features are then used to create 15 additional coordinates that are formed based on a random linear transformation of the first 5. A Gaussian noise N (0, 1) is added to each feature. Next, additional 480 nuisance coordinates are added in the same manner. These features have no effect on the class label. Finally, 1% of the labels are flipped. To be consistent with the XOR and two moons experiments, we use 1, 500 points from this dataset, and evaluate our proposed method in terms of its predictive power and ability to detect the informative features. We vary the regularizaton parameter λ in the range [0.01, 10] and evaluate the classification accuracy using 5 folds cross validation. In this example, we restrict our comparison to Random Forest and LASSO. We focus on Random Forest as it was the strongest competitor to our method in all of our experiments, while LASSO is evaluated because it is a widely used embedded feature selection method.

As evident from Fig. 8 (a), our method achieves the highest accuracy while using less features. Moreover, as depicted from this figure, peak performence occurs when selecting 5 features, thus, our method provides a clear indication to the true number of informative features. Both LASSO and RF on the other hand, do not provide a clear indication of the true number of relevant features. In Fig. 8 (b), we evaluate the effect of λ on the number and quality of selected features. As shown in this plot, there is a wide range of λ's such that our method only selects relevant features. Finally, as evident from the plato on the right hand side of the red plot, there is a range of λ's such that exactly 5 features are selected. An empirical evaluation of the effect the regularization parameter λ. The IFWR and the number of selected features are presented on both sides of the y-axis of this plot. For both plots, the mean is presented as a solid/dashed line, while the standard deviation is marked as a shaded color around the mean. (c) Box plots for the median rank of the 5 original informative features. Black line and dashed red represent the median and mean of each method. Optimal median rank in this experiment is 3.

J.2 TWO MOONS CLASSIFICATION WITH NUISANCE FEATURES:

In this experiment, we construct a dataset based on "two moons" shape classes, concatenated with noisy features. The first two coordinates x 1 , x 2 are generated by adding a Gaussian noise with zero mean and the variance of σ 2 r = 0.1 onto two nested half circles, as presented in Fig. 9 (a). Nuisance features x i , i = 3, ..., D, are drawn from a Gaussian distribution with zero mean and variance of σ 2 n = 1. We reserve the 70% as a test set, and use 10% of the remaining training set as a validation set. We follow the same hyperparameter tuning procedure as in the XOR experiment. The classification accuracy is in Fig. 9 (b). Based on the classification accuracies, it is evident that for a small number of nuisance dimensions all methods correctly identify the most relevant features. The proposed method (STG) and Random Forest (RF) are the only methods that achieve near perfect classification accuracy for a wide range of nuisance dimensions. The other NN based method (DFS) seem to converge to sub-optimal solutions. We note that the median rank for all the methods is 1.5. 

J.3 SPARSE HANDWRITTEN DIGITS CLASSIFICATION:

In the following toy example, we attempt to distinguish between images of handwritten digits of 3's and 8's using samples from MNIST (LeCun et al., 1998) . The orientation and location of the digits is more or less the same throughout this dataset, therefore for these two classes (3's and 8's), we expect that some of the left side features (pixels) would be sufficient for the separation. The experiment is performed as followed. We reserve 90% of the data as the test set, and train on the remaining 10%.

We then apply STG and evaluate the classification accuracy and the number of selected features. We use the architecture [200, 50, 10] with tanh activations. The experiment was repeated 10 times, the extracted features and accuracies were consistent over 20 trials. We noticed a relatively small number of selected features, which are positioned southwest and close to the center of the images, achieve very high classification accuracy. An example of 9 randomly selected samples overlaid with the weights of the selected features is presented in Fig. 10(a) . Furthermore, we also evaluate the effect of λ on the the number of selected features and accuracy of the method. We apply our method (STG) and its variant using the Hard-Concrete distribution to a randomly sampled training set of size N = 1500 and vary λ in the range of [0.001, 0.01]. In Fig. 10 (b) we present the accuracy and sparsity level vs. the λ parameter. This experiment demonstrates the improved performance of the proposed distribution compared to the Hard-Concrete (HC (Louizos et al., 2017) ), which was designed for neural net model compression. Not only that the overall accuracy is superior, but also it seems that the transition as a function of λ is smoother, which suggests that the method is less sensitive to the choice of λ.

J.4 GISETTE DATASET:

The Gisette is a handwritten dataset also appears in the NIPS 2003 feature selection challenge. The data consists of 13, 500 handwritten digits of '4' and '9'. The original digits have 28 × 28 pixels, which were modified for the feature selection challenge. Specifically, a random subset of the features are embedded in a 2500 dimensional space and additional 2500 irrelevant probes are added. The train, validation and test dimensions are 6000/1000/7500 respectively. We apply the proposed STG based feature selection method and compare its performance to Random Forests (RF) and Extremely Randomized Trees (Tree). As evident from Fig. 11 , we obtain a high accuracy even for a dramatic reduction in the feature size. LASSO is not presented in this experiment as its performance is dramatically inferior to the alternatives (accuracy < 0.6). where h 0 (t) is a baseline hazard function, which is often the exponential or Weibull distribution, and θ is the parameter of interests.

One of the difficulties in estimating θ in survival analysis is that a large portion of the available data is censored. However, in order to obtain estimates, Cox observed that it is sufficient to maximize the partial-likelihood, which is defined as follows: L(θ) = Tiuncensored e θ T x i T j ≥T i e θ T x j .

In (Katzman et al., 2018) , the authors propose DeepSurv, which uses a deep neural network model to replace the linear relations between the covariate x and θ, demonstrating improvements of survival time prediction over existing models such as CPH and the random survival forest (Ishwaran & Kogalur, 2007) , (Ishwaran et al., 2008) .

The Molecular Taxonomy of Breast Cancer International Consortium (METABRIC) dataset consists of gene expression data and clinical features for 1, 980 patients, and 57.72% have an observed death due to breast cancer with a median survival time of 116 months.

The METABRIC dataset involves 24,368 features (genes). Most genes are irrelevant for outcome prediction. To demonstrate the advantage of STG in the context of survival analysis, we selected 16 well-known genes relevant for survival (out of the 24,368 genes) that correspond to the Oncotype DX test, a gene panel used for treatment decision making. We also include five additional clinical features (hormone treatment indicator, radiotherapy indicator, chemotherapy indicator, ER-positive indicator, and age at diagnosis). We then added 200 additional irrelevant gene variables that we selected randomly from the remaining list of genes.

After we omit the null values, the number of samples is 1969. We reserve the 20% for test, and use the 20% of the remaining training set as validation. (That is, Train: 1260, Valid: 315, Test: 394 samples.)

Experimental Detail For DeepSurv, we manually select the architecture using the validation set so that we obtain a similar performance reported in (Katzman et al., 2018) . The learning rate decay is set to 1. The learning rate and the regularization parameter are optimized via Optuna using the validation set, where the search range is set LR : [1e -3, 1] and λ : [1e -3, 1]. The hyperparameters used in the experiment are the following: architecture :[60, 20, 3], activation: selu (as suggested by (Katzman et al., 2018) ), learning rate : 0.152, learning rate decay: 1.0, σ : 0.5, λ : 0.023, training epoch: 2000. Note that to see the effect of feature selection, we used the hyperparameters optimized for DeepSurv to test our method (STG-DeepSurv).

J.7 ADDITIONAL EXPERIMENTAL DETAILS:

Here we provide a full description of the procedures we have performed in the experimental parts of the paper.

For synthetic datasets are first split into train, validation and test. Validation is always 10% of the train, while the exact ratios between train and test is detailed for each experiment separately. All the neural network weights are initialized by drawing from N (0, 0.1) and bias terms are set to 0. All the batch sizes are equal to the number of training samples. In table 3, we detail the search range layer. For other methods we just used the feature relevance returned by the trained model. Finally, the LASSO's IFWR in the XOR experiment was omitted from the manuscript as it suffered from high variance.

Regarding the comparison performed in the two-moons and XOR problem, we believe that adding IFWR along with classification accuracy versus number of feature selected provides a complementary perspective in demonstrating the efficacy of feature selection techniques. We emphasize that our goal is not to just rank features but select features by assigning the weight of 0 to irrelevant features while simultaneously obtaining good predictive accuracy.

:

λ has a one-to-one correspondence to k in the convex setting via Lagrangian duality.

http://www.dcc.fc.up.pt/∼ltorgo/Regression/DataSets.html

Figures and Tables:

Figure 1: Probability of successfully recovering the support of a sparse signal β * ∈ R D vs. the number of observations N . Comparison between the proposed method, HC and LASSO under a linear additive noise model y = Xβ * + w using (a) D = 64 and (b) D = 128. Central lines are the means while shaded are represent the standard deviation.

Figure 2: (a) Classification accuracy (mean and standard deviation) vs. the number of irrelevant noisy dimension (D) for the XOR problem. (b) The Informative Features Weight Ratio (IFWR), central lines are the means while shaded are represent the standard deviation. IFWR is the sum of weights W d over the informative features divided by the sum over all weights. (c) Box plots for the median rank of the two informative features. Black line and dashed red represent the median and mean of each method. Optimal median rank in this experiment is 1.5.

Figure 3: Classification of T-cells sub-populations. Accuracy vs. number of features selected by each method. Comparison of the proposed method (STG and HC) to Random Forests (RF) and LASSO. 7.2 COX PROPORTIONAL HAZARD MODELS FOR SURVIVAL ANALYSIS

θ := θ -γ∇ θ L and µ := µ -γ∇ µ ( L + R) 16: end for 17: Algorithm 1: STG: Feature selection using stochastic gates We empirically observe that setting the number of Monte Carlo samples K = 1 and the standard deviation of the Gaussian distribution σ d = 0.5 for d = 1, ..., D suffices for feature selection in our experiments. See Section E in the Appendix for more details about the specific choice of σ d . After training, the set of indices for selected features is: {d : min(1, max(0, µ d + 0.5)) > 0}.

2 dThe effect of σ can be understood by looking at the value of ∂ ∂µ d E Z ||Z|| 0 . In the first iteration during the training, µ d is 0. Therefore, during the initial phase of training, it is close to λ √

Figure 4: The plot of ∂ ∂µ E Z ||Z|| 0 | µ=0 = 1 √ 2πσ 2 e -1 8σ 2 for σ = [0.001, 2].

Figure 5: Comparison between STG and Hard-Concrete on the Two-Moon dataset (a), XOR (b) and MADELON (c). The shaded area represents the standard deviation, calculated by running the same experiment 10 times with different random initializations.

Figure 6: Support recovery in the linear regression (see Section 6.1), the dimension D = 64. (a) Phase transition comparison between the STG, HC, LASSO and the deterministic non convex (DNC). (b) An example of the output gates, STG (z d ) and DNC (z d ) for N = 60. Each marker represents the gate value of the corresponding feature, of which the first 10 features are active. The deterministic gates are not pushed towards 1. (c) Demonstration of the "second chance" inherit to the STG. In the deterministic setting if a feature is eliminated its gradient vanishes for the rest of the training. The STG reevaluates this features and "recovers" after 4000 epochs.

Figure8: (a) Classification accuracy on the MADELON data sets. We evaluate performance using 5-fold cross validation for different number of selected features. In this dataset, only the first 20 coordinates are informative. In that regime the proposed method outperforms RF and LASSO. (b) An empirical evaluation of the effect the regularization parameter λ. The IFWR and the number of selected features are presented on both sides of the y-axis of this plot. For both plots, the mean is presented as a solid/dashed line, while the standard deviation is marked as a shaded color around the mean. (c) Box plots for the median rank of the 5 original informative features. Black line and dashed red represent the median and mean of each method. Optimal median rank in this experiment is 3.

Figure 9: (a) Realizations from the "Two moons" shaped binary classification class. X 1 and X 2 are the relevant features, X i , i = 3, ..., D are noisy features drawn from a Gaussian with zero mean and variance of 1. (b) Classification accuracy (mean and standard deviation based on 20 runs) vs. the number of irrelevant noisy dimension.

Figure 12: Classification of the version of the RCV1 textual dataset. The total number of feature is 47, 236. Here we compare the performance of the proposed approach to RF and LASSO. The lines represent a least squares polynomial fit plot of the accuracy vs. number of selected features.

. DNN represents a deep neural network without feature selection.

Table content: <table><tr><td>EXP</td><td>LASSO</td><td>RF</td><td>SG-L1-NN</td><td>SRFF</td><td>DNN</td><td>HC</td><td>STG</td></tr><tr><td>SE1</td><td>0.29 (0.01)</td><td>0.30 (0.01)</td><td>0.29 (0.01)</td><td colspan="3">0.27 (0.01) 0.29 (0.01) 0.29 (0.01)</td><td>0.29 (0.01)</td></tr><tr><td>SE2</td><td>2.22 (0.10)</td><td>2.34 (0.17)</td><td>2.35 (0.18)</td><td colspan="3">1.60 (0.10) 2.05 (0.11) 1.19 (0.31)</td><td>0.87 (0.15)</td></tr><tr><td>SE3</td><td>0.68 (0.002)</td><td>0.50(0.01)</td><td>0.68 (0.01)</td><td colspan="3">0.48 (0.03) 0.73 (0.02) 0.33 (0.05)</td><td>0.14 (0.10)</td></tr><tr><td>RCP</td><td>9.69 (0.71)</td><td>3.52 (0.12)</td><td>9.64 (0.65)</td><td colspan="3">2.52 (0.18) 2.89 (0.25) 2.74 (0.75)</td><td>2.44 (0.08)</td></tr><tr><td>REL</td><td>0.47 (0.01)</td><td>0.58 (0.01)</td><td>0.44 (0.01)</td><td colspan="4">0.31 (0.03) 0.61 (0.01) 0.33 (0.04) 0.27 (0.002)</td></tr><tr><td>RAI</td><td>0.43 (0.02)</td><td colspan="2">0.48 (0.003) 0.47(0.002)</td><td>0.41(0.02)</td><td colspan="2">0.44 (0.01) 0.41 (0.01)</td><td>0.39 (0.01)</td></tr></table>

Performance comparison of survival analysis on METABRIC. We run the same experiment 5 times with different train/test split and report the mean and the standard deviation on the test set. In(Katzman et al., 2018), it is reported that DeepSurv outperforms other existing survival analysis methods such as Random Survival Forest (RSF)(Ishwaran et al., 2008) and the original Cox Propotional Hazard Model.

Table content: <table><tr><td/><td>DEEPSURV</td><td>RSF</td><td>COX-LASSO</td><td>COX-HC</td><td>COX-STG</td></tr><tr><td>C-INDEX</td><td colspan="5">0.612 (0.009) 0.626 (0.006) 0.580 (0.003) 0.636 (0.007) 0.633 (0.005)</td></tr><tr><td># FEATURES</td><td>221 (ALL)</td><td>221 (ALL)</td><td>44 (0)</td><td>8 (0.89)</td><td>2 (0)</td></tr></table>

Peer review (iclr2020):
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.