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: In this work, authors propose the Counterfactual Maximum Likelihood Objective CMLE, where the parameters of the statistical model are inferred by maximizing the likelihood of the data under the interventional distribution, as opposed to the observational data. For the assumed model X -> Y <- T and X -> T, where X is the covariate, T is the treatment and Y is the outcome, the objective is to learn the parameters \theta that predict Y|X=x,T=t using just the causal link T->Y and not the spurious features in X=x. Ideally, if one could intervene on T and set do(T=t), it is possible to collect data where the outcome is observed upon the intervention done on T. Then maximizing the following likelihood: (CMLE objective) \argmax_\theta \sum_t E_{X|t}  E_{Y|x,t} p_\theta(Y|X=x,do(T=t)) would give us the MLE estimate of the parameter \theta under the set of interventional distributions do(T=t), thereby eliminating any effect of the spurious features in X on the prediction of the outcome label Y. The paper proposes two upper bounds for the above CMLE objective: i) Implicit MLE: focuses on obtaining better representations of the covariates X, such that the distribution of these representations is similar (by Wasserstein distance) for every value of T; and ii) Explicit MLE: where counterfactual examples are generated during training. Expectations in the theoretical derivations of the above upper bounds, can be estimated by Monte Carlo evaluations using only the observational data. Results are provided on two downstream tasks where models often rely upon spurious correlations: i) Natural Language Inference (NLI); and ii) Image Captioning (IC).