Title: Adversarial Learning for Feature Shift Detection and Correction
Venue: NeurIPS 2023
Accepted: yes
Keywords: feature shift detection

distribution shift

shift

data-centric AI
Abstract: 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.

Adversarial Learning for Feature Shift Detection and Correction

Abstract

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 .

1 Introduction

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

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.

2 Related Work

̸

Distribution Shift Detection and Localization. Distribution shift detection involves identifying if p = 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 KolmogorovSmirnov (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 Fisherdivergence 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], ANOV A-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].

3 Proposed Framework

Definition 1. [Feature Shift] We are given two sets of d -dimensional samples X = { x 1 , x 2 , ..., x N } and Y = { y 1 , y 2 , ..., y N } , with x i , y i ∈ R d , from distributions p and q , respectively. A feature shift between p and q occurs when D ( p, q ) > ε and D ( p S , q S ) ≤ ε , 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 ∪ C | = d , and p S and q S are the distributions restricted to S .

Wewill 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 ε = 0 , D ( p S , q S ) = 0 , D ( p, q ) ≫ 0 , and | S | > | C | ≥ 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 ) > ε , where p i and q i represent the marginal distribution of the i th dimension resulting from additive and non-linear

Filtered

reference

Iterative

Discriminator

Training

Reference

Query

Feature

Importances

Distribution

Distribution

Yes

Filtered

Query

Filtered

Query

No

Delete

features

Evaluation of

Feature Shift Detection

Classifier

Feature

removal policy

function

If

transformations; correlation shifts where D ( p, q ) > ε but D ( p i , q i ) ≤ ε for all i ; and correlation shifts where D ( p S , q S ) ≤ ε and D ( p C , q C ) ≤ ε but D ( p, q ) > ε . 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 C = S , that satisfies D ( p C , q C ) ≤ ε , that is C = argmin D ( p C ,q C ) ≤ ε | 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 , ..., y ′ N } such that y ′ i ∼ q ′ and D ( p, q ′ ) ≤ ε , 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 ′ = ϕ ( y ) , prediction or imputation methods y ′ C = ϕ ( y C ) , or optimization procedures or heuristics Y ′ = Φ( 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]:

