SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Josh Gardner
Zoran Popovic
Ludwig Schmidt<sup>a,b</sup>
<sup>a</sup> University of Washington <sup>b</sup> Allen Institute for AI {jpgard, zoran, schmidt}@cs.washington.edu
Robustness to distribution shift has become a growing concern for text and image models as they transition from research subjects to deployment in the real world. However, high-quality benchmarks for distribution shift in tabular machine learning tasks are still lacking despite the widespread real-world use of tabular data and differences in the models used for tabular data in comparison to text and images. As a consequence, the robustness of tabular models to distribution shift is poorly understood. To address this issue, we introduce TABLESHIFT, a distribution shift benchmark for tabular data. TABLESHIFT contains 15 binary classification tasks in total, each with an associated shift, and includes a diverse set of data sources, prediction targets, and distribution shifts. The benchmark covers domains including finance, education, public policy, healthcare, and civic participation, and is accessible using only a few lines of Python code via the TABLESHIFT API. We conduct a large-scale study comparing several state-of-the-art tabular data models alongside robust learning and domain generalization methods on the benchmark tasks. Our study demonstrates (1) a linear trend between in-distribution (ID) and out-of-distribution (OOD) accuracy; (2) domain robustness methods can reduce shift gaps but at the cost of reduced ID accuracy; (3) a strong relationship between shift gap (difference between ID and OOD performance) and shifts in the label distribution.<sup>1</sup>
Modern machine learning models have achieved near- or even super-human performance on many tasks. This has contributed to deployments of models across critical domains, including finance, public policy, and healthcare. However, in tandem with the growing deployment of machine learning models, researchers have also demonstrated concerning model performance drops under distribution shift – when the test/deployment data are not drawn from the same distribution as the training data. Analyses of these performance drops have primarily been confined to the domains of vision and language modeling (e.g. [38, 50], where effective benchmarks for distribution shift exist. Despite the widespread use of tabular data, the impact of distribution shift on tabular data has not been thoroughly investigated. While there are existing benchmarks for IID tabular classification, none of these focus on distribution shifts [36, 33].
This is particularly concerning in light of the known differences between tabular data and the modalities mentioned above (images, text, audio). First, in contrast to these modalities, where large neural models are the undisputed state-of-the-art, there is considerable debate about whether deep
<sup>1</sup>The benchmark data, Python package, model implementations, and more information about TABLESHIFT are available at https://github.com/mlfoundations/tableshift and https://tableshift.org.
Figure 1: In-domain (ID) and out-of-domain (OOD) accuracy show a linear trend across 15 TableShift tasks and 19 model types  $(\rho = 0.81)$ . ID accuracy  $(x$ -axis values) and change in the label distribution  $\Delta_y$  (color) together explain  $99\%$  of the variance in OOD accuracy  $(R^2 = 0.993)$ . For exact results see Section E.3.
learning models improve performance over non-neural baselines (e.g. XGBoost, LightGBM) for tabular data, even without the presence of distribution shift [46, 15, 78]. Second, tabular data tends to contain structured features extracted from raw data (e.g. counts of activities, coded responses to questions), as opposed to the raw signals (e.g. activity event streams, pixel values, audio of responses) where modern machine learning methods perform well and where previous studies of distribution shift have focused. Third, tabular data requires fundamentally different preprocessing procedures, and the impact of these decisions is not widely understood, despite being known to have empirical impact [31]. Finally, high-quality tabular datasets can be difficult to access [35]; for example, due to the personal nature of many tabular datasets, tabular data cannot simply be scraped at Internet scale as many text and image datasets are. This makes finding high-quality tabular distribution shift datasets particularly challenging.
Thus, the machine learning research community currently lacks not only (1) an empirical understanding of the impact of distribution shift on tabular data models, but also (2) a shared set of accessible and high-quality benchmarks to enable such investigations. We address both gaps in this work. Our main contributions are:
TableShift Benchmark: we introduce a curated, high-quality set of publicly-accessible benchmark tasks for (binary) tabular data classification under distribution shift. We describe the tasks in §3.1 and the API in §3.2. TableShift includes a set of real-world tabular datasets from domains including finance [30], public policy [24], healthcare [19, 47, 18, 74, 87], and civic participation [5]. We select these datasets to ensure a diversity of tasks, distribution shifts, and dataset sizes.
Large-scale empirical study of distribution shift in tabular data: We conduct a large-scale study in §4, including state-of-the-art tree-based tabular models, tabular neural networks, distributional robustness methods, domain generalization methods, and label shift robustness methods. Our findings show (1) a strong linear trend between in-distribution (ID) and out-of-distribution (OOD) accuracy across benchmark tasks and models that was not previously identified for tabular data; (2) that no model consistently outperforms baseline methods, and (3) a correlation between the shift gap and the shift in label distribution, which is not ameliorated by label shift robustness methods included in our study.
Accessible TableShift API and baselines: We release a Python API for constructing rich datasets directly from their raw public forms. The API provides built-in documentation of data types and feature codings, alongside standardized preprocessing and transformation pipelines, making the datasets accessible in multiple formats suitable for training tabular models (e.g. in scikit-learn and PyTorch). We also release the set of baseline model implementations (including both state of the art tabular data models, robust learning models, and domain generalization methods) and end-to-end training code in order to facilitate future research on distribution shift in tabular data.
2 Setup, Task, and Notation
2.1 Task and Setting
Consider a dataset composed of examples  $(x,y,d)\sim P_d$  where  $x$  is the input,  $y$  is the prediction target, and  $d$  the domain from which that example is drawn. All examples drawn from  $P_{d}$  have domain label  $d$ . We can view the overall data distribution as a mixture of domains  $\mathcal{D} = \{d_1,\dots ,d_D\}$ , where  $D\geq 2$ . Training examples are drawn from the training distribution  $P^{\mathrm{train}} = \sum_{d\in D}q_{d}^{\mathrm{train}}P_{d}$  and testing examples from  $P^{\mathrm{test}} = \sum_{d\in D}q_{d}^{\mathrm{test}}P_{d}$ , with domain weights  $q_{d}\in [0,1]$ . We can define the training and testing domains as  $\mathcal{D}^{\mathrm{train}} = \{d\in \mathcal{D}:q_d^{\mathrm{train}} > 0\}$  and  $\mathcal{D}^{\mathrm{test}} = d\in \mathcal{D}:q_d^{\mathrm{test}} > 0$ , respectively. We refer to cases where  $|\mathcal{D}^{\mathrm{train}}|\geq 2$  as "domain generalization" tasks, because domain generalization models require at least two subdomains in the training data.
In a standard (IID) setting, our goal is to learn a classifier  $f_{\theta}$  that accurately predicts  $y$  using examples from  $\mathcal{D}^{\mathrm{train}}$ . A distribution shift (or domain shift) occurs due to the fact that  $P^{\mathrm{train}} \neq P^{\mathrm{test}}$ . As a consequence of this shift, the joint distributions  $p^{\mathrm{train}}(x,y) \neq p^{\mathrm{test}}(x,y)$  differ in training and testing. This difference can be composed of one or more changes to the underlying data generating process. This includes covariate shift, where  $p(x)$  changes; label shift, where  $p(y)$  changes, and concept shift, where  $p(y|x)$  changes. In almost all real-world scenarios, distribution shifts are composed of an unknown mixture of all three forms of shift<sup>2</sup>. For a fixed classifier  $f_{\theta}$ , we refer to
$$
\Delta_ {\mathrm {A c c}} = \operatorname {A c c} \left(f _ {\theta}, \mathcal {D} ^ {\text {t e s t}}\right) - \operatorname {A c c} \left(f _ {\theta}, \mathcal {D} ^ {\text {t r a i n}}\right) \tag {1}
$$
as the "shift gap" (where both metrics are computed on examples not seen at training time). Note that the shift gap can be affected by changes in  $p(y)$ ,  $p(y|x)$ , and  $p(x)$ . While disentangling the effects of these forms of shift is not a focus of the current work, we provide initial exploratory results on the impact of changes in  $p(y)$ ,  $p(y|x)$ , and  $p(x)$  over the benchmark tasks in Sections 5 and E.
In our setting, we assume that no information about the target  $\mathcal{D}^{\mathrm{test}}$  is available - i.e., there is no knowledge of the change in  $p(y), p(y|x)$ , and  $p(x)$ , and no unlabeled data from the target domain.
2.2 Related Work
Here we provide a brief overview of related work necessary to contextualize our benchmark and main results. For a detailed overview of related work, see Section D.
Our work is closely related to the literature on distribution shift in machine learning. A series of recent works have demonstrated that even state-of-the-art models experience significant performance drops under distribution shift in tasks including vision, language modeling, and question answering [61, 62, 38, 50, 9]. This has led to the development of methods to mitigate susceptibility to such shifts [76, 53, 1, 6, 90, 89, 54, 46]. High-quality benchmarks, specifically tailored to distribution shift, have been essential in both measuring these gaps and assessing progress toward closing them [38, 50]. The use of tabular data is widespread in practice [15, 46, 78], including the use of sensitive personal data (race, gender, age) and for important tasks (credit scoring, medical diagnosis). However, the impact of distribution shift in the tabular domain has received little attention in the research literature. In particular, benchmarks containing tabular distribution shifts are lacking (one notable exception is Shifts [57] and Shifts 2.0 [57]), a multimodal benchmark of five tasks, two of which are tabular; for a more detailed overview of domain shift benchmarks and a comparison to TABLESHIFT, see Section G).
2We note that this is a slight abuse of the terminology, as e.g. "label shift" typically refers to the case where only  $p(y)$  changes.
3 Tableshift: A Distribution Shift Benchmark for Tabular Data
This work introduces the TABLESHIFT benchmark. TABLESHIFT contains a set of 15 curated tasks designed to be a rigorous, challenging, diverse, and reliable benchmarking suite for tabular data under distribution shift, and we encapsulate them within a Python API.
3.1 TableShift Benchmark Tasks
To select tasks for TABLESHIFT, we identified datasets meeting the following formal criteria:
Open source: datasets must be publicly available, including data dictionaries documenting the source of the data (i.e. conditions of its collection), definitions of variables, and any preprocessing applied.
Real-world: does not contain simulated data.
Table 1: Summary of TABLESHIFT tasks and their associated distribution shifts. For details on each task, see Section E. "Domain Generalization" indicates whether there are multiple training subdomains  $(|\mathcal{D}^{\mathrm{train}}| \geq 2)$  and thus whether domain generalization models can be applied to this task. "Baseline gap" gives the "shift gap"  $\Delta_{\mathrm{Acc}}$  (difference between ID and OOD test accuracy, see Equation (1)) of the tuned XGBoost or LightGBM model with the best validation accuracy after following our hyperparameter tuning procedure ( $\S 4.2$ ).
Task	Target	Shift	Domain Generalization	Baseline Gap ΔAcc	SE(ΔAcc)
ASSISTments	Next Answer Correct	School	✓	-34.49 %	0.011
College Scorecard	Low Degree Completion Rate	Institution Type	✓	-11.16 %	0.010
ICU Hospital Mortality	ICU patient expires in hospital during current visit	Insurance Type	✓	-6.30 %	0.008
Hospital Readmission	30-day readmission of diabetic hospital patients	Admission source	✓	-5.94 %	0.002
Diabetes	Diabetes diagnosis	Race	✓	-4.48 %	0.001
ICU Length of Stay	Length of stay >= 3 hrs in ICU	Insurance Type	✓	-3.39 %	0.015
Voting	Voted in U.S. presidential election	Geographic Region	✓	-2.58 %	0.016
Food Stamps	Food stamp reciprocity in past year for households with child	Geographic Region	✓	-2.39 %	0.002
Unemployment	Unemployment for non-social security-eligible adults	Education Level	✓	-1.28 %	0.001
Income	Income >= 56k for employed adults	Geographic Region	✓	-1.25 %	0.002
FICO	Repayment of Home Equity Line of Credit loan	Est. third-party risk level	-22.58 %	0.029
HELOC
Public Health Insurance	Coverage of non-Medicare eligible low-income individuals	Disability Status	-14.46 %	0.001
Sepsis	Sepsis onset within next 6hrs for hospital patients	Length of Stay	-6.05 %	0.001
Childhood Lead	Blood lead levels above CDC Blood Level Reference Value	Poverty level	-5.12 %	0.005
Hypertension	Hypertension diagnosis for high-risk age (50+)	BMI Category	-4.36 %	0.003
Sufficient dimensionality and size: contains at least three features (in all cases, our benchmark datasets contain many more than three features) and at least 1000 observations. In particular, having large test sets is critical for making reliable statistical comparisons between models.
Heterogeneous: contains features of mixed types.
Binary Classification: supports a meaningful binary classification task (regression tasks are not included).
Shift Gap: We explicitly select datasets where strong hyperparameter-tuned tabular baselines display a statistically significant shift gap ( $\Delta_{\mathrm{Acc}} \neq 0$ , see Eqn. (1)).
In addition to these criteria, we selected benchmark tasks and data sources that were diverse. TABLESHIFT includes tasks from many domains (finance, policy, civic participation, medical diagnosis) and from a variety of raw data sources (electronic health records, surveys/questionnaires, etc.) and with a diversity of shift gap  $(\Delta_{\mathrm{Acc}})$  magnitudes.
A summary of the benchmark tasks is shown in Table 1. We give a detailed overview of each task, including background and motivation, information on the data source, and distribution shifts, in Section B. Datasets and each individual feature of each task are also documented in the Python package. One important aspect of TableShift's diversity, shown in Table 1, is that not all real-world tasks support domain generalization (i.e. not all tasks have multiple training subdomains,  $|\mathcal{D}^{\mathrm{train}}\geq 2|)$  To reflect this, we include both types of tasks in the TableShift benchmark.
While the intended use of TableShift is for distribution shift, the package is also likely to be of high utility to all researchers studying tabular data modeling due to the data quality, detailed documentation, flexible preprocessing, and ease of use of the datasets in TableShift.
3.2 TableShift API
Successful existing benchmarks for distribution/domain shift in machine learning (e.g. WILDS, DomainBed) not only include high-quality datasets, but also make the data accessible by providing a high-quality API as an interface to the otherwise-disparate sources. This section describes the TableShift API. Providing this API for tabular data is particularly important, for several reasons.
First, the input and output of tabular data pipelines differ from other modalities: tabular datasets are stored in different formats from image and text datasets, and are used with a greater variety of machine learning tools (e.g. scikit-learn). Second, the preprocessing operations used in tabular data differ significantly from other data modalities. These preprocessing operations also require unique feature-level metadata such as data types (i.e. categorical vs. numeric; numeric values for categorical features are a common encoding scheme in practice) and codings for categorical variables. Finally, raw sources used to build tabular datasets can be difficult to access. Datasets are often scattered across hundreds or even thousands of files (e.g., the Sepsis task dataset is constructed from over  $40k$  data files; the Childhood Lead dataset is joined from nearly 100 files containing disjoint feature sets provided by the National Health and Nutrition Examination Survey (NHANES)).
The TableShift API addresses each of these issues. It defines a set of primitives which allow for the construction of data pipelines which go from raw data sources to preprocessed data of any TableShift benchmark task in a few lines of Python code<sup>3</sup>. The resulting data is documented - each feature in the benchmark includes metadata which describes the feature and any encodings. The API natively supports a set of common data transformations, including one-hot and label encoding for categorical data; scaling and binning of numeric data; and handling of missing values. TableShift provides native output in a variety of data formats, including PyTorch DataLoaders, Pandas DataFrames, and Ray Datasets. Finally, any dataset in the TableShift benchmark can be loaded with default preprocessing parameters with an identical call to the API, providing a unified interface.
We provide a detailed comparison between TableShift and related existing benchmarks in Section G. However, we emphasize that there is no existing benchmark suite for distribution shift in tabular data, and existing distribution shift benchmarks are incompatible with the unique constraints of tabular data discussed above.
<sup>3</sup>See https://tableshift.org and https://github.com/mlfoundations/tableshift
We conduct a set of experiments to demonstrate the potential insights to be gained from using TableShift. As previously mentioned, there has been considerable debate about whether tree-based models (XGBoost, LightGBM, etc.) or specialized deep learning-based models (i.e. ResNet- and Transformer-based architectures) are more effective for tabular data modeling. However, previous investigations have not explored how these models perform under distribution shift in tabular data. Additionally, many methods have been proposed for robust learning and domain generalization but also not rigorously evaluated on tabular data. We present a series of experiments to evaluate 19 distinct methods using the TABLESHIFT benchmark.
4.1 Tabular Data Classification Techniques in our Comparison
We train and evaluate a set of tabular data classifiers from several families. For each, we give additional details and description in Section F, and the full hyperparameter grids in Table 19. Implementations of these classifiers, including the hyperparameter tuning framework used to tune them, are available in the TABLESHIFT API. The classifiers compared in our experiments are:
**Baseline Models:** These models do not include any intervention for robustness to domain shift, but are generally effective for tabular data in the IID setting. We evaluate multilayer perceptrons (MLP), XGBoost [20], LightGBM [48], and CatBoost [25] as baseline methods. While we refer to these as "basielines" for convenience, we note that the methods based on gradient-boosted trees (XGBoost, LightGBM, CatBoost) are still considered state-of-the-art on many tasks [37].
Tabular Neural Networks: We also include a set of state-of-the-art methods for modeling tabular data. The models we use are SAINT [79], TabTransformer [43], NODE [70], FT-Transformer, and tabular ResNet (the latter two via [36]).
Domain Robustness Models: These models attempt to ensure good performance on distributions close to the training data. These models attempt to optimize an objective over a worst-case distribution with bounded distance from the training data. We evaluate distributionally robust optimization (DRO) with both  $\chi^2$  and CVaR geometry [53], and group DRO (where the groups are domains) [76]. Both the DRO and group DRO models are parameterized over MLPs, as in both original works.
Label Shift Robustness Models: These models attempt to ensure good performance when the label distribution  $P(y)$  changes. We evaluate Group DRO (where the groups are class labels) and the adversarial label robustness method of [92].
Domain Generalization Models: These are models designed with the goal of achieving low error rates on unseen test domains. In practice, this is achieved by achieving low error disparity across the subdomains in  $\mathcal{D}^{\mathrm{train}}$ . These methods require domain labels at training time, and training data drawn from multiple different domains ( $|\mathcal{D}^{\mathrm{train}}| \geq 2$ ). Domain generalization models in our study are: Domain-Adversarial Neural Networks (DANN) [1], Invariant Risk Minimization (IRM) [6], Domain MixUp [90, 89], Risk Extrapolation (VReX) [52], DeepCORAL [82] and MMD [54].
We note that our goal of the current study is not to propose novel methods for distributionally robust learning; it is to conduct a comprehensive comparison of a large set of existing methods, many of which have not been previously compared to each other, on a high-quality benchmark. For example, while domain generalization models have been applied to image and text classification tasks (e.g. [50, 38]), to our knowledge these methods have not been previously investigated for mitigating distribution shift in tabular data in a large-scale benchmarking study. Indeed, we are aware of no prior applications of many of these domain generalization methods to tabular data. As a result, it is not clear a priori how these methods might compare to existing robustness or baseline methods due to the aforementioned differences between tabular data and these other data modalities.
The experiments described above cover both model architectures (different functional forms for the predictor  $f_{\theta}$ ) and loss functions (different objective functions  $\mathcal{L}$  used to train the model by attempting to find  $\min_{\theta} \mathcal{L}(f_{\theta}(\mathcal{D}^{\mathrm{train}}))$ ). In order to train a classifier with gradient-based training, both are required. Except where noted otherwise, any method requiring gradient-based training (MLP, Tabular Neural Networks, Domain Generalization Models) is trained with standard empirical risk minimization and cross-entropy loss. Similarly, any method which itself is a loss function (i.e. all variants of DRO) is
trained with  $f$  parameterized as an MLP, as is standard in prior works implementing and comparing these methods (e.g. [53, 76, 33]).
4.2 Methods
For each task, we conduct the following procedure.
First, we split the full dataset into  $\mathcal{D}^{\mathrm{train}}$  and  $\mathcal{D}^{\mathrm{test}}$ . We summarize the domain splits in Tables 1,1 and describe the splitting for each task in detail, along with background and motivation for each task domain split, in Section B. Within each domain, we have both a validation and a test set. We use the same domain splits, data preprocessing, and train/validation/test splits for all models and training runs, except where explicitly noted.
Second, we then conduct a hyperparameter sweep for each model described in Section 4.1. We use HyperOpt [13] to sample from the model hyperparameter space, in accordance with previous works (e.g. [36, 46]) which largely use adaptive hyperparameter optimization due to the variability in effective hyperparameter settings between datasets. We only train on the training set, and use the in-domain validation accuracy for hyperparameter tuning. We give the complete grid for each model in §I. Each model is tuned for 100 trials.