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: This paper proposes a novel framework CAGE (causal probing of deep generative models) for inferring the causal-effect relationship in deep generative models. The treatment is implemented by moving linearly along the hyperplane normal. Then the attribute treatment effect can be quantified via a proxy classifier. The authors evaluated the CAGE on low-dimensional synthetic and high-dimensional datasets. Results show the ability to infer causal relationships of the proposed CAGE.