Distributionally robust self-supervised learning for tabular data

TL;DR

This work tackles distributional robustness for tabular data by introducing a two-stage self-supervised framework that learns robust representations during pre-training. The first stage uses Masked Language Modeling to obtain latent encodings, while the second stage applies JTT or DFR to develop feature-specific decoder heads, forming an ensemble that selects the most informative representation for each test instance. Empirical results on Bank and Census datasets demonstrate that Deep Feature Reweighting yields substantial improvements over ERM and Just Train Twice, with notable AUROC gains and better subgroup generalization. The approach offers a practical pathway to fairness-aware robustness in tabular data without requiring explicit group annotations during pre-training, and it provides open avenues for future enhancements in high-cardinality settings and causal bias mitigation.

Abstract

Machine learning (ML) models trained using Empirical Risk Minimization (ERM) often exhibit systematic errors on specific subpopulations of tabular data, known as error slices. Learning robust representation in presence of error slices is challenging, especially in self-supervised settings during the feature reconstruction phase, due to high cardinality features and the complexity of constructing error sets. Traditional robust representation learning methods are largely focused on improving worst group performance in supervised setting in computer vision, leaving a gap in approaches tailored for tabular data. We address this gap by developing a framework to learn robust representation in tabular data during self-supervised pre-training. Our approach utilizes an encoder-decoder model trained with Masked Language Modeling (MLM) loss to learn robust latent representations. This paper applies the Just Train Twice (JTT) and Deep Feature Reweighting (DFR) methods during the pre-training phase for tabular data. These methods fine-tune the ERM pre-trained model by up-weighting error-prone samples or creating balanced datasets for specific categorical features. This results in specialized models for each feature, which are then used in an ensemble approach to enhance downstream classification performance. This methodology improves robustness across slices, thus enhancing overall generalization performance. Extensive experiments across various datasets demonstrate the efficacy of our approach. The code is available: https://github.com/amazon-science/distributionally-robust-self-supervised-learning-for-tabular-data.

Figures (6)

• Figure 1: Schematic of our method. (a) Dataset construction of robust pretraining using JTT and DFR. (b) Robust pre-training strategy using MLM loss. We train the encoder $h$ and the reconstruction head of the $j^{th}$ feature $f_j$ for each sample. The embedding $z$ will be used for downstream tasks. We do this for all categorical features, obtaining a pre-trained model per category.
• Figure 2: Comparing overall performance of the downstream classifiers using ERM, JTT and DFR
• Figure 3: Ablation study on the Bank dataset comparing the performance of Just Train Twice (JTT) across different feature categories. The plot illustrates the impact of JTT on subgroup performance, highlighting how the model's accuracy changes when key features are ablated. Subgroups that were underrepresented or more challenging to classify show notable improvements in accuracy, underscoring the effectiveness of JTT in mitigating bias and enhancing model robustness.
• Figure 4: Ablation study on the Census dataset comparing the performance of Just Train Twice (JTT) across different feature categories. The plot illustrates the impact of JTT on subgroup performance, highlighting how the model's accuracy changes when key features are ablated. Subgroups that were underrepresented or more challenging to classify show notable improvements in accuracy, underscoring the effectiveness of JTT in mitigating bias and enhancing model robustness.
• Figure 5: Comparison of accuracies across different categorical features in the Bank dataset, evaluated with Empirical Risk Minimization (ERM), Just Train Twice (JTT), and Deep Feature Reweighting (DFR) for Bank dataset for positively labeled samples ($y=1$). Each subplot represents a distinct feature, and the x-axis indicates the category within each feature. The y-axis shows the accuracy for each method on that category. DFR consistently improves performance across most categories, particularly in underrepresented subgroups, highlighting its effectiveness in mitigating bias compared to ERM and JTT.