Adversarial Learning for Feature Shift Detection and Correction
Authors: Míriam Barrabés
Daniel Mas Montserrat
Margarita Geleta
Xavier Giró-i-Nieto
Alexander G Ioannidis
Venue: NeurIPS 2023 poster (2023)
Keywords: feature shift detection
distribution shift
shift
data-centric AI
TL;DR: We introduce a framework inspired in adversarial learning to detect and correct features originating a distribution shift between datasets.

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
Míriam Barrabés 12, Daniel Mas Montserrat 1.* Margarita Geleta
Xavier Giró-i-Nieto Alexander G. Ioannidis 11
Stanford University Universitat Politècnica de Catalunya University of California, Berkeley "Amazon
Abstract
Data shift is phenomenon present in many real- world applications, and while there are multiple methods attempting to detect shifts, the task localizing and correct- ing the features originating such shifts has been studied in depth. Feature shifts can occur in many datasets, including in multi-sensor data, where some sensors in tabular and structured data, including biomedical, finan- cial, and survey data, where faulty standardization and data processing pipelines can lead 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 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
malfunctioning,
com/ AI Al-sandbox/DataFF
1 Introduction
Distribution shifts in multi-dimensional data, caused by one more "corrupted" dimensions, are common in various real-world applications. Data streams from multi-senso environments in fields like medicine, industry, finance, and defense, can experience shifts due to faulty sensors [I]. 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 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 "Jocalization" of feature shifts is the task of finding which features (i.e., dimensions) of the dataset are causing the shift within multiple sets data. This step can be critical order to address the error
"Equal contribution.
I Work done prior joining Amazon.
I "ooreeppnndecee to A.G.I. [joannidis @ stanfford.edu].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 8J. 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 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 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 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] [1]
distribution alignment problem [12 [3),
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 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 [5 16 17]. In this work, we introduce "DataFixx", 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- -FF-Cctre"" which directly modifies the samples in order to decrease the distribution shift between two datasets. DF-Loc consists of an iterative process where 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 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 * y. 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 77 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 | applies hypothesis testing for concept drift detection. In 20), two-sample hypothesis testing is explored for shift detection, and comparison is made between multivariate hypothesis tests via Maximum- Mean Discrepancy (MMD) 244. univariate hypothesis tests with marginal Kolmogorov- Smimov (KS) tests and the Bonferroni correction (25,, and dimensionality reduction techniques, among others. However, these works focus on detecting 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-basee approaches: K-Nearest Neighbors with KS statistic KS), multivariate Gaussian with KS (MB-KS), multivariate Gaussian and

---

