SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Refinement stage. DF-Locate stores all intermediate steps, including the indexes of the predicted corrupted feature locations and the estimated  $\hat{D}$  at each iteration, in order to revisit the iterative filtering process and select the optimal stopping point. To determine the optimal iteration, the elbow or knee is found from a processed curve depicting the empirical total variation distance as a function of the total number of removed features (see Section E.3). This refinement stage effectively eliminates false positives and enhances the accuracy of feature shift localization.
Algorithm 1 DF-Locate
1: Inputs:  
X; ▷ Reference  
Y; ▷ Query  
τ; ▷ Feature Selection Threshold  
ε; ▷ Divergence Threshold
2:  $X^{(0)} = X$
3:  $Y^{(0)} = Y$
4:  $i = 0$
5:  $k^{(0)} = 0$
6: while  $\hat{D}_{\theta}(X^{(i)},Y^{(i)}) > \epsilon$  and  $k^{(i)} < \frac{|Y|}{2}$  and  $k^{(i)} - k^{(i-1)} > 0$  do
7:  $\theta \gets \operatorname{Train}(X^{(i)}, Y^{(i)})$ $\triangleright$  Train discriminator
8:  $\hat{D} \gets \hat{D}_{\theta}(X^{(i)}, Y^{(i)})$ $\triangleright$  Estimate divergence
9:  $\beta \leftarrow F_{\theta}(X^{(i)},Y^{(i)})$ $\triangleright$  Estimate feature importance
10:  $C_i\gets \psi_\tau (\beta ,\hat{D})$   Select corrupted features
11:  $X^{(i + 1)},Y^{(i + 1)}\gets X_{\overline{C_i}}^{(i)},Y_{\overline{C_i}}^{(i)}$  Remove detected features
12:  $k^{(i + 1)}\gets k^{(i)} + |C_i|\triangleright$  Update detected feature counter
13:  $i\gets i + 1$
14: end while
15:  $C' = \bigcup_{j=0}^{i-1} C_j$ $\triangleright$  Combine all detected features
16:  $C\gets \operatorname {Refine}(C^{\prime})$   Refine detected features
17: return C
Algorithm 2 DF-Correct
1: Inputs:  
X; ▷ Reference  
Y; ▷ Query  
C; ▷ Corrupted Features  
ε; ▷ Divergence Threshold
2:  $V = \{Y^0, Y^1, Y^2\} \gets \operatorname{Impute}(X, Y, C)$
3:  $Y^{\prime}\gets \mathrm{argmin}_{Y^{i}\in V}\hat{D}_{\theta}(X,Y^{i})$
4: if  $\hat{D}_{\theta}(X,Y^{\prime}) < \epsilon$  then
5: return  $Y^{\prime}$
6: end if
7: for epoch do
8:  $\theta \gets \operatorname{Train}(X, Y')$
9:  $L\gets \mathrm{DetectIncorrect}(\mathcal{D}_{\theta}(Y'))$ $\triangleright$  Detect samples that require feature correction
10:  $B\gets$  GenerateProposals(X,Y',V,C)
11: for  $i\in L$  do
12:  $b_{i}\gets \operatorname{argmax}_{b\in B}r_{\theta_{i}}(y_{i}^{(b)})\triangleright$  Find best proposal
13:  $Y^{\prime}\gets \mathrm{update}(Y^{\prime},y_{i}^{(b_{i})})$  Dataset with updated sample
14: end for
15: if  $\hat{D}_{\theta}(X,Y^{\prime}) < \epsilon$  then
16: return  $Y^{\prime}$
17: end if
18: end for
19: return  $Y^{\prime}$
5 Feature Shift Correction: DF-Correct
After the set  $C$  of features originating the shift has been detected by DF-Locate, DF-Correct (Figure 2, Algorithm 2) is applied to generate a new query dataset  $Y'$  that rectifies the distribution shift. Ideally, the objective is to find  $Y' \sim q'$  such that:
$$
Y ^ {\prime} = \underset {Y \sim q ^ {\prime}; \left| \left| Y _ {\overline {{C}}} - Y _ {\overline {{C}}} ^ {\prime} \right| \right| = 0} {\operatorname {a r g m i n}} D (p, q ^ {\prime}) \tag {6}
$$
where  $D$  is a valid statistical divergence. Because  $p$  and  $q^{\prime}$  are unknown, we approximate the optimization problem by using a discriminator  $\mathcal{D}_{\theta}$  to predict the empirical estimate of an  $f$ -divergence:
$$
Y ^ {\prime} = \underset {Y} {\operatorname {a r g m i n}} \underset {\theta} {\max } \hat {D} _ {\theta} (X, Y) \tag {7}
$$
This approach parallels the adversarial minimax optimization setting adopted in GANs, where samples are generated by a generator  $Y' = \mathcal{G}_{\omega}(u)$ , and a discriminator tries to accurately classify the generated samples as "fake". This training leads to the minimization of the Jenson-Shannon divergence, and the setting for GAN training can be generalized to any  $f$ -divergence [68]. Here, instead of training a generator, we directly update the corrupted values of the query dataset by trying to minimize the predicted likelihood of the updated samples coming from  $q$ , approximated with the discriminator  $\mathcal{D}_{\theta}$ . While neural network-based discriminators allow for direct optimization of the values of  $Y$  through backpropagation, these can have suboptimal performance when classifying tabular and similar structured data, which is the focus of our work. Hence, we employ tree-based
Query  $Y$
Feature Shift Correction
Figure 2: DF-Correct overview diagram.
Imputed query
techniques instead. Despite not being differentiable, derivative-free optimization and search heuristics can be used to update  $Y$ .
Two important aspects need to be taken into consideration: (a) the allowed search space for the values of  $Y$ , and (b) the frequency of training or updating the discriminator. In GANs, the space of possible values for generated samples is restricted by the complexity of the generator network, and, typically, the discriminator is updated in sync with each generator update. In our setting, if we allowed the search space to be the whole Euclidean space without updating or retraining the discriminator, we would be effectively conducting an adversarial attack [73, 74], where the empirical divergence would decrease while the true divergence would not. However, frequent retraining of the discriminator is computationally infeasible for tree-based techniques due to their lack of online training capabilities, necessitating retraining from scratch. Therefore, our discriminator is retrained only once at each iteration within our iterative process. Furthermore, we limit the search space for the possible values of the features in  $Y'$  to a set of proposal values  $B$ , generated from the reference dataset  $X$ .
Initial Imputation. DF-Correct starts by setting the features  $C$  of the query dataset, previously detected as corrupted by DF-Locate, as missing. Initial missing data imputation is then performed with three distinct techniques: KNN, linear regression, and random sampling from the reference dataset  $X_{C}$ . This yields three imputed query datasets. A set of discriminators consisting of binary CatBoost [75] classifiers are trained for each reference and imputed query pair, and the empirical total variation distance is estimated following the same procedure as in DF-Locate. The imputed query dataset providing the lowest empirical divergence is selected as a starting point for the iterative process of DF-Correct. If the initial empirical divergence is already lower than  $\varepsilon = 0.1$ , the correction process is finalized. If not, DF-Correct applies the iterative process described below.
Iterative Process. The reference, the imputed query, and augmented samples (see Section F) are used to train  $k = 2$  CatBoost binary classifiers to discriminate between reference and query datasets.  $k$ -fold splitting of the datasets is used to ensure that each classifier does not see the same sample during training and inference. The set of samples to be corrected  $L$  is obtained by selecting  $|Y| / 2$  samples from  $Y$  with the highest probability of being corrupted (or equivalently, the lowest probability of being from the reference distribution). Then, we generate a set of new feature value proposals  $B = \{b_1, \dots, b_N\}$ , with  $b_j \in \mathbb{R}^{|C|}$ . This set comprises all the feature values within positions  $C$  from the reference dataset  $X$  and the imputed query with linear regression, alongside random permutations of the reference values. In other words, each  $b \in B$  contains  $|C|$  values obtained from  $X$  (and imputed  $Y$ ) that can be a potential replacement for the corrupted features of each query sample  $y \in L$ . Note that the size of  $B$  is proportional to the size of  $X$ . Then, for every query sample  $y \in L$  classified as "corrupted" by the discriminator, we replace the shifted features by all candidate values in  $B$ , and select the  $b_i$  that provides the highest empirical likelihood of being a non-corrupted sample:
$$
b _ {i} = \underset {b \in B} {\operatorname {a r g m a x}} r _ {\theta_ {i}} \left(y _ {i} ^ {(b)}\right) \tag {8}
$$
where  $y_{i}^{(b)}$  is the query sample  $y_{i}$  with the corrupted features  $C$  replaced with the values of  $b$  such that  $y_{iC} = b$ , and  $r_{\theta_i}$  is the likelihood ratio function from the classifier used to process the  $i$ th sample. In other words, we evaluate each sample  $y \in L$  a total of  $|B|$  times with the discriminator and select the  $b$  providing the highest probability that  $y^{(b)}$  is from the reference distribution. Then, we update the current corrected query dataset  $Y^{\prime (k)}$  by replacing the corrupted sample  $y_{i}$  with  $y_{i}^{(b_{i})}$ , such that  $Y^{\prime (k + 1)} = (Y^{\prime (k)} \setminus y_i) \cup y_i^{(b)}$ . By replacing corrupted features with values from the Eq. 8, the empirical divergence is decreased such that  $\hat{D}_{\theta}(X,Y^{\prime (k + 1)}) \leq \hat{D}_{\theta}(X,Y^{\prime (k)})$ . After updating all the
"corrupted" features within the query dataset, the classifiers are retrained and the process is repeated until no divergence is detected or until the empirical divergence stops decreasing. The iterative process of replacing corrupted features, which reduces the empirical divergence, and retraining discriminators, which allows the estimation of more accurate likelihoods and divergences, approximates the minimax optimization process in Eq. 7. More details are in Section F.
6 Experimental Results
Real world datasets. We use multiple datasets including UCI datasets such as Gas [76], Energy [77], and Musk2 [78], OpenML datasets including Scene [79], MNIST [80], and Dilbert [81], datasets with DNA sequences such as Founders [82] and a private Dog DNA dataset (Canine), a subset of phenotypes from the UK Biobank (Phenotypes) [83], Covid-19 data [84], and simple datasets including values generated from Cosine and Polynomial functions. The datasets range from 8 to 198,473 features, and from 1,444 to 70,000 samples, including continuous and categorical datasets. We normalize each feature to have values from 0 to 1. We randomly divide each dataset into two equally sized subsets, corresponding to the reference  $X$  and query  $Y$  samples. Multiple manipulations are applied to randomly selected features of the query dataset. A detailed description of the datasets and the pre-processing applied is available in Section B.1.
Feature Shift Manipulations. We apply 10 different manipulations to the real world datasets in order to generate various distribution shifts. Table 1 describes each manipulation. Manipulations 1, 2, 4, 5, 6, and 7 distort the marginal distributions, with manipulation 4 leaving the mean approximately unchanged. Manipulations 4 and 6 have different levels of strength, controlled by parameters  $\alpha$  and  $\rho$ , respectively. Manipulations 3 and 6 shuffle the feature values across samples, leaving the marginal distributions unchanged ( $p_i = q_i$ ) but changing their correlations. Manipulation 3 performs a different random permutation at each feature, removing all correlation between features, while manipulation 6 performs the same permutation for all features, such that  $p_C = q_C$ , but  $p \neq q$ . Manipulations 9 and 10, applied to continuous and categorical variables respectively, replace the features with values predicted by k-nearest neighbor (KNN) trying to reconstruct the corrupted features. The nature of the shifts originated by KNN will depend on the given dataset and distribution. We discuss the nature of shifts originated by predictive and imputation models in Section D. Manipulations are applied to continuous features, categorical features, or both. In total, we apply 10 manipulations on continuous features and 8 manipulations on categorical features. Each manipulation is applied to  $5\%$ ,  $10\%$ , and  $25\%$  of the features in the query. This produces a reference dataset with 24 and 30 query datasets for categorical and continuous data, respectively. Each query dataset corresponds to a distinct transformation applied to a specific fraction of the features.
Table 1: Manipulation types applied to continuous and/or categorical features.
Type	Mapping	Description	Shift	Data
1	x ~ Uniform(0, 1)	Each value is substituted by a random number between 0 and 1.	pi ≠ qi	Cont.
2	1 - x	Each value is negated.	pi ≠ qi, E[qi] = 1 - E[pi]	Both
3	PiXi	Pi is a random permutation matrix applied to feature i.	pi = qi, pC ≠ qC, qC = ∏i∈C qi	Both
4.1-4.3	clamp0,1(x + ασ) σ ~ Rademacher(0.5)	Add constant noise with a random sign. α ∈ {0.02, 0.05, 0.1} for 4.1-4.3 respectively.	pi ≠ qi, E[pi] ≈ E[qi]	Cont.
5	round(x)	Values are binarized.	pi ≠ qi	Cont.
6.1-6.3	b(1 - x) + (1 - b)x b ~ Bernoulli(ρ)	Values are negated with probability ρ ∈ {0.2, 0.4, 0.6} for 6.1-6.3 respectively.	pi ≠ qi, E[qi] = ρ + (1 - 2ρ)E[pi]	Cat.
7	MLP(x)	Forward through an MLP with min-max normalization or binarization.	pi ≠ qi	Both
8	PXi	P is a random permutation matrix applied to all features simultaneously.	pi = qi, pC = qC, p ≠ q	Both
9	KNN(x)	Predict feature with KNN (Regressor).	-	Cont.
10	KNN(x)	Predict feature with KNN (Classifier).	-	Cat.
Probabilistic Simulations. We generate 15 datasets with 1000 features and 5000 samples, simulated from probabilistic distributions including multivariate Gaussians, multivariate Bernoulli distributions, Gaussian mixture models, and Bernoulli mixture models. By having full access to the generating distributions, the real distribution shift and divergence between distributions can be measured. The manipulations include marginal shifts in the mean and/or variance between distributions and shifts in
the feature correlations. Note that here we do not apply the shifts described in Table 1, but instead directly simulate datasets from distributions having a shift. See Section B.2 for more details.
Experimental Details. In both shift localization and correction, each method only has access to the reference dataset  $X$  and the corrupted query dataset  $Y$ , while ground truth information, such as actual corrupted feature locations  $C^*$  or the original (pre-shifted) query dataset  $Y^*$ , is not accessible. The localization task is evaluated by comparing the localized corrupted features  $C$  and the true locations  $C^*$  with the F-1 score, and the correction task is evaluated by comparing the corrected query dataset  $Y'$  and the reference dataset with non-parametric empirical divergence estimators. We use the simulated datasets to perform hyperparameter search for DataFix and all competing methods, while the real datasets are used as a hold-out testing set (see below and Section G for more details).
Feature Shift Localization. We evaluate our method and 8 competing techniques, including four feature shift localization methods (MB-SM, MB-KS, KNN-KS, and Deep-SM) and four feature selection methods (MI, selectKbest, MRMR, and Fast-CMIM) (Fig. 3). We evaluate MB-SM, MB-KS, KNN-KS, and Deep-SM with both their recommended configuration, without a priory specification of the number of corrupted features  $|C|$ , and with the hyperparameter configuration that yielded optimal results and includes the ground truth  $|C|$  (with *). The rest of the methods do not have access to the ground truth  $|C|$ . We measure the F-1 score of feature shift localization and average it across the percentage of manipulated features and manipulation types. Figure 3 (left) shows the median and mean F-1 scores across the real and simulated datasets, respectively. We make use of the median because some F-1 scores are missing as some techniques can not process the larger datasets such as "Founders" and "Canine" given the assigned time budget of 30h. DataFix outperforms all the competing techniques, in both real and simulated datasets, even when compared with MB-SM, MB-KS, KNN-KS, and Deep-SM that make use of the ground truth  $|C|$ .
Figure 3 (right) shows the median (left) and mean (right) F-1 scores divided by type of manipulation and datasets. DataFix outperforms all competing methods in all types of manipulations, except for type 10, where selectKbest provides a higher F-1 score. Techniques using univariate tests, MRMR and Fast-CMIM provide good results for manipulations causing marginal distribution shifts, but completely fail when facing manipulations affecting feature correlations (manipulations 3 and 8), while conditional testing-based techniques (MB-SM, MB-KS, KNN-KS, and Deep-SM) and DataFix are able to detect them. DataFix outperforms competing methods in most of the real datasets, with an overall lower F-1 score for datasets with a larger number of dimensions (phenotypes, founders, and canine).
Figure 3: (left) median and mean F-1 score of real and simulated datasets, and (right) median and mean F-1 score across manipulation types and datasets. Higher is better.
Feature Shift Correction. We evaluate our method and 13 competing techniques, including predictive models like KNN, linear regression (LR), and multilayer perceptron (MLP), imputation methods including HyperImpute, ICE, MIRACLE, MissForest, and SoftImpute, optimal transport methods such as Sinkhorn, adversarial learning methods like GAIN, and domain adaptation techniques including deep destructors (DD) and iterative alignment flows (INB). We perform a manipulation-agnostic evaluation, where we set the corrupted features to 0 (or missing) and treat all manipulation types equally. Both the reference and the query are provided to each method, and non-parametric statistical metrics computed between the reference and the updated (corrected) queries are used to evaluate each method's performance. Specifically, we use the Wasserstein distance  $(W_2^2)$  as in [48], the empirical estimator of the Henze-Penrose divergence  $(D_{hp})$  [85, 86, 87], and the empirical
estimator of the symmetric Kullback-Leibler divergence  $(D_{skl})$  [88]. We report the metrics after subtracting the "background" divergences computed with the reference and query datasets previous to any manipulation. Figure 4 shows the median and mean metrics for the real and simulated datasets respectively. DataFix is able to provide a corrected query dataset  $Y'$  with the lowest empirical divergences. Despite their simplicity, KNN and linear regression provide competitive results, followed by MLP, HyperImpute, ICE, INB, and Sinkhorn. See Section J for more details.
Figure 4:  $W_2^2$ ,  $D_{hp}$ , and  $D_{skl}$  of real and simulated datasets. Lower is better.
DataFix Analysis. Figure 5 shows the iterative process of DF-Locate before and after performing shift correction, in simulated dataset 12 with 200 corrupted features out of 1000 (see Section B.2 for more details). Fig. 5 (left) shows the total variation distance estimated by the random forest (blue) which lower-bounds its ground truth monte carlo estimate (black) as corrupted features are detected and removed. The F-1 detection score increases until all features are detected and the iterative process is stopped. Fig. 5 (right) shows the iterative process applied to the corrected query with different methods. ICE and DD provide an updated query that leads to a lower empirical divergence, while the other methods provide an updated query that increased the shift instead of reducing it. DF-Correct (Purple) provides an accurately corrected query, with no empirical divergence detected by DF-Locate.
Figure 5: DF-Locate iterative process before and after shift correction.
Extended experimental results are present in the appendix, including an analysis of the method's computational time (Section H), the effect of the classifiers used in DataFix (Section K), experimental results of using corrected datasets in downstream classification and regression (Section M), a detailed division of the quantitative results (Sections I and J), and a discussion of limitations (Section O).
In this paper, we introduced a new framework "DataFix" which makes use of tree-based classifiers, combined with iterative heuristics, to localize and correct feature shifts. The system, inspired by adversarial learning and feature selection frameworks, is able to accurately detect and correct a wide range of distribution shifts in many types of datasets, surpassing existing techniques.
This work was partially supported by a grant from the Stanford Institute for Human-Centered Artificial Intelligence (HAI) and by NIH under award R01HG010140. This research has been conducted using the UK Biobank Resource under Application Number 24983.
[103] J. Qian, Y. Tanigawa, W. Du, M. Aguirre, C. Chang, R. Tibshirani, M. A. Rivas, and T. Hastie, "A fast and scalable framework for large-scale and ultrahigh-dimensional sparse regression with application to the uk biobank," PLoS genetics, vol. 16, no. 10, p. e1009141, 2020.
[104] W. H. Press and S. A. Teukolsky, "Savitzky-golay smoothing filters," Computers in Physics, vol. 4, no. 6, pp. 669-672, 1990.
[105] V. Satopaa, J. Albrecht, D. Irwin, and B. Raghavan, "Finding a" kneedle" in a haystack: Detecting knee points in system behavior," in 2011 31st international conference on distributed computing systems workshops, pp. 166-171, IEEE, 2011.
A Extended Related Work
The task of feature shift localization and correction is connected to various areas of statistics and machine learning, including out-of-distribution (OOD) generalization, outlier detection, and Generative Adversarial Networks (GANs), among many others.
Out-of-distribution Generalization. Similar to feature shift localization and correction, OOD generalization addresses the issue of divergence between datasets or data sources. It does so by aiming to enhance the test (target) accuracy in settings where the test data significantly deviates from the training (source) data distribution. This divergence is often a consequence of the target domain being distinct from the source domain used for training the machine learning models. While both shift localization/correction and OOD generalization aim to reduce distribution shifts, OOD generalization has a primary focus on the task of domain adaptation, setting it apart from the primary objectives of feature shift localization and correction. Extensive work on OOD generalization includes studies into the conditions under which a classifier trained on source data can be expected to perform effectively on target data [89], learning invariant representations for domain adaptation [90], generalizing to unseen domains via distribution matching [91], and using neural networks to learn representations that possess discriminative qualities for the main learning task while remaining indiscriminate concerning the domain shift [92]. The latter goal partially aligns with the objective of feature shift localization and correction, excluding conditional shift, where class-conditional distributions of input features change between source and target domains.
Anomaly Detection. Anomaly detection finds observations that deviate noticeably from the data distribution. It has widespread applications in fault detection, particularly in industrial processes, where training samples with normal patterns are used to identify operations that deviate from those. The field of anomaly detection has seen extensive research, including machine learning approaches [93] and deep learning techniques [94, 95, 96], particularly auto-encoders [91, 1].
Generative Adversarial Networks. GANs have gained extensive traction in the imputation of missing data, facilitating the creation of realistic fake data through adversarial training. Among these, GAIN [49] stands out as one of the most renowned works, serving as the foundation for subsequent advancements. These improvements include data augmentation, a feature enhanced by MisGAN [97], architectural and loss modifications exemplified by GAMIN [98], SGAIN, WSGAIN-CP, and WSGAIN-GP [99], the incorporation of implicit label information as demonstrated by PC-GAIN [100], the ensemble of GANs by MI-GAN [101] and the capability to combine global and local information demonstrated by GAGIN [102].
B.1 Real Datasets