SubTab: Subsetting Features of Tabular Data for Self-Supervised Representation Learning

1 Introduction . In recent years , the self-supervised learning has successfully been used to learn meaningful representations of the data in natural language processing [ 34 , 41 , 11 , 28 , 10 , 21 , 9 ] . A similar success has been achieved in image and audio domains [ 7 , 15 , 37 , 5 , 17 , 13 , 8 ] . This progress is mainly enabled by taking advantage of spatial , semantic , or temporal structure in the data through data augmentation [ 7 ] , pretext task generation [ 11 ] and using inductive biases through architectural choices ( e.g . CNN for images ) . However , these methods can be less effective in the lack of such structures and biases in the tabular data commonly used in many fields such as healthcare , advertisement , finance , and law . And some augmentation methods such as cropping , rotation , color transformation etc . are domain specific , and not suitable for tabular setting . The difficulty in designing similarly effective methods tailored for tabular data is one of the reasons why self-supervised learning is under-studied in this domain [ 46 ] . The most common approach in tabular data is to corrupt data through adding noise [ 43 ] . An autoencoder maps corrupted examples of data to a latent space , from which it maps back to uncorrupted data . Through this process , it learns a representation robust to the noise in the input . This approach may not be as effective since it treats all features equally as if features are equally informative . However , perturbing uninformative features may not result in the intended goal of the corruption . A recent work takes advantage of self-supervised learning in tabular data setting by introducing a pretext task [ 46 ] , in which a de-noising autoencoder with a classifier attached to representation layer is trained on 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . corrupted data . The classifier ’ s task is to predict the location of corrupted features . However , this framework still relies on noisy data in the input . Additionally , training a classifier on an imbalanced binary mask for a high-dimensional data may not be ideal to learn meaningful representations . In this work , we turn the problem of learning representation from a single-view of the data into the one learnt from its multiple views by dividing the features into subsets , akin to cropping in image domain or feature bagging in ensemble learning , to generate different views of the data . Each subset can be considered a different view . We show that reconstructing data from the subset of its features forces the encoder to learn better representation than the ones learned through the existing methods such as adding noise . We train our model in a self-supervised setting and evaluate it on downstream tasks such as classification , and clustering . We use five different datasets ; MNIST in tabular format , the cancer genome atlas ( TCGA ) [ 42 ] , human gut metagen-omic samples of obesity cohorts ( Obesity ) [ 36 , 26 ] , UCI adult income ( Income ) [ 24 ] , and UCI BlogFeedback ( Blog ) [ 4 ] . SubTab can : i ) construct a better representation by using the aggregate of the representation of the subsets , a process that we refer as collaborative inference ii ) discover the regions of informative features by measuring predictive power of each subset , which is useful especially in high-dimensional data iii ) do training and inference in the presence of missing features by ignoring corresponding subsets and iv ) use smaller models by reducing input dimension , making it less prone to overfitting . 2 Method . The augmentation methods such as adding noise , rotation , cropping etc . are commonly used in image domain . Among them , the cropping is shown to be the most effective technique [ 7 ] . Inspired from this insight , we propose subsetting features of tabular data . Figure 1 presents SubTab framework , in which we have an encoder ( E ) , a decoder ( D ) , and an optional projection ( G ) . For the purpose of this paper , we will refer h as latent , or representation , z as projection , x̃ , and X̃ as the reconstruction of subset , and whole data respectively . Small letters are associated with subsets while capital latters are associated the whole set of features . Moreover , throughout this work , when we say that a representation is `` good '' , we refer to its performance in a classification task using a linear model . In SubTab framework , we divide tabular data to multiple subsets . Neighbouring subsets can have overlapping regions , defined as a percentage of a dimension of the subset . Each of the subsets is fed to the same encoder ( i.e . parameter sharing ) to get their corresponding latent representation . A shared decoder is used to reconstruct either the subset fed to the encoder , or full tabular data ( i.e . reconstructing all features from the subset of features ) . We chose the latter in our experiments since it is more effective in learning good representations . We should also note that , in the latter case , the autoencoder can not learn the identity , eliminating the constraint on the dimension of the bottleneck ( i.e . representation ) . We compute one reconstruction loss term per subset . Moreover , we can optionally add contrastive loss to our objective by using all combination of pairs of projections from all subsets . For example , given three subsets as in Figure 1 , there are three combinations of two : n k = 3 2 = 3 ! 2 ! ( 1 ) ! = 3 . For four subsets , it would be 6 pairs of combination , and so on . We can add one more loss term , referred as distance loss , to reduce the distance between the pairs of projections of the subsets by using a loss function such as mean squared error ( MSE ) . All three loss terms apply a pulling force on positive samples while contrastive loss also applies a push force between positive and negative samples as shown in Figure 2a . Once the dataset is divided into subsets in data preparation step , a process that is similar to feature bagging in ensemble learning , their location is fixed . Thus , we don ’ t change the relative order of features in a subset during training since standard neural network architectures are not permutation invariant . This is to ensure that same features are fed to the same input units of neural network . However , our method can be extended to permutation invariant setting as a next step . 2.1 Strategies for adding noise . Our framework is complementary to other augmentation techniques used in tabular data setting . Thus , we experimented with adding noise to randomly selected entries in each subset by using three types of noise : i ) adding Gaussian noise , N ( 0 , 2 ) , ii ) overwriting the value of a selected entry with another value randomly sampled from the same column , referred as swap-noise , iii ) zeroing-out randomly selected entries , referred as zero-out noise . Moreover , we use three different strategies when selecting the features to add noise to , as shown in Figure 2b : i ) a random block of neighboring columns ( NC ) , ii ) random columns ( RC ) iii ) random features per each sample ( RF ) . To add noise , we create a binomial mask , m , and a noise matrix , ✏ , with same shape as the subset , in which the entries of the mask is assigned to 1 with probability p , and to 0 otherwise . The corrupted version , x1c , of subset x1 is generated as following : x1c = ( 1 m ) x1 +m ✏ ( 1 ) 2.2 Training . Our objective function is : Lt = Lr + Lc + Ld , ( 2 ) where Lt , Lr , Lc and Ld are total , reconstruction , contrastive , and distance losses , respectively . i ) Reconstruction loss : Given a subset , denoted by xk , we can reconstruct either the same subset , x̃k or the entire feature space X̃k . Then , we can compute the reconstruction loss for kth subset by computing mean squared error using either ( xk , x̃k ) , or ( X , X̃k ) pair as shown in Figure 1 . We chose the latter since it was more effective . Overall reconstruction loss : Lr = 1 K KX k=1 sk , where sk = 1 N NX i=1 ⇣ X ( i ) X̃k ( i ) ⌘2 ( 3 ) where K is the total number of subsets , N is the size of the batch , sk is the reconstruction loss for kth subset , and Lr is the average of reconstruction loss over all subsets . ii ) Contrastive loss : If the dataset is rich in the number of classes such that chances of sampling negative samples are high , we can use a projection network ( G ) to get projections , z0s , of representations , h0s . Samples at the same rows of two subsets , z1 and z2 , can be considered as positive pairs while remaining rows in the subsets can be considered as negative to those samples.This allows us to compute the contrastive loss for each pair of projections using a loss function such as the normalized temperature-scaled cross entropy loss ( NT-Xent ) [ 7 ] . For three subsets , { x1 , x2 , x3 } , we can compute such a loss for every pair { za , zb } of total three pairs from the set S = { { z1 , z2 } , { z1 , z3 } , { z2 , z3 } } . Overall contrastive loss is : Lc = 1 J X { za , zb } 2S p ( za , zb ) , where p ( za , zb ) = 1 2N NX i=1 h l ( za ( i ) , zb ( i ) ) + l ( zb ( i ) , za ( i ) ) i ( 4 ) l ( za ( i ) , zb ( i ) ) = log exp ( sim ( za ( i ) , zb ( i ) ) /⌧ ) PN k=1 1k 6=i exp ( sim ( za ( i ) , zb ( k ) ) /⌧ ) ( 5 ) where J is the total number of pairs in set S , p ( za , zb ) is total contrastive loss for a pair of projection { za , zb } , l ( za ( i ) , zb ( i ) ) is the loss function for a corresponding positive pairs of examples ( za ( i ) , zb ( i ) ) in subsets { za , zb } , and Lc is the average of contrastive loss over all pairs . iii ) Distance loss : We can also add mean-squared error ( MSE ) loss for pairs of projections of subsets since the corresponding samples in subsets should be close to each other . Hence , we can compute an overall MSE loss as : Ld = 1 J X { za , zb } 2S p ( za , zb ) , where p ( za , zb ) = 1 N NX i=1 ⇣ z ( i ) a z ( i ) b ⌘2 ( 6 ) The pseudocode of algorithm can be found in Algorithm 1 in the Appendix . We should note that both Lc and Ld in equation ( 2 ) are optional , and we used them only in some experiments .

Summary: The paper proposes SubTab, a self-supervised (pre-training) method for learning good representations of tabular data. The method generates multiple views for each (unlabeled) example by selecting different (possibly overlapping) subsets of features. The model is trained with a three term loss, each averaged over all pairwise comparisons between the views -- a reconstruction term that teaches the model to reconstruct the full example from the cropped example, a contrastive term that encourages different crops of the same example to be closer together than crops of different examples (within the minibatch, infoNCE loss), and a distance term that encourages crops of the same example to be close together in L2 distance. Optionally, the input example can be corrupted by noise. After pre-training a linear classifier on the latent representation is trained using supervision. In this step, the latent embedding used is actually an average of those arising from feeding in several different random crops of the example. The authors test the method against VIME (Neurips 2020), context encoders, and denoising autoencoders on 5 datasets, which were used in the VIME paper.