disher-divergence test statistics (MB-SM), and deep density neural network models with Fisher- divergence test (Deep-SM). However, performing feature shift localization remains challenging.
Feature Selection. Feature selection methods can remove redundant features affecting efficiency or performance, and find the most relevant features, providing interpretability. Common techniques
include: filtering [26, [2]] wrapper [28] [9] and embedded methods [30] 311] Filtering methods,
including univariate Mutual Information statistics (MI) [32] ANOVA-F test (33 and Chi-square test
34] rank features based on data statistics. Minimum Redundancy Maximum Relevance (MRMR) 35
36 selects relevant features while minimizing redundancy with the selected ones. Fast- Cond-Condiiinna
Mutual Information Maximization (FAST- MMM [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 [9] Most methods select the features with the highest importance scores or all that exceed 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 [II], 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. Hyperlmpute 47 provides state-of-the -te-tart imputation by using an iterative process with MI. 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 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.
distance between distributions. The work in (48) applies a Sinkhorn-
Data-centric AL Data-centric AI focuses on systematically enriching data quality and quantity as means to boost AI and ML performance. covers strategies that impact every phase of the data
lifecycle, from data collection (50). labeling 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 statisties for automatic data processing 66 62 63), as well as learning- based automation techniques
that optimize data automation procedures, typically using machine learning [64 (5)666) 66). The detection
and removal of distribution shifts are becoming essential steps of the data-centric Al toolbox [17.
3 Proposed Framework
Definition 1. [Feature Shift] We are given two sets of d-dimensional samples X 121,
and Y = {y1, {2,,32, yy..... yN with Ti, y/ E RRd, from distributions p and q, respectively. A feature shift between p and q occurs when D(p, q) > and D(ps, qs) < where D) is a valid divergence or and C are the subsets of non-corrupted and corrupted features respectively, such that [SUC| = d, and ps and qs are the distributions restricted to S.
distance between distributions.
We will refer X and 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 detect and correct. Here we consider scenarios with E = 0, D(ps, 9s) = O, D(p g) > O, and [S] > |C| 2 1. That is, there are more non-corrupted than corrupted features, the divergence becomes 0 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, 4:) > E, where pt and q, represent the marginal distribution of the ith dimension resulting from additive and non-linea
3

---

Feature Shift Detection
Disciiiintoo Training
Fesure Ieposiance
Yo
benters polley fincion g
Refermes
If DIX, DIX
Diaiitaton
Damibutio
Distaation D(X.))
Qaay
Filiaed
Quary
Figure I: DF-Locate overview diagram.
transformations; correlation shifts where D(p, > E but D(pi, q:) S & for all i; and correlation shifts where (ps, qs) and D(pc,9c) but (p: q) > In the latter case, correlations 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 S, that satisfies D(Po. 90) S E, that is C argmin ( [C|.
The number of corrupted features |C] will be assumed to be unknown priori. Furthermore, because the distributions and are assumed to be unknown and only their samples are accessible, the task needs to be approximated, requiring careful consideration of trade-offs such falsely flagging non-corupted 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's 162-.. ..., y'y non-corrupted features unchanged, that is =s..
such that v/. ~ q' and D(p, d') S e, while keeping the
The correction of shifted features can be done through parametric models that perform sample- wise transformations y' = p(y), prediction or imputation methods v/ = 6(v)), or optimization procedures heuristics Y' ((X, Our proposed approach falls in the latter category.
Feature Shift Detection: DF-Locate
DF-Locate (Fig. [ Algorithm employs an iterative process detect the presence of shift and determine the features causing it. At each iteration, 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 corrupted are removed, and the process is repeated until no divergence is detected.
Shift Detection. distributions p and q by using an empirical approximation of the variational form 67 68 69
We detect whether shift is present by estimating an f-divergence between the
N
1
Dx(X,Y)=
f'ro(zz))
Ng (f(((o())
Σ
(1)
where X = {#1,-... (..,.... IN and Y (VI,.., with 2; P. yj ~ q. and (2) is function approximating the likelihood ratio between p and q, obtained by training binary classifier Dy Do(x) (a) such that ro(x)= Do(=)/1-Pp(1 We use 5-fold arain-evaluation such that 80% of the samples are used to train random forest binary classifier Da (z), and the 20% left used to estimate the empirical divergence with N₂ and N, testing samples of the reference and query datasets, respectively. The resulting estimates are averaged across folds to reduce variance. The function / is the generator function of the f-divergence, I' is its first order derivative, and f" is its Fenchel conjugate (70). By changing f. various divergences can be obtained such as the Kullhack-Leibe divergence, the Shannonn divergence used in GANs, or the total variation distance, among others.

---

The true f-divergence is recovered in expectation (p, q) = Eeeno.yrg Y)] when re ra(2) = r=(x) = p())/(()) the true likelihood ratio function, defining the Bayes decision rule that optimally separates y and q. In this work, we make use of the total variation distance 0 S Drv (p, g) S 1, defined by setting ['(m) = (f" o (u) = signn( - 1). The discrete nature of ('(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 Drv (X, Y) > O, or equívalently, the test balanced accuracy is larger than 0.5 (random chance), it might indicate that + 9. 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 tech- niques often involve approximating the following estimation problem, either implicitly or explicitly:
C argmax t2) OlSk
where C' is the feature subset of size up to k, which maximizes the mutual information / between the input sample 2, restricted the features in C, and the label t. By defining 2 as the samples coming from X and Y, and t as the binary label of origin, such that zle=o ~ p and 2 = ~ q. Fano's s inequality 77] and LeCam's 's method 722 can be used relate the mutual information between z and I and the total variation distance between p and q:
(Drv (2. +1 +1) S I(z:t) S (p.().g)
H2
(3)
2
where Hg is the entropy with base-2 logarithm. If the number of manipulated features is known |C| = k. and Dry (p, = Drv (pc, (c) = 1, then C = argmax|cigs ((=c;)) = argminn Drv (0 IC), thereby rendering the problem of feature selection and feature shift de- tection equivalent. See Section forr further discussion. In practice, the distributions are unknown making Eq. 2 intractable and most methods predict a score that approximates the amount of informa- tion relative to the label present at each feature: 3 = F(Z, T"), where 3 E IR 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 Dg as the feature importance scores.
Feature removal policy. We introduce the feature removal policy function 1a as an heuristic to approximate Eq. [ which utilizes the predicted feature importances 3 and the estimated total
variation distance D = Px, Y) to select likely corrupted features C = (3,, D). First, /3 is |,21 and sorted 8",, such that ((()) 2 (((1)) 2 ... 2 (d
1 Given the cumulative sum ((k) ΣØ (²O) the sorted features from 0 to J are selected as corrupted, where J is the smallest index such that (kk 2 dD (Eq. 4. and is a hyperparameter set by hyperparameter optimization. Finally, & returns the features from to with scores higher than a (Eq.
normalized B' Σ. o [ed
J(T, D,3)= argmin ((k)
(4)
c=(m- :0) J. )) (5)
kn(()2rD
Note that - acts as a threshold to select how many features are selected as corrupted. If the threshold is small, v simply returns the feature with the highest feature importance, while larger threshold will make ei return more features. By defining the threshold as the product of and D, we ensure that when the shift is small (low Đ), smaller number of features are flagged as corrupted.
Iterative Process. DF-Loe performs iterative process that starts with the reference (0) X and query (()) = Y datasets. At each iteration i, it trains set of discriminators g(i) argmax, DTV (x()) y()) consisting of random forest binary classifiers. The classifier predictions used to estimate ( Drv(((() (()) and, combined with the Gini importance from the random forests (()) F(X y()), ) set of corrupted features are localized using the feature removal policy C. = (8(()) D()). The selected features C, are removed from the dataset, such that
((+1) = (1) and (i-1) = ()) and the process is repeated while Dry (() > 0.02 or
5

---

until half of the features have been removed or no features are removed at the current iteration. After I iterations, the set of corrupted features is obtained C" U C.. Finally, refinement stage (see
below and Section E.3 predicts the final set C. More details are provided in Section
E
Refinement stage. Do-Loc stores all intermediate steps, including the indexes of the predicted corrupted feature locations and the estimated D each iteration, in order revisit the iterative filtering process and select the optimal stopping point. To determine the optimal iteration, the elbow or knee is found from processed curve depicting the empirical total variation distance function of the total number of removed features (see Section E.3). This refinement stage effectively eliminates false positives and enhances the accuracy of feature shift localization.
Algorithm 1 Do-Lte
Algorithm 2 DF-Correct
1: Inputs:
1: Inputs:
Reference > Query
X:
D. Reference
Query
D Feature Selection Threshold Divergence Threshold
C:
Corrupted Features
b Divergence Threshold 2: (r", vr, Impute(X, Y,C) y t- argmin Do(X,
2: x
3: y
4:
D. Y") < then
5: K
6: while (x(X(), v(0)
2 and k0
< [Y] and A(9)
return
6: end
A(1-1) > d
for epoch do
7:
- (1), (())
> Train discriminator
8:
0 Traint X, Y')
Train discriminatoes
8:
D - Do(X(1) y(
D Estimate divergence
9:
L +- Detecclncoorett (Da (Y"))
Detect samples that
9:
A (o( p(o)
Estimate feature importance
require feature conection
10:
C. - 1-11 D)
P Sclect Seloetcoomppted corrupted features
10:
B + Generutee
Y", V, C)
11:
X(++1), ((--))
& Remove detected
11:
12:
b +- argmax (v)
E
Find best proposal
12: k(+1)) e- 10 |Ca| Update detected feature counter
features
13:
y' +- update(Y' Y", .t.
b Dataset with updated
13: --i+1
14: end while
15:
14:
sample
end for
Dx(X, Y then
15: C' U', 3-0 16: C Refine(C")
U
Combine all detected features
16:
return y
17:
end if
Refine delected features
18:
end for
17: return C
19:
return y"
5 Feature Shift Correction: DF-Correct
After the set C' of features originating the shift has been detected by DF-Locete, DF-Correct (Figure 2. Algorithm 2 is applied generate new query dataset Y' that rectifies the distribution shift. Ideally, the objective is to find Y ~ g' such that:
Y' argmin Y Yo Yg |=
D(p.d')
(6)
where D is valid statistical divergence. Because p and of are unknown, we approximate the optimization problem by using discriminator Dy to predict the empirical estimate of an f-divergence:
Y' = argmin max max (X,
This approach parallels the adversarial minimax optimization setting adopted in GANs, where samples are generated by a generator Y' = (u), and discriminator tries to accurately classify the generated samples as "fake". This training leads to the minimization of the 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 minmmize the predicted likelihood of the updated samples coming from g, approximated with the discriminator Dn. 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, employ base

---

Feature Shift Correction
Select
argmin
M proposals
Classifiet
Query Y
probabilinies
Inggaeed query
Figure 2: DF-Correct overview diagram.
techniques instead. Despite not being differentiable, derivative-free optimization and search heuristics can be used to update Y.
Two important aspects need 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 the whole Fuclidean space without updating or retraining the discriminator, we would be effectively conducting an adversarial attack [73 [4], 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 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 then performed with three distinct techniques: KNN, linear regression, and random sampling from the reference dataset Xc. 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. the initial empirical divergence is already lower than e 0.1, 1, the correction process is finalized. If not, DF-Corred applies the iterative process described below.
Iterative Process. The reference, the imputed query, and augmented samples (see Section 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 T. is obtained by selecting IV/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 set of new feature value proposals B = [b..... by J. with by E [ This set comprises all the feature values within positions C' from the reference dataset and the imputed query with linear regression, alongside random permutations of the reference values. In other words, each b E B contains (CI values obtained from X (and imputed Y) that can be a potential replacement for the corrupted features of each query sample y E Note that the size of B is proportional to the size of X. Then, for every query sample y E L classified as "corrupted" by the discriminator, we replace the shifted features by all candidate values B, and select the b; that provides the highest empirical likelihood of being non-corrupted sample:
b, argmax To, (B, (b) BEB
(8)
where y: is the query sample with the corrupted features C' replaced with the values of b such that yic = b, and re, is the likelihood ratio function from the classifier used to process the ith sample. In other words, we evaluate each sample y E L total of B times with the discriminator and select the 6 providing the highest probability that (b) is from the reference distribution. Then, we update the current corrected query dataset (*) by replacing the corrupted sample y: with (b,). such that (*(+1) (Y*(k) \ wi) U w/ By replacing corrupted features with values from the Eq. 8 the
empirical divergence is decreased such that Do (X,,Y((( ((+1)) Do (1())). After updating all the
7

---

"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. More details are in Section
6 Experimental Results
Real world datasets. We use multiple datasets including UCI datasets such as Gas [76]. Energy tir
and Musk2 778., OpenML datasets including Scene [79], MNIST [80) and Dilbert [8]] datasets with DNA sequences such as Founders 82 and a private Dog DNA dataset (Canine), subset of phenotypes from the UK Biobank (Phenotypes) [83, Covid-19 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, 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 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.I
Feature Shift Manipulations. We apply 10 different manipulations to the real world datasets in order to generate various distribution shifts. Table I describbe each manipulation. Manipulations 4, 5, 6, and 7 distort the marginal distributions, with manipulation 4 leaving the mean approximately unchanged. Manipulations and 6 have different levels of strength, controlled by parameters a and p. respectively. Manipulations 3 and 6 shuffle the feature values across samples, leaving the marginal distributions unchanged (p: = q:) but changing their correlations. Manipulation 3 performs a different random permutation at each feature, removing all correlation between features, while manipulation performs the same permutation for all features, such that Pc = Ic, but p * q. Manipulations 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, apply 10 manipulations continuous features and 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 specific fraction of the features.
Table 1: Manipulation types applied to continuous and/or categorical features.
Type
Mapping
Description
Shift
Data
1
x Uniform(0, (0, 1)
Each value substituted by random number
Pt. # 4
Cont.
between and I.
2
1 1- P.X.
Each value is nogated.
# 4. E(y.) = 1- 1ppp]]
Both
P. is a random permutation matrix applied
p. 4. $ 40:
Both
feature s.
4C IIeo 9:
4.1-4.3
clampo, (x (x aa) (ademacher(0.55)
Add constant noise with a random sign. a E
Pc # q. Elp.l Big.]
Cont.
(0.02, 0.05, 0.1) for 4.1-4.3 respectively. Values binarized. Values are negated with probability o E (0.2, 0.1,0. 0.6) for 6.1-6.3 respectively.
5
round(:z)
q:
Cont.
6.1-6.3
b(1 - 2) + (1- b)2
Pi * qi
Cat.
Bernouli(p)) MLP(#)
E(g.] = p 2p)][pp]
7
Forward through an MLP with min-max normal- ization binarization. P is random permutation matrix applied to all features simallameously. Predict feature with KNN (Reyressor). Predict feature with KNN (Classifier).
Both
8
PX:
a. PC #
Both
9
KNN(.:) KNN(w)
Cont.
10
Cat.
Probabilistic Simulations. We generate 15 datasets with 1000 features and 5000 samples, simulated from probabilistic distributions including multivariate Gaussians, multivariate Bernoulli distributions, Gaussian mixture models, and Bernouli 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
8

---

the feature correlations. Note that here we do not apply the shifts described in Table [ 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 hold-out testing set (see below and Section 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, selectKhest, MRMR, and Fast-CMIM) (Fig. B} 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'l (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-MM, MB-KS KS, KNN-KS, and De 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- 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- besting-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).
Real Datasets
Simulated Datasets
Manipulation Type
Dataset musk2
1.0
1.0
0.
4,3
Scand
energs
130
06
05
0.5
6.1
poly
ga6
0.4
0.4
03
0.3
5,2
10
cosine
caninu
0.2
6.3
dibert
founders
phenotypes
Datatic
MRINI
MU-K5 KNNN-KS
NENT
KNN-K57
DataFx
NRMR
M4 SM
IOV
MI
FAST-CMMM
HB-KS-
Deep- 5014
14
FAST-CHIM
MS-33
Deep-SM!
selectibestt
ME SM4
Decp SM
selec|Chen?
Figure 3: (left) median and mean F-1 score of real and simulated datasets, and (right) median and mean F-1 score across manipulation types and datasets. Higher is better.
Feature Shift Correction. We evaluate our method and 13 competing techniques, including predictive models like KNN, linear regression (LR), and multilayer perceptron (MLP), imputation methods including HyperImpute, ICE, MIRACLE, MissForest, and SoftImpute, optimal transport methods such as Sinkhorn, adversarial learning methods like GAIN, and domain adaptation techniques including deep destructors (DD) and iterative alignment flows (INB). We perform a manipulation- agnostic evaluation, where we set the corrupted features to 0 (or missing) and treat all manipulation types equally. Both the reference and the query are provided to each method, and non-parametric statistical metrics computed between the reference and the updated (corrected) queries are used to evaluate each method's performance. Specifically, we use the Wasserstein distance (W?) as in 488, the empirical estimator of the Henze-Penrop divergence (D)p) [85] 86, 87 and the empirical

---

estimator of the symmetric Kullback-Leiber divergence (Deu)) 88 We report the metrics after subtracting the "background" divergences computed with the reference and query datasets previous to any manipulation. Figure shows the median and mean metrics for the real and simulated datasets respectively. DataFix able to provide corrected query dataset Y' with the lowest empirical divergences. Despite their simplicity, KNN and linear regression provide competitive results, followed by MLP, Hyperlmpute, ICE, INB, and Sinkhorn. See Section D for more details.
Simulated
W33
Dhp
Dat
W2
Dip
Dat
2
2.4
0.7
10?
2,2 10
0.6
10
0.6
10"
18x 10
0.5
0.5
166 109
10
10
0.2
x02x10
0.1
0.1
10
0.0
10
10
0.0
10"
KNN
INB MLP
Hyperimaute
Mean
MissForest
Softimpute
Figure 4: W2. Dup- and Dak! of real and simulated datasets. Lower is better.
DataFix Analysis. Figure shows 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 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-Lace
Before shift correction
After shift correction
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
50
100
150
200
250
0
50
100
150
200
250
Removed Features
Removed Features
TVD
DataFix
INB
GAIN
Mean
Sinkhorn
F1 Score
KNN
LR
Hyperimpute
MIRACLE
Softimpute
Original
DD
MLP
ICE
MissForest
Figure DF-Locate iterative process before and after shift correction.
Extended experimental results present in the appendix, including analysis of the method's
computational time (Section H., the effect
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) [nd 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, able accurately detect and correct wide range of distribution shifts in many types of datascts, surpassing existing techniques.
10

---

Acknowledgments and Disclosure of Funding
This work was partially supported by a grant from the Stanford Institute for Fluman-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
[1] 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. 1, 2022.
[2] 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.
[3] 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.
[4] 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.
[5] B. L. Browning and S. R. Browning. "Genotype imputation with millions reference samples," The American Journal of Human Geneties, vol. 98, pp. 116-126, 2016.
[6] T. Enamorado, B. Fifield, and K. Imai, "Using probabilistic model to assist merging of large- scale administrative records," American Political Science Review, vol. 113, no. 2, pp. 353-371, 2019.
[7] 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.
[8] S. Yu, X. Wang, and J. C. Príncipe, "Request-ann-reverffyy Hierarchical hypothesis testing for concept drift detection with expensive labels," arXiv preprint 1806.1180011111 2018.
[9] S. Kulinski, Bagchi, and 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.
[10] Van Buuren and K. Groothuis- Oudshoorn, "mice: Multivariate imputation by chained
equations in r," Journal of statistical software, vol. 45, pp. 1-67, 2011.
[11] D. J. Stekhoven and P. Bühlmann, mixed-type data," Bioinformatics, vol. 28, no. 1, pp. 112-118, 2012.
missing value imputation for
[12] Z. Zhou, Z. Gong, P. Ravikumar, and D. I. Inouye, "Iterative alignment flows," in International Conference on Artificial Intelligence and Statistics, pp. 1409-6444 PMLR, 2022.
[13] D. Inouye and P. Ravikumar, "Deep density destructors," in International Conference on Machine Learning, pp. 2167-2175, PMLR, 2018.
[14] L. Piano, 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.
[15] 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.
[16] T. Ginart, M. Zhang, and J. Zou, "Mldemon: Deployment monitoring for machine learning systems," in International Conference on Artificial Intelligence and Statistics, pp. 3962-3997, PMLR, 2022.
[17] D. Zha, Z. P. Bhat, K.-H. Lai, F. Yang, Z. Jiang, S. Zhong, and X. Hu, "Data-centric artificial intelligence: A survey," arXiv preprini arXiv:2303. 10158, 2023.
[18] J. Gama, I. /liobaite, 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.

---

[19] 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-2366 2018.
[20] 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.
[21] A. Liu, J. Lu, and G. Zhang, "Diverse instance- 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.
[22] V. Losing, B. Hammer, and H. Wersing, "Knn classifier with self adjusting memory for heterogencous concept drift," in 2016 IEEE 16th international conference on data mining (ICDM), pp. 291-300, IEEE, 2016.
[23] N. Lu, Lu, G. Zhang, and R. I.. De Mantaras, "A concept drift-tolerant case-base editing technique," Artificial Intelligence, vol. 230, pp. 108-133, 2016.
[24] 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. pp. 723-773, 2012.
[25] J. M. Bland and D. G. Altman, "Multiple significance tests: the bonferroni method," Bmj.
vol. 310, no. 6973, p. 170, 1995.
[26] I. M. Nasir, M. A. Khan, M. Yasmin, J. Shah, M. Gabryel, R. Scherer, and R. Damaševičius, "Pearson correlation- feature selection for document classification using balanced training." Sensors, vol. 20, no. 23, p. 6793, 2020.
[27] Hopf and Reifenrath, "Filter methods for feature selection supervised machine learning applications-review and benchmark," arXiv preprint arXiv:2 /2/40, 2021.
[28] 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.
[29] A. Mustaqeem, S. M. Anwar, M. Majid, and A. R. Khan, "Wrapper method for feature selection to classify cardiac arrhythmia," " in 20/7 39th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC), pp. 3656-365555 3659, IEEE, 2017.
[30] X. Huang, L. Zhang, B. Wang, F. Li, and Z. Zhang, "Feature clustering based support vector machine recursive feature climination for gene selection," " Applied Intelligence, vol. 48, no. 3,
pp. 594-607, 2018.
[31] B. Tran, M. Zhang, and B. Xue, "A pso based hybrid feature selection algorithm for high- dimensional classification," " in 2016 IEEE congress on evolutionary computation (CEC), pp. 3801-3808, IEEE, 2016.
[32] 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.
[33] 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.
[34] S. Bahassine, A. Madani, M. Al-Sarem, and M. Kissi, "Feature selection using improved chi-square for arabic text classification," " Journal of King Saud University- Computtrr and Information Sciences, vol. 32, no. 2. pp. 225-222-232 2351 2020.
[35] 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.
[36] J. Li, K. Cheng, S. Wang, Morstatter, R. P. Trevino, Tang, and H. Liu, "Feature selection: A data perspective," ACM Computing Surveys (CSUR), vol. 50, no. 6, p. 94, 2018.
[37] F. Fleuret, "Fast binary feature selection with conditional mutual information."" Journal of Muchine learning research, vol. 5, no. 9, 2004.
[38] 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.
[39] 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.
12

---

[40] A. Bommert, X. Sun, B. Bischl, J. Rahnenführer, and M. Lang, "Benchmark for filter methods for feature selection in sigh-dimensional classification data," Computational Statistics & Data Analysis, vol. 143, 106839, 2020.
[41] R. R. Andridge and R. J. Little, "A review of hot deck imputation for survey non non-response," International statistical review, vol. 78, no. 1. pp. 40-64, 2010.
[42] R. Mazumder, T. Hastie, and R. Tibshirani, "Spectral regularization algorithms for learning large incomplete matrices," " The Journal of Machine Learning Research, vol. 11, pp. 2287- 2322. 2010.
[43] 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 Leurning Research, vol. 16, no. 1, pp. 3367-3402, 2015.
[44] P. J. Garcfa-Laencina, J.-L. Sancho-Gómez, and A. R. Figueiras-Vidaa "Pattern classification with missing data: a review," Neural Computing and Applications, vol. 19, no. 2, pp. 263-282, 2010.
[45] P.-.. Mattei and J. Frellsen, "Miwae: Deep generative modelling and imputation of incomplete data sets," in International conference on muchine learning, pp. 4413-4423, PMLR, 2019.
[46] 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.
[47] 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.
[48] B. Muzellec, Josse, C. Boyer, and M. Cuturi, "Missing data imputation using optimal transport." in International Conference Machine Learning, pp- 7130-7140, PMLR, 2020.
[49] 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.
[50] A. Bogatu, A. A. Fernandes, N. W. Paton, and N. Konstantinou, "Dataset discovery in data lakes," 2020 IEEE 36th International Conference on Data Engineering (ICDE), pp. 709-720, IEEE, 2020.
[51] O. Dekel, "Collecting high-quality labels from a crowd," in Proc. 22nd Annual Conference on Learning Theory (COLT), 2009, 2009.
[52] B. Boecking, W. Neiswanger, E. Xing, and A. Dubrawski, "Interactive weak supervision: Learning useful heuristics for data labeling," arXiv preprint arXiv:2012.06004 2020.
[53] 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.
[54] D. M. Montserrat, Q. Lin, J. Allebach, and E. J. Delp, "Training object detection and recogni- tion enn models using data augmentation." Electronic Imaging, vol. 2017, no. 10, pp. 27-36, 2017.
[55] M. Geleta, D. Mas Montserrat, X. Giro-i Nieto, and A. G. loannidis, "Deep variational autoencoders for population genetics," biorxiv, pp. 2023-09, 2023.
[56] M. Lenzerini, "Data integration: theoretical perspective," in Proceedings of the twenty-first ACM SIGMODD-SGACTT-SIGA. 2002.
symposium on Principles of database systems, pp. 233 24-4
[57] S. Krishnan and E. A. Wu, "Automatic generation of data cleaning pipelines," arXiv preprint arXiv:/904. /127, 2019.
[58] Z. Zhang, "Missing data imputation: focusing on single imputation," " Annals of translational medicine, vol. 4, no. 1, 2016.
[59] A. 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.
[60] P. J. M. Ali, R. H. Faraj, Koya, P. J. M. Ali, and R. Faraj. "Data normalization and standardization: technical report," Mach Learn Tech Rep, vol. 1, no. 1, pp. 1-6, 2014.
13

---

[61] B. Azhagusundari, A. S. Thanamani, et al., "Feature selection based on information gain," International Journal of Innovative Technology and Exploring Engineering (IJITEE), vol.
no. 2. pp. 18-21, 2013.
[62] 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.
[63] 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.
[64] 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.
[65] 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.
[66] P. Xanthopoulos, M. Pardalos, T. B. Trafalis, P. Xanthopoulos, P. M. Pardalos, and T. B. Trafalis, "Linear discriminant analysis," Robust data mining, pp. 27-33, 2013.
[67] X. Nguyen, M. 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.
[68] 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.
[69] G. Huang, H. Berard, A. Touati. G. Gidel, Vincent, and S. Lacoste-Julien, "Parametric adver- sarial divergences are good losses for generative modeling," arXiv preprint arXiv:1708.0255 2017.
[70] J.-B. Hiriart- Urruty and C. Lemaréchal, Fundamentals of convex analysis. Springer Science Business Media, 2004.
[71] 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.
[72] M. J. Wainwright, High-dimensione statistics: A non-asymptotic viewpoint, vol. 48. Cam- bridge university press, 2019.
[73] D. M. Montserrat and A. G. Ioannidis, "Adversarial attacks on genotype sequences." in ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Process- ing (ICASSP), pp. 1-5, IEEE, 2023.
[74] P.-Y. Chen, H. Zhang, Y. Sharma, J. Yi, and C.-J. Hsieh, "Zoo: Zeroth order optimization based black-box attacks deep neural networks without training substitute models," in Proceedings of the 10th ACM workshop on artificial intelligence and security. pp. 15-26, 2017.
[75] A. V. Dorogush, V. Ershov, and A. Gulin, "Catboost: gradient boosting with categorical features support," arXiv preprint arXiv: 1eXiv:1810.1 1363, 2018.
[76] R. Huerta, T. Mosqueiro. Fonollosa, N. Rulkov, and Rodriguez-Lujan, "Online decorre- lation of humidity and temperature in chemical sensors for continuous monitoring," Chemo- metrics and Intelligent Laboratory Systems, vol. 157, pp. 169-176, 2016
[77] 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.
[78] C. Blake, "Uci repository of machine learning databases," http://www. ics. uci. edu/ mleam//MLRepository html, 1998.
[79] M. R. Boutell, J. Luo, X. Shen, and C. M. Brown, "Learning multi-labe scene classification." Pattern recognition, vol. 37, no. 9, pp. 1757-1771, 2004.
[80] L.. Deng, "The mnist database of handwritten digit images for machine learning research," IEFE Signal Processing Magazine, vol. 29, no. 6. pp. 141-142, 2012.
[81] J. Vanschoren, 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.
14

---

[82] M. Perera, D. M. Montserrat, M. Barrabes, M. Geleta, X. Giró-i Nieto, and A. G. Ioanni- dis, "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.
[83] C. Bycroft. C. Freeman, D. Petkova, G. Band, L. T. Elliott, K. Sharp. A. Motyer, D. Vukcevic, O. Delaneau, J. O' O'Conelll, el al., "The uk biobank resource with deep phenotyping and genomic data," Nature, vol. 562, no. 7726, pp. 203-209, 2018.
[84] C. T. Force, "United states covid-19 19 cases and deaths by state over time," 2022.
[85] S. Ghanem, H. Krim, H. S. Clouse, and W. Sakla, "Metric driven classification: non- parametric approach based on the
test statistic," IEEE Transactions on Image Processing, vol. 27, no. 12, pp. 5947-5956, 2018.
[86] S. Y. Sekeh, B. Oselio, and A. O. Hero, "A dimension-indepeedden diseriminant between distributions." in 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4419-4423, IEEE, 2018.
[87] 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.
[88] Q. Wang. S. R. Kulkarni, and S. Verdú, "Divergence estimation for multidimensional densities via --eearsst-neihhbor distances," IEEE Transactions on Information Theory, vol. 55, no. 5,
pp. 2392-2405, 2009.
[89] 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.
[90] 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.
[91] I. Albuquerque. J. Monteiro, M. Darvishi, T. H. Falk, and I. Mitliagkas, "Generalizing to unseen domains via distribution matching," arXiv preprint arXiv:1911.00800 2019.
[92] 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. I, pp. 2096-2030, 2016.
[93] 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.
[94] 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.
[95] G. Li and J. J. Jung, "Deep learning for anomaly detection in multivariate time series: Ap- proaches, applications, and challenges," " Information Fusion, vol. 91, pp. 93-102, 2023.
[96] 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.
[97] S. C.-X. Li, B. Jiang, and B. Marlin, "Misgan: Learning from incomplete data with generative adversarial networks." " arXiv preprint /rXiv:/902.00. 999999999 2019.
[98] S. Yoon and S. Sull, "Gamin: Generative adversarial multiple imputation network for highly missing data," in Proceedings of the IEEE/VV Conference on Computer Vision and Pattern Recognition, pp. 8456-8464, 2020.
[99] D. T. Neves, M. G. Naik, and A. Proença, "Sgain, wsgain-cp and wsgain-gp: Novel gan methods for missing data imputation," International Conference on Computational Science, pp. 18-113. Springer, 2021.
[100] Y. Wang, D. Li, X. Li, and M. Yang. "Po-gain: Pseudo-label conditional generative adversarial imputation networks for incomplete data," Neural Networks, vol. 141, pp. 395-403, 2021.
[101] Z. Dai, Z. Bu, and Q. Long, "Multiple imputation via generative adversarial network for high- dimensional blockwise missing value problems," " in 2021 20th IEEE International Conference on Machine Learning and Applications (ICMLA), pp. 791-798, IEEE, 2021.
[102] W. Wang, Y. Chai, and Y. Li, "Gagin: generative adversarial guider imputation network for missing data," Neural Computing and Applications, pp. 1-14, 2022.
15

---

[103] J. Qian, Y. Tanigawa, W. Du, M. Aguirre, C. Chang, R. Tibshirani, M. A. Rivas, and T. Hastie, "A fast and scalable framework for large-scale and altrahigh-dimensione sparse regression with application to the uk biobank," PLoS genetics, vol. 16, no. 10, p. e1009141, 2020.
[104] W. H. Press and S. A. Teukolsky, ""avitzky-golay smoothing filters," " Computers in Physics, vol. 4, no. 6, pp- 669-672, 1990.
[105] V. Satopaa, J. Albrecht, D. Irwin, and B. Raghavan, "Finding a" kneedle" in haystack: Detecting knee points in system behavior," in 2011 31st international conference on distributed computing systems workshops, pp. 166-171, IEEE, 2011.
16

---

A Extended Related Work
The task of feature shift localization and correction is connected to various areas of statistics and ma- chine learning, including out-of-distributio (OOD) generalization, outlier detection, and Generative Adversarial Networks (GANs), among many others.
Generalization. Similar to feature shift localization and correction, OOD generalization addresses the issue of divergence between datasets or data sources. It does 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 consequence of the target domain being distinct from the source domain used for training the machine learning models. While both shift localization/corrcttoon 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 [9I]. 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 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 [..
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 197, architectural and loss modifications exemplified by GAMIN 198). SGAIN, WSGAIN-CP, and
WSGAIN-GI [99]. the incorporation of implicit label information as demonstrated by PC-GAIN [000, the ensemble of GANs by MI-GAN [101] and the capability to combine global and local
information demonstrated by GAGIN 102
17

---

B Datasets
B.1 Real Datasets
In this work, we use a total of continuous and 3 categorical datasets, with their dimensions and data types shown in Table 2 The Gas, Covid, and Energy datasets are the same 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 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 subset of categorical phenotypes from the UK Biobank, following the same pre- rre-processing in [03]. The Canine and Founders datasets comprise binary-coder 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² + bx + The parameters a, b, and e are fixed for each sample. The feature values are derived by evaluating the polynomial function for the different values within sample. The values for
b, and are uniformly sampled from the range [-10, 10]. The Cosine dataset consists of samples with features following a cos( bz + c). Similar to the Polynomial dataset, b, and
fixed for each sample, and each feature value obtained by evaluating the cosine function for the different x values within sample. Here, the values of b and uniformly sampled within the range [- n..
cosine function f(x)
Table 2: Datasets used to evaluate DataFix.
Dataset
No. of attributes
No. of samples
Data
Gas
12,815
Cont.
Covid
10
9,889
Cont.
Energy
26
19,735
Cont.
Musk2
166
6,598
Cont.
Scene
294
2,407
Cont.
MNIST
784
70,000
Cont.
Phenotypes
1,227
31,424
Cat.
Dilbert
2,000
10,000
Cont.
Founders
10,000
4,144
Cat.
Canine
198,473
1,444
Cat.
Cosine
1,000
10,000
Cont.
Polynomial
1,000
10,000
Cont.
B.2 Simulated Probabilistic Datasets
We generate 15 simulated datasets containing 1000 features and 5000 samples by sampling from pre- defined 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 each dataset. For every dataset, two distributions, p and q. are defined, so that D(p. 4) > 0 but D(p 4 O. In fact, for all datasets except for dataset 8, we also have that D(Pc. qc) > 0. A further discussion of the effect of shifts with D(pc, gc) > 0 and with D(pc, ((Pc.qc) = 0. and its relation to the equivalence of feature shift localization and feature selection, provided the following sections. Because we have access to p and q, we can compute real divergence between the distributions. In practice, we make use of a Monte Carlo estimate, as shown in Figures[Jand]] because the divergences might not have a dosed-form solution or can be computationally intractable.
18

---

Datasets 1-3 are based on multivariate Gaussian, with diagonal covariance used in datasets and 2. and covariance Σ used in dataset 3. The covariance matrix Σ is defined by Gaussian kernel such that the ijj component is Σnj = exp -لزجااا where acts scale parameter, and 0 S Σij I 1 with Σn = In practice, when constructing the covariance matrix, we perform a shuffle of the feature order 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 Eg in dataset 3. Datasets 4 and follow lognormal distribution, defined X exp V) with V ~ ((, Σ) and ~ Lognormal( (u. )) We use s = 0.05 and & = 0.002 to define 2ig in datasets 4 and 5, respectively. Datasets 6-8 follow logit-normal distribution defined X (VV with the sigmoid transformation. We use $ = 0.05, 0.002, and 0.002 to define Eij in datasets 6, and 8, respectively. Datasets 9-12 follow a multivariate Bernoulli with independent features. Each feature has a frequency of fi. where (O 2/)). NVO I). and (-)o,1 clampo,1 (-) min(maxx- 0), 1) is the clamping function to ensure that the frequencies are between O and 1. Dataset 13 follows a Gaussian Mixture Model distribution, with 3 mixtures of equal weights, P ~ A(O 0.01I), and Σ, defined with 0.3. Datasets 14 and 15 follow a Bernoulli Mixture Model distribution with 3 mixtures and ft Uniform (0, 1).
V ~ N(,, Σ) and X ~ P(N(H, Σ)), where
Datasets 1, 3, 4, 6, 9, 10, 11, 12, 13, and 15 apply a shift to the marginal means, such that for all i E C. Elpc, * E[ Such datasets include marginal shifts of similar nature the shifts generated by manipulation types 1, 2, and 6 applied to real datasets. Dataset performs a shift of the marginal standard deviation while maintaining their mean such that for all E C, E[pc;] = Elgc,] but pc, * qc, and Var[pc,] + Var[gc;], leading to a marginal shift similar to the one applied by manipulation type 4 used in the real datasets. Datasets 13 and 15 apply 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[pc,] * E[gc,]. Datasets 5, 7, and 14 apply distribution shift consisting of removing the correlation between features, equivalently to manipulation type 3 applied in real datasets, such that 9c = [iec qi and pC # qc. Dataset also applies shift originating from modifying the correlation
between features, but in this case, E' defined
Mn -0 (
i. if C with j E C, or E with
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 shift equivalent to manipulation type 8 applied in real datasets, such that PC = 9c. Pa = 9r. but p # q.
19

---

Table Probabilistic datasets used to evaluate DataFix.
pc
qc
Description
N(O, I)
N(O.51, I)
Multivariate Gaussians with diagonal variance and a shifted mean.
2
N(0, I)
N(0, 1.57)
Multivariate Gaussians with diagonal co- variance and a shifted scale.
((0 ))
W/0055, Σ)
Multivariate Gaussians with diagonal covariance a shifted mean. Multivariate lognormal with non- diagonal covariance and a shifted mean. Multivariate lognormal with diagonal (pc) and diagonal (gc) covariance.
(0g 2)
Lognormal(0.51, 2)
5
LognormalO,, 2))
Lognormal I)
6
P(N(O, x))
P(N(0.57, E))
Multivariate logit-normal with non- diagonal covariance and a shifted mean. Multivariate logit-normal with diagonal (pc) and diagonal (ςc) covari- ance.
POV(O, )))
P(N(0, f))
8
PON(O, Σ))
P(N(0, E\))
Multivariate nogitnnormal with different non non- diagonal covariances.
(emoulli(.)
Bemoulli((f + 0.05e)v.1)
Multivariate independent Berouulii with shifted mean.
10
Bermoulli(J)
Bernoulii(S + 0.1e)o..)
Multivariate independent Bernoulli with shifted mean.
11
Bernouli(J))
Bernuuli(( 0.5c)o.1)
Multivariate independent Bernoulli with shifted mean.
12
Bernoulli(J)
Sermouli(({+++ 1.0c)o.1)
Multivariate independent Bernoulli with shifted mean.
!E /Vies 2x)
1 D .V Mx)
Gaussian Mixture Model with one ture shifted such that μí i + 10, 21 112, and μí Bernoulli Mixture Model with different means fattaa
14
BMM(G f2.
BMM(J',f.
fas
BMM(((fi + 0.2)o.. fal)
Bernouli Mixture Model with one mix- ture shifted.
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:
C = argmax [(zc; ClSk
(10)
where C is the feature subset size up to kt which maximizes the mutual information between the input sample restricted to the features in C, and the label t. In practice, the distributions of t are unknown, making Eq. intrrcctaee Instead, most methods predict a score that approximates the amount of information relative the label present at each feature: B F(Z, T),, where 3 Rd and F(.) is a function mapping samples Z and labels T to an importance score for each feature. Informally, these scores provide ordering such that if features i and j have |33|| |33|, then T( f) T(2;
Section shows that by defining as the samples coming from X and Y. and as the binary label of origin, such that 2 z p and 2|1=1 72 can used to relate the mutual information between and and the total variation distance between p and q. Namely, when the number of manipulated features is known |C| = k. and the
then Fano's inequality [7I] and LeCam's 's method
divergence the corrupted features is larger than 0, such that Drv (p. q) = Drv (pc( 9c) 1, and Drv (rr ) 0, then feature selection and feature shift localization are equivalent:
20

---

C argmax ( 1
argmin Drv (Pm0r)=0
|C|
(11)
Even if |C| is unknown, if only marginal shifts present, can perform feature shift detection by iteratively solving Eq with 1, and removing the detected features at each iteration from & as long [(cctt) t) > O, making the iterative feature selection and feature shift localization equivalent problems:
C argmax ((
(12)
|C
where C = Ú fact, the approach presented in section Do-Locate, can be seen as an i=1 approximation of this iterative process, where one more features selected each step by the feature removal policy function.
However, the equivalence of both tasks breaks down for distribution shifts such the applied in the manipulation type 8 for real datasets, and in dataset probabilistic simulations, where PC qc and Pe 9c. but p * q. That is, when considering only the corrupted features C or the non-corrupted features C in isolation, the shift is impossible detect unless all features 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 [(zer; i) > 0 because I (=c; () and I(z;; () = Furthermore, if |C] = |C| the problem of feature shift detection becomes unsolvable. This is because, even if technique was able to properly identify a subset of features G C, would not be possible to know if the detected subset G contains corrupted or non-corrupted features, that if G 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 the probabilistic datasets, shows that dataset is the one providing the lowest F-1 scores. While the low F-1 score in dataset 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 in datasets 5 and 7, where correlation shifts (with pc * qc) applied.
21

---

D Feature Shift from Imputation Methods
Imputation and supervised methods trained to reduce the expected mean square error (MSE) between the predicted Ie and real Ie features of given sample z can lead to distribution shifts. Note that the optimal function minimizing Ex=pll {(>²) is the expected value of Ic conditioned E|ze| |-)l 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 êc = f (zz), with a distribution where ((tc) for &c = E[rc| and P(tc) 0 everywhere else. If Var Var(sclæd because Var(ăc|az) [ا = 0.
in 2τ., that is f (*7)
0, then D(P((c)) ((x))
For example, given a dataset of samples T1, 22.,.... EN, with शंद Zic and Fic # Zjc for all i and j. in other words, Var[ac] > 0, the optimal regression model, in terms of MSE, will predict Ric = Blsic/aacl for all i, such that Varloc|rc where the features C and C are independent, such that PCPC, and pC Ni,,, I), then tc E[ac
0. Similarly, consider dataset
[[xr]]] = μ., where μ E [R|C| is constant vector. The simulated probabilistic dataset I follows this form, with PC AO,, I) and the lowest MSE will be the one predicting &c E[ " = - 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.
A00.5,, I). where the method providing
200
200
200
500
100
100
150
200
100
50-
-2 -1
-3
5000
5000
4000
4000
300
300
300
200
200-
100
1000
1000
-2 -1
-3 -2
-2 -1
5000
250
5000
3000
1200
200
4000-
1000
3800
3000
3000
2000
2000
2000-
250
0
Reference
MLP
ICE
MissForest
DataFix
INB
GAIN
Mean
Sinkhom
KNN
LR
Hyperlmpute
MIRACLE
Saftimpute
Figure 6: Histograms of first feature values before and after shift correction with various techniques for probabilistic dataset I.
22

---

E DF-Locate
Loup from
Features importances
Label
g
F(
Normatzation
Rundor forest thoos
Reference
|s90|
X (1 R
Update
3t. p
X
x((
X 911
Sample 2
Inttaliaation
y
(6)
x
X
Fool samples
Fllendd Query
y-10) y
y
V1.... 9/
Feerumn nemiraa poilyy function
C. v1309 B())
Counry
Extinttot wal wwwwenncoeeegoooo
(300) 84 g
D Dxv(X (())
Figure DF-Locate diagram.
DF-Locate (Section 4 Figures [anrr [ and Algorithm D the proposed method within DataFix that localizes the features originating the distribution shift by performing feature selection in iterative way. First, starting with i 0, and reference (() = X and query (0)) Y datasets, set of discriminators are trained 0 argmax (x( (()). The discriminators are used to predict the empirical total variation distance (TVD) between distributions D (x)(), y()),, and feature importance score for each feature 3 F& (x ( The divergence and feature importances used to select potentially corrupted features with the feature removal policy function (see section 4
C. = u5 D). Then, the detected features removed from X and Y. such that ( X C,
and (i+1) - The process is repeated as long as the estimated divergence is smaller than threshold
> e, less than half of the features of the dataset are removed k:00) < [ or
least one feature selected by the feature removal policy function each step k(() 66-1) 0.
In this work we use 0.02 as the stopping threshold. After the iterative process is stopped, refinement step, described below, is applied to remove features that might have been incorrectly selected as corrupted.
E.1 Shift Detection
We detect whether shift present by estimating an f-divergence between the distributions p and g. Specifically, we make use of the variational form 67
68 69
Dy (p. 2 E. R. [mm=p(((ro(()) Ey-- (5" (f"
(13)
where (z) is a function approximating the likelihood ratio between and q. The inequality becomes an equality when ro(2) - (x:) p(e)/o(a)) is the true likelihood ratio function, defining the Bayes decision rule that optimally separates and q- f is the generator function of the f-divergence, f' is its first order derivative and I' is its Fenchel conjugate [70].
In practice, estimating (p q) is intractable, as we do not have direct access to and g. Hence, we use empirical estimator of the divergence by training a binary classifier Do (x) to discriminate the samples coming from X = { IN} } and Y {11 ..., yN with x; ~ p and yj ~ q. The likelihood ratio can easily estimated as r=(x) = Da(=)/I-D, (m) = expo (D)(a)), 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 train a random forest classifier Dy (x). and the 20% left is used to estimate the empirical divergence:
23

---

1 N Σ (((((rr(i))) N₂
XY, = N (((o(()))
(14)
where N₂ and Ny are the number of testing samples of the reference and query datasets respectively, and Dy(p,g) q) 2 Pewp.rgg Y)j. By training the random forest, we are approximately finding
the discriminator that provides a maximal divergence 0" argmax DX,, within the random forest hypothesis space.
E.2 Feature Removal Policy Function
The feature removal policy function v introduced in section 4 makes use of the predicted feature importances 3 = F& (z( and the estimated total variation distance ò = Dry (X, ()) to select the corrupted features C² (3, D). This function utilizes the sorted normalized absolute value of the importance scores :94/ Σラ [&d1
1 with (0)) 2 ... 2 2 At each iteration, the first
features selected as corrupted, where J is dynamically selected at each iteration (equation 4 as the smallest index & such that r(k) = - 2 +D (Eq. 4 with selected by cross-validation. The features selected corrupted C' include the features from 0 to J with scores higher than 1 (Eq. 5).
The product DD specifies a threshold that defines how many features are selected as corrupted at each iteration. When r D small, to returns the feature with the highest feature importance, while larger DD leads to more features flagged as corrupted. The value of T is selected through hyperparameter optimization using the simulated datasets and maintained fixed afterward, and î computed at each iteration with values that are, in expectation, non-increasing. When the empirical shift is low, the product 7D is small, leading to smaller number of features being flagged corrupted. Figure presents a comparative analysis of the sorted normalized feature importance scores 3 in the first, middle, and last iterations of the DF-Locate iterative process at the first, second, and third figure column respectively. Different thresholds T (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 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 small number of remaining corrupted features and most features exhibiting similar low importance scores. Larger thresholds T 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 an increased number of false positives, where many non corrupted features are incorrectly selected as corrupted, leading to low F- 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 itcration, we store the indexes of the features selected as corrupted, their corresponding estimated Đ, 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 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 processed curve depicting the empirical total variation distance (TVD) as a function
24

---

First teration
Middle heracion
Last Iteration
0.0015
0.003
0.004
0.0010
0.002
0.001
0.0000
0.000
0.000
500
1000
1500
2000
500
1300
1500
500
1000
1500
0.004
0.0015
0.002
0.002
0.0005
0.0000
0.000
0.000
500
1000
1500
2000
500
1000
1500
500
1000
1500
0.004
0.0015
0.034
0.0010
0.002
0.002
0.0005
0.0000
0.000
0.000
1000
1500
1000
0.004
0.0010
0.002
0.035
0.0005
0.0000
0.000
0.000
500
1000
1500
2000
1500
1000 1500 importancas
importances
importancas
Compupted
Nancorrupted
Figure 8: Sorted and normalized feature importance scores of the discriminators in the first, middle, and last iterations of DF-Locate applied to the Dilbert dataset with manipulation type 8. Each row represents a different threshold used to localize and filter the corrupted features each iteration.
0.05
Removed Features
0.3
150
100
Iteration
Figure 9: Number features removed each iteration for different threshold 7 values in DF-Locate applied to the Dilbert dataset with manipulation type 8.
of the total number removed features (Figure T0p. Once such optimal iteration is determined, all selected features up that iteration flagged corrupted. The curve the true divergence would always exhibit convex non- non-increasing shape but due to the intrinsic randomness in the training and evaluation process of the discriminator, this ideal shape locating the knee challenging. In order to make the curve smooth and non- non-increasing, we make use of the Savitzky-Goll Gollay filter [104] due its ability 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 left neighbor. Furthermore, we process initial iterations of the curve to enforce strict decreasing behavior. After processing the empirical divergence curve, we employ the knee locator method introduced by [05] identify the optimal stopping criterion.
always achieved, making the task of
25

---

The Savitzky-Goon 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 + 1), where is chosen from the set {1, 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 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 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
Smoothed
Processed TVD
F1 Score
0.0
Removed Features
Figure 10: Knee location based on distribution shift. Number of removed Features, TVD, and F-1 score in shift localization task.
26

---

DF-Correct
Set cornupted features
Imputation with linear regression
Inilial candidales
{Y Y Y2}
Reference
Select
X {01, (Z1,...,
Y arg min miny yiey (XX
3z ~ p
No
Yes
Do(X,Y') <E
y
CatBoost classifiers
Update
proposals
Update yr
Query
B {bi,.... by
Y (yI...., ]NN
yi
M proposals
Figure 11: DFF-Correct diagram.
DF-Correct (Section 5 Figures 2 and I1 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 iterative approach. Given the set of corrupted features C. DF-Correct tries to generate new query dataset Y', such that Yg Y while D(X,Y")
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
completely independent of each other. The set of the 3 initial candidates, V = {Y", Y Y²). J. is evaluated by computing the empirical total variation distance with a classifier, and the providing the lowest empirical divergence is selected the initial corrected query y' argminyiev Dx Y*). If YY) c, the correction process is finished and Y' is returned. Otherwise, if (x(X,Y') > e, an iterative process where some samples of y" are modified performed.
The iterative process tries find new values of Y {1/ Vly,.., .., vvv,} } that reduce the empirical total variation distance:
N (x,Y') = E g(ro(()))
N g(ro(1())
Σ
(15)
N
N
with g(u) = 2 sign(u - -1). Because the values of X are left untouched, this becomes equivalent to solving:
27

---

Y' argmin max
g(ro()))) = argmax max Σ g(ro(vg))
(16)
In order to perform such an optimization process, set of classifiers are trained using X, Y', and data augmentation consisting 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 find which samples need to be corrected. Note that when p 9. we have E (X, (X,Y)) 0, and:
B Σ s(r(()))) E JEN,
Σ g(ro(um))) MEN+
- E Σ g(ro(y)))] =0
(17)
nEN
where N+ {i: ro()) > } are the indices of samples classified as positive by the discriminator, and N = {i ro < } are the indices of samples classified as negative. Furthermore, both sets have, in expectation, the same size B+] [=[ |] E | when p 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 [Nxl samples with the highest probability of being corrupted:
== {i: ro, (1s) r0.41 (11+1) < <1}
(18)
with |Z| < [ull features. This proposal set 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 E L. each of the proposals
2- Next, we construct set of feature value proposals B that will replace the corrupted
B placed as an alternative to
the corrupted features, generating a sample y (b) where yt b. and y, (b) с yig. The proposal providing the highest probability of being "non-corrupted" selected:
b. argmax re, (1; (0)
(19)
BE
Finally, the updated sample y/u) is placed inside the corrected query Y". After updating all samples L, the divergence is computed again, and if Dx Y') > e, the process is repeated for a number of epochs. Typically, the number of epochs is set to or 2, 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- 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 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-Corred is able to reduce the divergence between datasets by comparing it with KNN imputation. A KNN imputer fills the missing feature values with combination of the top-k closest samples with respect to L2 distance on non- 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 sense, the KNN training samples are replaced by our feature proposals, and the KNN 1.2 distance replaced by the prediction of discriminator, making the correction task similar to divergence-reducing imputation problem.
28

---

Cormupted
Worst Proposal
Intermediate Propesal
Intermediete Proposal
Best Proposal
5
5
2
2
2_
2_
Figure 12:
Comparison of worst, intermediate, and best proposals from a single iteration of DF- Correct on two samples of the MNIST dataset against their corrupted version with manipulation type 2.
G Experimental Details
G.1 Hyperparameter Search
In this section, we present the experimental details of our hyperparameter optimization process for DataFix, along with the benchmarking methods used for feature shift localization and correction. The hyperparameter optimization involves conducting a grid search on the simulated datasets, with the provision that experiments exceeding 30-hour runtime 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 their respective optimal value. Similarly, Table provides insight into the search space for each parameter tuned in each correction benchmarking method, 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 discriminator, detailed in Section K Additionally, we conduct hyperparameter search for DF-Locate, focusing on variables such as window length and polyorder, which is further elaborated upon Section E33
Table 4: Search space and optimal values for tuned parameters in detection benchmarking methods.
Method
Tuned Parameter
Search Space
Optimal
MI
threshold
[0.008, 0.01, 0.02, 0.05, 0.1]
0.02
selectKbest
significance_leve level [0.008, 0.01, 0.02, 0.05, 0.1]
0.01
MB-SM, MB-KS
expectation
[30, 100, 500, 10000]
30
KNN-KS, Deep-SM
n bootstrap_runs
[1, 5, 10, 25, 50, 250]
250
G.2 Evaluation Setting
In both shift localization and correction, each method only has access to the reference dataset X and the corrupted query dataset Y 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*, not used by DataFix any step when predicting the localization of corrupted features or when correcting the query dataset.
29

---

Table 5: Search space and optimal values for tuned parameters in correction benchmarking methods.
Method
Tuned Parameter
Search Space
Optimal
KNN
[10, 25, 50, 100]
DD
[10, 50, 100, 200]
100
INB
[50, 100, 200, 500]
200
ndim
[30, 60, 90]
90
MLP
activation
['relu', "tanh"]
'relu'
hidden layer_sizes sizes alpha
[100, 1000]
100
[0.0001, 0.001]
0.0001
batch size
['auto', 64]
"auto"
learning rate init
[0.0001, 0.001]
0.001
GAIN
batch size
[16, 64, 256]
16
o_epochs
[1000, 50000]
50000
hint_rate rate__rte____
[0.8, 0.9]
0.9
Hyperlmpute
baseline imputer
[0. 1, 2] [1000, 5000]
optimize thresh
5000
inner iter
[40, 80]
40
ICE
max iter initial strategy
[1000, 2000]
1000
[0. 1, 1,
MIRACLE
Ir batch_sit size hidden reg la__aa reg beta window
[0.0001, 0.001]
0.0001
[1024, 512]
1024
[16, 32, 64]
16
[0.1, 1, 10]
10
[1,
[10, 20]
20
MissForest
estimators max max_iter iter
[10, 20, 50]
10
[100, 500, 1000]
100
Sinkhom
[0.001, 0.01]
0.01
n_epochs
[500, 1000)
1000
batch_size noise scaling
[256, 512]
512
[0.001, 0.01]
0.001
[0.9]
0.9
SoftImpute
maxit
[1000, 2000]
1000
convergence_threntt threshold
[0.00001, 0.0001]
0.0001
max_rank
[2.
shrink_lambda
[0.5, 0]
30

---

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 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 be reasonably efficient method, exhibiting good scalability as the dataset size increases, while providing the best localization performance.
10a
10a
101
100
Logso (Hours)
10-1
10
10
10
10-5
10
10
8.c
5.1e 11s
1e'
dilbert 8.
4e'
5.4e!
29c
covid. gas
energy
scene
musk2
cosine, pelynomial
phenotypes. founders mnist
canine
DataFix
selectKbest
FAST-CMIM
MR-K5*
Deep-SM*
MI
MRMR
MB-SM*
KNN-KS*
Figure 13: Computational time for shift localization methods based on real dataset size.
In Figure [4] we conduct comparative analysis of the computational time for shift correction methods. Although DF-Correct is not faster than simpler techniques like median or linear regression, its speed surpasses several competing methods, such as MIRACLE and HyperImpute. Furthermore, DF- Correct exhibits reasonable runtime even for the largest high- high-dimensional Canine dataset, successfully correcting the distribution shift within the 30-hour time limit. Many 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 terms of distribution shifts, while still providing competitive or faster speeds than competing methods.
All experiments were done with Intel Xeon Gold with 12 CPU cores.
31

---

10
10
Logiu (Hours)
10
10-2
10
10
10-5
3.cf
S1e' 1.1e
dilbert 3
Covid. gas
energy
scene
musk2
cosine, palynomial
phenotypes, founders mnist
canine
DataFix
INB
GAIN
Mean
Sinkhorn
KNN
LR
Hyperlmpute
MIRACLE
Softimpute
MLP
ICE
MissForest
Figure 14: Computational time for shift correction methods based on real dataset size.
I Extended Feature Shift Localization Results
We evaluate the performance of DF-Locate across different fractions of manipulated features: 5%, 10%, and 25%. The F-1 scores, depicted in Figure 15 were computed by averaging across manipula- tion 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 small percentage of ma- nipulated features present, a smaller amount of features need 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 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 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 provide accurate results while scaling large datasets. DF-Loc consistently exhibits superior performance compared to all benchmarking methods across most datasets. The only exception observed is with the Phenotypes dataset, where selectK Best slightly outperforms DF-Locate in locating the corrupted features. Nevertheless, for the remaining datasets, DF-Locate surpasses competing approaches and is only slightly surpassed or equaled in performance few occasions by Deep-SM (+) or KNN-KS (*), both of which use the ground truth |C|. Additionally, is worth noting that MI always outperforms selectKbest on datasets with continuous features, while the opposite holds true for datasets with categorical features.
32

---

1.0
0.9
0.8
DataFix
0.7
MI
SelectKbest
0.6
Sccoe
MRMR
0.5
FAST-CMIM
MB-SM*
0.4
MB-K5*
KNN-K5*
0.3
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
Fraction of Manipulated Features
Figure 15: Median F-1 scores of shift localization methods by fraction of manipulated features real datasets.
1.0
0.9
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0.0
gas
covid
energy
musk2
scene
mnist
1.0
0.9
0.8
0.7
0.6
0.5
[ 0.4
0.3
0.2
0.1
0.0
XXXX
polynomial
cosine
dilbert
phenotypes
founders
canine
DataFix
selectKbest
FAST-CMIM
MB-KS*
Deep-SM*
MI
MRMR
MB-SM*
KNN-KS*
Figure 16: Mcan F-1 scores of shift localization methods by real datasets. 'x' indicates missing experiment. Higher is better.
Figure 17 shows the median F-1 scores of DF-Locate and other competing methods, categorized by the feature manipulation type applied the real datasets. The average F-1 score computed across fractions of manipulated features, followed by the computation of the median F-1 score across different datasets, The symbol "x' is used 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 these three datascts within the specified time constraint of 30 hours. Manipulations involving shifts caused
33

---

by the correlation between features (manipulations 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, the most challenging to detect by DataFix. Note that such manipulation can lead, in some scenarios, small or undetectable distribution shifts, KNN provides perfect predictions when its training dataset size goes to infinity.
1.0
0.9
0.8
0.7
Scott
500
0.6
0.5
W 0.4
0.3
0.2
0.0
3
4.1
4.2
5
1.0
0.9
0.8
0.7
Sco18
0.6
0.5
0.4
0.3
0.2
0.1
0.0
6.1
6.2
6.3
8
10
DataFix
selectKbest MRMR
FAST-CMIM MB-SM*
MB-K5
Deep-SM*
MI
KNN-KS*
Figure 17: Median F-1 scores of shift localization methods by feature manipulation type on real datasets. 'x' indicates missing experiment. Higher is better.
Figure 18 displays the mean F-1 scores of shift localization methods across the simulated datasets. As observed in the results presented for real datasets Figure [6] DF-Locate consistently outperforms or matches all competing methods, except for datasets 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 fair comparison where [C] is not used, MB-SM and Deep-SM perform poorly (see Figure Do-Loc obtains lower scores in simulated datasets 5, and 8 compared to the other datasets, which involve shifts originated by mismatching correlation between distributions.
34

---

1.0
0.9
0.8
0.7
5
0.6
0.5
W 0.4
0.3
0.2
0.1
0.0
3
5
1.0
0.9
0.8
0.7
Score
0.6
0.5
0.4
0.3
0.2
0.1
0.0
8
9
10
11
12
13
14
15
DataFix
selectKbest MRMR
FAST-CMIM MB-SM*
MB-KS* KNN-K5*
Deep-SM*
MI
Figure 18:
Mean F-1 scores of shift localization methods on simulated datasets. Higher is better.
35

---

Extended Feature Shift Correction Results
Figures 19 and 20 provide comprehensive evaluation of the performance of DF-Correct and competing shift correction methods across real datasets, by using the Wasserstein distance (W2) as [48] the empirical estimator the Henze-Penrose divergence (Dup)
85 86 87], and the empirical estimator of the symmetric Kullback-Leiber divergence (DaRt) 88 W Dap- and D are pa--mmamttticc estimators of distances between the underlying distributions p and that make use of the reference 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 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 no query datasets prior to manipulation, second reference dataset used to obtain background empirical divergences.
As shown in Figures 19 and 20 DataFix outperforms all other methods by significant margin for most datasets, demonstrating its high effectiveness 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, Hyperlmpute, INB, Sinkhorn, and MLP yield favorable results for specific datasets. Furthermore, DataFix consistently achieves the lowest W values across most datasets, with exceptions being MNIST and Phenotypes. Specifically, for the MNIST dataset, KNN, MLP. GAIN, Hyperlmpute, and ICE surpass DataFix achieving the best W value of O. Note that for high-dimensional datasets such Dilbert, Phenotypes, Founders, and Canine, the W3 metric saturates to 0 for multiple methods, making the other metrics a better alternative to compare the quality between methods.
10
w2
10
gas
covid
energy
musk2
mnist
0.8
0.6
0.4
0.2
0.0
gas
covid
energy
musk2
scene
mnist
102
10
10
gas
covid
energy
musk2
scene
mnist
DataFix
INB
GAIN
Mean
Sinkhorn
KNN
LR
Hyperlmpute
MIRACLE
SoftImpute
MLP
ICE
MissForest
Figure 19: W3, Dap, and Dant of shift correction methods on real datasets. Lower is better. (Part 1)
36

---

102
10
10
JooK
XX1 300 XX
polynomial
cosine
dilbert
phenotypes
founders
canine
0.8
0.6
Dhp
0.4
0.2
0.0
book
XX vox o0
polynomial
cosine
dilbert
phenotypes
founders
canine
10
0
10"
20000
XX1 1000000 00
polynomial
cosine
dilbert
phenotypes
founders
canine
DataFix
INB
GAIN
Mean
Sinkhorn Softimpute
KNN
LR
Hyperlmpute
MIRACLE
DD
MLP
ICE
MissForest
Figure 20: W2,, Dap. and DakI of shift correction methods on real datasets. Lower is better. (Part 2)
37

---

K Classifier Analysis for Localization and Correction
Figure 21 provides the F-1 score results for DF-Locate when using different classifiers as discrimina- tors within the iterative process 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 provided much faster training times while being highly competitive in localization accuracy.
Simulated Datasets
Heal Datasets
0.9
0.8
0.8
0.6
0.5
0.5
2
0.4
0.3
0.2
0.
0.0
RF
CatBoost LogReg ExtraTree LGBM
SVC
RF
CatBoost Loq ExtraTree LGBM
SVC
Figure 21: Mean F-1 scores by DF-Locate using different classifiers simulated (left) and real (right) datasets. Higher is better.
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 DF- Locate, free-based methods provide similar results, surpassing the linear methods in most metrics. We use Dap as a metric to select our method because provides an estimate that tightly bounds the total variation distance. Namely, we select CatBoost as discriminator as it provides competitive performance with the other tree-based methods in the simulated datasets, and clearly outperforms others in the real datasets.
Simulated Datesets Dhp
W?
D
10"
3 10
0.1
2 10
Datasets
W3
Dhp
Daxl
10a5
0.2
10a0
10a
10
0.1
10
10a
10a
CatBanst
LagReg
ExtraTree
LGBM
SVC
Figure 22: Mean WJ. Dap, and Dau by DF-Correct using different classifiers on simulated (top) and real (bottom) datasets. Lower is better.
38

---

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 given dataset, namely the Energy dataset.
In Figure 23 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 W2, Dhp and DaAS 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.
present the variability of detection methods, conveyed by the and standard
1.0
0.9
0.8
0.7-
0.66
0.4
0.3
0.2
0.0
DataFix
selectKbest
FAST-CMIM
MB-KS KNN-KS
Dnep-SM SM M8-SM
MB KS* KNN-KS KS
Deep SM
MI
MRMR
ME-SM
Figure 23: Variability of detection methods on the Energy dataset. Mean and standard deviation of F-1 score across five different seeds, for 25% corrupted features and all manipulation types. Higher is better.
W3
Dip
Dat
0.6
0.5
0.4
10:
11
0.2
0.1
10a
0.0
10a
DataFix
DD
LR
MIRACLE
Sinkhorn Softimpute
INB
MLP
Hyporimaute
MissForest
Figure 24: Variability of correction methods on the Energy dataset. Mean and standard deviation of W Daps and Dek across five different seeds. Lower is better.
39

---

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 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 classifiet/regressoor 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 [ and the regression results are reported in Table [7] using balanced accuracy for classification and mean square error (RMSE) for regression. Methods are sorted by their performance without non-corrupted features (0%). Notably, the classification/reress// downstream performance when using DataFix surpasses the competing methods 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 Datalix is small.
Table 6: Downstream classification task with Musk2 dataset, with all corrected features and all non- corrupted features (100%), with half of the non-corrupted features (50%), and without non-corrupted features (0%).
Method
Balanced Ace. (0%)
Balanced Acc. (50%)
Balanced Acc. (100%
Original Query
0.951
0.960
0.964
DataFix
0.902
0.953
0.957
LR
0.871
0.932
0,938
Hyperlmpute
0.863
0.932
0.945
ICE
0.848
0.937
0,937
MIRACLE
0.828
0.920
0,937
Sinkhorn
0,818
0,924
0,936
MLP
0.817
0.907
0.911
KNN
0.815
0.925
0.934
INB
0.812
0.950
0.954
MissForest
0.779
0.908
0.944
DD
0.576
0.950
0.956
GAIN
0.574
0.910
0.939
Mean
0.500
0.952
0.961
SoftImpute
0.493
0.620
0.836
40

---

Table Downstream regression task with Energy dataset, with all corrected features and all non- corrupted features (100%), with half of the non-corrupted features (50%), and without non-corrupted features (0%).
Method
RMSE (0%)
RMSE (50 %))
RMSE (100 %
DataFix
0.089
0.082
0.080
Original Query
0.093
0.084
0.082
KNN
0.105
0.094
0.091
Sinkhorn
0.107
0.090
0.085
INB
0,114
0.094
0.092
Mean
0.122
0.088
0.083
GAIN
0.127
0.200
0.279
Hyperlmpute
0.129
0.124
0.123
DD
0.137
0.120
0.116
MissForest
0.147
0,155
0.156
MLP
0.175
0.187
0,200
SoftImpute
0.180
0.182
0.252
LR
0.191
0.189
0.205
ICE
0.217
0.225
0.234
MIRACLE
0.311
0.335
0.343
End-to-end DataFix Analysis
Figure 25 provides five additional examples that illustrate the iterative process of DF-Locate before and after shift correction, simulated datasets 1, 2, 3, 6, and 10. each dataset, there are 200 corrupted features out of 1000. Similar Figure 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 Figure 25 showcases the iterative process applied the corrected query using different methods.
In simulated datasets and 2, the methods DDD, INB, MIRACLE and MLP all yield updated query that results lower empirical divergence. Notably, DD stands out by providing a corrected query with 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 DF-Correct, and simulated dataset 6, INB and DD yield an updated query that reduces the empirical divergence, with both techniques achieving 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.
able lower the empirical divergence for simulated dataset 10. Moving to
DataFix Limitations
DataFix is specifically tailored for tabular datasets and will not work optimally when applied to other data types such as images, vidcos, 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 generally outperforms competing methods in terms of computational efficiency.
41

---

Before shift correction
After shift correction
0.8
0.6
0.4
0.4
0.2
0.2
0.0
0.0
50
100
150
200
250
50
100
150
200
250
100
150
200
250
100
150
200
250
1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
50
100
150
200
100
150
200
250
1.0
1.0
0.8
0.6
0.4
0.2
0.2
0.0
0.0
50
100
150
200
250
50
100
150
200
250
1.0
100
150
100
200
Removed
TVD
DataFix
INB
GAIN
Mean
Sinkhom
Hyperlmpute
Softimpute
Original
DD
MLP
ICE
MissForest
Figure 25: DF-Locate iterative process before (left) and after (right) shift correction. Each row corresponds to different simulated dataset.
42


Reviews:
[{"review_id": "nCZtfOAjSF", "reviewer": "Reviewer_dDZx", "summary": "This 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\n", "questions": "\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", "limitations": "NA. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "it makes an interesting connection between distributional shift and GAN. ", "weaknesses": "Second part sounds a bit hackish. I also have a few questions (see below). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear Reviewer,**\n\n**Thank you for the thoughtful review! We agree that some of the introduced heuristics can feel a “bit hackish”, however, they have proved to be more accurate than previous works and than our attempts with more conventional supervised learning methods. One way to get a better intuition for the correction system is to compare it with a k-NN imputer: a k-NN imputer fills the missing feature values with a combination of the top-k closest samples with respect to L2 distance on non-missing features.** **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 correction task is a divergence-reducing imputation problem.** **We are including a more in-depth discussion, as well as additional figures (see attached pdf), to provide better intuition to the reader and to further motivate our approach.**\n\n&nbsp;\n\n*“statistical distinguishable does not imply computational distinguishable”*\n\n**Definitely, thank you for catching this error, we have re-phrased the text to avoid making this mistake.**\n\n&nbsp;\n\n*“Divergence measures are asymmetric.”* \n\n\n**While many f-Divergences are asymmetric (e.g. KL-Divergence), there are a few that are in fact symmetric, such as the “Hellinger distance” or the “Total Variation” used in this work.**\n\n\n&nbsp;\n\n*“When 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 \n**Deciding to use corrected features or simply discarding them will depend largely on the application at hand. For example, this method (using corrections) is currently being deployed for biomedical applications (for instance genomics), where simple statistics and correlations are desired from the features for scientific discovery. Using corrected features allows one to use all of the data, while not distorting the statistical estimates as severely as using the corrupted data would.**\n"}, {"review_id": "UzSw0CAe0m", "reviewer": "Reviewer_kRpA", "summary": "This 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 )", "questions": "- 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.", "limitations": "The 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear Reviewer,**\n\n**Thank you for the highly constructive suggestions. We want to address each of the questions and concerns:**\n\n\n*“the presentation of the work can be improved…  lot of discussion in the main paper is not essential … moving parts of appendix D and E to the main paper and maybe moving some of the discussion on f-divergences and the relationship to the appendix. I would be happy to revise my score...”*\n\n**We completely agree that many important details of the method are in the appendix and that non-essential theoretical information could be moved from the main text. Therefore, we are moving material from sections 4 and 5 to the appendix and adding a more detailed description of the methods (currently in the Appendix) into the main text. Furthermore, with the extra available page upon acceptance, we will try to place the algorithm boxes into the main text.**\n\n\n*“how the proposal values are constructed, can be a bit confusing ... would request the authors to provide a small example in the paper”*\n\n**We agree that it can be confusing to follow the details of the methods, therefore, we are including multiple figures (see attached pdf) in the appendix to provide a better intuition of the inner workings of localization (attached fig 1) and correction methods (attached fig 2). We are including two additional diagrams describing the methods, which do not fit in the attached rebuttal pdf.**\n\n\n*“The authors propose using divergence measures to evaluate …, I think the purpose of correcting for feature shifts is to ensure the performance of a classifier is the same on the reference as on the corrected dataset … small overall divergence between, not guarantees that class conditioned divergence is small. I will be happy to raise my score if the authors could elaborate on this”*\n\n**While in localization/detection tasks there are well-stablished evaluation metrics, in the distribution shift removal problem, the evaluation procedures are not so standardized. We selected the given evaluation metrics because (1) we wanted application-independent metrics, as the corrected datasets might be used in many scenarios (statistical analysis, visualization, machine learning). For example, the proposed algorithm is already being used to homogenize biomedical data in which statistics and correlations are estimated for scientific discovery. Furthermore, some of the datasets do not even have semantic labels, making an evaluation with downstream classification/regression tasks impossible.  (2) As both data imputation and shift correction are closely related tasks (each tries to replace missing/corrupted values with new ones), we wanted to adopt metrics from the imputation literature. Recent papers [a] have used W2 as an imputation evaluation metric, and we included estimates of the symmetric KL divergence and the HP-divergence as they are computationally tractable, and appeared to be sensible estimates of the quality of the shift correction results. Note that we don’t include MSE as it doesn’t properly characterize divergence between distributions (see Appendix C). Finally, note that in imputation literature, application-independent metrics are used.**\n\n*“under different fraction of shifted features Data-Fix performs much better. How would the performance scale? When could break down?”*\n\n**It depends a lot on the type of manipulation. For example, if the mean of just one feature is shifted largely, DataFix can detect it without problem, despite only one feature being present. For more challenging manipulations, such as removing the correlation between features, a small number of features can make DataFix fail. The true underlying divergence between distributions (which is typically unknown) will partly indicate when DataFix will fail. Figure 18,19 in the appendix, which includes an example of how the true and predicted divergences evolve through the filtering process, might help to get a better intuition.**\n\n*“GANs are not used.  But some of the datasets used are data modalities where NNs can do well. Do you think a GAN based adversarial approach for feature correction could perform better?”*\n\n**We do include a GAN-based method for feature correction: GAIN, which combines reconstruction and adversarial losses. However, the method did not provide competitive results and did not scale well for more large dimensional datasets.**\n\n*“A lot feature shift scenarios were implemented. are there feature shift scenarios which are  encountered more often? where these feature shifts naturally occur?”*\n\n**We are encountering many of the proposed shifts when using the system in a biomedical application. Features collapsing to 0 or being negated are not uncommon, especially within genomic sequences. Furthermore, bad standardization (e.g., incorrectly encoding features with the metric or imperial system) can lead to mean shifts. We are open to suggestions for other shifts so we can incorporate them in future works.**\n\n*“Are there details for the computing infrastructure in the appendix?”*\n\n**We forgot to add the hardware details in the appendix, thanks for noticing that! All experiments were done with the same compute resources which included an Intel Xeon Gold with 12 CPU cores. We are adding all this information in the appendix.**\n\n*“The authors mention in the main paper that limitations are discussed in the appendix, but I wasn't able to find this”*\n\n**The description of the results includes some discussion of the limitations (e.g., which manipulation types tend to be detected incorrectly). However, we agree that a section in the appendix dedicated to discussing the limitations of the method would benefit the paper. We are including such a section.**\n\n**We hope that by addressing your concerns you would consider raising the score. Thank you for helping us improve the quality of the manuscript!**\n\n**[a] ICML, 2020, Jarrett, Daniel, et al. Hyperimpute: Generalized iterative imputation with automatic model selection.**\n"}, {"review_id": "LN0shLrhsp", "reviewer": "Reviewer_MQtc", "summary": "The 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.", "questions": "The 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?", "limitations": "-", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We want to thank the reviewer for the insightful comments. We want to address each of the questions and concerns:**\n\n*“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.”*\n\n**We believe that the reviewer has a perfectly reasonable misunderstanding. In the task of feature shift correction, there is no need for a train/test split. The evaluation process is similar to data imputation evaluation: algorithms take as input a matrix (concatenated reference+query) with some missing (or corrupted) values, and the algorithm outputs the same matrix with the missing/corrupted values replaced. The evaluation is done by taking the imputed/corrected matrix (the output matrix) and comparing it with the matrix before any corruption/missingness was present (the ground truth matrix). As long as the algorithm does not see the “ground truth matrix” during the prediction of the “output matrix”, and only sees the “corrupted input matrix”, the evaluation setting is perfectly valid, and it is the common setting adopted in missing data imputation tasks. Both shift detection and correction tasks follow a similar setting, where the training/testing split is not required. As our algorithms train classifiers internally, as part of the filtering and correction process, we acknowledge that there can be confusion regarding the lack of common train/test split used in regular supervised classification/regression tasks, therefore, we will add some clarifying text in both the main text and appendix to avoid causing any confusion to the readers.**\n\n&nbsp;\n\n*“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?”*\n\n**We selected the optimal hyperparameters by doing a grid search for all methods in both localization and correction tasks. We acknowledge that this was not properly explained in the paper so we are including this information in the main text, and all the details of the search in the appendix. Note that in the localization task, some methods are marked with “*” indicating that the methods make use of the number of corrupted features as extra input, making the comparison between the other methods not fair.**\n\n&nbsp;\n\n*“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. 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.”*\n\n**Thanks for the comment! We agree that a lot of the theoretical part could be moved in the appendix and more details could be included in the main text. We are restructuring the paper to include a more detailed description of the methods in the main text.** \n\n&nbsp;\n\n *“what are the k classifiers and the CatBoost classifiers trained to predict?”*\n\n**All the classifiers trained in localization and correction tasks are trained to classify the samples as “corrupted” vs “non-corrupted” (i.e. reference vs query).**\n\n&nbsp;\n\n*“How are the values of B generated?”*\n\n**B is a set that consists of the feature values within the reference, the feature values after performing imputation with linear regression and nearest neighbor, and a shuffle of feature values from the reference panel. We are adding an additional description of its construction in the appendix.**\n\n&nbsp;\n\n*“How is the objective of equation 9 optimized?”*\n\n**Equation 9 is a simple combinatorial problem. We use the discriminator to evaluate every sample with each proposal b in B, and simply select the b that provides the highest probability of being non-corrupted.**\n\n&nbsp;\n\n*“Error bars of the results are not present, contrary to what has been specified by authors”*\n\n**Thank you for pointing this out; we did not provide a plot with error bars. While performing the whole evaluation benchmark for both localization and correction multiple times is infeasible and expensive (as it would require around 2 months of computing resources), we are including a plot when performing the evaluation of both localization and detection with multiple random seeds for a given dataset, showcasing the variability between runs for each method. The plot is added in the appendix and provided in the attached pdf (attached figures 3 and 4).**\n\n&nbsp;\n\n*“What is the motivation behind definining the threshold as kD (lines 213-214 in text)? Is it not possible for a single feature to cause high divergence? “*\n\n**After normalizing and sorting the features, we select the features with the highest importance until their cumulative sum surpasses tD. Where 0<“t”<1 is a fixed hyperparameter and “D” is the empirical divergence (which is typically non-decreasing). You can imagine “D” as acting as a “learning rate scheduler” where, as the filtering process evolves, forces that fewer features are filtered at each step, to avoid incorrectly removing incorrect features, and “t” acts as a scaling factor. It is completely possible to have a single feature causing high divergence. If that is the case, it is likely that the normalized sorted features will look like a 1 followed by 0s (Note that the normalized and sorted features add to 1), because tD < 1, the first feature will be selected but not the following ones, leading to a correct detection.**\n\n&nbsp;\n\n**We hope that we have addressed all your doubts and concerns, that we have provided enough arguments showing the validity of the adopted evaluation setting, and that you will consider increasing the score.**\n"}, {"review_id": "9TxdUAyEIn", "reviewer": "Reviewer_5LeU", "summary": "This 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.", "questions": "- 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?", "limitations": "Cf. 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.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Dear reviewer,**\n\n**Thank you for your constructive suggestions. We believe there has been misunderstanding regarding the goals and scope of the paper. We introduce two systems: the first localizes faulty features, which can be used to detect incorrectly standardized or processed values or to localize malfunctioning sensors in multi-sensor environments. The second system proposes new values for the corrupted data, by performing divergence-reducing “imputation” of the data. All classifiers described in the paper are discriminators trained to classify between real and corrupted as part of the iterative processes of localization and correction.** \n\n**Thus, your summary line “... and then classifiers trained on samples from p generalize to samples from q” is not performed at any point throughout the paper. While the proposed method could be used as a way to provide more homogenized datasets to train ML systems, training classifiers with out-of-distribution generalization is out of the scope of this paper and not directly comparable to our approaches. In fact, many of the datasets used do not have semantic labels, so training a classifier/regressor with them would not in any case be possible. Feature shift localization (and correction) is useful in many applications and settings and is an independent problem from OOD generalization. The paper [NeurIPS 2020] provides a similar motivation to our one, and we use it as our benchmark baseline.**\n\n&nbsp;\n\n*“... Is the setting under consideration discussed in the OOD generalization literature? If so, it should be highlighted ...”*\n\n**There has been a misunderstanding with respect to the problem that we are solving; OOD generalization is not directly comparable to our work. Our setting provides a mechanism to “detect” data errors and “fix” them, in an application-independent manner.**\n \n**As feature shift localization (and correction) is related to multiple fields of statistics and ML, we are including an extended related work section in the appendix, where we are mentioning OOD generalization, outlier detection, and GANs among others.**\n\n&nbsp;\n\n*“... extra assumptions over data conditional label distributions are required since resulting pruned features are domain invariant. In other words, p(y|x∼p)=p(y|x∼q), which is not stated.”*\n\n**Our setting does not include any classification tasks, and there is no training of supervised ML methods for prediction. The “y” (label) is not part of our framework, in fact, most of the datasets do not have semantic labels.**\n\n&nbsp;\n\n*“Empirical assessment limited to tabular data, and limited to a particular pair of distributions.”*\n\n**Our work focused on tabular data as it is where the problem of feature shift localization most commonly arises. For example, this system is being used to homogenize tabular datasets for medical statistical analysis. We adopt a similar framework as in [NeurIPS 2020] where only two distributions are considered. The proposed method can be applied to more distributions/data sources by applying the system iteratively.**\n\n&nbsp;\n\n*“The proposal and the evaluation ignore representation learning techniques... Invariance-based approaches to out-of-distribution robustness typically project data onto a space where domains cannot be discriminated ...”*\n\n**We do not consider out-of-distribution robustness / representation learning as they are not directly comparable. Those techniques learn a new space that provides better generalization across domains for classification. Here we try to localize and fix erroneous features within datasets.** \n\n&nbsp;\n\n*“Are the resulting data any useful/meaningful after the correction approach is performed?”*\n\n**The correction approach can be seen as a “divergence-removing” imputation. The usefulness and meaningfulness of the corrected dataset will vary largely depending on the application and type of data, in a similar way to data imputation.** \n\n&nbsp;\n\n*“How does the proposal assume labels shift across distributions/domains?”*\n\n**We don’t solve a classification problem and there are no labels.**\n\n&nbsp;\n\n*“How does features selection compare with invariant representation learning?”*\n\n**While our system can generate datasets that could provide more invariance for downstream classification, our approach differs from invariant rep. Learning, as this learns a representation of the input so it has invariance between domains. Here we try to localize and correct a dataset with corrupted features.**\n\n&nbsp;\n\n*“Does inducing invariance across a pair of domains p and q influence invariance across other domains? Under which conditions?”*\n\n**We are trying to detect faulty features and provide corrected datasets; we are not learning an invariant representation.**\n\n&nbsp;\n\n*“How would the feature selection approach fare in more complex structured data, such as images or text for instance?”*\n\n**Applying feature selection methods directly to natural images could lead to poor results. However, the problem we are trying to solve arises most commonly in tabular data.**\n\n&nbsp;\n\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\n**Similar to data imputation, the usefulness of the corrected features will depend on the nature of the data and the application at hand. For example, in order to compute statistics and correlations between features, using a corrected dataset can help to reduce some bias that might appear if the corrupted dataset is used.**\n\n&nbsp;\n\n**We hope that we have been able to clarify the goal and scope of the paper and you will consider increasing the score.**\n\n\n**[NeurIPS 2020] Kulinski, Sean, Saurabh Bagchi, and David I. Inouye. \"Feature shift detection: Localizing which features have shifted via conditional distribution tests.\"** \n"}]

Author rebuttal:
**We would like to thank all reviewers for their constructive comments,**

**First, we are restructuring the text to improve the readability of the paper. As suggested by the reviewers, we are including more details about the proposed methods in the main text, and moving some theoretical results to the appendix. Additional figures (fig 1,2,5 in the attached document) are included in the appendix providing better intuition for the design details of the methods. Furthermore, we are including in both the main text and appendix more details regarding the evaluation setting, hyperparameter search, and computational resources, as well as an additional figure with error bars showing the variability between random seeds (see fig 3,4 in the attached pdf).**

**We want to clarify the motivation for the proposed work: while presented jointly, both the localization and correction work can be useful by themselves. The shift localization can be used to detect incorrectly formatted features, errors arising during data processing, or in multi-dimensional sensor applications, detecting malfunctioning sensors. The work presented in [Neurips 2020] provides a similar setting and set of motivations and acts as our benchmark baseline. The shift correction system can be used in any application where simple imputation would be used, providing a replacement for corrupted/missing values that minimize an empirical divergence. Deciding to use one or both systems is completely dependent on the application. We are modifying the text accordingly to properly depict the motivation of the work.**

**While the proposed methods can be used to extend training datasets for training other ML models (e.g., supervised classification or regression), they are not limited to such tasks, and in fact, many of the datasets used do not have labels, making an evaluation using downstream ML classifiers impossible. In fact, our system is currently being deployed in a biomedical setting in order to homogenize data that is later used to compute statistics and correlations between features for scientific discovery. Therefore, we adopt a more appropriate evaluation setting, independent of the application where DataFix is applied, as is commonly done in related tasks such as data imputation [ICML 2020]. Finally, we want to note that the evaluation setting adopted is correct, no train/test split is required (as in data imputation tasks), and we have a set of datasets (simulated data) for hyperparameter tuning and a set for purely testing (real datasets). We are modifying the text accordingly to properly justify the adopted evaluation setting.**

**The attached pdf with figures include:**
- Figure 1: examples of feature importance during the iterative localization process with different parameters.
- Figure 2: number of features filtered at each iteration, with different hyperparameters.
- Figure 3: error bars for shift correction.
- Figure 4: error bars for shift localization.
- Figure 5: example of proposals in MNIST dataset for shift correction.

**Two extra diagrams providing more details of our proposed methods are included in the appendix.**

**We hope that we have addressed all of your concerns. As the method provides accurate results, it is useful in many applications (in fact, it is already being deployed), and your suggestions have helped us to largely improve the text and structure of the manuscript, we believe that our work is a strong submission to NeurIPS, and we hope that you will consider raising the score.**

**[NeurIPS 2020] Kulinski, Sean, Saurabh Bagchi, and David I. Inouye. "Feature shift detection: Localizing which features have shifted via conditional distribution tests."** 

**[ICML 2020], Jarrett, Daniel, et al. Hyperimpute: Generalized iterative imputation with automatic model selection.**


Meta-review: My recommendation is to accept the paper.

After active discussion with the authors, reviewers agreed that the paper makes a valuable contribution. It does seem like some edits are in order for the camera ready to clarify points that reviewers found confusing during the review and discussion periods, but these seem very feasible post-review. Please look over the reviewer suggestions carefully and incorporate them (some of this is already promised in the rebuttal).
Decision: Accept (poster)