Source-linked AI summary

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

Talip Ucar, Ehsan Hajiramezanali, Lindsay Edwards

arXiv:2110.04361v2cs.LGstat.ML

TL;DR

Tabular data lacks the spatial, temporal, or semantic structure that supports effective self-supervised augmentation in images, audio, and text. SubTab addresses this by dividing features into multiple views, reconstructing data from subsets, and aggregating their representations at test time. It reaches state-of-the-art performance on tabular MNIST and reports benefits on other tabular datasets, while its broader extensions remain future work.

  • Problem

    Self-supervised representation learning has progressed less for tabular data because commonly useful augmentation structures and inductive biases may be absent.

  • Method

    SubTab divides features into subsets, reconstructs the full data from each subset, and aggregates subset representations at test time.

  • Results

    SubTab achieves state-of-the-art performance on tabular MNIST, on par with CNN-based state-of-the-art image models, and surpasses baselines on three other real-world datasets.

  • Takeaways & Limitations

    Subset reconstruction and aggregation provide the main reported performance gains and support representation learning from partial feature sets.

  • Takeaways & Limitations

    The framework is computationally quadratic during training when contrastive or distance losses require projection combinations, and extensions beyond tabular data remain future work.

Abstract

from arXiv · show

Self-supervised learning has been shown to be very effective in learning useful representations, and yet much of the success is achieved in data types such as images, audio, and text. The success is mainly enabled by taking advantage of spatial, temporal, or semantic structure in the data through augmentation. However, such structure may not exist in tabular datasets commonly used in fields such as healthcare, making it difficult to design an effective augmentation method, and hindering a similar progress in tabular data setting. In this paper, we introduce a new framework, Subsetting features of Tabular data (SubTab), that turns the task of learning from tabular data into a multi-view representation learning problem by dividing the input features to multiple subsets. We argue that reconstructing the data from the subset of its features rather than its corrupted version in an autoencoder setting can better capture its underlying latent representation. In this framework, the joint representation can be expressed as the aggregate of latent variables of the subsets at test time, which we refer to as collaborative inference. Our experiments show that the SubTab achieves the state of the art (SOTA) performance of 98.31% on MNIST in tabular setting, on par with CNN-based SOTA models, and surpasses existing baselines on three other real-world datasets by a significant margin.

1 Introduction

SubTab addresses the difficulty of learning representations from tabular data, where spatial, semantic, and temporal structures supporting augmentation may be absent. It creates multiple feature-based views and reconstructs from subsets to improve representation learning.

  • Motivation: Tabular data lacks the structures and inductive biases that enable effective augmentation in image, audio, and language representation learning.This limitation is relevant to tabular datasets used in healthcare, advertising, finance, and law.
  • Motivation: Noise-based autoencoders treat all features equally, so perturbing uninformative features may fail to produce the intended corruption.A related classifier-based method also relies on noisy inputs and predicts corrupted-feature locations using an imbalanced high-dimensional mask.
  • SubTab framework: SubTab divides features into subsets, treating each subset as a distinct view for self-supervised representation learning.The approach is analogous to cropping in images or feature bagging in ensemble learning.
  • SubTab framework: Reconstructing the full data from feature subsets is presented as a way to force better representations than adding noise.The framework is evaluated on classification and clustering across MNIST, TCGA, Obesity, Income, and Blog datasets.
  • SubTab framework: SubTab aggregates subset representations through collaborative inference and can identify informative feature regions, handle missing features, and reduce input dimensionality.These properties are described as useful for high-dimensional data and for reducing overfitting risk.

2 Method

SubTab forms multiple feature subsets, encodes them with shared components, reconstructs the full feature space, and optionally aligns projections with contrastive and distance losses. At test time, it aggregates available subset representations into a joint representation.

  • Feature subsets: SubTab divides tabular features into multiple, potentially overlapping subsets and feeds each subset through a shared encoder.Subset locations remain fixed because standard neural networks are not permutation invariant.
  • Reconstruction: The shared decoder reconstructs the full feature space from each subset rather than reconstructing only the subset itself.The authors report that full-space reconstruction is more effective and prevents the autoencoder from learning the identity mapping.
  • Auxiliary losses: SubTab optionally adds contrastive loss over projection pairs and distance loss that brings corresponding subset projections closer together.Contrastive loss uses same-row samples as positives and other rows as negatives when the dataset provides enough classes.
  • Noise augmentation: Training can add Gaussian, swap, or zero-out noise using neighboring-column, random-column, or random-feature selection strategies.These augmentation choices are complementary to the subset-based framework.
  • Training objective: The objective combines total, reconstruction, contrastive, and distance losses, with contrastive and distance terms optional in some experiments.The reconstruction loss averages per-subset reconstruction losses across the batch and subsets.
  • Test-time inference: At test time, SubTab aggregates representations from available subsets using mean, sum, min, max, or another aggregation function.Mean aggregation is used in the experiments, while weighted aggregation is proposed as a possible extension.

