Counterfactual Maximum Likelihood Estimation for Training Deep Networks

1 Introduction . Deep neural networks have been tremendously successful across a variety of tasks and domains in recent years . However , studies have shown that deep learning models trained with traditional supervised learning framework tend to learn spurious correlations as predictive clues [ 1–4 ] . For example , in computer vision , deep learning models can rely on surface-level textures [ 1 , 5 ] or background environment [ 2 , 6 ] instead of the object of interest in images . In natural language processing ( NLP ) , question-answering models are insensitive to the choice of question [ 7 ] and natural language inference models are surprisingly accurate at predicting the logical relationship between a pair of sentences from just one of them [ 3 ] . In image captioning , a phenomenon called object hallucination is observed where models tend to include nonexistent objects in a caption based on their common association with other objects that are actually present in the image [ 4 ] . Another related problem in the current associative learning framework , like maximum likelihood estimation ( MLE ) , is that neural networks can achieve almost perfect performance on one dataset but dramatically fail to generalize on another because of a naturally occurring or adversarially enforced distribution shift [ 8–11 ] . One explanation is that the model learns ‘ fake ’ features induced by spurious correlations instead of invariant features that would hold in different domains for the same task [ 12 ] ; this kind of invariance can be explained by the underlying causal mechanism of the task [ 13 ] . 1Our code is released at https : //github.com/WANGXinyiLinda/CMLE . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . It is usually impractical to identify and disentangle all of the potential sources of spurious correlations . Various training algorithms have been proposed to reduce the influence of spurious correlations by learning the underlying causal mechanism . One approach for discovering the underlying causal mechanism is by invariant prediction based on the invariance of causal mechanisms across different environments [ 13 , 12 , 14 ] . However , constructing such diverse environments is usually unrealistic . Another approach is counterfactual data augmentation [ 15–19 ] , which directly modifies the part of the input that causes the target variable , but usually involves expensive human efforts in generating the counterfactual examples . We propose Counterfactual Maximum Likelihood Estimation ( CMLE ) , which tries to perform MLE on an interventional distribution instead of the observational distribution . Here , obervational distribution means the distribution that the train data is sampled from . Interventional distribution means the distribution with one or some of the variables intentionally intervened ( set to some fixed value ) . More specifically , we aim to remove observable confounders and reduce spurious correlations by directly intervening on the label variable , without requiring human annotators or the curation of diverse environments . In this paper , we focus on the setting of predicting the outcome Y of a certain action T on X , with the underlying causal model as X → Y ← T and an observed train dataset in the form of ( X , Y , T ) . Such a setting enables us to assume that there exist some observed confounders in X that influence both Y and T . This means there is potentially a false causal relation from X to T , indicated by a red arrow in Figure 1a . In our framework , instead of trying to identify all the confounders , we try to rule out the effect of the confounders by considering the interventional distribution that directly deleting the false causal edge from X to T with a resulting causal graph as shown in Figure 1d . This causal graph can be viewed as a Randomized Controlled Trial ( RCT ) , which is the gold standard approach for estimating the treatment effect [ 20 ] . Our method is inspired by individual treatment effect ( ITE ) prediction [ 21 , 22 ] , which tries to predict the expected effect E [ Y1 − Y0|X = x ] ( e.g . differnce in blood pressure ) of a treatment T ∈ { 0 , 1 } ( e.g . drug/surgery ) on a individual unit X = x ( e.g . a patient ) . In this case , the observed confounder in X can be explained by selection bias [ 22 ] , which means the treatment T applied to an individual X is dependent on X . For example , young patients are more likely to be treated by surgery , while elder patients are more likely to be treated by drugs . There are two broad kinds of tasks to which our proposed CMLE framework can apply : relationship prediction tasks that predicts the relation T given two inputsX and Y ( e.g . natural language inference [ 23 , 24 ] , paraphrase identification [ 25 , 26 ] , natural language for visual reasoning [ 27 ] , etc . ) , and conditional generation tasks that generate Y given preconditionX and constraint T ( e.g . style transfer [ 28 , 29 ] , controllable image/text generation [ 30–32 ] , controllable image captioning [ 33 , 34 ] etc. ) . For conditional generation tasks , the causal relation of X , Y , T fits in our setting . For relation prediction tasks , we first assume there is an underlying conditional generation process of Y for a given pair of X and T . This is a natural assumption as this is usually how the datasets for this type of task is generated . Then we augment the original dataset using the counterfactual examples generated by our method . In this paper , we try to reduce the spurious correlations contained in both kinds of tasks . Since we generally can not observe the interventional distribution , we derive two different upper bounds of the expected interventional negative log-likelihood using only the observational distribution : one is Implicit CMLE , as it does not involve any explicit generation of counterfactual examples ; while the other is Explicit CMLE , as it explicitly generates counterfactual examples during the training process . We test our framework with deep neural networks on a simulated dataset and two real-world tasks with well-known spurious correlations : natural language inference [ 23 ] and image captioning [ 35 ] . Compared to regular MLE , we improve the out-of-domain accuracy of NLI on a hard adversarial dataset by 2.9 % relatively and beat baselines on human preference evaluations by more than 10 % , while maintaining comparable performance on automated evaluations . Our results show that our learning framework can better capture the underlying causal mechanism and reduce spurious correlations without degrading performance . 2 Related Work . A growing body of work has investigated algorithmic improvements for machine learning model training by leveraging underlying causal mechanisms . One line of work tries to utilize the invariance of causal mechanisms across different environments . Invariant Causal Prediction ( ICP ) [ 36 ] aims to X Y T ( a ) Observed pO ( X , Y , T ) X Y1 Y2 Ym T=1 T=2 T=m T … YT … ( b ) Expand all possible outcomes X Y1 Y2 Ym T=1 T=2 T=m … ... ( c ) Intervene on T X Y T ( d ) Intervened pI ( X , Y , T ) Figure 1 : The Structural Causal Model ( SCM ) assumed by Counterfactual MLE , where ( a ) and ( b ) have the same distribution , ( c ) is the resulting causal graph of intervening on T . ( d ) is the intervened distribution by only removing the X → T edge from ( a ) . Shaded nodes means being observed and red arrows denote the causal relation causing the spurious correlation . learn a linear invariant causal predictor to identify the causal parents of given target variables with data obtained from different environments . Further refinements on this line of work include nonlinear and nonparametric ICP assessment methods [ 37 ] and Invariant Risk Minimization ( IRM ) where out-of-distribution generalization is enabled by learning a nonlinear invariant causal predictor shared across different environments [ 12 ] . However , these methods generally require a set of environments that are both sufficiently diverse to eliminate confounders while maintaining the underlying causal structure . Though this can be achieved by intervening in the non-causal variables in a synthetic or interactive setting , it is very difficult to obtain or even define such environments with complex high-dimensional observational data like text or images . Our proposed CMLE framework only requires a single observational dataset , making it much more practical . Another line of work approaches this problem with data augmentation . Methods including programmatically editing text to address gender bias [ 16–18 ] , employing crowd workers to augment training data to capture potential unmeasured variables [ 38 ] , and manually editing minimal text spans to flip the predicted label [ 15 ] have all been proposed to produce counterfactually augmented data ( CAD ) for improving out-of-domain generalization . Conceptually , in contrast to the first approach that involves interventions on a potentially large number of non-causal variables , CAD directly intervenes on a target variable [ 39 ] , which is more practical for a specific task . Refinements on this work attempt to better utilize the manually edited or synthetically generated counterfactual data by pairing the factual and counterfactual examples together [ 40 , 41 ] , and leverage local causal structures to generate counterfactual experience in training reinforcement learning algorithms [ 19 ] . While human annotators can generate high-quality counterfactual examples that are effective in reducing spurious correlations , they are expensive to obtain and do not easily scale . Meanwhile , our proposed CMLE framework is fully automated , requiring no human annotations . Our CMLE framework is inspired by a technique for predicting the Individual Treatment Effect ( ITE ) from observational data [ 22 ] . This method only observes one possible outcome of an individual X = x with one binary treatment T = 0 or T = 1 , and it assumes there are no hidden confounders but potentially observed ones . Assuming a similar underlying causal model as ours ( m = 2 ) gives an error bound as the sum of the standard generalization error and the distance between the control and treated distributions in a representation space . Improvements on this work include disentangling the observable confounders in ITE using representation learning [ 42 ] , and estimating the causal effects of linguistic properties [ 43 ] . While in this work we do not aim to separate the confounders from the causal factors in X , we try to reduce the spurious correlations caused by the possible confounders contained in X using our proposed CMLE training framework . 3 Method . In this section , we first introduce our problem setting and the goal of Counterfactual Maximum Likelihood Estimation ( CMLE ) . As we can not observe the counterfactual outcomes in observational data , we propose two different ways of estimating the counterfactual likelihood by deriving two different upper bounds : Implicit CMLE , which is faster in training , and Explicit CMLE , which is more flexible . We also briefly introduce the deep learning architecture we use in our experiments . 3.1 Problem Setting . We consider the scenario with one input variable X ∈ X that contain all the confounders , one discrete treatment variable T ∈ { 1 , 2 , ... , m } where m ≥ 2 , and one outcome variable Y ∈ Y , as shown in Figure 1a . Note that X and Y can be very complicated high dimensional space like text or image and all three variables are observed in the dataset . Suppose the data we observed is D = { ( x1 , t1 , y1 ) , ( x2 , t2 , y2 ) , ... , ( xn , tn , yn ) } and only xi ’ s are sampled i.i.d . In this paper , we adopt the causal notations from [ 44 ] . Below we give a formal definition of Structural Causal Model ( SCM ) [ 44 ] that we use to model our framework : Definition 1 . A Structural Causal Model ( SCM ) is a pair 〈M , p ( u ) 〉 , where M is a triple 〈U , V , F 〉 and p ( u ) is a probability function defined over the domain of U . U = { U1 , U2 , ... , Un } is a set of background variables ( exogenous ) that are determined by factors outside the model . V = { V1 , V2 , ... , Vn } is a set of variables ( endogenous ) and F = { f1 , f2 , ... , fn } is a set of functions s.t . Vi = fi ( Pa ( Vi ) , Ui ) , i = 1 , 2 , ... , n , where Pa ( Vi ) ∈ V \Vi are the parents of Vi . We say that “ Vj is a direct cause of Vi '' if Vj ∈ Pa ( Vi ) and illustrate this relation as Vj → Vi in a causal graph . In Figure 1 , we omit the exogenous and only show the causal relations between endogenous variables . The shaded nodes indicate that the corresponding variables are observed . Intervention on a specific variable is formalized as fixing the value of this variable to a specific value . This is known as the do-operation , changing the causal graph by removing all edges incident to the intervened variable . We denote the outcome of Y corresponding to do ( T = t ) as Yt . We expand all the possible outcomes of Y in Figure 1b while we actually can only observe one outcome for each data point in the observational dataset . We also make the following standard assumption [ 22 , 45–49 ] : Assumption 1 . ( Strong Ignorability Condition ) ( Y1 , Y2 , ... , Ym ) ⊥ T |X and 0 < p ( T |X ) < 1 . Under this condition , the counterfactual outcomes are identifiable since p ( Yt = y|X = x ) = p ( Y = y|X = x , T = t ) . i.e . The counterfactual likelihood can be written as observable conditional probabilities estimated from data . Our goal is to eliminate the false causal relation from X to T represented by the red arrow in Figure 1a . To achieve this goal , we define our objective over the interventional distribution pI ( X , Y , T ) as shown in Figure 1d , instead of over the observed data distribution pO ( X , Y , T ) as shown in Figure 1a . To construct the interventional distribution , we only delete the causal edge between X and T and keep all the other causal relationships unchanged . In this distribution , T is uniformly sampled from { 1 , 2 , ... , m } . Note that pI ( Y |X , T ) = pI ( YT |X ) = pO ( YT |X ) = pO ( Y |X , T ) . As p ( YT |X ) is the same for both observational distribution and interventional distribution , we omit the superscript for simplicity . Then for predicting Y , the objective of counterfactual maximum likelihood estimation ( CMLE ) would be : EpI ( X , Y , T ) [ − log pIθ ( Y |X , T ) ] = EX [ 1 m m∑ i=1 EYi|X [ − log pθ ( Yi|X ) ] ] ( 1 ) Similarly , for predicting T , we have : EpI ( X , Y , T ) [ − log pIθ ( T |X , Y ) ] = EX [ 1 m m∑ i=1 EYi|X [ − log p I θ ( T = i|X , YT ) ] ] ( 2 ) Here θ is the learnable parameters . In the remaining part of paper , the probabilities without superscripts are default to be observational probabilities . In the following sections , we first focus on deriving two different upper bounds of the CMLE objective for predicting Y : Implicit CMLE and Explicit CMLE . In Implicit CMLE , we add a regularizer to balance the distribution of X given T in the representation space . For Explicit CMLE , we directly generate counterfactual examples during training . To predict T , we train a classifier using normal MLE on a augmented dataset with the generated counterfactual examples to estimate the interventional distribution pI ( X , Y , T ) . We define a counterfactual example as follows : Definition 2 . A counterfactual example is an example ( x , t , y ) with y sampled from p ( Yt|X = x ) s.t . ( x , t , y ) /∈ D and there is an example ( x , t′ , y′ ) ∈ D where t′ 6= t. We can rewrite our objective function in Equation 1 as Ex [ ∑m t=1 Lθ ( x , t ) ] , by defining the expected loss function of a specific input and treatment as : Definition 3 . We denote the loss function for predicting Y as Lθ ( x , t , y ) = − log pθ ( Yt = y|X = x ) . Then the expected loss function given X = x and T = t is Lθ ( x , t ) = EYt|X=x [ Lθ ( x , t , Yt ) ] . Because of the Strong Ignorability Assumption , we have Lθ ( x , t ) > 0 . Then we define an expected factual loss T=tF and expected counterfactual loss T=t CF when observing a specific treatment T = t : Definition 4 . The expected factual loss and expected counterfactual loss given T = t are T=tF = EX|T=t [ Lθ ( X , t ) ] and T=tCF = EX|T 6=t [ Lθ ( X , t ) ] , respectively . Here T=tF measures how well the model performs on the observational distribution while T=t CF measures how well the model would perform on an alternative world where a different treatment is chosen . Then we can rewrite our CMLE objective as : Ex [ m∑ t=1 Lθ ( x , t ) ] = m∑ t=1 [ p ( T = t ) T=tF + p ( T 6= t ) T=tCF ] ( 3 )

Summary: The paper tackles the problem of spurious correlations caused by observed confounders and offers two potential solutions: implicit and explicit counterfactual maximum likelihood estimation. Specifically the paper deals with the setup of predicting the outcome Y of some action T on X and considers X to potentially be a confounder. The paper compares the performance of the proposed methods to on a natural language inference and an image captioning task and shows improvements in terms of human evaluation but little difference in terms of automatic evaluation.