CAGE: Probing Causal Relationships in Deep Generative Models

1 INTRODUCTION . Causal models play a fundamental role in analyzing real-world datasets , as many of the questions underlying scientific research require reasoning beyond simply associations ( Pearl , 2009 ; Rubin , 2005 ) . Fundamentally , a causal model characterizes the generative process of the observed data via establishing cause-effect relationships between attributes ; for example , a disease causes a symptom effect ( and not vice versa ) . Consequently , a causal analysis is concerned with data generated and collected in the real world via natural processes and biological agents ( Glymour et al. , 2019 ) . In contrast , a major series of advancements in machine learning in the last decade concerns deep generative models that are learned via statistical objectives to artificially generate high-dimensional data . The widespread use of these models across the artificial intelligence spectrum , from vision ( Karras et al. , 2020 ) and language ( Brown et al. , 2020 ) to scientific discovery ( SanchezLengeling & Aspuru-Guzik , 2018 ) and decision making ( Chen et al. , 2021 ) , raises an important question : What are the causal mechanisms governing a deep generative model ? With a statistical training objective , it is unclear to what extent ( if any ) the causal relationships from the observed dataset are retained within a generative model and consequently , the outputs expected from the model under interventions and counterfactuals . In addition to controlling and interpreting these black-box models , such a fine-grained understanding can be paramount for formally characterizing their biases and generalization capabilities prior to high-stake deployments . Present Work . We propose CAGE , a framework for inferring implicit cause-effect relationships in the latent space of deep generative models . Our framework builds off the potential outcomes framework of Neyman-Rubin causal model ( Neyman , 1923 ; Rubin , 1974 ) and defines a new notion of generative average treatment effects ( GATE ) . A fundamental problem of causal inference is that , by construction , we can not observe the potential outcomes under all treatments ( Holland , 1986 ) i.e. , at any given time , any individual can be assigned only one treatment ( a.k.a . the factual ) but not both . However , when studying treatment effects for deep generative models , CAGE exploits the generative nature of such models to overcome this challenge and explicitly generate the counterfactual . In Figure 1 : Counterfactual generation of Blonde Males using a deep generative model ( Karras et al. , 2020 ) . CAGE infers the implicit causal structure as Gender → Blonde Hair . Top : The left panel contains the original image of a blond female , other panels display counterfactual samples where gender ( treatment ) is changed to male . Bottom : Equivalent visualizations starting from a nonblond male ( left panel ) and counterfactually changing hair color ( treatment ) to blond . Generating counterfactual samples that respect the implicit causal structure yields subjectively better samples . particular , our approach utilizes probabilistic classifiers for characterizing the axis of variations for relevant attributes in the latent space of a deep generative model . Given the axis for the treatment attribute classifier , we develop a geometric manipulation scheme for simulating the assignment of specific treatments ( in particular , the counterfactual ) to an individual . Thereafter , we use an outcome attribute classifier to quantify the difference in outcomes for the factual and counterfactual generations and thus , estimate GATE . We refer to our overall framework for causal probing of deep generative models as CAGE . Figure 1 shows illustrative samples generated from our approach . While there is a body of existing work on enforcing a causal structure on deep generative models e.g. , ( Kocaoglu et al. , 2017 ; Yang et al. , 2021 ; Suter et al. , 2019 ) or using generative models as tools to aid causal inference e.g. , ( Louizos et al. , 2017 ; Yoon et al. , 2018 ; Sauer & Geiger , 2021 ) , our work stands out in extracting the causal structure inherent in existing generative models . As such , CAGE applies readily to any pretrained latent variable model , without any changes to their training procedure that can introduce any undesirable trade-offs in learning ( Locatello et al. , 2019 ) . We further leverage the treatment effects estimated via CAGE to define a natural measure of causal direction over a pair of variables associated with a deep generative model . Our score is simply defined as the difference in magnitudes of GATE when considering each variable as a treatment . The intuition behind such a score is that the ATE when considering the effect as a treatment should be approximately zero . In contrast , we expect the magnitude of the ATE when considering the cause as a treatment to be large . We study the consistency of the inferred causal directions for two highdimensional image datasets , MorphoMNIST ( Pawlowski et al. , 2020 ) and CelebaHQ ( Karras et al. , 2017 ) , with known or biologically guessed prior cause-effect relationships . Finally , we use these inferred causal directions to guide counterfactual generation for data augmentations . Data augmentations from generative models are immensely useful for dealing with class and subgroup imbalance ( Ramaswamy et al. , 2021 ; Xu et al. , 2018 ; Sattigeri et al. , 2018 ) . We show that augmentations generated by respecting the causal directions inferred via CAGE improve classification accuracy by over 2.7 % over the baseline and 11.7 % over the anti causal direction on the CelebaHQ dataset for the worst subgroup in the challenging group DRO task . Our main contributions are summarized as follows : 1 . We introduce the notion of generative average treatment effects ( GATE ) for probing causeeffect relationships in the latent space of deep generative models . 2 . We propose CAGE , a GATE estimation framework for deep generative models based on latent counterfactual manipulations . 3 . Finally , as a practical application , we use the ATEs to infer causal direction over pairs of attributes within deep generative models and employ CAGE for generating counterfactual data augmentations with an application to distributionally robust optimization . 2 BACKGROUND AND PRELIMINARIES . 2.1 THE POTENTIAL OUTCOMES FRAMEWORK . We briefly overview the potential outcomes framework of Neyman ( 1923 ) and Rubin ( 1974 ) , upon which we base our proposed method in the next section . In its simplest form , this framework considers the causal effect of assigning a binary treatment , T ∈ { 0 , 1 } . Such a framework posits the existence of potential outcomes for the ith individual both whilst receiving treatment , YT=1 ( i ) , and when treatment is withheld , YT=0 ( i ) . The causal effect for a given individual , i , is defined as the difference in these potential outcomes . The “ fundamental problem of causal inference ” is that whilst we may posit the existence of both potential outcomes , YT=1 ( i ) and YT=0 ( i ) , we only ever observe the outcome under one treatment ( Holland , 1986 ) . For this reason , causal inference is often performed at the population level , by considering quantities such as the average treatment effect ( ATE ) : τ = Ei [ YT=1 ( i ) − YT=0 ( i ) ] . ( 1 ) Equation ( 1 ) is the expected difference in potential outcomes of individuals receiving treatment T = 1 and T = 0 . Under standard assumptions of unconfoundedness , positivity , and stable unit treatment value assumption , the ATE is identifiable and can be estimated by a statistical estimate of the associational difference in expectations ( Rubin , 1978 ) . 2.2 DEEP LATENT VARIABLE GENERATIVE MODELS . The goal of generative modeling is to learn an underlying data distribution given a training dataset . For high-dimensional data , this can be done effectively via deep neural networks by learning a mapping G from unobserved factors of variation Z ( i.e. , latent variables ) to observed variables X ( e.g. , raw pixels ) . The mapping G is parameterized using deep neural networks , and can be learned via a variety of training objectives , such as adversarial training ( e.g. , generative adversarial networks ( GAN ; Goodfellow et al . ( 2014 ) ) , exact maximum likelihood estimation ( e.g. , Normalizing Flows ( Dinh et al. , 2017 ) ) , and variational Bayes ( e.g. , Variational Autoencoders ( VAE ; Kingma & Welling ( 2013 ) ; Rezende & Mohamed ( 2015 ) ) . For our current work , we can use any such latent variable generative model . Our experiments will exhibit this generality by considering state-of-theart flows ( Papamakarios et al. , 2017 ; Song et al. , 2019 ) and GANs ( Karras et al. , 2020 ) . In addition to G ( i.e. , the decoder ) , we will assume the existence of an encoder ENC : X → Z that can map any observed datapoint to its latent vector representation . For instance , for flows , the encoder is given by the inverse of the generator mapping , ENC = G−1 . For VAEs , the encoder is the variational posterior distribution trained alongside G. For GANs , either the encoder can be trained during training ( e.g. , Donahue et al . ( 2016 ) ; Dumoulin et al . ( 2016 ) ) , after training ( e.g. , Richardson et al . ( 2021 ) ) , or implicitly learned by backpropogating through G to find a latent vector z ∈ Z that best represents a target input x , as measured via a reconstruction loss ( e.g. , Lipton & Tripathi ( 2017 ) ; Bora et al . ( 2017 ) ) or a pretrained domain classifier ( e.g. , ( Karras et al. , 2020 ) ) . 3 ESTIMATING CAUSAL EFFECTS IN DEEP GENERATIVE MODELS . Our goal is to interrogate a deep generative model with a view to better understanding which—if any—causal associations have been implicitly encoded . We therefore look to answer the question : In the context of a generator , G , does considering the counterfactual over a given “ causal ” attribute ( e.g. , gender ) lead to significant changes in a “ effect ” attribute ( e.g. , hair color ) ? Formally , we assume whitebox access to a pretrained latent variable deep generative model G : Z → X . In addition , we further assume a finite dataset of n annotated observations , D = ⋃n i=1 { xD ( i ) , m1 ( i ) , m2 ( i ) } . Here , the examples { xD ( 1 ) . . . , xD ( N ) } are drawn from the model ’ s training distribution and we annotate each example xD ( i ) with binary metadata for two variables m1 ( i ) ∈ { 0 , 1 } and m2 ( i ) ∈ { 0 , 1 } . These two annotated variables can be drawn from a larger unobserved ( w.r.t . G ) causal process . For example , xD ( i ) could correspond to a highdimensional image and m1 ( i ) and m2 ( i ) could be any two attributes of the image , such as gender expression and hair color in Figure 1 . Apriori , we do not know which of the two variables ( if any ) is cause or effect but through D and G , we seek to identify the nature of their relationship within G. We note that any causal conclusions obtained under our proposed framework only reflect implicit properties of a generator G , as opposed to properties of the data on which the generator was trained . That is to say , it is perfectly permitted , forG to have learned a causal relationship that is inconsistent with the observational data used for training . As we shall see in §4 , shedding light on the implicit causal relationship learned by G is abundantly useful for pinpointing shortcomings in the model and informing best practices for counterfactual generations . As a result , this work—by design— is in sharp contrast with the majority of causal inference literature that focuses on identifying the generative mechanisms of a fixed dataset as opposed to our focus on the generative model itself .

Summary: The goal of this paper is to inspect causal relationships learned by a generative model. In particular, the paper studies how moving in the latent space with respect to a putative cause variable (e.g. age) affects the probability of the putative effect variable (e.g. baldness). This is done by training two linear classifiers in the latent space, one for each variable. The counterfactual for a given cause variable and datapoint is defined by changing the location of the data point in the latent space in the direction perpendicular to the linear classification boundary for the cause variable. Then, the paper inspects what happens to the putative effect variable, using its respective decision boundary. A null distribution is defined using permutation. The usefulness of techniques is demonstrated for inferring the causal relationships (including direction) for pairs of variables and for generating data augmentations for robust classification.