3 Experiments

Experiments evaluate SubTab across diverse tabular datasets, examine subset configuration and collaborative inference, and report strong performance and robustness to missing features. Ablations further assess losses, aggregation, architecture depth, and feature ordering.

  • Experimental setup: SubTab was evaluated on MNIST, TCGA, Obesity, Income, and BlogFeedback using classification performance from logistic regression on learned representations.For SubTab, mean aggregation of subset embeddings forms the joint representation for training and test evaluation.
  • MNIST results: K = 4 with 75% overlap and K = 7 with 50% overlap performed best on MNIST, while swap noise further improved SubTab performance.The base model without noise already outperformed autoencoder and other self-supervised baselines with the same architecture.
  • Collaborative inference: 87.5% test accuracy was achieved using one subset, and accuracy increased as additional subset representations were aggregated, reaching its highest level with all subsets.This indicates useful performance can be obtained when only a small subset of features is available.
  • Missing subsets: Aggregating all seven subsets improved results even when the model was trained on only one subset, and performance generally increased as more subsets were used.Similar behavior occurred with an untrained model, while individual subset informativeness varied across image regions.
  • Missing features: SubTab accommodated missing features at test time and remained effective when trained with missing subsets, with predictions improving as more feature information was added.The experiments model settings such as healthcare sites having access to different feature sets.
  • Real-world datasets: On Obesity, SubTab with Gaussian noise achieved 71.13 ± 4.08% accuracy, above the 64.82% baseline and the reported 66 ± 3.2% DAE result.The passage also reports that SubTab outperformed other self-supervised models on both TCGA and Obesity.
  • Ablation study: Ablations found reconstruction loss effective on its own, with further gains from noise, additional losses, and a larger representation dimension; mean aggregation outperformed concatenation.Shuffling MNIST features did not substantially change performance, and shallow architectures significantly improved results on MNIST and TCGA.

4 Related works

Prior tabular self-supervised methods mainly reconstruct corrupted data or predict corruption masks, but high-dimensional sparse masks create learning challenges. Related work also includes shallow and deep SubTab architecture comparisons.

  • Tabular self-supervised learning: Recent tabular self-supervised methods commonly use denoising autoencoders or pretext tasks that recover corrupted data or predict masked-feature locations.These approaches include VIME-self, TabNet, and TaBERT.
  • SubTab architecture comparisons: Table 3 compares shallow and deep SubTab architectures.
  • Tabular self-supervised learning: Predicting high-dimensional, sparse, imbalanced binary masks can make representation learning challenging.

5 Conclusion

SubTab learns tabular representations by reconstructing all features from feature subsets and aggregating subset embeddings. The framework also supports missing-feature inference and has computational trade-offs that depend on the training losses used.

  • Conclusion: SubTab’s main performance gains come from reconstructing all features from feature subsets and aggregating their embeddings into a joint representation.
  • Conclusion: SubTab can support inference with missing features by ignoring corresponding subsets, and it may enable distributed training across feature subsets.
  • Conclusion: Using contrastive or distance losses makes training computationally quadratic in the number of subsets, although test-time aggregation remains linear.
  • Conclusion: Extending SubTab beyond tabular data to images, audio, and text remains future work.

6 Broader Impact

Tabular data is widely used in healthcare, finance, law, and other fields, motivating representation learning methods that can broaden its applications. The paper’s algorithm divides each minibatch into subsets, encodes and reconstructs them, and optionally applies pairwise contrastive and distance losses.

  • Broader impact: Tabular data is widely used in healthcare, finance, law, and many other fields, while unsupervised representation learning has focused more on images, text, and audio.The paper connects progress in tabular representation learning to potential transfer, distributed, and multi-view applications, while noting bias and privacy concerns.
  • Learning procedure: The learning algorithm divides each minibatch into K feature subsets and optionally adds noise to each subset.
  • Learning procedure: Each subset is passed through an encoder, projection head, and decoder to produce latent representations, projections, and reconstructions.
  • Optional losses: The algorithm can compute contrastive and distance losses over all pairs of subset projections.It averages these losses over the total number of pairs before updating network parameters.
  • Optional losses: The total loss combines reconstruction loss with averaged contrastive and distance losses when those optional objectives are applied.

B.1 Adult Income Dataset

