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 aims to use self-supervision to identify the most useful features for downstream tasks particularly in the context of correlated features. This is an important aspect that is lacking from many feature selection approaches commonly used within the highly correlated datasets of healthcare. This approach - using both feature vector reconstruction and gate vector estimation appears to provide better results for both true structure estimation in a  dataset with known truth as well as downstream predictive performance in real world data sets.