Self-Supervision Enhanced Feature Selection with Correlated Gates

1 INTRODUCTION . High-dimensional datasets are increasingly prevalent in a variety of fields , including critical domains such as medicine and biology . Discovering features responsible for the target variable is an essential but challenging step in the analysis of such data . For example , while next generation sequencing can detect the expression of tens of thousands of genes per sample ( e.g. , RNA-seq ) , many genetic disorders stem from the variation in only a few groups of related genes ( Jackson et al. , 2018 ) . Identification of such disease-related factors is crucial for the design of therapeutic treatments . Furthermore , feature selection has additional benefits , including reduced experimental costs , greater interpretability , and improved model generalization ( Min et al. , 2014 ; Ribeiro et al. , 2016 ) . However , the high-dimensional and often noisy nature of such data prevents relevant features from being readily discovered . Moreover , in many domains there is a scarcity of label information due to cost , privacy reasons , or experimental study design . While deep learning approaches have resulted in improvements in feature selection ( e.g. , Yamada et al . ( 2020 ) ; Lemhadri et al . ( 2021 ) , see Related Work ) , such methods typically assume access to many labeled samples . With low sample size , feature selection methods , in particular those employing deep learning , have shown a propensity to overfit high-dimensional data ( Kuncheva et al. , 2020 ) . This issue is further exacerbated by the inherent structure of such data . Many ( high-dimensional ) datasets exhibit substantial inter-correlations or multicollinearity – i.e. , there exist features that are ( highly ) correlated among themselves – which impacts the performance of feature selection methods ( Chong & Jun , 2005 ; Katrutsa & Strijov , 2017 ; Belsley et al. , 2005 ) . In particular , this structure can cause parameter estimates to be unstable ( Dobson & Barnett , 2018 ) and , in the extreme , can prevent variable effects from being separated , confounding the problem of feature selection ( Meloun et al. , 2002 ) . Existing deep learning-based methods for feature selection do not ( explicitly ) consider the correlations between features which will often result in the selection of redundant features . ∗Equal contribution When labeled samples are scarce , large numbers of unlabeled samples are often available ( e.g. , Perez-Riverol et al . ( 2019 ) ) . However , the importance of unlabeled data has been overlooked in feature selection despite its potential to prevent the model selection process from overfitting by allowing informative feature representations to be learned . Therefore , the development of methods for feature selection that can exploit both labeled and unlabeled data is of great practical importance . Contributions . We propose a novel method for feature selection that addresses two key challenges in real-world scenarios : limited labeled data and correlated features . We use a self-supervised approach to train an encoder using unlabeled data via two pretext tasks : feature vector reconstruction and gate vector estimation . This pre-conditions the encoder to learn informative representations from partial feature sets , aligning the self-supervision with the model selection process of the downstream feature selection task . In addition , we introduce a novel gating procedure that accounts for the correlation structure of the input features . This ensures the pretext tasks remain challenging by preventing the model from memorizing trivial relations between features . Moreover , unlike previous deep learning-based feature selection methods , the correlated gate vectors encourage our method to select the most relevant features by making multiple correlated features compete against each other . We validate our approach through experiments on synthetic and multiple real-world datasets , including clinical and omics , where only a small number of labeled samples are available . Our model discovers relevant features that provide superior prediction performance compared to state-of-the-art benchmarks , and we corroborate these features with supporting medical and scientific literature . 2 RELATED WORK . Feature Selection Methods . Feature selection is a well-studied problem , with a number of proposed solutions including wrapper ( Kohavi & John , 1997 ) and filter ( Liu & Setiono , 1996 ; Kira & Rendell , 1992 ) methods . Recent advances in deep learning provide an elegant way of training embedded feature selection methods by jointly learning to perform feature selection while training a prediction network ( Huang et al. , 2020 ) . These methods learn to perform the non-differentiable process of selecting feature subsets by approximating it either with Lasso or elastic net penalization ( Li et al. , 2016 ) , using an MCMC sampling approach ( Liang et al. , 2018 ) , or more recently with continuous relaxation using independent Gaussian random variables ( Yamada et al. , 2020 ) . However , supervised feature selection methods can fail to identify relevant features when limited labeled samples are available due to overfitting ( Kuncheva et al. , 2020 ) , impacting their suitability in many real-world scenarios . Moreover , these methods do not consider the underlying correlation structure of the input features which can be problematic when selecting relevant features among ( highly ) correlated ones . Relatively few feature selection methods use unlabeled samples . Abid et al . ( 2019 ) used autoencoders to identify a pre-specified number of features that are sufficient for reconstructing the data . Lindenbaum et al . ( 2020 ) improved the well-known Laplacian score ( He et al. , 2005 ) by selecting feature subsets that better capture the “ local structure ” of the data . However , both approaches are fully unsupervised and , without the guidance of label information , can fail to identify features relevant to the target outcome . While a number of semi-supervised feature selection methods have been proposed ( Sheikhpour et al. , 2017 ) , they have typically been extensions of traditional methods , such as a Laplacian score with modified affinity scores using labels ( Zhao et al. , 2008 ) or manifold regularization based on linear SVMs ( Dai et al. , 2013 ) . To the best of our knowledge , this is the first deep learning framework that fully utilizes both labeled and unlabeled samples for feature selection . Self-Supervised Learning . Self-supervised learning methods create ( weak ) supervised signals from unlabeled data , employing contrastive learning or pretext task ( s ) to provide surrogate labels . While self-supervised learning has found success in computer vision ( Chen et al. , 2020 ) and natural language processing ( Devlin et al. , 2018 ) , the tabular domain , which is most relevant in feature selection , has been largely neglected . Methods for tabular data seek to reconstruct data either based on a corrupted sample alone ( Vincent et al. , 2008 ; Arik & Pfister , 2019 ; Yin et al. , 2020 ) or with knowledge of which entries have been corrupted ( Pathak et al. , 2016 ) . Recently , Yoon et al . ( 2020 ) jointly learn to recover the original sample and predict the mask vector used to corrupt the data . While our pretext tasks are also reconstruction-based , we propose a novel method for generating the gate vector used to produce the input feature vector . This is particularly important when there exists substantial correlation between features ( see Section 5 ) . All existing methods have used indepen- dent , uniform sampling ; however , this allows highly correlated features to be readily reconstructed since it is likely that not all such features will be corrupted . We prevent this by incorporating the correlation structure into the gating procedure . In this work , we employ pretext tasks that mirror the goal of the feature selection process : the input to the encoder is a subset of the features selected at random . Employing partial feature sets for self-supervised learning encourages the encoder to learn better representations of these partial feature sets . This in turn benefits learning both the model selection function and the feature selection step since they are also trained using partial features sets . 3 PROBLEM FORMULATION . Let X = ( X1 , · · · , Xp ) ∈ X p and Y ∈ Y be random variables for the high-dimensional input features ( e.g. , gene expressions ) and the target outcome ( e.g. , disease traits ) whose realizations are denoted as x = ( x1 , · · · , xp ) and y , respectively . Throughout the paper , we will often use lowercase letters to denote realizations of random variables . Embedded feature selection aims to select a subset , S ⊂ [ p ] , of features that are relevant for predicting the target as part of the model selection process . Denote ∗ be any point not in X and define XS = ( X∪ { ∗ } ) p.1 Then , given X ∈ X p , the selected subset of features can be denoted as XS ∈ XS where xS , k = xk if k ∈ S and xk = ∗ if k /∈ S. Let f : XS → Y be a function in F that takes as input the subset XS and outputs Y . Then , selecting relevant features can be achieved by solving the following optimization problem : minimize f∈F , S⊂ [ p ] Ex , y∼pXY [ ` Y ( y , f ( xS ) ) ] subject to |S| ≤ δ ( 1 ) where δ constrains the number of selected features and ` Y ( y , y′ ) = − ∑C c=1 yc log y ′ c for C-way classification tasks ( i.e. , Y = { 1 , · · · , C } ) and ` Y ( y , y′ ) = ‖y − y′‖22 for regression tasks ( i.e. , Y = R ) .2 Unfortunately , the combinatorial problem in ( 1 ) becomes intractable for high-dimensional data as the search space increases exponentially with p. Hence , we instead focus on a relaxation by converting the combinatorial search in ( 1 ) into a search over the space of possibly correlated binary random variables . It is worth highlighting that unlike existing work ( e.g. , Yamada et al . ( 2020 ) ; Yoon et al . ( 2019 ) ) , we do not assume independence among these random variables . Let M = ( M1 , · · · , Mp ) ∈ { 0 , 1 } p be binary random variables governed by distribution pM , whose realization m is referred to as the gate vector , for indicating selection of the corresponding features . Then , the selected features given gate vector m can be written as x̃ , m x + ( 1−m ) x̄ ( 2 ) where x̄ = Ex∼pX [ x ] and indicates element-wise multiplication . Here , we replace not-selected features with their means to resolve any issue when a feature having zero value ( e.g. , turned-off genes ) has specific meaning . Finally , we can approximately achieve ( 1 ) by jointly learning the model selection f and the gate vector distribution pM based on the following optimization problem : minimize f , pM Ex , y∼pXY Em∼pM [ ` Y ( y , f ( x̃ ) ) + β‖m‖0 ] ( 3 ) where f is implemented as a neural network and β is a balancing coefficient that controls the number of features to be selected . Challenges . In practice ( especially in the healthcare domain ) , there are two main challenges that confound selecting relevant features for predicting the target : ( i ) inter-correlation or multicollinearity , which is the existence of features that are ( highly ) correlated among themselves , and ( ii ) the absence of sufficient labeled samples . These challenges make embedded feature selection vulnerable to overfitting . More specifically , model selection ( i.e. , learning f ) and feature selection ( i.e. , learning pM ) are conducted jointly . Either one being overfit to correlated or noisy irrelevant features will end up discovering spurious relations and poor prediction . For instance , if the network is biased toward irrelevant features , the selection probability ( i.e. , importance ) of those features will be increased as if those features were “ discriminative ” or “ predictive ” , and vice versa . 1To enable neural networks to be trained with varying feature subsets , we set ∗ = x̄ . 2For Y = { 1 , · · · , C } , we will occasionally abuse notation and write yc to denote the c-th element of the one-hot encoding of y when clear in the context .

Summary: This work presents some additional tweaks and tricks in using neural networks for feature selection. The main contributions include an inclusion of a parameterized masking function that selects the feature and a self-supervised component for pretraining. The overall proposal makes sense and there are strong experimental results that justifies the merits.