SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Míriam Barrabés  $^{1,2,*}$  Daniel Mas Montserrat  $^{1,*}$  Margarita Geleta  $^{3}$
Xavier Giró-i-Nieto  $^{4,\dagger}$  Alexander G. Ioannidis  $^{1\dagger}$
$^{1}$ Stanford University  $^{2}$ Universitat Politecnica de Catalunya
<sup>3</sup>University of California, Berkeley <sup>4</sup>Amazon
Data shift is a phenomenon present in many real-world applications, and while there are multiple methods attempting to detect shifts, the task of localizing and correcting the features originating such shifts has not been studied in depth. Feature shifts can occur in many datasets, including in multi-sensor data, where some sensors are malfunctioning, or in tabular and structured data, including biomedical, financial, and survey data, where faulty standardization and data processing pipelines can lead to erroneous features. In this work, we explore using the principles of adversarial learning, where the information from several discriminators trained to distinguish between two distributions is used to both detect the corrupted features and fix them in order to remove the distribution shift between datasets. We show that mainstream supervised classifiers, such as random forest or gradient boosting trees, combined with simple iterative heuristics, can localize and correct feature shifts, outperforming current statistical and neural network-based techniques. The code is available at https://github.com/AI-sandbox/DataFix.
Distribution shifts in multi-dimensional data, caused by one or more "corrupted" dimensions, are common in various real-world applications. Data streams from multi-sensor environments in fields like medicine, industry, finance, and defense, can experience shifts due to faulty sensors [1]. Tabular and structured data used in domains such as economics, biology, genomics, and social sciences, can encounter distribution shifts caused by improper standardization, erroneous data processing, data collection procedures, or human entry errors [2, 3]. Combining databases from diverse sources is common in many fields: various hospitals collect different phenotypic characteristics from patients, and governmental institutions often capture different socio-economic indicators of citizens. However, proper standardization of data is seldom applied across data sources, and merging these poses challenges that often require data-dependent and domain-specific techniques [4]. For example, efficient computational techniques are used to process high-dimensional genomic sequences [5], and domain knowledge is incorporated in processing data from social sciences [6]. This process typically involves changing, imputing, merging, and removing both features and samples, potentially leading to complex pipelines that can originate distribution shifts between datasets. These shifts can create erroneous scientific results if present in datasets used for medical or social studies, or indicate infrastructure failure if present in sensory data in industrial applications. Therefore, detecting, localizing, and correcting such shifts is an important task that remains a relatively unsolved problem.
The "localization" of feature shifts is the task of finding which features (i.e., dimensions) of the dataset are causing the shift within multiple sets of data. This step can be critical in order to address the error
$^{\dagger}$ Work done prior to joining Amazon.
$^{\ddagger}$ Correspondence to A.G.I. [ioannidis@stanford.edu].
source, by physical intervention in the multi-sensory scenario, or by data removal or fixing in tabular data-based applications. There has been extensive work on distribution shift detection and anomaly detection [7, 8], however, such methods focus mainly on detecting if two sets of data follow the same distribution, or on detecting outliers, and the task of identifying the exact components generating the shift remains partially unexplored. Recently, work combining machine learning techniques with statistical testing provided state-of-the-art results showing that localizing corrupted features can be done accurately [9]. In this work, we explore "feature selection" methods as a mechanism to detect potential feature shifts. The "correction" of feature shifts consists of replacing corrupted feature values with new ones, ensuring that the updated dataset follows a distribution that is equal to, or more similar to, the distribution without corruptions. This can be useful in data homogenization and quality control pipelines, enabling the elimination of shifts between multiple data sources and their combination for downstream applications, including data discovery or machine learning training. Feature shift correction can be framed as a supervised problem, employing regression or classification methods to predict the new feature values; as a missing data imputation problem, by treating corrupted features as missing values [10, 11]; or as a distribution alignment problem [12, 13], where mappings between distributions are learned using adversarial learning, optimal transport, or statistical divergences. Effective techniques for correcting distribution shifts, either by learning parametric models to predict new values for corrupted features, or directly updating values without explicit model learning, can be very useful to avoid incorrect or biased results in analysis based on data containing feature shifts. We provide a more formal definition of the feature shift detection and correction tasks in the following sections.
Tools for automated shift detection, data monitoring, and feature analysis and correction are becoming more common in data quality control, data homogenization, data processing pipelines, data-centric AI, MLOps, and deployment monitoring for ML Systems [14, 15, 16, 17]. In this work, we introduce "DataFix", a framework that makes use of discriminators trained to distinguish samples coming from two different distributions, in order to estimate distribution shifts, localize which features are causing them, and modify the samples in order to reduce or remove the shift. DataFix is composed of two systems: "DF-Locate", which locates the features causing the distribution shift, and "DF-Correct", which directly modifies the samples in order to decrease the distribution shift between two datasets. DF-Locate consists of an iterative process where a discriminator combined with a feature importance estimator is used to identify the most discriminative features between samples coming from two different distributions. The detected features are removed from the dataset, a new discriminator is trained, new feature importances are estimated, and the process is repeated until no distribution shift is detected. DF-Correct replaces the values of the corrupted features detected by DF-Locate with values that exhibit a minimal probability of being corrupted when processed by a discriminator.
Our contributions are as follows: (1) we motivate, define, and formalize the problem of feature shift detection and correction and relate it to feature selection and adversarial learning frameworks; (2) we propose an iterative algorithm that makes use of feature selection techniques from discriminators to accurately detect manipulated features; (3) we propose an iterative algorithm that makes use of a discriminator to guide the correction of a corrupted dataset in order to decrease its distribution shift; (4) we provide an in-depth experimental evaluation with multiple manipulation types and datasets.
Distribution Shift Detection and Localization. Distribution shift detection involves identifying if  $p \neq q$ , where  $p$  and  $q$  are the reference and query distributions respectively. While there is extensive research on detecting distribution shifts in univariate distributions [18, 19, 7], exploration in the multivariate setting is relatively limited [20]. Multivariate distributions can exhibit various types of shifts, including covariate, label, concept, or marginal shifts, among others [21, 22, 23]. The work in [8] applies hypothesis testing for concept drift detection. In [20], two-sample hypothesis testing is explored for shift detection, and a comparison is made between multivariate hypothesis tests via Maximum-Mean Discrepancy (MMD) [24], univariate hypothesis tests with marginal Kolmogorov-Smirnov (KS) tests and the Bonferroni correction [25], and dimensionality reduction techniques, among others. However, these works focus on detecting if a distribution shift is present but do not localize which features are causing it. Recently, a conditional test method was able to identify the features originating the shift [9] with model-free and model-based approaches: K-Nearest Neighbors with KS statistic (KNN-KS), multivariate Gaussian with KS (MB-KS), multivariate Gaussian and
Fisher-divergence test statistics (MB-SM), and deep density neural network models with Fisher-divergence test (Deep-SM). However, performing feature shift localization remains challenging.
Feature Selection. Feature selection methods can remove redundant features affecting efficiency or performance, and find the most relevant features, providing interpretability. Common techniques include: filtering [26, 27], wrapper [28, 29] and embedded methods [30, 31]. Filtering methods, including univariate Mutual Information statistics (MI) [32], ANOVA-F test [33], and Chi-square test [34], rank features based on data statistics. Minimum Redundancy Maximum Relevance (MRMR) [35, 36] selects relevant features while minimizing redundancy with the selected ones. Fast-Conditional Mutual Information Maximization (FAST-CMIM) [37] selects features that maximize MI, conditional to previously selected features. Wrapper methods select feature subsets by training models on them and adding and removing features through a search process that can be computationally intensive. Embedded methods use the inherent scores computed by predictive models, such as logistic regression weights [38] or the mean decrease of impurity (Gini index) in random forests [39]. Most methods select the features with the highest importance scores or all that exceed a given threshold [40].
Missing Data Imputation. Missing data imputation methods can remove distribution shifts by considering the shifted features as missing and re-constructing them. Imputation methods range from simple record deletion, zero imputation, mean imputation, and deck imputation [41], to machine learning-based methods that apply regression and classification to reconstruct missing data including MICE [10], MissForest [11], and Matrix Completion [42, 43]. Some generative approaches include Expectation Maximization algorithms [44] and deep learning-based methods such as the MIWAE autoencoder [45]. Causality-based techniques such as MIRACLE have proven to be highly accurate [46]. HyperImpute [47] provides state-of-the-art imputation by using an iterative process with ML methods, including gradient boosting machines and neural networks trained to impute missing data.
Optimal Transport, Distribution Alignment, and Adversarial Learning. Feature shift correction can also be performed with distribution alignment methods, optimal transport, and adversarial learning. Distribution alignment and representation learning methods learn a mapping between distributions and can be used to remove shifts by projecting the shifted samples into the non-shifted distribution. Iterative Alignment Flows [12] and Deep Density Destructors [13] are neural network examples that map samples between distributions. Optimal transport-based methods reduce the Wasserstein distance between distributions. The work in [48] applies a Sinkhorn-based optimization process to perform imputation. Adversarial learning has been explored in generative adversarial networks such as GAIN [49], where the loss of a discriminator is used to improve imputation accuracies.
Data-centric AI. Data-centric AI focuses on systematically enriching data quality and quantity as a means to boost AI and ML performance. It covers strategies that impact every phase of the data lifecycle, from data collection [50], labeling [51, 52], augmentation [53, 54, 55], and integration [56], to the crucial processes of data cleaning [57, 58], feature extraction [59], and transformation [60]. These include programmatic automation strategies, which rely on programs guided by heuristics and statistics for automatic data processing [61, 62, 63], as well as learning-based automation techniques that optimize data automation procedures, typically using machine learning [64, 65, 66]. The detection and removal of distribution shifts are becoming essential steps of the data-centric AI toolbox [17].
Definition 1. [Feature Shift] We are given two sets of  $d$ -dimensional samples  $X = \{x_{1}, x_{2}, \ldots, x_{N}\}$  and  $Y = \{y_{1}, y_{2}, \ldots, y_{N}\}$ , with  $x_{i}, y_{i} \in \mathbb{R}^{d}$ , from distributions  $p$  and  $q$ , respectively. A feature shift between  $p$  and  $q$  occurs when  $D(p, q) > \varepsilon$  and  $D(p_{S}, q_{S}) \leq \varepsilon$ , where  $D$  is a valid divergence or distance between distributions,  $S$  and  $C$  are the subsets of non-corrupted and corrupted features respectively, such that  $|S \cup C| = d$ , and  $p_{S}$  and  $q_{S}$  are the distributions restricted to  $S$ .
We will refer to  $X$  and  $Y$  as the "reference" and "query" datasets, and to  $p$  and  $q$  as the "reference" and "query" distributions, respectively. We will assume that the reference contains only "non-corrupted" features, while the query contains one or more "corrupted" dimensions that we will want to detect and correct. Here we consider scenarios with  $\varepsilon = 0$ ,  $D(p_{S},q_{S}) = 0$ ,  $D(p,q)\gg 0$ , and  $|S| > |C|\geq 1$ . That is, there are more non-corrupted than corrupted features, the divergence becomes 0 if the corrupted features are removed, and is large enough to be empirically detected otherwise. We will consider multiple types of distribution shifts: marginal shifts with  $D(p_{i},q_{i}) > \varepsilon$ , where  $p_i$  and  $q_{i}$  represent the marginal distribution of the  $i$ th dimension resulting from additive and non-linear
Figure 1: DF-Locate overview diagram.
transformations; correlation shifts where  $D(p,q) > \varepsilon$  but  $D(p_i,q_i)\leq \varepsilon$  for all  $i$ ; and correlation shifts where  $D(p_{S},q_{S})\leq \varepsilon$  and  $D(p_{C},q_{C})\leq \varepsilon$  but  $D(p,q) > \varepsilon$ . In the latter case, correlations are maintained locally, but a shift is present when considering all features simultaneously.
Definition 2. [Feature Shift Localization Task] The task of localizing a feature shift consists of finding the smallest subset of features  $C$ , with  $\overline{C} = S$ , that satisfies  $D(p_{\overline{C}}, q_{\overline{C}}) \leq \varepsilon$ , that is  $C = \operatorname*{argmin}_{D(p_{\overline{C}}, q_{\overline{C}}) \leq \varepsilon} |C|$ .
The number of corrupted features  $|C|$  will be assumed to be unknown a priori. Furthermore, because the distributions  $p$  and  $q$  are assumed to be unknown and only their samples are accessible, the task needs to be approximated, requiring careful consideration of trade-offs such as falsely flagging non-corrupted features (false positives) and failing to detect corrupted ones (false negatives).
Definition 3. [Feature Shift Correction Task] The task of correcting a feature shift consists of finding a new matrix  $Y' = \{y_1', y_2', \dots, y_N'\}$  such that  $y_i' \sim q'$  and  $D(p, q') \leq \varepsilon$ , while keeping the non-corrupted features unchanged, that is  $Y_S' = Y_S$ .
The correction of shifted features can be done through parametric models that perform samplewise transformations  $y' = \phi(y)$ , prediction or imputation methods  $y_C' = \phi(y_{\overline{C}})$ , or optimization procedures or heuristics  $Y' = \Phi(X, Y)$ . Our proposed approach falls in the latter category.
4 Feature Shift Detection: DF-Locate
DF-Locate (Fig. 1, Algorithm 1) employs an iterative process to detect the presence of a shift and determine the features causing it. At each iteration, a classifier is trained to detect the origin (binary label indicating reference vs query) of each sample, the output predictions are used to estimate the divergence between distributions, and the feature importance scores provided by the binary classifier are used to locate the features originating the shift. At the end of each iteration, the features detected as corrupted are removed, and the process is repeated until no divergence is detected.
Shift Detection. We detect whether a shift is present by estimating an  $f$ -divergence between the distributions  $p$  and  $q$  by using an empirical approximation of the variational form [67, 68, 69]:
$$
\hat {D} _ {\theta} (X, Y) = \frac {1}{N _ {x}} \sum_ {i = 1} ^ {N _ {x}} f ^ {\prime} \left(r _ {\theta} \left(x _ {i}\right)\right) - \frac {1}{N _ {y}} \sum_ {j = 1} ^ {N _ {y}} f ^ {*} \left(f ^ {\prime} \left(r _ {\theta} \left(y _ {j}\right)\right)\right) \tag {1}
$$
where  $X = \{x_{1},\dots,x_{N}\}$  and  $Y = \{y_{1},\dots,y_{N}\}$ , with  $x_{i}\sim p$ ,  $y_{j}\sim q$ , and  $r_{\theta}(x)$  is a function approximating the likelihood ratio between  $p$  and  $q$ , obtained by training a binary classifier  $\mathcal{D}_{\theta}(x)$  such that  $r_{\theta}(x) = \mathcal{D}_{\theta}(x) / 1 - \mathcal{D}_{\theta}(x)$ . We use 5-fold train-evaluation such that  $80\%$  of the samples are used to train a random forest binary classifier  $\mathcal{D}_{\theta}(x)$ , and the  $20\%$  left is used to estimate the empirical divergence with  $N_{x}$  and  $N_{y}$  testing samples of the reference and query datasets, respectively. The resulting estimates are averaged across folds to reduce variance. The function  $f$  is the generator function of the  $f$ -divergence,  $f'$  is its first order derivative, and  $f^{*}$  is its Fenchel conjugate [70]. By changing  $f$ , various divergences can be obtained such as the Kullback-Leibler divergence, the Jensen-Shannon divergence used in GANs, or the total variation distance, among others.
The true  $f$ -divergence is recovered in expectation  $D_{f}(p,q) = \mathbb{E}_{x\sim p,y\sim q}[\hat{D}_{\theta}(X,Y)]$  when  $r_{\theta}(x) = r^{*}(x) = p^{(x)} / q(x)$  is the true likelihood ratio function, defining the Bayes decision rule that optimally separates  $p$  and  $q$ . In this work, we make use of the total variation distance  $0 \leq D_{TV}(p,q) \leq 1$ , defined by setting  $f^{\prime}(u) = (f^{*} \circ f^{\prime})(u) = \frac{1}{2}\mathrm{sign}(u - 1)$ . The discrete nature of  $f^{\prime}(u)$  makes the empirical estimator robust to poorly calibrated classifiers. Furthermore, its value is proportional to the empirical balanced accuracy on the evaluation test. Intuitively, if  $\hat{D}_{TV}(X,Y) > 0$ , or equivalently, the test balanced accuracy is larger than 0.5 (random chance), it might indicate that  $p \neq q$ . After detecting the presence of a shift, the next step is to localize the features originating it.
Feature Selection for Shift Localization. Our proposed method employs feature importance and feature selection techniques in order to detect the features originating a distribution shift. These techniques often involve approximating the following estimation problem, either implicitly or explicitly:
$$
C = \underset {| C | \leq k} {\operatorname {a r g m a x}} I \left(z _ {C}; t\right) \tag {2}
$$
where  $C$  is the feature subset of size up to  $k$ , which maximizes the mutual information  $I$  between the input sample  $z$ , restricted to the features in  $C$ , and the label  $t$ . By defining  $z$  as the samples coming from  $X$  and  $Y$ , and  $t$  as the binary label of origin, such that  $z|_{t=0} \sim p$  and  $z|_{t=1} \sim q$ , Fano's inequality [71] and LeCam's method [72] can be used to relate the mutual information between  $z$  and  $t$  and the total variation distance between  $p$  and  $q$ :
$$
1 - H _ {2} \left(\frac {D _ {T V} (p , q) + 1}{2}\right) \leq I (z; t) \leq D _ {T V} (p, q) \tag {3}
$$
where  $H_{2}$  is the entropy with base-2 logarithm. If the number of manipulated features is known  $|C| = k$ , and  $D_{TV}(p,q) = D_{TV}(p_C,q_C) = 1$ , then  $C = \operatorname{argmax}_{|C|\leq k}I(z_C;t) = \operatorname{argmin}_{D_{TV}(p_{\overline{C}},q_{\overline{C}}) = 0}|C|$ , thereby rendering the problem of feature selection and feature shift detection equivalent. See Section C for further discussion. In practice, the distributions are unknown making Eq. 2 intractable and most methods predict a score that approximates the amount of information relative to the label present at each feature:  $\beta = F(Z,T)$ , where  $\beta \in \mathbb{R}^d$  and  $F(\cdot)$  is a function mapping samples  $Z$  and labels  $T$  to a feature-wise importance score. Here, we use the mean decrease of impurity from the random forest classifier  $\mathcal{D}_{\theta}$  as the feature importance scores.
Feature removal policy. We introduce the feature removal policy function  $\psi$  as an heuristic to approximate Eq. 2, which utilizes the predicted feature importances  $\beta$  and the estimated total variation distance  $\hat{D} = \hat{D}_{TV}(X,Y)$  to select likely corrupted features  $C = \psi (\beta ,\hat{D})$ . First,  $\beta$  is normalized  $\beta^{\prime} = \frac{|\beta|}{\sum_{i = 0}^{d - 1}|\beta_i|}$  and sorted  $\beta_{\pi}^{\prime}$ , such that  $\beta_{\pi (0)}^{\prime}\geq \beta_{\pi (1)}^{\prime}\geq \ldots \geq \beta_{\pi (d - 1)}^{\prime}$ . Given the cumulative sum  $\gamma (k) = \sum_{j = 0}^{k}\beta_{\pi (j)}^{\prime}$ , the sorted features from 0 to  $J$  are selected as corrupted, where  $J$  is the smallest index such that  $\gamma (k)\geq \tau \hat{D}$  (Eq. 4), and  $\tau$  is a hyperparameter set by hyperparameter optimization. Finally,  $\psi$  returns the features from 0 to  $J$  with scores higher than  $\frac{1}{d}$  (Eq. 5):
$$
J (\tau , \hat {D}, \beta) = \underset {k; \gamma (k) \geq \tau \hat {D}} {\operatorname {a r g m i n}} \gamma (k) \quad (4) \quad C = \left\{\pi^ {- 1} (j): 0 \leq j \leq J, \beta_ {\pi (j)} ^ {\prime} > \frac {1}{d} \right\} \tag {5}
$$
Note that  $\tau \hat{D}$  acts as a threshold to select how many features are selected as corrupted. If the threshold is small,  $\psi$  simply returns the feature with the highest feature importance, while a larger threshold will make  $\psi$  return more features. By defining the threshold as the product of  $\tau$  and  $\hat{D}$ , we ensure that when the shift is small (low  $\hat{D}$ ), a smaller number of features are flagged as corrupted.
Iterative Process. DF-Locate performs an iterative process that starts with the reference  $X^{(0)} = X$  and query  $Y^{(0)} = Y$  datasets. At each iteration  $i$ , it trains a set of discriminators  $\theta^{(i)} = \operatorname{argmax}_{\theta} \hat{D}_{TV}(X^{(i)},Y^{(i)})$  consisting of random forest binary classifiers. The classifier predictions are used to estimate  $\hat{D}^{(i)} = \hat{D}_{TV}(X^{(i)},Y^{(i)})$  and, combined with the Gini importance from the random forests  $\beta^{(i)} = F(X^{(i)},Y^{(i)})$ , a set of corrupted features are localized using the feature removal policy  $C_i = \psi(\beta^{(i)},\hat{D}^{(i)})$ . The selected features  $C_i$  are removed from the dataset, such that  $X^{(i+1)} = X_{C_i}^{(i)}$  and  $Y^{(i+1)} = Y_{C_i}^{(i)}$ , and the process is repeated while  $\hat{D}_{TV}(X^{(i)},Y^{(i)}) > 0.02$  or
until half of the features have been removed or no features are removed at the current iteration. After  $l$  iterations, the set of corrupted features is obtained as  $C' = \bigcup_{i=1}^{l} C_i$ . Finally, a refinement stage (see below and Section E.3) predicts the final set  $C$ . More details are provided in Section E.