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 paper is motivated by choosing the causal mechanism for the so-called "Gumbel-max" structural causal models (SCM). The authors study constructing a coupling between two categorical distributions through a Gumbel-max-style sampler such that the coupling is maximal. The authors seek to learn such a sampler through a smooth parametrization and optimization with gradient-based methods. Two such architectures are developed. The performance of the proposal is studied numerically.