SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In order to perform the refinement step, we store all intermediate steps so that we can revisit the iterative filtering process and select the optimal stopping point. At each iteration, we store the indexes of the features selected as corrupted, their corresponding estimated  $\hat{D}$ , and the number of removed features. While the ideal stopping iteration would be the iteration providing the highest F-1 score, this requires access to the ground truth, which is unavailable in practice. However, there exists a point at which the cost of removing additional features, some potentially non-corrupted, outweighs the decrease in the empirical divergence. To determine the optimal iteration, we locate the elbow or knee from a processed curve depicting the empirical total variation distance (TVD) as a function
Figure 8: Sorted and normalized feature importance scores of the discriminators in the first, middle, and last iterations of DF-Locate applied to the Dilbert dataset with manipulation type 8. Each row represents a different threshold used to localize and filter the corrupted features at each iteration.
Figure 9: Number of features removed at each iteration for different threshold  $\tau$  values in DF-Locate applied to the Dilbert dataset with manipulation type 8.
of the total number of removed features (Figure 10). Once such optimal iteration is determined, all selected features up to that iteration are flagged as corrupted. The curve of the true divergence would always exhibit a convex non-increasing shape but due to the intrinsic randomness in the training and evaluation process of the discriminator, this ideal shape is not always achieved, making the task of locating the knee challenging. In order to make the curve smooth and non-increasing, we make use of the Savitzky-Golay filter [104] due to its ability to remove noise without distorting the underlying signal. Additionally, we apply an opening operation to eliminate any local maxima in the curve, ensuring that each point is equal to or smaller than its left neighbor. Furthermore, we process the initial iterations of the curve to enforce a strict decreasing behavior. After processing the empirical divergence curve, we employ the knee locator method introduced by [105] to identify the optimal stopping criterion.
The Savitzky-Golay filter relies on two main parameters: the length of the filter window and the polyorder used for fitting the samples. We define the window length as  $\max(5, 2[\zeta \delta / 2] + 1)$ , where  $\zeta$  is chosen from the set  $\{1, 2, 3, 5, 7\}$  through cross-validation, and  $\delta$  represents the average number of removed features at each iteration. We explore different polyorders from the set  $\{3, 4\}$ . Based on experiments with simulated datasets, we select  $\zeta = 2$ , and the best polyorder as 4. The knee locator involves two primary parameters: sensitivity  $(S)$  and online mode. The sensitivity parameter determines the number of "flat" points we anticipate encountering in the original data curve before identifying a knee, while the online mode enables the correction of previous knee values. We explore different values for  $S$  from the set  $\{1, 3, 5, 7\}$ . Our experimentation reveals that the optimal values are  $S = 5$  and using the offline mode.
Figure 10 shows an example of the knee location, used to refine the selection of features. The F-1 score can serve as a way to evaluate the selected stopping point, showing the trade-off between reduced TVD and the number of removed features.
Figure 10: Knee location based on distribution shift. Number of removed Features, TVD, and F-1 score in shift localization task.
F DF-Correct
Figure 11: DF-Correct diagram.
DF-Correct (Section 5, Figures 2 and 11, and Algorithm 2) is the proposed method within DataFix that corrects the feature shifts by changing the values of the features in  $Y$  originating the distribution shift through an iterative approach. Given the set of corrupted features  $C$ , DF-Correct tries to generate a new query dataset  $Y'$ , such that  $Y_{\overline{C}} = Y_{\overline{C}}'$ , while  $D(X,Y') < D(X,Y)$ . DF-Correct starts by obtaining an initial candidate of  $Y'$  by setting the values within the subset  $C$  of  $Y$  as missing and performing imputation with linear regression and KNN. Furthermore, a naive initial candidate is generated by replacing the corrupted values of  $Y$  with randomly selected values of  $X$  (restricted to the features subset  $C$ ). Note that this naive initialization already fixes distribution shifts where features are completely independent of each other. The set of the 3 initial candidates,  $V = \{Y^0,Y^1,Y^2\}$ , is evaluated by computing the empirical total variation distance with a classifier, and the one providing the lowest empirical divergence is selected as the initial corrected query  $Y' = \operatorname*{argmin}_{Y^i\in V}\hat{D}_\theta (X,Y^i)$ . If  $\hat{D}_{\theta}(X,Y') < \epsilon$ , the correction process is finished and  $Y'$  is returned. Otherwise, if  $\hat{D}_{\theta}(X,Y') > \epsilon$ , an iterative process where some samples of  $Y'$  are modified is performed.
The iterative process tries to find new values of  $Y' = \{y_1', y_2', \dots, y_{N_y}'\}$  that reduce the empirical total variation distance:
$$
\hat {D} _ {\theta} ^ {T V} \left(X, Y ^ {\prime}\right) = \frac {1}{N _ {x}} \sum_ {i = 1} ^ {N _ {x}} g \left(r _ {\theta} \left(x _ {i}\right)\right) - \frac {1}{N _ {y}} \sum_ {j = 1} ^ {N _ {y}} g \left(r _ {\theta} \left(y _ {j} ^ {\prime}\right)\right) \tag {15}
$$
with  $g(u) = \frac{1}{2} \mathrm{sign}(u - 1)$ . Because the values of  $X$  are left untouched, this becomes equivalent to solving:
$$
Y ^ {\prime} = \underset {Y} {\operatorname {a r g m i n}} \underset {\theta} {\max } \sum_ {j = 1} ^ {N _ {y}} - g \left(r _ {\theta} \left(y _ {j} ^ {\prime}\right)\right) = \underset {Y} {\operatorname {a r g m a x}} \underset {\theta} {\max } \sum_ {j = 1} ^ {N _ {y}} g \left(r _ {\theta} \left(y _ {j} ^ {\prime}\right)\right) \tag {16}
$$
In order to perform such an optimization process, a set of classifiers are trained using  $X$ ,  $Y'$ , and data augmentation consisting of performing random permutations within the features of the reference  $X$  dataset, in order to generate extra samples of "corrupted" sequences. After training the classifiers, the next step is to find which samples need to be corrected. Note that when  $p = q$ , we have  $\mathbb{E}[\hat{D}_{\theta}^{TV}(X,Y)] = 0$ , and:
$$
\mathbb {E} \left[ \sum_ {j \in N _ {y}} g \left(r _ {\theta} \left(y _ {j}\right)\right) \right] = \mathbb {E} \left[ \sum_ {m \in N ^ {+}} g \left(r _ {\theta} \left(y _ {m}\right)\right) \right] - \mathbb {E} \left[ \sum_ {n \in N ^ {-}} g \left(r _ {\theta} \left(y _ {n}\right)\right) \right] = 0 \tag {17}
$$
where  $N^{+} = \{i : r_{\theta}(y_i) > 1\}$  are the indices of samples classified as positive by the discriminator, and  $N^{-} = \{i : r_{\theta}(y_i) < 1\}$  are the indices of samples classified as negative. Furthermore, both sets have, in expectation, the same size  $\mathbb{E}[|N^{+}|] = \mathbb{E}[|N^{-}|] = \mathbb{E}\left[\frac{|N_y|}{2}\right]$  when  $p = q$ . This indicates that when both distributions are equal, a discriminator will approximately classify half of the samples as positive and half as negative. Therefore we only correct the set of samples  $L$ , including up to  $\frac{|N_y|}{2}$  samples with the highest probability of being corrupted:
$$
L = \left\{i: r _ {\theta_ {i}} \left(y _ {i}\right) <   r _ {\theta_ {i + 1}} \left(y _ {i + 1}\right) <   1 \right\} \tag {18}
$$
with  $|L| \leq \frac{|N_y|}{2}$ . Next, we construct a set of feature value proposals  $B$  that will replace the corrupted features. This proposal set is constructed by including within  $B$  all the feature values of the reference  $X$ , of the initial imputed candidates  $V$ , of the current corrected query  $Y'$ , and random permutations of  $X$ . Then, for every sample  $i \in L$ , each of the proposals  $b \in B$  is placed as an alternative to the corrupted features, generating a sample  $y_i^{(b)}$ , where  $y_i^{(b)}_{C} = b$ , and  $y_i^{(b)}_{\overline{C}} = y_i \overline{C}$ . The proposal providing the highest probability of being "non-corrupted" is selected:
$$
b _ {i} = \underset {b \in B} {\operatorname {a r g m a x}} r _ {\theta_ {i}} \left(y _ {i} ^ {(b)}\right) \tag {19}
$$
Finally, the updated sample  $y_{i}^{(b)}$  is placed inside the corrected query  $Y'$ . After updating all samples in  $L$ , the divergence is computed again, and if  $\hat{D}_{\theta}(X,Y') > \epsilon$ , the process is repeated for a number of epochs. Typically, the number of epochs is set to 1 or 2, as the correction process can become computationally intensive for large datasets (see following sections).
F.1 MNIST Example Proposals
In Figure 12, we provide a visual comparison of proposals generated from a single iteration of DF-Correct using two image samples from the MNIST dataset, categorized from worst to best based on their likelihood of being corrupted. The comparison is made against their manipulated counterparts with manipulation type 2. The findings reveal that the top proposals generated by DF-Correct excel at correcting the distorted values. The outcomes closely resemble unaltered MNIST images, eliminating any distortions effectively.
F.2 Relationship between DF-Correct and KNN Imputation
We can provide more insight into how DF-Correct is able to reduce the divergence between datasets by comparing it with KNN imputation. A KNN imputer fills the missing feature values with a combination of the top-k closest samples with respect to L2 distance on non-missing features. Instead, our correction system replaces the corrupted feature values with the top-1 candidate with respect to the probability estimated by the discriminator. In a sense, the KNN training samples are replaced by our feature proposals, and the KNN L2 distance is replaced by the prediction of a discriminator, making the correction task similar to a divergence-reducing imputation problem.
Figure 12: Comparison of worst, intermediate, and best proposals from a single iteration of DF-Correct on two samples of the MNIST dataset against their corrupted version with manipulation type 2.
G Experimental Details
G.1 Hyperparameter Search
In this section, we present the experimental details of our hyperparameter optimization process for DataFix, along with the benchmarking methods used for feature shift localization and correction. The hyperparameter optimization involves conducting a grid search on the simulated datasets, with the provision that experiments exceeding a 30-hour runtime are excluded. For parameters left unspecified, we initialize them with their default values.
Table 4 outlines the search space for each parameter tuned in each detection benchmarking method, as well as their respective optimal value. Similarly, Table 5 provides insight into the search space for each parameter tuned in each correction benchmarking method, as well as the values that yielded optimal results.
DF-Locate and DF-Correct both undergo optimization with respect to the choice of classifier employed as a discriminator, as detailed in Section K. Additionally, we conduct a hyperparameter search for DF-Locate, focusing on variables such as window length and polyorder, which is further elaborated upon Section E.3.
Table 4: Search space and optimal values for tuned parameters in detection benchmarking methods.
Method	Tuned Parameter	Search Space	Optimal
MI	threshold	[0.008, 0.01, 0.02, 0.05, 0.1]	0.02
selectKbest	significance_level	[0.008, 0.01, 0.02, 0.05, 0.1]	0.01
MB-SM, MB-KS	n Expectation	[30, 100, 500, 10000]	30
KNN-KS, Deep-SM	niestream_run	[1, 5, 10, 25, 50, 250]	250
G.2 Evaluation Setting
In both shift localization and correction, each method only has access to the reference dataset  $X$  and the corrupted query dataset  $Y$  as input for the iterative process of training discriminators and computing heuristics. Note that the ground truth information, including the actual corrupted feature locations  $C^*$  and the original (pre-shifted) query dataset  $Y^*$ , are not used by DataFix at any step when predicting the localization of corrupted features or when correcting the query dataset.
Table 5: Search space and optimal values for tuned parameters in correction benchmarking methods.
Method	Tuned Parameter	Search Space	Optimal
KNN	k	[10, 25, 50, 100]	10
DD	n	[10, 50, 100, 200]	100
INB	n	[50, 100, 200, 500]	200
ndim	[30, 60, 90]	90
MLP	activation	['relu', 'tanh']	'relu'
hidden_layer_sizes	[100, 1000]	100
alpha	[0.0001, 0.001]	0.0001
batch_size	['auto', 64]	'auto'
learning_rate_init	[0.0001, 0.001]	0.001
GAIN	batch_size	[16, 64, 256]	16
n_epochs	[1000, 50000]	50000
hint_rate	[0.8, 0.9]	0.9
HyperImpute	baseline_imputer	[0, 1, 2]	1
optimize thresh	[1000, 5000]	5000
n_inner_iter	[40, 80]	40
ICE	max_iter	[1000, 2000]	1000
initial_strategy	[0, 1, 2]	1
MIRACLE	lr	[0.0001, 0.001]	0.0001
batch_size	[1024, 512]	1024
n Hidden	[16, 32, 64]	16
reg_lambda	[0.1, 1, 10]	10
reg_beta	[1, 3]	3
window	[10, 20]	20
MissForest	n_estimators	[10, 20, 50]	10
max_iter	[100, 500, 1000]	100
Sinkhorn	lr	[0.001, 0.01]	0.01
n_epochs	[500, 1000]	1000
batch_size	[256, 512]	512
noise	[0.001, 0.01]	0.001
scaling	[0.9]	0.9
SoftImpute	maxit	[1000, 2000]	1000
convergence_threshold	[0.00001, 0.0001]	0.0001
max_rank	[2, 3]	3
shrink_lambda	[0.5, 0]	0
H Computational Time
Large and high-dimensional datasets are becoming the norm, therefore, methods that detect and correct feature shifts should be able to properly scale with respect to the number of samples and features.
Figure 13 presents the average computational time of feature shift localization benchmarking methods as a function of the product between the number of samples and features for each dataset. MI and selectKbest stand out as the fastest methods (using the Chi-square test for categorical datasets and ANOVA-F test for continuous datasets). MRMR and FAST-CMIM, although performing adequately in terms of speed for small datasets, encounter challenges in scaling with larger dataset sizes. Consequently, they fail to produce results within the 30-hour time limit for Founders and Canine datasets. Furthermore, the feature-shift detection techniques KNN-KS and, particularly MB-KS, demonstrate significantly slower performance, rendering them incapable of delivering results within the given time constraint for Founders, Canine, Dilbert, and Phenotypes datasets. DF-Locate proves to be a reasonably efficient method, exhibiting good scalability as the dataset size increases, while providing the best localization performance.
Figure 13: Computational time for shift localization methods based on real dataset size.
In Figure 14, we conduct a comparative analysis of the computational time for shift correction methods. Although DF-Correct is not faster than simpler techniques like median or linear regression, its speed surpasses several competing methods, such as MIRACLE and HyperImpute. Furthermore, DF-Correct exhibits reasonable runtime even for the largest high-dimensional Canine dataset, successfully correcting the distribution shift within the 30-hour time limit. Many of the competing methods were unable to provide results for the Canine dataset due to their excessive time complexity and/or memory requirements. Therefore, DF-Correct provides the best correction in terms of distribution shifts, while still providing competitive or faster speeds than competing methods.
All experiments were done with an Intel Xeon Gold with 12 CPU cores.
Figure 14: Computational time for shift correction methods based on real dataset size.
I Extended Feature Shift Localization Results
We evaluate the performance of DF-Locate across different fractions of manipulated features:  $5\%$ ,  $10\%$ , and  $25\%$ . The F-1 scores, depicted in Figure 15, were computed by averaging across manipulation types and taking the median across real datasets for each feature shift localization method. Our findings show that DF-Locate obtains consistently higher performance, irrespective of the fraction of manipulated features involved. However, it is important to note that this is not the case for other methods, such as MRMR and FAST-CMIM, which exhibit limitations in their localization capabilities, particularly when there are only a few corrupted features. Note that when a small percentage of manipulated features is present, a smaller amount of features need to be localized as corrupted. However, this could also lead, in some cases, to lower distribution shifts, making the detection of such shifts more challenging. On the other hand, the presence of a larger percentage of manipulated features can make the localization task more challenging, while the empirical detection of the presence of the shift can be, in some cases, easier.
Figure 16 shows the mean F-1 score of DF-Locate and various shift localization methods applied to the real datasets. The symbol  $\mathbf{\dot{x}}$  denotes experiments that are missing due to exceeding the time limit of 30 hours. Note that most methods fail to process high-dimensional datasets such as Founders and Canine, whereas DF-Locate is able to provide accurate results while scaling to large datasets. DF-Locate consistently exhibits superior performance compared to all benchmarking methods across most datasets. The only exception observed is with the Phenotypes dataset, where selectKBest slightly outperforms DF-Locate in locating the corrupted features. Nevertheless, for the remaining datasets, DF-Locate surpasses all competing approaches and is only slightly surpassed or equaled in performance on a few occasions by Deep-SM  $(^{*})$  or KNN-KS  $(^{*})$ , both of which use the ground truth  $|C|$ . Additionally, it is worth noting that MI always outperforms selectKbest on datasets with continuous features, while the opposite holds true for datasets with categorical features.
Figure 15: Median F-1 scores of shift localization methods by fraction of manipulated features on real datasets.
Figure 16: Mean F-1 scores of shift localization methods by real datasets. 'x' indicates missing experiment. Higher is better.
Figure 17 shows the median F-1 scores of DF-Locate and other competing methods, categorized by the feature manipulation type applied to the real datasets. The average F-1 score is computed across fractions of manipulated features, followed by the computation of the median F-1 score across different datasets. The symbol  $\mathbf{\dot{x}}^{\prime}$  is used to indicate missing experiments for MB-KS and manipulation types 6.1-6.3 and 10. These manipulations are applied to categorical datasets only (Phenotypes, Founders, Canine), and the MB-KS method did not yield results for any of these three datasets within the specified time constraint of 30 hours. Manipulations involving shifts caused
by the correlation between features (manipulations 3 and 8) are not detected by methods such as MI, selectKbest, MRMR and Fast-CMIM, while being accurately detected by MB-SM, KNN-KS, Deep-SM, and our proposed method. Note that while manipulation type 8 breaks the theoretical equivalence between feature selection and feature shift localization problem (see previous sections), it is still accurately localized by DataFix. Manipulation 10, consisting of replacing feature values with the ones predicted with a KNN, is the most challenging to detect by DataFix. Note that such manipulation can lead, in some scenarios, to small or undetectable distribution shifts, as KNN provides perfect predictions when its training dataset size goes to infinity.
Figure 17: Median F-1 scores of shift localization methods by feature manipulation type on real datasets. 'x' indicates missing experiment. Higher is better.
Figure 18 displays the mean F-1 scores of shift localization methods across the simulated datasets. As observed in the results presented for real datasets in Figure 16, DF-Locate consistently outperforms or matches all competing methods, except for datasets 7 and 8, where MB-SM and Deep-SM obtain a higher F-1 localization score when using  $|C|$  as extra ground truth information. Note that in a fair comparison where  $|C|$  is not used, MB-SM and Deep-SM perform poorly (see Figure 3). DF-Locate obtains lower F-1 scores in simulated datasets 5, 7, and 8 compared to the other datasets, which involve shifts originated by a mismatching correlation between distributions.
Figure 18: Mean F-1 scores of shift localization methods on simulated datasets. Higher is better.
J Extended Feature Shift Correction Results
Figures 19 and 20 provide a comprehensive evaluation of the performance of DF-Correct and competing shift correction methods across real datasets, by using 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].  $W_2^2$ ,  $D_{hp}$ , and  $D_{skl}$  are non-parametric estimators of distances between the underlying distributions  $p$  and  $q$  that make use of the reference  $X$  and query  $Y$  datasets. By computing statistics based on distances between samples within and between datasets, these metrics estimate the distance between distributions. The reported metrics provide empirical estimates of the divergences between the corrected query datasets and the reference datasets. Note that first, the background empirical divergences between the reference and query dataset (prior to any manipulation) are computed and subtracted from the reported metrics. For simulated datasets where there are no query datasets prior to manipulation, a second reference dataset is used to obtain background empirical divergences.
As shown in Figures 19 and 20, DataFix outperforms all other methods by a significant margin for most datasets, demonstrating its high effectiveness in providing corrected query datasets that are close to the reference distribution. Following DataFix, simpler methods like KNN and linear regression demonstrate competitive performance for most datasets, while techniques such as ICE, HyperImpute, INB, Sinkhorn, and MLP yield favorable results for specific datasets. Furthermore, DataFix consistently achieves the lowest  $W_2^2$  values across most datasets, with exceptions being MNIST and Phenotypes. Specifically, for the MNIST dataset, KNN, MLP, GAIN, HyperImpute, and ICE surpass DataFix achieving the best  $W_2^2$  value of 0. Note that for high-dimensional datasets such as Dilbert, Phenotypes, Founders, and Canine, the  $W_2^2$  metric saturates to 0 for multiple methods, making the other metrics a better alternative to compare the quality between methods.
Figure 19:  $W_2^2$ ,  $D_{hp}$ , and  $D_{skl}$  of shift correction methods on real datasets. Lower is better. (Part 1)
Figure 20:  $W_{2}^{2}$ ,  $D_{hp}$ , and  $D_{skl}$  of shift correction methods on real datasets. Lower is better. (Part 2)
K Classifier Analysis for Localization and Correction
Figure 21 provides the F-1 score results for DF-Locate when using different classifiers as discriminators within the iterative process on both real and simulated datasets. Tree-based methods including Random Forest (RF), CatBoost, ExtraTree, and LightGBM (LGBM) provide highly similar results, with high F-1 scores, surpassing linear models such as logistic regression (LogReg) and a support vector classifier (SVC). We selected RF as our discriminator as it provided much faster training times while being highly competitive in localization accuracy.
Figure 21: Mean F-1 scores by DF-Locate using different classifiers on simulated (left) and real (right) datasets. Higher is better.