where X = { x 1 , ..., x N } and Y = { y 1 , ..., y N } , with x i ∼ p , y j ∼ q , and r θ ( x ) is a function approximating the likelihood ratio between p and q , obtained by training a binary classifier D θ ( x ) such that r θ ( x ) = D θ ( x ) / 1 -D θ ( x ) . We use 5 -fold train-evaluation such that 80% of the samples are used to train a random forest binary classifier D θ ( 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 ) = E x ∼ p,y ∼ q [ ˆ D θ ( X,Y )] when r θ ( 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 ≤ D TV ( p, q ) ≤ 1 , defined by setting f ′ ( u ) = ( f ∗ ◦ f ′ )( u ) = 1 2 sign ( u -1) . The discrete nature of f ′ ( 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 ˆ D TV ( X,Y ) > 0 , or equivalently, the test balanced accuracy is larger than 0 . 5 (random chance), it might indicate that p = 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:

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 ∼ p and z | t =1 ∼ 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 :

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 = argmax | C |≤ k I ( z C ; t ) = argmin D TV ( p C ,q 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: β = F ( Z, T ) , where β ∈ R d and F ( · ) 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 D θ as the feature importance scores.

Feature removal policy. We introduce the feature removal policy function ψ as an heuristic to approximate Eq. 2, which utilizes the predicted feature importances β and the estimated total variation distance ˆ D = ˆ D TV ( X,Y ) to select likely corrupted features C = ψ ( β, ˆ D ) . First, β is normalized β ′ = | β | ∑ d -1 i =0 | β i | and sorted β ′ π , such that β ′ π (0) ≥ β ′ π (1) ≥ ... ≥ β ′ π ( d -1) . Given the cumulative sum γ ( k ) = ∑ k j =0 β ′ π ( j ) , the sorted features from 0 to J are selected as corrupted, where J is the smallest index such that γ ( k ) ≥ τ ˆ D (Eq. 4), and τ is a hyperparameter set by hyperparameter optimization. Finally, ψ returns the features from 0 to J with scores higher than 1 d (Eq. 5):

Note that τ ˆ D acts as a threshold to select how many features are selected as corrupted. If the threshold is small, ψ simply returns the feature with the highest feature importance, while a larger threshold will make ψ return more features. By defining the threshold as the product of τ and ˆ D , we ensure that when the shift is small (low ˆ 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 θ ( i ) = argmax θ ˆ D TV ( X ( i ) , Y ( i ) ) consisting of random forest binary classifiers. The classifier predictions are used to estimate ˆ D ( i ) = ˆ D TV ( X ( i ) , Y ( i ) ) and, combined with the Gini importance from the random forests β ( i ) = F ( X ( i ) , Y ( i ) ) , a set of corrupted features are localized using the feature removal policy C i = ψ ( β ( i ) , ˆ D ( i ) ) . The selected features C i are removed from the dataset, such that X ( i +1) = X ( i ) C i and Y ( i +1) = Y ( i ) C i , and the process is repeated while ˆ 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 ′ = l ⋃ i =1 C i . Finally, a refinement stage (see below and Section E.3) predicts the final set C . More details are provided in Section E.

Refinement stage. DF-Locate stores all intermediate steps, including the indexes of the predicted corrupted feature locations and the estimated ˆ 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 2 DF-Correct

1:

Inputs:

X

;

▷

Reference

Y

;

▷

Query

C

;

▷

Corrupted Features

ϵ

;

▷

Divergence Threshold

2:

V

=

{

Y

0

, Y

1

, Y

2

Impute

(

X,Y,C

)

3:

Y

′

argmin

Y i

∈

V

ˆ

D

θ

(

X,Y

i

)

4:

if

ˆ

D

θ

(

X,Y

′

)

< ϵ

then

5:

return

Y

′

6:

end if

7:

for

do

8:

θ

Train(

X,Y

′

)

▷

Train discriminators

9:

L

DetectIncorrect

(

D

θ

(

Y

′

))

▷

Detect samples that

require feature correction

10:

B

GenerateProposals(

X,Y

′

, V, C

)

11:

for

i

∈

L

do

12:

b

i

argmax

b

∈

B

r

θ i

(

y

(

b

)

i

)

▷

Find best proposal

13:

Y

′

update

(

Y

′

, y

(

b i

)

i

)

▷

Dataset with updated

sample

14:

end for

15:

if

ˆ

D

θ

(

X,Y

′

)

< ϵ

then

16:

return

Y

′

17:

end if

18:

end for

19:

return

Y

′

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 ′ ∼ q ′ such that:

where D is a valid statistical divergence. Because p and q ′ are unknown, we approximate the optimization problem by using a discriminator D θ to predict the empirical estimate of an f -divergence:

This approach parallels the adversarial minimax optimization setting adopted in GANs, where samples are generated by a generator Y ′ = G ω ( 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 D θ . 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

Feature Shift Correction

proposals

Classifier

Select

argmin

Imputed query

⋮

⋮

⋮

probabilities

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 ε = 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 , ..., b N } , with b j ∈ 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 ∈ 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 ∈ L . Note that the size of B is proportional to the size of X . Then, for every query sample y ∈ 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:

where y ( b ) i is the query sample y i with the corrupted features C replaced with the values of b such that y i C = b , and r θ i is the likelihood ratio function from the classifier used to process the i th sample. In other words, we evaluate each sample y ∈ 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 ′ ( k ) by replacing the corrupted sample y i with y ( b i ) i , such that Y ′ ( k +1) = ( Y ′ ( k ) \ y i ) ∪ y ( b ) i . By replacing corrupted features with values from the Eq. 8, the empirical divergence is decreased such that ˆ D θ ( X,Y ′ ( k +1) ) ≤ ˆ D θ ( X,Y ′ ( 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 α and ρ , 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 = 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.

̸

̸

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).

RealDatasets

SimulatedDatasets

ManipulationType

Dataset

1.0

1.0

4.2

4.1

0.9

0.9

4.3

3

scene

energy

0.8

0.8

0.7

0.7

5

2

mnist

covid

0.6

0.6

1.0

1.0

0.5

0.5

6.1

0.2

0.40.60.8

1 poly.

0.2

0.4

gas

0.4

0.4

0.3

0.3

6.2

10

cosine

canine

0.2

0.2

0.1

0.1

6.3

9

dilbert

founders

0.0

0.0

7

8

phenotypes

DataFix

MRMR

MB-KS

MB-SM*

KNN-KS*

DataFix

MRMR

MB-SM*

KNN-KS*

MI

FAST-CMIM

KNN-KS

MB-KS*

Deep-SM*

MI

FAST-CMIM

MB-KS*

Deep-SM*

selectKbest

MB-SM

Deep-SM

selectKbest

Feature Shift Correction. Weevaluate 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 manipulationagnostic 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-Leiber 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.

RealDatasets

SimulatedDatasets

Dhp

Dsk1

Dhp

0.8

0.8

2.6 ×100

2.4 ×100

0.7

102

0.7-

2.2 ×100

0.6

102

0.6

102

2 ×100

0.5

0.5

1.8 × 100

1.6 x 100

0.4

101

0.4

1.4 × 100

0.3

101

0.3

101

0.2

0.2

1.2 × 100

0.1

0.1

100

0.0

1001

100

0.0

100

DataFix

DD

LR

GAIN

ICE

MIRACLE

Sinkhorn

KNN

INB

MLP

Hyperlmpute

Mean

MissForest

Softlmpute

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.

Beforeshiftcorrection

Aftershiftcorrection

1.0

1.0

0.9

0.9

0.8

0.8

0.7

0.7

0.6

0.6

0.5

0.5

0.4

0.4

0.3

0.3

0.2

0.2

0.1

0.1

0.0

0.0

0

50

100

150

200

250

0

50

100

150

200

250

RemovedFeatures

RemovedFeatures

TVD

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

Original

DD

MLP

ICE

MissForest

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).

7 Conclusions

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.

Acknowledgments and Disclosure of Funding

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.

References

J. Qian, Z. Song, Y. Yao, Z. Zhu, and X. Zhang, 'A review on autoencoder based representation learning for fault detection and diagnosis in industrial processes,' Chemometrics and Intelligent Laboratory Systems , p. 104711, 2022.

T. Dai, H. Hu, Y. Wan, Q. Chen, and Y. Wang, 'A data quality management and control framework and model for health decision support,' in 2015 12th International Conference on Fuzzy Systems and Knowledge Discovery (FSKD) , pp. 1792-1796, 2015.

K. A. Barchard and L. A. Pace, 'Preventing human error: The impact of data entry methods on data accuracy and statistical results,' Computers in Human Behavior , vol. 27, no. 5, pp. 1834-1839, 2011.

S. B. Lim, S. J. Tan, W.-T. Lim, and C. T. Lim, 'A merged lung cancer transcriptome dataset for clinical predictive modeling,' Scientific data , vol. 5, no. 1, pp. 1-8, 2018.

B. L. Browning and S. R. Browning, 'Genotype imputation with millions of reference samples,' The American Journal of Human Genetics , vol. 98, no. 1, pp. 116-126, 2016.

T. Enamorado, B. Fifield, and K. Imai, 'Using a probabilistic model to assist merging of largescale administrative records,' American Political Science Review , vol. 113, no. 2, pp. 353-371, 2019.

J. Pan, V. Pham, M. Dorairaj, H. Chen, and J.-Y. Lee, 'Adversarial validation approach to concept drift problem in user targeting automation systems at uber,' arXiv preprint arXiv:2004.03045 , 2020.

S. Yu, X. Wang, and J. C. Príncipe, 'Request-and-reverify: Hierarchical hypothesis testing for concept drift detection with expensive labels,' arXiv preprint arXiv:1806.10131 , 2018.

S. Kulinski, S. Bagchi, and D. I. Inouye, 'Feature shift detection: Localizing which features have shifted via conditional distribution tests,' Advances in Neural Information Processing Systems , vol. 33, pp. 19523-19533, 2020.

S. Van Buuren and K. Groothuis-Oudshoorn, 'mice: Multivariate imputation by chained equations in r,' Journal of statistical software , vol. 45, pp. 1-67, 2011.

D. J. Stekhoven and P. Bühlmann, 'Missforest-non-parametric missing value imputation for mixed-type data,' Bioinformatics , vol. 28, no. 1, pp. 112-118, 2012.

Z. Zhou, Z. Gong, P. Ravikumar, and D. I. Inouye, 'Iterative alignment flows,' in International Conference on Artificial Intelligence and Statistics , pp. 6409-6444, PMLR, 2022.

D. Inouye and P. Ravikumar, 'Deep density destructors,' in International Conference on Machine Learning , pp. 2167-2175, PMLR, 2018.

L. Piano, F. Garcea, V. Gatteschi, F. Lamberti, and L. Morra, 'Detecting drift in deep learning: A methodology primer,' IT Professional , vol. 24, no. 5, pp. 53-60, 2022.

V. Subasri, A. Krishnan, A. Dhalla, D. Pandya, D. Malkin, F. Razak, A. Verma, A. Goldenberg, and E. Dolatabadi, 'Diagnosing and remediating harmful data shifts for the responsible deployment of clinical ai models,' medRxiv , pp. 2023-03, 2023.

T. Ginart, M. J. Zhang, and J. Zou, 'Mldemon: Deployment monitoring for machine learning systems,' in International Conference on Artificial Intelligence and Statistics , pp. 3962-3997, PMLR, 2022.

D. Zha, Z. P. Bhat, K.-H. Lai, F. Yang, Z. Jiang, S. Zhong, and X. Hu, 'Data-centric artificial intelligence: A survey,' arXiv preprint arXiv:2303.10158 , 2023.

J. Gama, I. Žliobait˙ e, A. Bifet, M. Pechenizkiy, and A. Bouchachia, 'A survey on concept drift adaptation,' ACM computing surveys (CSUR) , vol. 46, no. 4, pp. 1-37, 2014.

J. Lu, A. Liu, F. Dong, F. Gu, J. Gama, and G. Zhang, 'Learning under concept drift: A review,' IEEE transactions on knowledge and data engineering , vol. 31, no. 12, pp. 2346-2363, 2018.

S. Rabanser, S. Günnemann, and Z. Lipton, 'Failing loudly: An empirical study of methods for detecting dataset shift,' Advances in Neural Information Processing Systems , vol. 32, 2019.

A. Liu, J. Lu, and G. Zhang, 'Diverse instance-weighting ensemble based on region drift disagreement for concept drift adaptation,' IEEE transactions on neural networks and learning systems , vol. 32, no. 1, pp. 293-307, 2020.

V. Losing, B. Hammer, and H. Wersing, 'Knn classifier with self adjusting memory for heterogeneous concept drift,' in 2016 IEEE 16th international conference on data mining (ICDM) , pp. 291-300, IEEE, 2016.

N. Lu, J. Lu, G. Zhang, and R. L. De Mantaras, 'A concept drift-tolerant case-base editing technique,' Artificial Intelligence , vol. 230, pp. 108-133, 2016.

A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola, 'A kernel two-sample test,' The Journal of Machine Learning Research , vol. 13, no. 1, pp. 723-773, 2012.

J. M. Bland and D. G. Altman, 'Multiple significance tests: the bonferroni method,' Bmj , vol. 310, no. 6973, p. 170, 1995.

I. M. Nasir, M. A. Khan, M. Yasmin, J. H. Shah, M. Gabryel, R. Scherer, and R. Damaševiˇ cius, 'Pearson correlation-based feature selection for document classification using balanced training,' Sensors , vol. 20, no. 23, p. 6793, 2020.

K. Hopf and S. Reifenrath, 'Filter methods for feature selection in supervised machine learning applications-review and benchmark,' arXiv preprint arXiv:2111.12140 , 2021.

S. Maldonado and R. Weber, 'A wrapper method for feature selection using support vector machines,' Information Sciences , vol. 179, no. 13, pp. 2208-2217, 2009.

A. Mustaqeem, S. M. Anwar, M. Majid, and A. R. Khan, 'Wrapper method for feature selection to classify cardiac arrhythmia,' in 2017 39th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC) , pp. 3656-3659, IEEE, 2017.

X. Huang, L. Zhang, B. Wang, F. Li, and Z. Zhang, 'Feature clustering based support vector machine recursive feature elimination for gene selection,' Applied Intelligence , vol. 48, no. 3, pp. 594-607, 2018.

B. Tran, M. Zhang, and B. Xue, 'A pso based hybrid feature selection algorithm for highdimensional classification,' in 2016 IEEE congress on evolutionary computation (CEC) , pp. 3801-3808, IEEE, 2016.

R. Battiti, 'Using mutual information for selecting features in supervised neural net learning,' IEEE Transactions on neural networks , vol. 5, no. 4, pp. 537-550, 1994.

N. O. F. Elssied, O. Ibrahim, and A. H. Osman, 'A novel feature selection based on one-way anova f-test for e-mail spam classification,' Research Journal of Applied Sciences, Engineering and Technology , vol. 7, no. 3, pp. 625-638, 2014.

S. Bahassine, A. Madani, M. Al-Sarem, and M. Kissi, 'Feature selection using an improved chi-square for arabic text classification,' Journal of King Saud University-Computer and Information Sciences , vol. 32, no. 2, pp. 225-231, 2020.

C. Ding and H. Peng, 'Minimum redundancy feature selection from microarray gene expression data,' Journal of bioinformatics and computational biology , vol. 3, no. 02, pp. 185-205, 2005.

J. Li, K. Cheng, S. Wang, F. Morstatter, R. P. Trevino, J. Tang, and H. Liu, 'Feature selection: A data perspective,' ACM Computing Surveys (CSUR) , vol. 50, no. 6, p. 94, 2018.

F. Fleuret, 'Fast binary feature selection with conditional mutual information.,' Journal of Machine learning research , vol. 5, no. 9, 2004.

Q. Cheng, P. K. Varshney, and M. K. Arora, 'Logistic regression for feature selection and soft classification of remote sensing data,' IEEE Geoscience and Remote Sensing Letters , vol. 3, no. 4, pp. 491-494, 2006.

E. V. Sylvester, P. Bentzen, I. R. Bradbury, M. Clément, J. Pearce, J. Horne, and R. G. Beiko, 'Applications of random forest feature selection for fine-scale genetic population assignment,' Evolutionary applications , vol. 11, no. 2, pp. 153-165, 2018.

A. Bommert, X. Sun, B. Bischl, J. Rahnenführer, and M. Lang, 'Benchmark for filter methods for feature selection in high-dimensional classification data,' Computational Statistics & Data Analysis , vol. 143, p. 106839, 2020.

R. R. Andridge and R. J. Little, 'A review of hot deck imputation for survey non-response,' International statistical review , vol. 78, no. 1, pp. 40-64, 2010.

R. Mazumder, T. Hastie, and R. Tibshirani, 'Spectral regularization algorithms for learning large incomplete matrices,' The Journal of Machine Learning Research , vol. 11, pp. 22872322, 2010.

T. Hastie, R. Mazumder, J. D. Lee, and R. Zadeh, 'Matrix completion and low-rank svd via fast alternating least squares,' The Journal of Machine Learning Research , vol. 16, no. 1, pp. 3367-3402, 2015.

P. J. García-Laencina, J.-L. Sancho-Gómez, and A. R. Figueiras-Vidal, 'Pattern classification with missing data: a review,' Neural Computing and Applications , vol. 19, no. 2, pp. 263-282, 2010.

P.-A. Mattei and J. Frellsen, 'Miwae: Deep generative modelling and imputation of incomplete data sets,' in International conference on machine learning , pp. 4413-4423, PMLR, 2019.

T. Kyono, Y. Zhang, A. Bellot, and M. van der Schaar, 'Miracle: Causally-aware imputation via learning missing data mechanisms,' Advances in Neural Information Processing Systems , vol. 34, pp. 23806-23817, 2021.

D. Jarrett, B. C. Cebere, T. Liu, A. Curth, and M. van der Schaar, 'Hyperimpute: Generalized iterative imputation with automatic model selection,' in International Conference on Machine Learning , pp. 9916-9937, PMLR, 2022.

B. Muzellec, J. Josse, C. Boyer, and M. Cuturi, 'Missing data imputation using optimal transport,' in International Conference on Machine Learning , pp. 7130-7140, PMLR, 2020.

J. Yoon, J. Jordon, and M. Schaar, 'Gain: Missing data imputation using generative adversarial nets,' in International conference on machine learning , pp. 5689-5698, PMLR, 2018.

A. Bogatu, A. A. Fernandes, N. W. Paton, and N. Konstantinou, 'Dataset discovery in data lakes,' in 2020 IEEE 36th International Conference on Data Engineering (ICDE) , pp. 709-720, IEEE, 2020.

O. Dekel, 'Collecting high-quality labels from a crowd,' in Proc. 22nd Annual Conference on Learning Theory (COLT), 2009 , 2009.

B. Boecking, W. Neiswanger, E. Xing, and A. Dubrawski, 'Interactive weak supervision: Learning useful heuristics for data labeling,' arXiv preprint arXiv:2012.06046 , 2020.

N. V. Chawla, K. W. Bowyer, L. O. Hall, and W. P. Kegelmeyer, 'Smote: synthetic minority over-sampling technique,' Journal of artificial intelligence research , vol. 16, pp. 321-357, 2002.

D. M. Montserrat, Q. Lin, J. Allebach, and E. J. Delp, 'Training object detection and recognition cnn models using data augmentation,' Electronic Imaging , vol. 2017, no. 10, pp. 27-36, 2017.

M. Geleta, D. Mas Montserrat, X. Giro-i Nieto, and A. G. Ioannidis, 'Deep variational autoencoders for population genetics,' biorxiv , pp. 2023-09, 2023.

M. Lenzerini, 'Data integration: A theoretical perspective,' in Proceedings of the twenty-first ACM SIGMOD-SIGACT-SIGART symposium on Principles of database systems , pp. 233-246, 2002.

S. Krishnan and E. A. Wu, 'Automatic generation of data cleaning pipelines,' arXiv preprint arXiv:1904.11827 , 2019.

Z. Zhang, 'Missing data imputation: focusing on single imputation,' Annals of translational medicine , vol. 4, no. 1, 2016.

A. O. Salau and S. Jain, 'Feature extraction: a survey of the types, techniques, applications,' in 2019 international conference on signal processing and communication (ICSC) , pp. 158-164, IEEE, 2019.

P. J. M. Ali, R. H. Faraj, E. Koya, P. J. M. Ali, and R. H. Faraj, 'Data normalization and standardization: a technical report,' Mach Learn Tech Rep , vol. 1, no. 1, pp. 1-6, 2014.

B. Azhagusundari, A. S. Thanamani, et al. , 'Feature selection based on information gain,' International Journal of Innovative Technology and Exploring Engineering (IJITEE) , vol. 2, no. 2, pp. 18-21, 2013.

J. C. Riquelme, J. S. Aguilar-Ruiz, and M. Toro, 'Finding representative patterns with ordered projections,' pattern recognition , vol. 36, no. 4, pp. 1009-1018, 2003.

M. Stonebraker, D. Bruckner, I. F. Ilyas, G. Beskales, M. Cherniack, S. B. Zdonik, A. Pagan, and S. Xu, 'Data curation at scale: the data tamer system.,' in Cidr , vol. 2013, 2013.

M. Stonebraker, I. F. Ilyas, et al. , 'Data integration: The current status and the way forward.,' IEEE Data Eng. Bull. , vol. 41, no. 2, pp. 3-9, 2018.

K. Yan and D. Zhang, 'Feature selection and analysis on correlated gas sensor data with recursive feature elimination,' Sensors and Actuators B: Chemical , vol. 212, pp. 353-363, 2015.

P. Xanthopoulos, P. M. Pardalos, T. B. Trafalis, P. Xanthopoulos, P. M. Pardalos, and T. B. Trafalis, 'Linear discriminant analysis,' Robust data mining , pp. 27-33, 2013.

X. Nguyen, M. J. Wainwright, and M. I. Jordan, 'Estimating divergence functionals and the likelihood ratio by convex risk minimization,' IEEE Transactions on Information Theory , vol. 56, no. 11, pp. 5847-5861, 2010.

S. Nowozin, B. Cseke, and R. Tomioka, 'f-gan: Training generative neural samplers using variational divergence minimization,' Advances in neural information processing systems , vol. 29, 2016.

G. Huang, H. Berard, A. Touati, G. Gidel, P. Vincent, and S. Lacoste-Julien, 'Parametric adversarial divergences are good losses for generative modeling,' arXiv preprint arXiv:1708.02511 , 2017.

J.-B. Hiriart-Urruty and C. Lemaréchal, Fundamentals of convex analysis . Springer Science & Business Media, 2004.

M.-J. Zhao, N. Edakunni, A. Pocock, and G. Brown, 'Beyond fano's inequality: Bounds on the optimal f-score, ber, and cost-sensitive risk and their implications,' The Journal of Machine Learning Research , vol. 14, no. 1, pp. 1033-1090, 2013.

M. J. Wainwright, High-dimensional statistics: A non-asymptotic viewpoint , vol. 48. Cambridge university press, 2019.

D. M. Montserrat and A. G. Ioannidis, 'Adversarial attacks on genotype sequences,' in ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pp. 1-5, IEEE, 2023.

P.-Y. Chen, H. Zhang, Y. Sharma, J. Yi, and C.-J. Hsieh, 'Zoo: Zeroth order optimization based black-box attacks to deep neural networks without training substitute models,' in Proceedings of the 10th ACM workshop on artificial intelligence and security , pp. 15-26, 2017.

A. V. Dorogush, V. Ershov, and A. Gulin, 'Catboost: gradient boosting with categorical features support,' arXiv preprint arXiv:1810.11363 , 2018.

R. Huerta, T. Mosqueiro, J. Fonollosa, N. F. Rulkov, and I. Rodriguez-Lujan, 'Online decorrelation of humidity and temperature in chemical sensors for continuous monitoring,' Chemometrics and Intelligent Laboratory Systems , vol. 157, pp. 169-176, 2016.

L. M. Candanedo, V. Feldheim, and D. Deramaix, 'Data driven prediction models of energy use of appliances in a low-energy house,' Energy and buildings , vol. 140, pp. 81-97, 2017.

C. Blake, 'Uci repository of machine learning databases,' http://www. ics. uci. edu/˜ mlearn/MLRepository. html , 1998.

M. R. Boutell, J. Luo, X. Shen, and C. M. Brown, 'Learning multi-label scene classification,' Pattern recognition , vol. 37, no. 9, pp. 1757-1771, 2004.

L. Deng, 'The mnist database of handwritten digit images for machine learning research,' IEEE Signal Processing Magazine , vol. 29, no. 6, pp. 141-142, 2012.

J. Vanschoren, J. N. Van Rijn, B. Bischl, and L. Torgo, 'Openml: networked science in machine learning,' ACM SIGKDD Explorations Newsletter , vol. 15, no. 2, pp. 49-60, 2014.

M. Perera, D. M. Montserrat, M. Barrabes, M. Geleta, X. Giró-i Nieto, and A. G. Ioannidis, 'Generative moment matching networks for genotype simulation,' in 2022 44th Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC) , pp. 1379-1383, IEEE, 2022.

C. Bycroft, C. Freeman, D. Petkova, G. Band, L. T. Elliott, K. Sharp, A. Motyer, D. Vukcevic, O. Delaneau, J. O'Connell, et al. , 'The uk biobank resource with deep phenotyping and genomic data,' Nature , vol. 562, no. 7726, pp. 203-209, 2018.

C. C. T. Force, 'United states covid-19 cases and deaths by state over time,' 2022.

S. Ghanem, H. Krim, H. S. Clouse, and W. Sakla, 'Metric driven classification: A nonparametric approach based on the henze-penrose test statistic,' IEEE Transactions on Image Processing , vol. 27, no. 12, pp. 5947-5956, 2018.

S. Y. Sekeh, B. Oselio, and A. O. Hero, 'A dimension-independent discriminant between distributions,' in 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pp. 4419-4423, IEEE, 2018.

S. Y. Sekeh, M. Noshad, K. R. Moon, and A. O. Hero, 'Convergence rates for empirical estimation of binary classification bounds,' Entropy , vol. 21, no. 12, p. 1144, 2019.

Q. Wang, S. R. Kulkarni, and S. Verdú, 'Divergence estimation for multidimensional densities via k -nearest-neighbor distances,' IEEE Transactions on Information Theory , vol. 55, no. 5, pp. 2392-2405, 2009.

S. Ben-David, J. Blitzer, K. Crammer, A. Kulesza, F. Pereira, and J. W. Vaughan, 'A theory of learning from different domains,' Machine learning , vol. 79, pp. 151-175, 2010.

H. Zhao, R. T. Des Combes, K. Zhang, and G. Gordon, 'On learning invariant representations for domain adaptation,' in International conference on machine learning , pp. 7523-7532, PMLR, 2019.

I. Albuquerque, J. Monteiro, M. Darvishi, T. H. Falk, and I. Mitliagkas, 'Generalizing to unseen domains via distribution matching,' arXiv preprint arXiv:1911.00804 , 2019.

Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and V. Lempitsky, 'Domain-adversarial training of neural networks,' The journal of machine learning research , vol. 17, no. 1, pp. 2096-2030, 2016.

S. Omar, A. Ngadi, and H. H. Jebur, 'Machine learning techniques for anomaly detection: an overview,' International Journal of Computer Applications , vol. 79, no. 2, 2013.

G. Pang, C. Shen, L. Cao, and A. V. D. Hengel, 'Deep learning for anomaly detection: A review,' ACM computing surveys (CSUR) , vol. 54, no. 2, pp. 1-38, 2021.

G. Li and J. J. Jung, 'Deep learning for anomaly detection in multivariate time series: Approaches, applications, and challenges,' Information Fusion , vol. 91, pp. 93-102, 2023.

G. Pang, C. Shen, H. Jin, and A. van den Hengel, 'Deep weakly-supervised anomaly detection,' in Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pp. 1795-1807, 2023.

S. C.-X. Li, B. Jiang, and B. Marlin, 'Misgan: Learning from incomplete data with generative adversarial networks,' arXiv preprint arXiv:1902.09599 , 2019.

S. Yoon and S. Sull, 'Gamin: Generative adversarial multiple imputation network for highly missing data,' in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 8456-8464, 2020.

D. T. Neves, M. G. Naik, and A. Proença, 'Sgain, wsgain-cp and wsgain-gp: Novel gan methods for missing data imputation,' in International Conference on Computational Science , pp. 98-113, Springer, 2021.

Y. Wang, D. Li, X. Li, and M. Yang, 'Pc-gain: Pseudo-label conditional generative adversarial imputation networks for incomplete data,' Neural Networks , vol. 141, pp. 395-403, 2021.

Z. Dai, Z. Bu, and Q. Long, 'Multiple imputation via generative adversarial network for highdimensional blockwise missing value problems,' in 2021 20th IEEE International Conference on Machine Learning and Applications (ICMLA) , pp. 791-798, IEEE, 2021.

W. Wang, Y. Chai, and Y. Li, 'Gagin: generative adversarial guider imputation network for missing data,' Neural Computing and Applications , pp. 1-14, 2022.

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.

W. H. Press and S. A. Teukolsky, 'Savitzky-golay smoothing filters,' Computers in Physics , vol. 4, no. 6, pp. 669-672, 1990.

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 Datasets

B.1 Real Datasets

In this work, we use a total of 9 continuous and 3 categorical datasets, with their dimensions and data types shown in Table 2. The Gas, Covid, and Energy datasets are the same as those used in [9], and we apply identical preprocessing procedures. The Musk2 dataset provides information on musk and non-musk molecules. The Scene dataset describes image characteristics. The MNIST dataset consists of images depicting handwritten digits. The Dilbert dataset is an image recognition dataset of pictures of objects rotated from various orientations. We additionally include 3 internal datasets of biomedical data: Phenotypes, Canine, and Founders datasets. The Phenotypes dataset contains a subset of categorical phenotypes from the UK Biobank, following the same pre-processing as in [103]. The Canine and Founders datasets comprise binary-coded sequences of DNA including Single Nucleotide Polymorphisms (SNPs), representing data for multiple dog breeds and human populations, respectively.

The Polynomial and Cosine are two internally generated datasets consisting of values obtained from deterministic simulations. The Polynomial dataset includes samples where each feature value is obtained by evaluating a second-degree polynomial function f ( x ) = ax 2 + bx + c . The parameters a , b , and c are fixed for each sample. The feature values are derived by evaluating the polynomial function for the different x values within a sample. The values for x , a , b , and c are uniformly sampled from the range [ -10 , 10] . The Cosine dataset consists of samples with features following a cosine function f ( x ) = a · cos( bx + c ) . Similar to the Polynomial dataset, a , b , and c are fixed for each sample, and each feature value is obtained by evaluating the cosine function for the different x values within a sample. Here, the values of b and c are uniformly sampled within the range [ -π, π ] .

B.2 Simulated Probabilistic Datasets

We generate 15 simulated datasets containing 1000 features and 5000 samples by sampling from predefined probabilistic distributions, including multivariate Gaussians, with and without transformations, multivariate Bernoulli distributions, Gaussian mixture models, and Bernoulli mixture models. A total of 200 features are shifted in each dataset such that | C | = 200 and | C | = 800 . Table 3 describes the distributions used in each dataset. For every dataset, two distributions, p and q , are defined, so that D ( p, q ) > 0 but D ( p C , q C ) = 0 . In fact, for all datasets except for dataset 8, we also have that D ( p C , q C ) > 0 . A further discussion of the effect of shifts with D ( p C , q C ) > 0 and with D ( p C , q C ) = 0 , and its relation to the equivalence of feature shift localization and feature selection, is provided in the following sections. Because we have access to p and q , we can compute the real divergence between the distributions. In practice, we make use of a Monte Carlo estimate, as shown in Figures 5 and 25, because the divergences might not have a closed-form solution or can be computationally intractable.

Datasets 1-3 are based on a multivariate Gaussian, with diagonal covariance used in datasets 1 and 2, and a covariance Σ used in dataset 3. The covariance matrix Σ is defined by a Gaussian kernel such that the ij component is Σ ij = exp -|| i -j || 2 s , where s acts as a scale parameter, and 0 ≤ Σ ij ≤ 1 with Σ ii = 1 . In practice, when constructing the covariance matrix, we perform a shuffle of the feature order to better depict tabular data, where, in many cases, the correlation between features does not follow any specific ordering (opposed to images or audio). We use s = 0 . 05 to define Σ ij in dataset 3. Datasets 4 and 5 follow a lognormal distribution, defined as X = exp( V ) with V ∼ N ( µ, Σ) and X ∼ Lognormal ( µ, Σ) . We use s = 0 . 05 and s = 0 . 002 to define Σ ij in datasets 4 and 5, respectively. Datasets 6-8 follow a logit-normal distribution defined as X = σ ( V ) with V ∼ N ( µ, Σ) and X ∼ P ( N ( µ, Σ)) , where σ is the sigmoid transformation. We use s = 0 . 05 , s = 0 . 002 , and s = 0 . 002 to define Σ ij in datasets 6, 7, and 8, respectively. Datasets 9-12 follow a multivariate Bernoulli with independent features. Each feature i has a frequency of f i , where f ∼ P ( N (0 , 2 I )) , ϵ ∼ N (0 , I ) , and ( · ) 0 , 1 = clamp 0 , 1 ( · ) = min(max( · , 0) , 1) is the clamping function to ensure that the frequencies are between 0 and 1. Dataset 13 follows a Gaussian Mixture Model distribution, with 3 mixtures of equal weights, µ i ∼ N (0 , 0 . 01 I ) , and Σ i defined with s = 0 . 3 . Datasets 14 and 15 follow a Bernoulli Mixture Model distribution with 3 mixtures and f i ∼ Uniform d (0 , 1) .

̸

Datasets 1, 3, 4, 6, 9, 10, 11, 12, 13, and 15 apply a shift to the marginal means, such that for all i ∈ C , E [ p C i ] = E [ q C i ] . Such datasets include marginal shifts of a similar nature as the shifts generated by manipulation types 1, 2, and 6 applied to real datasets. Dataset 2 performs a shift of the marginal standard deviation while maintaining their mean such that for all i ∈ C , E [ p C i ] = E [ q C i ] but p C i = q C i and Var[ p C i ] = Var[ q C i ] , leading to a marginal shift similar to the one applied by manipulation type 4 used in the real datasets. Datasets 13 and 15 apply a shift to the mean of just one mixture of the mixture model, leading to only 1 / 3 of the samples being shifted, while still ensuring that E [ p C i ] = E [ q C i ] . Datasets 5, 7, and 14 apply a distribution shift consisting of removing the correlation between features, equivalently to manipulation type 3 applied in real datasets, such that q C = ∏ i ∈ C q i and p C = q C . Dataset 8 also applies a shift originating from modifying the correlation between features, but in this case, Σ ′ is defined as:

̸

̸

where the correlation of the features within C and within C are maintained, but the cross-correlations between the C and C are lost, which leads to a shift equivalent to manipulation type 8 applied in real datasets, such that p C = q C , p C = q C , but p = q .

̸

̸

̸

C Feature Selection and Feature Shift Localization Equivalence

DF-Locate combines feature importance and feature selection techniques to localize the feature subset originating the distribution shift. Feature selection techniques often solve, either implicitly or explicitly, the following problem:

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 . In practice, the distributions of z or t are unknown, making Eq. 2 intractable. Instead, most methods predict a score that approximates the amount of information relative to the label present at each feature: β = F ( Z, T ) , where β ∈ R d and F ( · ) is a function mapping samples Z and labels T to an importance score for each feature. Informally, these scores provide an ordering such that if features i and j have | β i | > | β j | , then I ( z i ; t ) ⪅ I ( z j ; t ) .

Section 4 shows that by defining z as the samples coming from X and Y , and t as the binary label of origin, such that z | t =0 ∼ p and z | t =1 ∼ q , then 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 . Namely, when the number of manipulated features is known | C | = k , and the divergence of the corrupted features is larger than 0, such that D TV ( p, q ) = D TV ( p C , q C ) = 1 , and D TV ( p C , q C ) = 0 , then feature selection and feature shift localization are equivalent:

Even if | C | is unknown, if only marginal shifts are present, one can perform feature shift detection by iteratively solving Eq.2 with k = 1 , and removing the detected features at each iteration from z as long as I ( z C ; t ) > 0 , making the iterative feature selection and feature shift localization equivalent problems:

where C = l ⋃ i =1 C i . In fact, the approach presented in section 4, DF-Locate, can be seen as an approximation of this iterative process, where one or more features are selected at each step by the feature removal policy function.

̸

However, the equivalence of both tasks breaks down for distribution shifts such as the one applied in the manipulation type 8 for real datasets, and in dataset 8 of probabilistic simulations, where p C = q C and p C = q C , but p = q . That is, when considering only the corrupted features C or the non-corrupted features C in isolation, the shift is impossible to detect unless all features are considered jointly. Therefore, any feature selection technique that approximates either implicitly or explicitly equation 2, will need a subset of features G containing features from both C and C in order to obtain I ( z G ; t ) > 0 because I ( z C ; t ) = 0 and I ( z C ; t ) = 0 . Furthermore, if | C | = | C | the problem of feature shift detection becomes unsolvable. This is because, even if a technique was able to properly identify a subset of features G = C , it would not be possible to know if the detected subset G contains corrupted or non-corrupted features, that is if G = C or G = C , making the assumption of | C | < | C | a necessary condition.

̸

Figure 3 (right) and Figure 17 show the F-1 score for each manipulation type, indicating that DataFix is able to correctly localize manipulated features with manipulation type 8 on real datasets, despite breaking the equivalence between feature selection and feature shift localization approaches. In contrast, Figure 18, which illustrates the average F-1 on the probabilistic datasets, shows that dataset 8 is the one providing the lowest F-1 scores. While the low F-1 score in dataset 8 can be partly caused by the mismatch between feature selection and feature shift localization problems, it can also originate from the difficulty of detecting shifts caused by mismatching correlations, as it provides similar performance as in datasets 5 and 7, where correlation shifts (with p C = q C ) are applied.

D Feature Shift from Imputation Methods

Imputation and supervised methods trained to reduce the expected mean square error (MSE) between the predicted ˆ x C and real x C features of a given sample x can lead to distribution shifts. Note that the optimal function minimizing E x ∼ P [ || x C -f ( x C ) || 2 ] is the expected value of x C conditioned in x C , that is f ∗ ( x C ) = E [ x C | x C ] . Therefore, the method that predicts the corrupted (or missing) features C given the non-corrupted (non-missing) features C , which provides the lowest MSE, will generate predicted samples ˆ x C = f ∗ ( x C ) , with a distribution where P (ˆ x C ) = 1 for ˆ x C = E [ x C | x C ] , and P (ˆ x C ) = 0 everywhere else. If Var[ x C | x C ] > 0 , then D ( P (ˆ x C ) , P ( x C )) > 0 because Var[ˆ x C | x C ] = 0 .

̸

For example, given a dataset of samples x 1 , x 2 , ..., x N , with x i C = x j C and x iC = x j C for all i and j , in other words, Var[ x C | x C ] > 0 , the optimal regression model, in terms of MSE, will predict ˆ x iC = E [ x iC | x i C ] for all i , such that Var[ˆ x C | x C ] = 0 . Similarly, consider a dataset where the features C and C are independent, such that p = p C p C , and p C = N ( µ, I ) , then ˆ x C = E [ x C | x C ] = E [ x C ] = µ , where µ ∈ R | C | is a constant vector. The simulated probabilistic dataset 1 follows this form, with p C = N (0 , I ) and q C = N (0 . 5 I, I ) , where the method providing the lowest MSE will be the one predicting ˆ x C = E [ x C | x C ] = µ = 0 for all samples. Figure 6 shows the histogram of the first feature values before and after performing feature correction with multiple methods. Most techniques, especially imputation-based techniques, output the mean or values highly close to it, which while producing minimum MSE, does not reflect the real distribution, removing or reducing its variance and leading to a distribution shift.

F000

200

200-

009

200

250

500

150

150-

150

200

400-

100

100-

FO00

100

150

200-

100

50-

50

50

100

50

0

-3-2-10123

0

0

-3-2-10123

Q

3-2-1012

-0

-3-2-10123

500

5000

5000-

500-

400

400-

4000

4000-

400

300

000

-0000

0000

F000

200-

200-

2000-

2000

200

100-

100-

1000-

1000-

100-

-3 -2 -1 01 23

0

-3-2-10123

0

-3 -2 -1 0 1 2 3

5000

250

5000

5000

1250-

4000-

200-

4000

4000

1000

F0000

150

F0000

750-

2000-

100

2000

500

2000

1000-

50

1000-

250

1000-

-3-2-10123

-3-2-1

-3-2-10i

2

-3-2-101

Reference

DD

MLP

ICE

MissForest

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

E DF-Locate

Initialization

Reference

Query

Pool of samples

Sample

Random

forest

classifier

Label

Features importances

Normalization

Sort

Estimated total

variation divergence

Feature removal

policy function

Features to remove

Update

Loop from

Filtered

outputs

to

Filtered

Query

DF-Locate (Section 4, Figures 1 and 7, and Algorithm 1) is the proposed method within DataFix that localizes the features originating the distribution shift by performing feature selection in an iterative way. First, starting with i = 0 , and a reference X (0) = X and query Y (0) = Y datasets, a set of discriminators are trained θ = argmax θ ˆ D θ ( X ( i ) , Y ( i ) ) . The discriminators are used to predict the empirical total variation distance (TVD) between distributions ˆ D = ˆ D θ ( X ( i ) , Y ( i ) ) , and a feature importance score for each feature β = F θ ( X ( i ) , Y ( i ) ) . The divergence and feature importances are used to select potentially corrupted features with the feature removal policy function (see section 4) C i = ψ τ ( β, ˆ D ) . Then, the detected features are removed from X and Y , such that X ( i +1) = X ( i ) C i and Y ( i +1) = Y ( i ) C i . The process is repeated as long as the estimated divergence is smaller than a threshold ˆ D θ ( X ( i ) , Y ( i ) ) > ϵ , less than half of the features of the dataset are removed k ( i ) < | Y | 2 , or at least one feature is selected by the feature removal policy function at each step k ( i ) -k ( i -1) > 0 . In this work we use ϵ = 0 . 02 as the stopping threshold. After the iterative process is stopped, a refinement step, described below, is applied to remove features that might have been incorrectly selected as corrupted.

E.1 Shift Detection

We detect whether a shift is present by estimating an f -divergence between the distributions p and q . Specifically, we make use of the variational form [67, 68, 69]:

where r θ ( x ) is a function approximating the likelihood ratio between p and q . The inequality becomes an equality when r θ ( x ) = r ∗ ( x ) = p ( x ) / q ( x ) is the true likelihood ratio function, defining the Bayes decision rule that optimally separates p and q . f is the generator function of the f -divergence, f ′ is its first order derivative and f ∗ is its Fenchel conjugate [70].

In practice, estimating D f ( p, q ) is intractable, as we do not have direct access to p and q . Hence, we use an empirical estimator of the divergence by training a binary classifier D θ ( x ) to discriminate the samples coming from X = { x 1 , ..., x N } and Y = { y 1 , ..., y N } , with x i ∼ p and y j ∼ q . The likelihood ratio can be easily estimated as r θ ( x ) = D θ ( x ) / 1 -D θ ( x ) = exp σ -1 ( D θ ( x )) , where σ -1 is the inverse sigmoid. We use k -fold train-evaluation split, where the dataset is divided into k = 5 subsets. At each iteration, 80% of the samples are used to train a random forest classifier D θ ( x ) , and the 20% left is used to estimate the empirical divergence:

where N x and N y are the number of testing samples of the reference and query datasets respectively, and D f ( p, q ) ≥ E x ∼ p,y ∼ q [ ˆ D θ ( X,Y )] . By training the random forest, we are approximately finding the discriminator that provides a maximal divergence θ ∗ = argmax θ ˆ D θ ( X,Y ) within the random forest hypothesis space.

E.2 Feature Removal Policy Function

The feature removal policy function ψ , introduced in section 4, makes use of the predicted feature importances β = F θ ( Z, T ) and the estimated total variation distance ˆ D = ˆ D TV ( X,Y ) to select the corrupted features C = ψ ( β, ˆ D ) . This function utilizes the sorted normalized absolute value of the importance scores β ′ π = | β | π ∑ d -1 i =0 | β i | with β ′ π (0) ≥ ... ≥ β ′ π ( d -1) . At each iteration, the first J features are selected as corrupted, where J is dynamically selected at each iteration (equation 4) as the smallest index k such that γ ( k ) = ∑ k j =0 β ′ π ( j ) ≥ τ ˆ D (Eq. 4), with τ selected by cross-validation. The features selected as corrupted C include the features from 0 to J with scores higher than 1 d (Eq. 5).

The product τ ˆ D specifies a threshold that defines how many features are selected as corrupted at each iteration. When τ ˆ D is small, ψ returns the feature with the highest feature importance, while a larger τ ˆ D leads to more features flagged as corrupted. The value of τ is selected through hyperparameter optimization using the simulated datasets and maintained fixed afterward, and ˆ D is computed at each iteration with values that are, in expectation, non-increasing. When the empirical shift is low, the product τ ˆ D is small, leading to a smaller number of features being flagged as corrupted. Figure 8 presents a comparative analysis of the sorted normalized feature importance scores β ′ π in the first, middle, and last iterations of the DF-Locate iterative process at the first, second, and third figure column respectively. Different thresholds τ (0.05, 0.1, 0.2, and 0.3) are used to localize and filter the corrupted features at each iteration, with each row representing a specific threshold value. The features flagged as corrupted (as part of C ) by the feature removal policy function ψ are indicated in red color. As the number of iterations increases, ˆ D decreases and fewer features are flagged as corrupted, with the last iteration containing only a small number of remaining corrupted features and most features exhibiting similar low importance scores. Larger thresholds τ lead to more features selected at each iteration. However, for thresholds higher than 0.1, some of the selected features can be falsely classified as corrupted, indicating the need for lower threshold values.

Figure 9 shows the number of removed features at each iteration with different threshold values τ . Across all threshold values, the number of removed features decreases as the number of iterations increases, and in turn, the empirical divergence decreases. Note that higher threshold values lead to more features removed at each iteration, resulting in a smaller total number of iterations and faster computational times. However, higher thresholds also lead to an increased number of false positives, where many non-corrupted features are incorrectly selected as corrupted, leading to low F-1 scores. Therefore, the τ threshold provides a trade-off mechanism between feature shift localization speed and accuracy.

E.3 Refinement Stage

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 ˆ 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

FirstIteration

MiddleIteration

LastIteration

0.0015

0.003

0.004

°0.0010-

0.002

0.002

0.0005

0.001

0.0000

0.000

0.000

0

500

1000

1500

2000

0

500

1000

1500

0

500

1000

1500

0.004

0.0015

0.004

0.0010

0

0.002

0.002-

0.0005

0.0000

0.000-

0.000-

0

500

1000

1500

2000

0

500

1000

1500

0

500

1000

1500

0.004

0.0015

0.004

20.0010

0.002

0

0.002

0.0005

0.0000

0.000-

0.000

0

500

1000

1500

2000

0

500

1000

1500

0

500

1000

1500

0.010-

0.0015

0.004

0

0.002-

0.005

0.0005

0.0000

0.000

0.000-

0

500

1000

1500

2000

0

500

1000

1500

0

500

1000

1500

Sorted Feature Importances

Sorted FeatureImportances

Sorted Feature Importances

CorruptedFeatures

Noncorrupted Features

0.05

250

0.1

0.2

Features

S

200

0.3

150

RemovedF

100

50

0

0

5

10

15

20

25

Iteration

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 ⌊ ζδ/ 2 ⌋ +1) , where ζ is chosen from the set { 1 , 2 , 3 , 5 , 7 } through cross-validation, and δ 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 ζ = 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.

1.0

0.9

0.8

0.7

0.6

0.5

0.4

0.3

TVD

0.2

SmoothedTVD

一

Processed TVD

0.1

0.0

Knee

0

10

20

30

40

RemovedFeatures

F DF-Correct

Set of corrupted

features

Imputation with KNN

+ linear regression

Initial candidates

Select

Yes

CatBoost classifiers

Train

Update

Generate proposals

Update

Reference

Query

No

proposals

⋮

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 C = Y ′ 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 ′ = argmin Y i ∈ V ˆ D θ ( X,Y i ) . If ˆ D θ ( X,Y ′ ) < ϵ , the correction process is finished and Y ′ is returned. Otherwise, if ˆ D θ ( X,Y ′ ) > ϵ , 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 , ..., y ′ N y } that reduce the empirical total variation distance:

with g ( u ) = 1 2 sign ( u -1) . Because the values of X are left untouched, this becomes equivalent to solving:

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 E [ ˆ D TV θ ( X,Y )] = 0 , and:

where N + = { i : r θ ( y i ) > 1 } are the indices of samples classified as positive by the discriminator, and N -= { i : r θ ( y i ) < 1 } are the indices of samples classified as negative. Furthermore, both sets have, in expectation, the same size E [ | N + | ] = E [ | N -| ] = E [ | N y | 2 ] 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 | N y | 2 samples with the highest probability of being corrupted:

with | L | ≤ | 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 ∈ L , each of the proposals b ∈ B is placed as an alternative to the corrupted features, generating a sample y ( b ) i , where y ( b ) i C = b , and y ( b ) i C = y i C . The proposal providing the highest probability of being "non-corrupted" is selected:

Finally, the updated sample y ( b ) i is placed inside the corrected query Y ′ . After updating all samples in L , the divergence is computed again, and if ˆ D θ ( X,Y ′ ) > ϵ , 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.

Corrupted

WorstProposal

IntermediateProposal1

IntermediateProposal2

BestProposal

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.

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.

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.

103.

102.

101.

100

(SInOH)

10-1.

10

-2

01607

10

-3

10

-4

10-5.

10-6.

10-7.

9.9e4

5.1e57.1e5

1.1e6

1e7

2e7

4e7

5.4e7

2.9e8

covid,

gas

energy

scene

cosine,

polynomial

dilbert

phenotypes,

founders

mnist

canine

DataFix

selectkbest

FAST-CMIM

MB-KS*

+

Deep-SM*

MI

MRMR

MB-SM*

KNN-KS*

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, DFCorrect 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.

101

100

(SInoH) 01607

10

-1.

10-2-

10-3

10-4-

10-5-

10-6-

9.9e4

5.1e57.1e5

1.1e6

1e7

2e7

4e7

5.4e7

2.9e8

covid,

se6

energy

scene

polynomial

dilbert

phenotypes,

founders

mnist

canine

cosine,

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

DD

MLP

ICE

MissForest

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 '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.

1.0

0.9

0.8

DataFix

0.7

MI

selectKbest

0.6

Score

MRMR

0.5

FAST-CMIM

MB-SM*

0.4

MB-KS*

0.3

KNN-KS*

Deep-SM*

0.2

0.1

0.0

0.050

0.075

0.100

0.125

0.150

0.175

0.200

0.225

0.250

FractionofManipulatedFeatures

1.0-

0.9-

0.8

0.7-

0.5

0.4

0.3

0.2-

0.1

0.0

gas

covid

energy

scene

mnist

1.0-

0.9

0.8

0.7-

core

0.6

0.4

0.3

0.2

0.1-

0.0-

xxxx

XXXXX

polynomial

cosine

dilbert

phenotypes

founders

canine

DataFix

selectKbest

FAST-CMIM

MB-KS*

Deep-SM*

MI

MRMR

MB-SM*

KNN-KS*

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 'x' 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.

1.0

0.9

0.8

0.7

e

0.6-

0.5

C

F0.4

0.3

0.2

0.1-

0.0

1

2

3

4.1

4.2

4.3

5

1.0

0.9

0.8

0.7-

core

0.6

00.5

0.4

0.3

0.2

0.1

0.0

6.1

6.2

6.3

7

8

9

10

DataFix

selectKbest

FAST-CMIM

MB-KS*

Deep-SM*

MI

MRMR

MB-SM*

KNN-KS*

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.

1.0-

0.9

0.8

0.7

0.6

0.5

F0.4

0.3

0.2

0.1-

一

0.0

1

2

3

4

5

6

7

1.0-

0.9

0.8

0.7

core

0.6

0.5

F0.4

0.3

0.2

0.1-

0.0

8

9

10

11

12

13

14

15

DataFix

selectKbest

FAST-CMIM

MB-KS*

Deep-SM*

MI

MRMR

MB-SM*

KNN-KS*

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-Leiber 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.

101.

WV

100.

gas

covid

energy

scene

mnist

0.8

0.6

0.4

D

0.2

0.0

gas

covid

energy

scene

mnist

102.

101.

100

gas

covid

energy

scene

mnist

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

DD

MLP

ICE

MissForest

102.

W

101.

100.

xxxxxx

polynomial

cosine

dilbert

phenotypes

founders

canine

0.8

0.6

0.4

D

0.2

0.0

双

X

polynomial

cosine

dilbert

phenotypes

founders

canine

102

D

100

XX

xxXx

polynomial

cosine

dilbert

phenotypes

founders

canine

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

DD

MLP

ICE

MissForest

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.

SimulatedDatasets

RealDatasets

1.0

1.0

0.9

0.9

0.8

0.8

0.7

0.7

0.6

S

0.5

S

0.5

0.4

0.4

0.3

0.3

0.2

0.2

0.1

0.1

0.0

0.0

RF

CatBoostLogReg ExtraTreeLGBM

SVC

RF

CatBoostLogReg ExtraTreeLGBM

SVC

Figure 22 provides the feature shift correction metrics for DF-Correct when using different classifiers as discriminators within the iterative process on both real and simulated datasets. Similar to DFLocate, tree-based methods provide similar results, surpassing the linear methods in most metrics. We use D hp as a metric to select our method because it provides an estimate that tightly bounds the total variation distance. Namely, we select CatBoost as our discriminator as it provides competitive performance with the other tree-based methods in the simulated datasets, and clearly outperforms the others in the real datasets.

Simulated Datasets

Dhp

Dsk1

0.2

4×100

3×100

101

0.1-

2×100

100

0.0-

100

RealDatasets

Dnp

Dsk1

1025.

0.2

1020

102

1015

0.1-

1010

101

105

100

0.0-

100.

CatBoost

RF

LogReg

ExtraTree

SVC

L Variability Analysis for Location and Correction

Conducting the whole evaluation benchmark for both localization and correction multiple times is infeasible and expensive. However, to provide insight into the variation among runs for each method, we present plots showcasing the variability of detection and correction methods with multiple random seeds for a given dataset, namely the Energy dataset.

In Figure 23, we present the variability of detection methods, conveyed by the mean and standard deviation of the F-1 score across five different seeds on the Energy dataset. This evaluation covers scenarios involving 25% corrupted features and all manipulation types. Similarly, Figure 24 provides insights on the variability of correction methods, conveyed by the mean and standard deviation of W 2 2 , D hp , and D skl computed across five different seeds on the Energy dataset. All detection and correction methods, particularly DF-Locate and DF-Correct, exhibit a low variability across different seeds.

1.0

0.9

0.8

0.7

Score

0.6

0.5

0.4

0.3

0.2

0.1

0.0

DataFix

selectKbest

FAST-CMIM

MB-KS

Deep-SM

MB-KS*

Deep-SM*

MI

MRMR

MB-SM

KNN-KS

A

MB-SM*

KNN-KS*

Dhp

1.5 × 100

0.8

102

0.7

1.4 × 100

0.6

1.3 × 100

0.5

0.4

101

1.2 × 100

0.3

1.1 × 100

0.2

0.1

100

0.0

100

DataFix

DD

LR

GAIN

ICE

MIRACLE

Sinkhorn

KNN

INB

MLP

Hyperlmpute

Mean

MissForest

Softlmpute

M Evaluation of Downstream Classification and Regression Tasks after Correction

In this section, we include an evaluation of downstream classification and regression tasks for the Musk2 and Energy datasets, which include categorical and continuous labels, respectively.

We train the downstream models with the corrected query datasets and evaluate their performance on the reference dataset. We perform the experiment using each query dataset version, which includes the original pre-corrupted dataset, and the correct datasets with each correction method. The evaluation of the classifier/regressor is carried out under three different scenarios: when all features are used (100%), when only 50% of the non-corrupted features are included, and when 0% of non-corrupted features are included (only the corrected features are used). The classification results are reported in Table 6, and the regression results are reported in Table 7, using balanced accuracy for classification and root mean square error (RMSE) for regression. Methods are sorted by their performance without non-corrupted features (0%). Notably, the classification/regression downstream performance when using DataFix surpasses the competing methods in most settings, in some cases even surpassing the performance of directly using the original query dataset (ground truth pre-corruption). This observation demonstrates that the class-conditional divergence between the reference and the corrected query with DataFix is small.

N End-to-end DataFix Analysis

Figure 25 provides five additional examples that illustrate the iterative process of DF-Locate before and after shift correction, in simulated datasets 1, 2, 3, 6, and 10. In each dataset, there are 200 corrupted features out of 1000. Similar to Figure 5, the left column in Figure 25 displays the TVD estimated by the random forest (blue), which provides a lower bound for its ground truth Monte Carlo estimate (black), as the iterative process detects and removes corrupted features. The F-1 detection score progressively increases until all corrupted features are identified, leading to the termination of the iterative process. The right column in Figure 25 showcases the iterative process applied to the corrected query using different methods.

In simulated datasets 1 and 2, the methods DD, INB, MIRACLE and MLP all yield an updated query that results in a lower empirical divergence. Notably, DD stands out by providing a corrected query with no empirical divergence detected by DF-Locate. However, DD produces queries with considerably higher empirical divergence in the case of the other simulated datasets when compared to DF-Correct, and is not able to lower the empirical divergence for simulated dataset 10. Moving to simulated dataset 6, INB and DD yield an updated query that reduces the empirical divergence, with both techniques achieving an empirical divergence that is almost negligible. The other shift correction benchmarking methods generate an updated query that increases the shift instead of reducing it. In simulated dataset 10, none of the benchmarking methods are able to generate an updated query that leads to a lower empirical divergence. Remarkably, DF-Correct (Purple) provides a precisely corrected query with no empirical divergence detected by DF-Locate for all datasets.

O DataFix Limitations

DataFix is specifically tailored for tabular datasets and will not work optimally when applied to other data types such as images, videos, audio, speech, or textual data. DataFix also exhibits limitations due to its computational cost, preventing scalability for online or streaming scenarios, and potentially resulting in reduced speed with very large datasets, although it generally outperforms competing methods in terms of computational efficiency.

Before shift correction

After shift correction

1.0

1.0

0.8

0.8

0.6-

0.6

1

0.4

0.4

0.2

一

0.2

00

0.0

0

50

100

150

200

250

0

50

100

150

200

250

1.0

1.0

0.8

0.8

0.6-

0.6

2

0.4

0.4

0.2-

0.2

0.0七

0.0

0

50

100

150

200

250

0

50

100

150

200

250

1.0

1.0

0.8-

0.8

0.6

0.6

3

0.4

0.4

0.2

一

0.2

00

0.0

0

50

100

150

200

250

0

50

100

150

200

250

1.0

1.0

0.8-

0.8

0.6

-0.6

6

0.4

一

0.4

0.2

一

0.2

100

0.0

0

50

100

150

200

250

0

50

100

T

150

200

250

1.0-

1.0

0.8-

0.8

0.6

0.6

0.4

-0.4

0.2-

一

0.2

10:0

0.0

0

50

100

150

200

250

0

50

100

150

200

250

Removed Features

Removed Features

TVD

DataFix

INB

GAIN

Mean

Sinkhorn

KNN

LR

Hyperlmpute

MIRACLE

Softlmpute

Original

DD

MLP

ICE

MissForest

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces an invariance-based approach to out-of-distribution generalization where, for a given pair of distributions or domains $p$ a $q$, a subset of the input dimensions is filtered away so as to minimize some divergence between the filtered data, and then classifiers trained on samples from $p$ generalize to samples from $q$. The proposal leverages feature selection approaches to find a subset of the input space dimensions that minimizes a divergence between modified $p$ and $q$, as estimated by a domain discriminator. In addition, an adversarial approach is further introduced where features deemed responsible for the shift are corrected.\n\nSTRENGTHS:\n- A very relevant problem is tackled with potential to practical relevance.\n- The manuscript is very clearly written and easy to follow.\n- The proposal applies well established approaches to features selection and estimation of distribution shifts.\n\nWEAKNESSES:\n- Lack of contextualization with out-out-distribution generalization literature. Is the setting under consideration discussed in the OOD generalization literature? If so, it should be highlighted. It seems the type of shift under consideration is the standard covariate shift case, where data marginal distributions shift (in a particular way since only a subset of the input dimension shifts), and the divergence used for feature selecting is the standard H-divergence estimator (prediction accuracy of a domain discriminator) [1].\n- Potentially hidden assumptions: extra assumptions over data conditional label distributions are required since resulting pruned features are domain invariant. In other words, $p(y|x \\sim p) = p(y|x \\sim q)$, which is not stated. Refer to [1,2] for discussion on settings and corresponding assumptions for OOD generalization cases.\n- Empirical assessment limited to tabular data, and limited to a particular pair of distributions. That is, the resulting set of features is only invariant across the pair of domains used for selection, and little is known about what would happen had a new shifted distributed been observed.\n- The proposal and the evaluation ignore representation learning techniques, which are an alternative to feature selection. Invariance-based approaches to out-of-distribution robustness typically project data onto a space where domains cannot be discriminated, with domain-adversarial approaches being a popular example of such an approach [3, 4].\n- The motivation to correction approach is unclear to me. Are the resulting data any useful/meaningful after the correction approach is performed?\n\n[1] Ben-David, Shai, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. \"A theory of learning from different domains.\" Machine learning 79 (2010): 151-175.\n\n[2] Zhao, Han, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. \"On learning invariant representations for domain adaptation.\" In International conference on machine learning, pp. 7523-7532. PMLR, 2019.\n\n[3] Ajakan, Hana, Pascal Germain, Hugo Larochelle, François Laviolette, and Mario Marchand. \"Domain-adversarial neural networks.\" arXiv preprint arXiv:1412.4446 (2014).\n\n[4] Albuquerque, Isabela, João Monteiro, Mohammad Darvishi, Tiago H. Falk, and Ioannis Mitliagkas. \"Generalizing to unseen domains via distribution matching.\" arXiv preprint arXiv:1911.00804 (2019).\n\nQUESTIONS:\n- How does the proposal assume labels shift across distributions/domains?\n- How does features selection compare with invariant representation learning?\n- Does inducing invariance across a pair of domains $p$ and $q$ influence invariance across other domains? Under which conditions?\n- How would the feature selection approach fare in more complex structured data, such as images or text for instance? On could first encode the data with a pre-trained model and apply the proposal on top of resulting features.\n- Could the authors clarify a bit the motivation on the correction approach? One would modify their data so it's \"less different\" than what is observed from a different data source, but then are the resulting corrected data more useful than the invariant subset of the features?\n\nLIMITATIONS:\nCf. weaknesses section for details. Main limitations are lack of contextualization with previous work on OOD generalization, and a limited evaluation. It's also unclear what would be the motivation for a correction approach such as the one proposed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle the problems of feature shift localization and correction, i.e., identifying columns leading to divergence between two distributions, and imputing new values in their place, leading to lower divergence. For the first task they employ a random forest classifier trained to predict between samples coming from the reference or query distributions. The corrupted features are identified based on the feature importance scores of the random forest models.\n\nFor the correction task, a set of proposed corrections is generated per-sample, and the correction yielding the lowest probability under a classifier trained to distinguish between the distributions is selected as the imputation.\n\nSTRENGTHS:\n1. The proposed method is outperforming all the baselines\n2. the underlying models are relatively simple and computationally efficient\n3. The authors provided the full code in the supplementary\n\nWEAKNESSES:\n1. Most importantly, how were train/test splits performed? The text implies performance is measured directly on the train tests. Analyzing the code for the feature correction task it seems that the models were indeed evaluated directly on the training data.\n2. Several experimental details are missing. Why was the optimal hyperparameter configuration selected only for a subset of the baselines in the feature localization tasks? How were hyperparameters selected for the baselines for the feature correction tasks?  \n3. The paper could be written in a clearer way at times. Especially the description of both methods (Sections 4 and 5) are somewhat convoluted, in particular the last paragraph of Section 5. what are the **k** classifiers and the CatBoost classifiers trained to predict? How are the values of B generated? How is the objective of equation 9 optimized? The method could have been described clearer in the main text, instead of describing e.g., the total variance distance for the case not applicable to practice.\n4. Error bars of the results are not present, contrary to what has been specified by authors\n\nPlease note that my rating of the paper is mostly influenced by points 1 and 2, as they seem to invalidate the empirical results presented by authors.\n\nQUESTIONS:\nThe main weakness of the paper are the missing experimental details, i.e., points 1. and 2. of Weaknesses.\n\nFor 1., it could very well be that I misunderstood the code and/or the text. If so, please indicate the fragments of attached codebase, were the train/test splits are indeed performed. In my understanding, judging for example by lines 461-470 in correct/src/scripts/run_benchmarks.py, the same dataset is used for model fitting as well as evaluation.\nFor 2., authors should clarifiy how they selected hyperparameters for all baseline methods.\n\nOther questions:\n\n1. What is the motivation behind definining the threshold as $\\kappa D$ (lines 213-214 in text)? Is it not possible for a single feature to cause high divergence?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method called Datafix that tries to identify and correct feature shifts in datasets.\n\nThis method is composed of two distinct algorithms\nDataFix-locate which first identifies what features have shifted between datasets and Datafix-Correct which tries to correct for these features shifts.\n\n\nDatafix-locate learns a random forest discriminator which is trained to classify reference data from query data. Here the reference data is the non-shifted data while the query data is the data to be tested and corrected for feature shifts.\n This discriminator classifier is used to obtain an approximation to the likelihood ratio between the reference and query data distributions. This approximation is then used to estimate Total Variation (TV) distance between the reference and query data. \nThe mean decrease of impurity for the random forest classifier is used to estimate importance scores for all features within the data. The top $K$ important features (the top $K$ features whose cumulative importance scores are greater than the scaled TV-distance between the reference and the query datasets) are removed, and the entire procedure is iterated for a set number of times.\nAll the features throughout this iterative process are the \"located features\" which have shifted between the reference and the query datasets.\n\nDatfix correct then tries to fix these corrupted or discarded features by first imputing these features through different methods (linear classifiers ,knn,random replacement from reference data) and computes the TV distance between the query data after imputation and t reference data. If this divergence is less than a set threshold, the imputed values are selected as the final corrected values. Otherwise, different \"proposals\" values for these corrupted features are constructed (which consist of different feature values from reference data as well as previously imputed queries. The proposal value which leads  to the corresponding query having the highest probability of  being a non corrupted sample by the discriminant classifier. The selected proposal values are then chosen to replace the corrupted features. This process is repeated until the TV distance between the reference data and the corrected data is less than a threshold.\n\nExtensive experiments are performed to validate this method. The proposed method is able to outperform baselines on both locating the  corrected  manipulated features (measured through F1 score) and correcting these features (measured through divergence measures )\n\nSTRENGTHS:\n- This paper addresses an important problem which is of great significance to the community.\nFeatureshift detection is an important problem and this paper proposes a promising method to both identify these shifts and correct for these shifts.\n- The method is validated on numerous datasets under numerous feature shift settings (though most of these settings are enforced)\n- The experimental setting and results also demonstrate that DataFix can even correct for correlated shifts in features\n- The proposed method performs much better than existing baselines on correctly identifying and correcting feature shifts\n\nWEAKNESSES:\n- I think the presentation of the work can be improved. Particularly the layout of the paper. There is a lot of discussion in the main paper on material that is not essential to the paper. This particularly includes  section 4 (page 4 line 161 ), such as the discussion on f-divergences (yes, granted there are relevant but not sure how much relevant given only TV is used throughout the paper), and the discussion on mutual information (page 5, line 189 onwards), equation (4).  \nThis causes a lot of important material to be pushed to the appendix, particularly how the proposals are generated when correcting for corrupted features, or how  the algorithms for Feature selection and feature correction work (as provided in the details on the appendix).\nThis, in my opinion, hurts the readability of the paper  and makes the proposed methods appear much more complicated than it is. It's very difficult to understand the paper without going through the details in the appendix. I would encourage authors to consider improving the readability of the paper by moving parts of appendix D and E to the main paper and maybe moving some of the discussion on f-divergences and the relationship between feature selection  and shift localization to the appendix.\n I would be happy to revise my score if there important reasons for or against this suggestion and I look forward to your response.\n\n- Some of the details, such as how the proposal values are constructed, can be a bit confusing In their current form and would request the authors to provide a small example in the paper (or the appendix) to help explain this. \n\n- Metric for evaluating feature correctness. The authors propose using divergence measures to evaluate how close the corrected query values are to the reference values. This might be fair but from a practical standpoint, I think the purpose of correcting for feature shifts is to ensure the performance of a classifier (or some other ML model) is the same on the reference dataset as that on the corrected dataset. I think such classifier results which evaluate the difference between performance before and after correction would have been helpful.\nThe corrected features might have  small overall divergence between, but I am not sure this guarantees that class conditioned divergence between the reference and the query data is small. I will be happy to raise my score if the authors could elaborate on this choice of evaluation metric for feature correctness.\n\nQUESTIONS:\n- The experiments show that that under different fraction of shifted features (0 to 0.20), Data-Fix performs much better than other methods. How would the performance scale even under more sever shifted features ? When do you think Data-Fix could break down?\n-  The paper mentions that GANs are not used to produce the corrected features as they are not suitable for tablular data. But some of the datasets used such as MNIST (image data) and simulated cosine and ploynomial (functional data) are data modalities where NNs can do well. Do you think a GAN based adversarial approach for feature correction could perhaps perform better in these scenarios?\n- A lot feature shift scenarios were implemented within the paper. Though, are there feature shift scenarios which are more realistic and are perhaps encountered more often in real world settings? Are there existing datasets where these feature shifts naturally occur (rather than generating feature shifts as done in the paper)?\n- The appendix provides details on computational running times for Data-fix. Are there details for the computing infrastructure used for running these experiments in the appendix? I think these details would be helpful to better understand these computational runtimes.\n\nLIMITATIONS:\nThe authors mention in the main paper that limitations are discussed in the appendix, but I wasn't able to find discussion on this.\nI think any discussion on the limits of DataFix and when it could fail could be helpful for readers and potential users of DataFix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the problem when a subset of coordinates in the features have shifts in distributions. It specifically proposes algorithms to detect the shifts, and methods to “repair” the shifted coordinates. \n\nImo, the most interesting part of the paper is that it relates certain inequalities in divergence measures with GAN via likelihood ratio test. Under this interpretation, the problem of identifying distributional shifts becomes tuning some classifiers. Their experiments also confirm that this approach is effective. \n\nThen the second part of repairing the corrupted coordinates smell a bit hackish, e.g., it proposes a few candidate solutions, then it runs some classification algorithms to determine which candidates to use. Nevertheless, it appears to me that the central idea is to use unchanged features to “complete” the changed ones, which I believe is plausible. In general, this part feels heavy lifting and I do not have much intuition on that. \n\nOverall, I believe this is a quite solid contribution to neurips.\n\nSTRENGTHS:\nit makes an interesting connection between distributional shift and GAN.\n\nWEAKNESSES:\nSecond part sounds a bit hackish. I also have a few questions (see below).\n\nQUESTIONS:\nI have a few major concerns/questions. \n\nThis work seems not very careful on the expository text. For example, statistical distinguishable does not imply computational distinguishable (foundation of modern/any cryptography?). So D(p, q) > \\epsilon does not always mean that they can be detected. \nDivergence measures are asymmetric. Can you comment about this, e.g., do we have two thresholds or always put the original distribution to the left? \nWhen we can “complete” the features with shifted distributions for training and/or inference purpose, does that mean these features are not important and can be excluded anyway?\n\nLIMITATIONS:\nNA.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}