Learning Generalized Gumbel-max Causal Mechanisms

1 Introduction . Pearl [ 2009 ] presents a “ ladder of causation ” that distinguishes three levels of causal concepts : associational ( level 1 ) , interventional ( level 2 ) , and counterfactual ( level 3 ) . As an illustrative example , suppose we wish to compare two treatments for a patient in a hospital . Level 1 corresponds to information learnable from passive observation , e.g . correlations between treatments given in the past and their outcomes . Level 2 coresponds to active intervention , e.g . choosing which treatment to give to a new patient , and measuring the distribution of outcomes it causes ( called an interventional distribution ) . Level 3 corresponds to reasoning about hypothetical interventions given that some other outcome actually occurred ( called a counterfactual distribution ) : given that the patient recovered after receiving a specific treatment , what would have happened if they received a different one ? The three levels are distinct in the sense that it is generally not possible to uniquely determine higher level models from lower level information [ Bareinboim et al. , 2020 ] . In particular , although we can determine level 2 information ( such as the average effect of each treatment ) by actively intervening ⇤Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in the world , we can not determine level 3 information in this way ( such as the effect two different treatments would have had for a single situation ) . Nevertheless , we still desire to reason about counterfactuals . First , counterfactual reasoning is fundamental to human cognition , e.g. , in assigning credit or blame , and as a mechanism for assessing potential alternative past behaviors in order to update policies governing future behavior . Second , counterfactual reasoning is computationally useful , for example allowing us to shift measurements from an observed policy to an alternative policy in an off-policy manner [ Buesing et al. , 2018 , Oberst and Sontag , 2019 ] . Doing counterfactual reasoning thus requires us to make an assumption about the causal mechanism of the world , which specifies how particular choices lead to particular outcomes while holding “ everything else ” fixed . Different assumptions , however , lead to different counterfactual distributions . One approach , as exemplified by Oberst and Sontag [ 2019 ] , is axiomatic . There , an intuitive requirement of counterfactuals is presented , and then a causal mechanism is chosen that provably satisfies the requirement . The resulting proposal is Gumbel-max SCMs which assume causal mechanisms are governed by the Gumbel-max trick . In this work , we instead view the choice of causal mechanism as an optimization problem , and ask what causal mechanism ( that is consistent with level 2 observations ) yields the most desirable behavior under some statistical or computational criterion . For example , what mechanism leads to the lowest variance estimates of treatment effects in a counterfactual setting ? If we are ultimately interested in estimating a level 2 property , choosing a level 3 mechanism that minimizes the variance of our estimate can lead to algorithms that converge faster . Alternatively , we can view minimizing variance as a kind of stability assumption on the treatment effect : we are interested in the causal mechanism for which the treatment effect is as “ evenly divided ” as possible across realizations of the exogenous noise . More generally , casting the problem in terms of optimization gives additional flexibility to choose a causal mechanism that is specifically-tuned to a distribution of observations and interventions of interest , and in terms of a loss function that measures the quality of a counterfactual sample . A key insight to lay the foundation for specifically-tuned causal mechanisms is to view the average treatment effects ( or other measure of interest ) from the perspective of a coupling between interventional and counterfactual distributions . We begin by drawing connections between causal mechanisms and couplings , and show that defining a level 3 structural causal model consistent with level 2 observations is equivalent to defining an implicit coupling between interventional distributions . Next , to motivate the need for specificallytuned causal mechanisms , we prove limitations of non-tuned mechanisms ( including Gumbel-max ) and the power of tuned mechanisms by drawing on connections to literature on couplings , optimal transport , and common random numbers . We then introduce a continuously parameterized family of causal mechanisms whose members are identical when used in a level 2 context but different when used in a level 3 context . The families contain Gumbel-max , but a wide variety of other mechanisms can be learned by using gradient-based optimization over the family of mechanisms . Empirically we show that the mechanisms can be learned using a variant of Gumbel-softmax relaxation [ Maddison et al. , 2017 , Jang et al. , 2017 ] , and that the resulting mechanisms improve over Gumbel-max and other fixed mechanisms . Further , we show that the learned mechanisms generalize , in the sense that we can learn a causal mechanism from a training set of observed outcomes and counterfactual queries and have it generalize to a test set of observed outcomes and counterfactual queries that were not seen at training time . 2 Background . Structural Causal Models . Here we briefly summarize the framework of structural causal models ( SCMs ) , which enable us to ask and answer counterfactual questions . SCMs divide variables into exogenous background ( or noise ) variables and endogenous variables that are modeled explicitly . Each endogenous variable vi has a set of parent endogenous variables pai , an associated exogenous variable ui , and a function fi . Intuitively , the function fi specifies how the value of vi depends on the other variables pai of interest , and ui represents “ everything else ” that may influence the value of vi . Putting these together along with a prior distribution over ui ’ s defines the causal mechanism governing vi , as vi = fi ( pai , ui ) ; we emphasize that this mapping is deterministic , as all of the randomness in vi is captured by either pai or ui . Marginalizing over the exogenous variable ui yields the interventional distribution p ( vi|pai ) , which specifies how vi is affected by modifying its parents . Counterfactual reasoning , on the other hand , amounts to holding ui fixed but considering multiple values for pai and consequently for vi . In other words , if v ( 1 ) i = fi ( pa ( 1 ) i , ui ) and v ( 2 ) i = fi ( pa ( 2 ) i , ui ) , the counterfactual distribution is p ( v ( 2 ) i |pa ( 2 ) i , v ( 1 ) i , pa ( 1 ) i ) = R ui p ( v ( 2 ) i |pa ( 2 ) i , ui ) p ( ui|v ( 1 ) i , pa ( 1 ) i ) . Gumbel-max SCMs The Gumbel-max trick is a method for drawing a sample from a categorical distribution defined by logits l 2 RK , e.g . p ( X = k ) / exp lk . It is based on the fact that , if k ⇠ Gumbel ( 0 ) for k 2 { 1 , . . . , K } , then2 P h x = argmaxk [ lk + k ] i = exp ( lx ) P k exp ( lk ) . ( 1 ) Oberst and Sontag [ 2019 ] use this as the basis for their Gumbel-max SCM , which uses vectors of Gumbel random variables 2 RK as the exogenous noise , and specifies the causal mechanism for each variable x with parents pa as x = f ( pa , ) = argmaxk [ lk + k ] , ( 2 ) where lk = log pk = log p ( X = k|pa ) is the interventional distribution of x given the choice of pa. Oberst and Sontag [ 2019 ] motivate their Gumbel-max SCMs by appealing to a property known as counterfactual stability : if we observe X ( 1 ) = i under intervention p ( 1 ) = p , then the only way we can observe X ( 2 ) = j 6= i under intervention p ( 2 ) = q is if the probability of j relative to i has increased , i.e . if qjpj > qi pi or equivalently qjqi > pj pi . This property generalizes the well-known monotonicity assumption for binary random variables [ Pearl , 1999 ] to the larger class of unordered categorical random variables , and corresponds to the intuitive idea that the outcome should only change in a counterfactual scenario if there is a reason for it to change ( i.e. , some other outcome ’ s relative probability increased more than the observed outcome ’ s ) . Oberst and Sontag [ 2019 ] show that Gumbel-max SCMs are counterfactually stable . Gumbel-max SCMs have interesting properties due to the Gumbel-max trick introducing statistical independence between the max-value and the argmax , cf . Maddison et al . [ 2014 ] . Oberst and Sontag [ 2019 ] exploit this to sample from counterfactual distributions by sampling the exogenous variables from the posterior over Gumbel noise : conditioned on an observed outcome x = argmaxk [ lk + k ] we can sample x ⇠ Gumbel ( lx ) and then sample the other i from truncated Gumbel distributions . We observe that sharing the exogenous noise between two different logit vectors l ( 1 ) and l ( 2 ) yields a joint distribution over pairs of outcomes : ⇡gm ( x , y ) = P h x = argmaxk h l ( 1 ) k + k i and y = argmaxk h l ( 2 ) k + k i i . ( 3 ) We call this a Gumbel-max coupling . Couplings . A coupling between categorical distributions p ( x ) and q ( y ) is a joint distribution ⇡ ( x , y ) such that X x ⇡ ( x , y ) = q ( y ) and X y ⇡ ( x , y ) = p ( x ) . ( 4 ) The set of all couplings between p and q is written C ( p , q ) . A core problem of coupling theory is find a coupling ⇡ 2 C ( p , q ) in order to estimate Ex⇠p [ h1 ( x ) ] Ey⇠q [ h2 ( y ) ] = Ex , y⇠⇡ [ h1 ( x ) h2 ( y ) ] for some real-valued cost functions h1 , h2 , with minimal variance Varx , y⇠⇡ [ h1 ( x ) h2 ( y ) ] . Interestingly , whenever h1 ( x ) , h2 ( y ) are monotone , such a coupling can be attained by ( F 1X ( U ) , F 1 Y ( U ) ) ) , when FX ( t ) , FY ( t ) are the cumulative distribution functions ( CDFs ) of X , Y and U is a uniform random variable over the unit interval . When h1 , h2 are not monotone , there are clearly cases where CDF inversion produces suboptimal couplings . Couplings are also used to define the Wasserstein distance W1 ( p , q ; d ) between two distributions p and q ( with respect to a metric d between samples ) : W1 ( p , q ; d ) = inf ⇡2C ( p , q ) Ex , y⇠⇡d ( x , y ) , ( 5 ) When d ( x , y ) = 1x 6=y , then a coupling that attains this infimum is known as a maximal coupling ; such a coupling maximizes the probability that X and Y are the same . 2Samples of Gumbel ( 0 ) random variables can be generated as log ( log ( u ) ) , where u ⇠ Uniform ( [ 0 , 1 ] ) Causal mechanisms as implicit couplings . Any causal mechanism for a variable vi defines a coupling between outcomes under two counterfactual interventions . In other words , for any two interventions pa ( 1 ) i and pa ( 2 ) i on the parent nodes pai , sharing the exogenous noise ui yields a coupling ⇡SCM between the interventional distributions p ( v ( 1 ) i |do ( pa ( 1 ) i ) ) and p ( v ( 2 ) i |do ( pa ( 2 ) i ) ) : ⇡SCM ⇣ v ( 1 ) i , v ( 2 ) i ⌘ = Pui h v ( 1 ) i = fi ( pa ( 1 ) i , ui ) and v ( 2 ) i = fi ( pa ( 2 ) i , ui ) i . ( 6 ) We call this an implicit coupling because fi ( · , u ) is not directly defined with respect to a particular pair of marginal distributions p , q , but instead arises from running the same causal mechanism forward with shared noise but different inputs , representing either p or q . This connection between SCMs and couplings enables us to translate ideas between the two domains . For instance , suppose we are interested in estimating Ex⇠p [ h1 ( x ) ] Ey⇠q [ h2 ( y ) ] between observed outcomes from p / exp l ( 1 ) and counterfactual outcomes from q / exp l ( 2 ) . We might do so using counterfactual reasoning in the Gumbel-max SCM : Ex , y⇠⇡ [ h1 ( x ) h2 ( y ) ] = X y q ( y ) E ⇠ ( G2|y ) h h1 ( argmaxk [ l ( 1 ) k l ( 2 ) k + k ] ) h2 ( argmaxk [ k ] ) i . Here G2|y is the Gumbel distribution with location log l ( 2 ) k conditioned on the event that y is the maximal argument ; the proof of this equality appears in Appendix B . However , if we are interested in minimizing the variance of a Monte Carlo estimate of the expectation , this may not be optimal .

Summary: The authors propose a method to parametrize the functions of a certain class of structural causal models (SCM) called Generalized Gumbel-max SCMs. The aim is to train the model while minimizing the variance of treatment (counterfactual?) effects computed from them.  The paper introduces two different families of functions (gadgets) that can be learned through an optimization process aiming to find a maximal coupling of a pair of logits vectors defining a discrete distributions.  The experimental section compares the two gadgets with the Gumbel-max representation from Oberst and Sontag, and a cumulative distribution functions (CDF) inversion method. The experimental results seem to suggest that the proposed gadgets improve may achieve lower variance.