The Adult Income dataset uses separate training and test sets, with an 80–20% training split for validation-based hyperparameter search. It contains mixed categorical and continuous attributes and an imbalanced target.

  • Train-validation-test split: The provided training data are split 80–20% into training and validation sets for hyperparameter search, then the final model uses the full training set.
  • Features: The dataset has 14 attributes: 8 categorical and 6 continuous features.Rows with missing values are dropped, categorical features are one-hot encoded, and features are standardized using training-set statistics.
  • Class imbalance: Only 25% of samples are positive, making the dataset imbalanced.
  • Train-validation-test split: The original dataset also includes one training set and 60 small test sets, which are combined into one test set in the described preprocessing.
  • Additional dataset description: The BlogFeedback dataset contains 280 features and a binary target indicating whether a blog post received comments in the next 24 hours.Its features are normalized with min-max scaling.
  • Additional dataset description: For BlogFeedback, approximately 36% of training samples and 30% of test samples are positive.

C.1 Model architectures and hyper-parameters

The appendix specifies the architectures, losses, evaluation procedure, and experimental setups used to assess SubTab representations and subset information content.

  • The reported architectures and hyper-parameters are summarized separately for the main and Shallow SubTab results.
  • SubTab uses reconstruction loss by default, with one MNIST ablation using only contrastive loss.
  • Representation quality is evaluated with a linear classifier on embeddings from training and test sets, using joint embeddings for SubTab.
  • The experiments divide MNIST into seven subsets and compare models trained on different subset combinations, including an untrained model.
  • A second experiment measures each subset’s information content by applying a linear classifier to its embeddings.
  • The appendix reports that over-complete first hidden layers perform best, while combining reconstruction and contrastive losses improves contrastive-loss stability for small batches.

D.2 Comments

The comments describe extensions and practical uses of SubTab, while identifying representation collapse as a limitation of dividing features into smaller subsets.

  • SubTab is presented as applicable to transfer learning, few-shot generalisation, domain adaptation, multi-task learning, and continual learning.
  • SubTab can be extended beyond tabular data to images, text, time series, and audio by using random subspaces.
  • Dividing features into smaller chunks can cause representation collapse, where subsets from very different samples obtain similar representations.
  • The framework includes contrastive-loss variants using either noisy copies of full data or subsets, and it can alternatively use separate encoders for subsets.
  • PCA and t-SNE plots compare representations as one through four overlapping subsets are progressively combined.
  • MNIST rows can be treated as subsets whose test accuracy measures the information content of individual image regions.
  • Figure A9 varies the number of subsets and overlap levels across Blog, Income, and TCGA datasets.

F.3 Sensitivity analysis for masking ratio (p) and initialization

The sensitivity analyses examine masking ratio, initialization, aggregation, and synthetic-data settings to characterize robustness across experimental choices.

  • Masking ratio: A masking ratio p in the range [0.1-0.3] is usually effective, and performance is generally robust across different p values on MNIST.
  • Sensitivity to initialization: Across 10 random initializations, the variation in test accuracy is very small, indicating that informative-subset discovery is not sensitive to initialization.
  • Aggregation functions: Mean aggregation generally performs better, while concatenating subset representations reduces MNIST classification performance; overall results are robust to aggregation choices.
  • Synthetic datasets: Synthetic datasets contain 10 classes and 10k samples, with 10% reserved for testing and deliberately difficult cluster separation.
  • Synthetic datasets: The synthetic-data setup uses two subsets with 25% overlap, masking ratio p = 0.2, Gaussian noise σ = 0.1, reconstruction loss, and mean aggregation.
  • Synthetic datasets: The described setup is reported to work well for most tabular datasets, and SubTab improves results on all three synthetic datasets.
  • Dataset selection: The OpenML-CC18 table summarizes the datasets used for the broader experiments, with the first eight selected for evaluation.

G.3 Evaluation

The evaluation compares logistic regression on raw features with logistic regression on SubTab embeddings, after preprocessing and dataset selection decisions.

  • SubTab embeddings improve test results over raw features on all three synthetic datasets, with gains reaching 100% on the most difficult dataset.
  • The evaluation excludes datasets with fewer than 10 features or fewer than 1,000 samples, except Electricity because of its relatively large sample size.
  • The datasets are cleaned by removing missing-data rows or user-ID features, min-max scaled, and split into 70-10-20% training, validation, and test sets.

H.2 Models and Evaluation

The evaluation compares six models across eight datasets, using a shared neural-network architecture and repeated-seed embedding evaluation. SubTab performs better on datasets suited to representation learning, while simple logistic regression is preferable on trivial datasets.

  • Six models were compared: logistic regression, autoencoder variants, VIME-self, and SubTab variants.
  • The evaluation covered eight datasets.
  • All neural networks used a four-layer [256, 256, 256, 256] encoder, with dropout models using p=0.04.SubTab used two non-overlapping subsets and reconstruction loss; neural-network models were trained and evaluated with 10 random seeds.
  • SubTab tended to outperform the other approaches, including VIME-self, on datasets suitable for pre-training or representation learning.This pattern was reported for datasets [1] [2] [5].
  • Logistic regression was preferable on trivial datasets where it already achieved very decent performance.This was reported for the Texture, DNA, and Climate datasets.
Loading 2110.04361v2…