{'fraction_non_alphanumeric': 0.03163136700378101, 'fraction_numerical': 0.009006793520027052, 'mean_word_length': 5.0457164095892955, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Counterfactual confounding adjustment for feature representations learned by deep models: with an application to image classification tasks

Causal modeling has been recognized as a potential solution to many challenging problems in machine learning (ML). While counterfactual thinking has been leveraged in ML tasks that aim to predict the consequences of actions/interventions, it has not yet been applied to more traditional/static supervised learning tasks, such as the prediction of labels in image classification tasks. Here, we propose a counterfactual approach to remove/reduce the influence of confounders from the predictions generated a deep neural network (DNN). The idea is to remove confounding from the feature representations learned by DNNs in anticausal prediction tasks. By training an accurate DNN using softmax activation at the classification layer, and then adopting the representation learned by the last layer prior to the output layer as our features, we have that, by construction, the learned features will fit well a (multi-class) logistic regression model, and will be linearly associated with the labels. Then, in order to generate classifiers that are free from the influence of the observed confounders we: (i) use linear models to regress each learned feature on the labels and on the confounders and estimate the respective regression coefficients and model residuals; (ii) generate new counterfactual features by adding back to the estimated residuals to a linear predictor which no longer includes the confounder variables; and (iii) train and evaluate a logistic classifier using the counterfactual features as inputs. We validate the proposed methodology using colored versions of the MNIST and fashion-MNIST datasets, and show how the approach can effectively combat confounding and improve generalization in the context of dataset shift. Finally, we also describe how to use conditional independence tests to evaluate if the counterfactual approach has effectively removed the confounder signals from the predictions.


Introduction
DNN models are able to learn complex prediction rules from biased data, tarnished by confounding factors and selection biases. Training DNN models in biased datasets allows the models to carelessly leverage all correlations in the data, including spurious correlations contributed by these data biases, which are unrelated to the causal mechanisms of interest. The unwanted consequence is a potential decrease in generalization performance under dataset shift. As a hypothetical example, consider a slight variant of the though experiment described in [2], where the goal is to classify cows versus camels. Suppose that the training set is comprised of pictures taken outdoors under natural settings where 99% of pictures of cows were taken in green pastures while 99% of pictures of camels are taken in deserts. In this situation it is possible that a convolutional neural network (CNN) might learn to exploit landscape color to classify cows versus camels, so that the CNN might easily misclassify pictures of cows taken, for instance, in a beach. If this is the case, and we want to evaluate the classifier performance on a test set where only 80% of the cow and camel pictures are taken in green pastures and beige deserts, respectively, then we might see degraded generalization performance.
In situations where the context where the picture is taken (e.g., green pasture versus desert in the above example) is available, we can use this observed confounding variable to try to remove, or at least reduce, the impact of the confounder on the predictive performance of the classifier. Confounding adjustment is an active research area in ML, with balancing approaches being commonly adopted remedies to prevent DNNs from leveraging spurious associations induced by confounders.
In this paper, we propose an alternative approach inspired by the counterfactual confounding adjustment described in [5]. While the approach in reference [5] is only directly applicable to linear structural causal models, here we show how it can be easily adapted to remove confounding biases from the learned representations generated by DNNs. The key insight is that by training a highly accurate DNN using softmax activation at the classification layer, and adopting the feature representation learned by the last layer prior to the output layer (as illustrated in Figure 1) as the features of a logistic regression model, we have that, by construction, our features will fit well a logistic regression model, since the softmax classification used to generate the output of the DNN is essentially performing logistic regression classification. Therefore, even though the associations between the raw inputs and the labels can be highly non-linear, the associations between the learned representation/features and the labels can be well captured by a linear classifier. As a consequence, we can use a simple linear regression model to adjust for confounding, using the counterfactual confounding adjustment proposed by [5]. Throughout the text we denote this process of learning a feature representation that fits well a linear classifier as the "linear modeling step" performed by DNNs.
The counterfactual confounding adjustment is presented in detail in Algorithm 1 in Section 4. In a nutshell, the approach is implemented by regressing each learned feature on the confounder and label variables in order to estimate regression coefficients and residuals, and then generate new counterfactual features by adding back the estimated residuals to a linear predictor which excludes the confounder variables. By doing so, we generate counterfactual features whose association with the confounder variable is exclusively mediated by the labels, and then use these counterfactual features as the inputs of logistic regression classifiers, which no longer leverage the associations with the confounder in their

predictions.
We validate the proposed counterfactual confounding adjustment using a variant of the "colored MNIST" dataset described in reference [2], which corresponds to a synthetic binary classification dataset derived from the MNIST data. Following [2], we assign label "0" to digits {0, 1, 2, 3, 4}, and label "1" to digits {5, 6, 7, 8, 9}, and color each digit image as either red or green, in a way that most of the training set images labeled as "0" are assigned the green color and most of the training set images labeled as "1" are assigned the red color. Using this dataset we evaluate the effectiveness of the approach in a series of experiments in both the: (i) absence of dataset shift, where the test set is colored in exactly the same manner as the training set; and (ii) presence of increasing amounts of dataset shift, where we decrease the proportion of green images labeled as "0" and red images labeled as "1" in the test set, in order to generate a shift in the joint distribution of the confounder (red or green color) and label ("0" or "1") variables. In order to validate our approach in a more challenging dataset, we also repeat our experiments using the "colored fashion-MNIST" dataset, generated by coloring the fashion-MNIST dataset [19] in the exact same way as the MNIST data.
Our experiments show that while increasing amounts of dataset shift lead to increasing degradation of the generalization performance for classifiers trained without any adjustment, the classifiers built with counterfactually adjusted features still achieved strong generalization performance across all levels of dataset shift. We also show that in the absence of dataset shift, a situation where the color variable improves the predictive performance, the counterfactual adjustment leads to a decrease in predictive performance, as it prevents the classifier from exploiting the image color for the predictions. Comparison against the SMOTE [7] balancing approach shows the proposed approach achieves comparable predictive performances, while still having the additional advantage of not requiring the re-training of DNNs using balanced training data.
The remaining of this paper is organized as follows. Section 2 provides notation and background required for the other sections. Section 3 describes the causal model underlying image databases. Section 4 describes in detail the proposed counterfactual adjustment. In Section 5 we show how to use conditional independence tests proposed in [4] to evaluate the effectiveness of the counterfactual adjustment. Sections 6 and 7 describe the colored MNIST and colored fashion-MNIST experiments, respectively. Section 8 presents the comparison against the SMOTE balancing approach. Section 9 shows how the adoption of a linear classifier is essential part of the approach. Section 10 describes related previous work. Finally, Section 11 presents final remarks and puts our work in a broader context of applications outside image classification.

Notation and background
Throughout the text, we let Y , A, and P represent, respectively, the classification labels, the confounder variable, and the raw image data (pixels). We use X to denote a feature representation learned by a DNN, and denote the set of learned features by X = {X 1 , X 2 , . . . , X k }, where k represents the number of units in the hidden layer prior to the output layer (e.g., k = 3 in the cartoon DNN in Figure 1). The counterfactual versions of the learned features are represented byX. Selection variables are represented S. We let the superscripts tr and ts represent, respectively, the training and test sets.
In this work, we subscribe to Pearl's mechanism-based approach to causation [11] where the joint probability distribution of a set of random variables is supplemented by a directed acyclic graph (DAG), describing the qualitative assumptions about the causal relation between the variables, also denoted as a causal diagram/graph. In this framework, the nodes on the causal graph represent the random variables, and the directed edges represent causal influences of one variable on another. We represent observed variables by round frames, and variables that have been assigned an specific value (i.e., variables that have been conditioned on) by squared frames.
A path is defined as any unbroken, nonintersecting sequence of edges in a DAG, which may go along or against the direction of the arrows. A path is said to be d-separated or blocked [11] by a set of nodes W if and only if: Otherwise, the path is said to be d-connected or open. A joint probability distribution over a set of variables is faithful [17,11] to a causal diagram if no conditional independence relations, other than the ones implied by the d-separation criterion are present. We adopt the notation V 1 ⊥ ⊥ V 2 and V 1 ⊥ ⊥ V 2 to represent marginal statistical dependence and independence, respectively. Conditional dependencies and independencies of V 1 and V 2 given V 3 are represented, respectively, by

The causal model underlying image classification databases
We first describe the "nature causal mechanisms" generating the observed variables, Y , A, and P , in a image classification database. While the causal model is general, for concreteness we describe it in terms of the cows versus camels classification example, and let the nature variable Y = {cow, camel} represent the labels and A = {green, beige} represent a surrogate variable for the environment (pasture or desert) where these animals live under natural circumstances. We let S represent a selection variable (described in more detail below). Figure 2 presents the causal diagram. The arrow Y → P indicates that the real world cow or camel causes the observed patterns and intensities of pixels in the image. Note that, in terms of the nature causal mechanisms, the cow or camel precedes the picture and when we take a picture of a cow or a camel the shape and color of these animals is imprinted as a pattern of pixels with varying intensities. Clearly, the observed pixels in the picture cannot cause a cow or a camel in the real world 1 . Similarly, the arrow A → P indicates that while the pasture or desert environments also influence the observed intensities and patterns of pixels in the image, the observed pixels cannot "cause" a pasture or desert environment in the natural world. Finally, note that the arrows pointing from both Y and A towards the square-framed S variable indicates the presence of a selection mechanism generating an association between Y and A in the database. Here, the squared frame around S indicates that we are conditioning on the S variable.
As an example, consider a database (denoted database 1) where the cow pictures were taken in organic livestock farms in France, while the camel pictures were taken in Middle East deserts. For this database, most of the cow pictures will be associated with green pasture backgrounds, while most of the camel pictures will be associated with beige desert backgrounds. Hence, conditional 1 One can argue that looking at the picture causes the human cognition to label it as a cow or a camel. However, this "human cognition causal process" is different from the "nature causal process" generating the observed variables. This distinction between nature causal mechanisms and human cognition causal processes was first presented in Figure 6 of reference [2]. Because we are interested in modeling the causal processes underlying a machine learning classification task, in our work we will focus on the nature causal mechanism giving rise to the observed data.
on the fact that database 1 was generated from pictures taken in France and in the Middle East, we observe a strong association between Y and A. In this example, we define the selection variable S as a binary variable assuming value 1 when a picture was taken in France or in the Middle East, and value 0 otherwise. Because any ML models generated from database 1 are conditional on the fact that that the data came from France or the Middle East, i.e., S = 1, we represent S using a square frame. (Note that by conditioning on S = 1 the A and Y variables become d-connected.) Now, consider a second database (denoted database 2) where the cow pictures were taken in natural grassland farms in semi-arid African countries and the camel pictures were taken in the Sahara desert. In this second dataset, the association between the green/beige background color and the cow/camel labels is expected to be weaker than in database 1, because except for the cow pictures taken during the rainy season, when the natural grasslands are green, the pictures taken during the dry season will also show beige backgrounds of dried grasslands. For this second dataset, we have that the S represents a binary selection variable assuming value 1 when a picture was taken in Africa, and 0 otherwise. Due to the different selection mechanisms, there is a clear drift in the joint distribution of {A, Y } in database 2 relative to database 1.
The causal diagram in Figure 2 shows that A represents a confounder of the relationship between the labels, Y , and the pixels, P . Figure 3 describes two ways we can potentially remove the spurious associations between Y and P generated by A. Panel a shows that we could attempt to remove the association between Y and A by balancing the data, while panel b shows that we could try to remove the causal influence of A on P by using a counterfactual approach. In this paper, we take the latter route. However, instead of attempting to directly model the highly non-linear causal relations Y → P and A → P , we leverage the "linear modeling step" described before and work with the learned representation, X, as described in the next section.

Counterfactual confounding adjustment for learned representations
Using the "linear modeling step" we generate a set of learned features, X.
Training a logistic regression classifier using X as inputs is essentially equivalent as training a DNN on the raw input data. Clearly, because X simply corresponds to a transformation of the raw data P , we can still represent the data generation process associated with Y , A, and X by the causal diagram in Figure 4a, where P is replaced by X.
However, training a logistic regression classifier on X will still generate predictions biased by A. In order to remove/reduce the influence of A on the predictive performance of the classifier, we apply the counterfactual adjustment proposed in [5] to generate counterfactual features,X, according to the causal diagram in Figure 4b. The key idea is to apply a modification of Pearl's three step approach for generating deterministic counterfactuals [11] where we: (i) regress each separate feature on the labels and the confounders; (ii) estimate the regression coefficients and model residuals; and (iii) generate new counterfactual features by adding back the estimated residuals to a linear predictor that no longer includes the confounder variable. Mechanistically, the approach is implemented according to Algorithm 1: Note that we do not make use of the test set labels in the estimation of the counterfactual test set features. Instead, we use the predicted test set labels (line 2 of the Algorithm 1) and the training set regression coefficients (line 5) in the computations.
Once the counterfactual features have been generated by Algorithm 1, we can then useX tr to train a logistic regression classifier, and then useX ts to generate predictions that are no longer biased by the confounder (or at least impacted by a lesser degree). In the next section, we describe how to empirically evaluate if the adjustment is working well in practice.

Empirical evaluation of the adjustment's effectiveness
Following [4], we describe here a simple approach to evaluate if the counterfactual adjustment is indeed removing the spurious associations generated by the confounders from the predictions generated by the logistic regression classifier.
The key idea is to consider the causal diagram underlying the classification task based on the counterfactual features (Figure 5a), and evaluate if the conditional independence (CI) relations between Y ts and the classifier prediction scores, R ts , and A ts hold in the test set data. If the counterfactual adjustment is effective (and faithfulness holds) then, from the application of d-separation [11] to the simplified causal graph in Figure  5b, we would expect to see the following pattern of CI relations in the data: That is, we would expectR ts and A ts to be conditionally independent given Y ts , since conditioning on Y ts blocks the path A ts → Y ts → R ts in Figure 5b. On the other hand, if the adjustment has failed (or if no • Using the training set, estimate regression coefficients and residuals from the linear regression model, and then compute the respective counterfactual feature,X tr j , by adding back the estimated residuals,ˆ tr A tr i , to the quantitŷ µ tr j +β tr X j Y Y tr (which represents the linear predictor obtained by excluding the confounder variables). That is, estimate the counterfactual feature as, • Using the test set, compute the counterfactual feature, Result: Counterfactual training features,X tr , and counterfactual test features,X ts . adjustment was performed), we would expect to see the conditional association R ts ⊥ ⊥ A ts | Y ts in the data. In the next sections we illustrate the application of these sanity checks in our experiments with the colored MNIST and colored fashion-MNIST data.

Colored MNIST experiments
We validated the proposed confounding adjustment using a variant of the "colored MNIST" dataset described in reference [2]. For each grey-scale image in the MNIST data we generate a synthetic RGB colored image by assigning the original grey-scale image pixels to either the red channel or the green channel of the synthetic colored image, depending on whether the image was colored as red or green. Similarly to [2], we assign label "0" to digits {0, 1, 2, 3, 4}, and label "1" to digits {5, 6, 7, 8, 9}, and color each digit image as either red or green, in a way that 98% of the training set images labeled as "0" are assigned the green color and 98% of the training set images labeled as "1" are assigned the red color. Using this training dataset, we evaluate the effectiveness of the approach in a series of separate experiments, with increasing amounts of dataset shift in the test set relative to the training set. Explicitly, we conducted the following 6 experiments: (i) "no shift experiment" where the joint distribution of labels and colors in the test set was the same as in the training set; (ii) "shift experiment 1" where 90% of the "0" labels were colored as green and 90% of the "1" labels were colored red; (iii) "shift experiment 2" where 70% of the "0" labels were colored as green and 70% of the "1" labels were colored red; (iv) "shift experiment 3" where 50% of the "0" labels were colored as green and 50% of the "1" labels were colored red; (v) "shift experiment 4" where 30% of the "0" labels were colored as green and 30% of the "1" labels were colored red; and (vi) "shift experiment 5" where 10% of the "0" labels were colored as green and 10% of the "1" labels were colored red. Figure 6 presents mosaic plots comparing the joint distribution of labels and color in the training set against the test set joint distributions in each of the 6 experiments. In all experiments we trained convolutional neural networks (CNNs) for 10 epochs using the RMSprop optimizer with learning rate 0.001, and mini-batches of 128 images according with the following architecture: Input → Conv2d(32, (3,3), relu) → Conv2d(32, (3,3), relu) → MaxPooling2d((2, 2)) → Dropout(0.25) → Conv2d(32, (3,3), relu) → Conv2d(32, (3,3), relu) → MaxPooling2d((2, 2)) → Dropout(0.25) → Dense(16, relu) → Dropout(0.5) → Dense(2, softmax), where: Conv2d(f , s, a) represents a convolutional layer with number of filters given by f , strides given by s, and activation given by a; MaxPooling2d(p) represents the max-pooling operation with pool-size p; Dropout(r) represents a dropout layer with dropout rate r; and Dense(u, a) represents a dense layer with number of units given by u, and activation function given by a. (Note that the learned representation used in these illustrations corresponds to the values extracted from the 16 units in the Dense(16, relu) layer.) All experiments were performed using the keras [1] R package [13] (a R wrapper of Keras[20]). Figure 7 reports the results. The boxplots represent the results over 100 independent replications of each experiment, where for each replication a different set of randomly selected images were colored as red or green (in both the training and test sets). Panel a reports the test set accuracies obtained by logistic regression classifiers trained with the CNN's 16 learned features (without any adjustment) across the 6 experiments and shows a pronounced reduction in generalization performance with larger amounts of dataset shift. These results illustrate that the feature representations learned by CNNs trained on the highly unbalanced training set were able to effectively capture the color signal. Panel b shows the test set accuracies obtained by logistic regression classifiers trained and evaluated with the counterfactual feature representations. Note the strong improvements in generalization performance in the experiments involving strong amounts of dataset shift. (Observe, as well, that the drop in accuracy in the "no shift experiment", and to a lesser degree in the "shift experiment 1", is expected since under no to small amounts of dataset shift the color actually helps the predictions, but the counterfactual adjustment removes its contribution.) Finally, panels c and d present the results of the conditional independence "sanity checks" for the "no shift experiment" (panel c) and "shift experiment 5" (panel d), for both the "no adjustment" (orange boxplots) and counterfactual adjustment (blue boxplots) cases. In both experiments, the blue boxplots show that the CI relations are consistent with the causal graph in Figure 5b, showing that the counterfactual approach effectively removed the contribution of the color confounder to the predictive performance (as that the partial correlation betweenR ts and A ts given Y ts is close to 0), while the orange boxplots provide evidence that the color confounder is still contributing to the predictive performance for classifiers trained with the unadjusted learned features (i.e., R ts ⊥ ⊥ A ts | Y ts ). The results for the other experiments are similar, and are reported in Supplementary Figure S1.

Colored fashion-MNIST experiments
We also repeated our experiments using the fashion-MNIST dataset [19]. This dataset shares the same image size, number of classes, and training and test set sample sizes as the original MNIST data, and is often used as a direct drop-in replacement for the MNIST dataset. It contains grey-scale images of clothing items such as shirts, coats, shoes, sandals, etc, and represents a more challenging classification task. We colored the fashion-MNIST images and performed the exact same experiments as we did for the colored MNIST experiments. The results were largely consistent, although the predictive performances were slightly but consistently lower and noisier than in the colored MNIST experiments. The results are reported in Supplementary Figure S2.

Comparison with SMOTE
We compared the proposed counterfactual confounding method against the SMOTE (Synthetic Minority Over-sampling TEchnique) balancing adjustment [7] using the colored MNIST data. The basic idea behind SMOTE is to balance the training set by under-sampling examples from the majority categories and over-sampling synthetically modified examples from the minority categories.
In our experiments, the training set (Figure 6a) contained the following four {label, color} categories (at the following proportions): {"0", "green"} (49%); {"1", "green"} (1%); {"0", "red"} (1%); and {"1", "red"} (49%). Application of the SMOTE adjustment generated balanced datasets with 25% of examples in each of the 4 categories where we randomly under-sampled images from the {"0", "green"} and {"1", "red"} categories, and over-sampled synthetically modified versions of images in the {"0", "red"} and {"1", "green"} categories. (Namely, for each one of the 600 images in each of these minority categories we generated 24 additional randomly rotated versions of the original image.) We repeated the same 6 experiments reported in Section 6 using the SMOTE adjustment (where we applied SMOTE to the same 100 distinct training sets used in the counterfactual adjustment experiments reported in Figure 7), and compared the SMOTE performance to the proposed counterfactual approach. Figure 8 reports the results and shows the approaches are on average comparable with SMOTE (panel a) performing slightly better under the experiments based on small amounts of dataset shift, while the counterfactual approach (panel b) outperformed SMOTE in the experiments with stronger amounts of dataset shift. It is important to point out, however, that the counterfactual approach has the additional practical advantage that it does not require the re-training of the DNN using balanced data. This is particularly advantageous in applications involving massive datasets and large DNNs that have already been trained on biased data.

On the need of linear classifiers
Here, we illustrate how the adoption of a linear classifier is an essential part of the proposed approach. Using random forest classifiers, we show how non-linear classifiers trained on the counterfactual features can still exploit the confounder signal and suffer from decreased generalization performance under dataset shift conditions. To this end, we repeated the same 6 experiments reported in Section 6 (colored MNIST) comparing the generalization performance obtained by adopting random forest and penalized logistic classifiers (based on ridgeregression and lasso penalties) as the shallow classifiers trained on the learned feature representations. Figure 9 reports the results. The top panels show the results for the random forest classifier, where we see that the counterfactual adjustment (panel b) is still showing an accentuated decrease in generalization performance for the experiments with strong shift. Panel c reports the partial correlation between the classifier prediction and the color confounder conditional on the label variable, for classifiers trained on the original/unadjusted learned features (orange boxplots), as well as, on the counterfactual features (blue boxplots). The fact that cor(R, A | Y ) is removed from 0 illustrates that the random forest prediction is still leveraging the color signal (even when trained on the counterfactual features). The middle and bottom row panels show the results based on penalized logistic regression models (using the exact same learned feature data as the random forest), and illustrate how these linear classifiers achieve better generalization when trained on counterfactual features (panels e and h).
(For completeness, Supplementary Figures S3, S4, and S5 show the CI patterns for each of these 3 models).

Related work
Published work in the intersection between causality and ML can be roughly split into two different categories. The first, largely focus on the estimation of causal effects, and only uses supervised ML techniques as a tool to improve the estimation of causal effects [9]. The second, uses causality to improve the robustness and generalizability of ML predictions. Here, we focus on this second category. The potential benefits of causal modeling in challenging ML tasks has been well articulated in [12,14]. Causal approaches based on counterfactual thinking have been used to predict the outcomes of different actions, policies, and interventions using observational data [3,18,8,15], as well as, to generate causality-aware predictions in static ML tasks based on linear structural causal models [5]. An interesting line of research has also explored the close connection between invariance and causation aiming to learn invariant associations across different environments in order to improve generalization under dataset shifts [2]. Related work in the particular context of image classification include [6,10,16].

Discussion
In real life applications, the effective use of counterfactual adjustments require the specification of a model that fits well the observed data. For instance, attempting to fit a linear regression model to variables showing non-linear associations may lead to the estimation of regression coefficients that fail to capture the associations between outputs and inputs, so that these associations might still be leaked to the residuals. The main insight in this paper is the realization that the "linear modeling step" performed by DNNs allows for effective counterfactual adjustment in anticausal classification tasks where we can train an accurate DNN. (Anticausal tasks correspond to prediction problems where the output variable has a causal influence on the inputs.) While we have only illustrated the approach using CNNs in a simple binary classification task influenced by a single discrete confounder variable, the approach might also be potentially used with alternative DNN models, in multiclass tasks, and accounting for multiple discrete and/or continuous confounders. However, empirical performance on these alternative settings is yet to be evaluated.
In addition to the image classification applications used as illustrations in this paper, the approach can be potentially applied outside computer vision problems. For instance, diagnostic tasks in health applications, where the goal is to predict disease status based on input variables that measure disease symptoms, represent another important class of anticausal classification tasks that are often plagued by confounding biases.
Finally, it is important to point out that the proposed approach is trivial to implement (as the generation of the counterfactual features only require the computation of linear regression fits), and that it does not require re-training DNNs (an important practical advantage in applications involving very large neural networks trained on massive datasets).  The results for the "no shift experiment" and "shift experiment 5" are presented in Figure 7.