Learning Generalized Gumbel-max Causal Mechanisms

Abstract: To perform counterfactual reasoning in Structural Causal Models (SCMs), one needs to know the causal mechanisms, which provide factorizations of conditional distributions into noise sources and deterministic functions mapping realizations of noise to samples. Unfortunately, the causal mechanism is not uniquely identified by data that can be gathered by observing and interacting with the world, so there remains the question of how to choose causal mechanisms. In recent work, Oberst & Sontag (2019) propose Gumbel-max SCMs, which use Gumbel-max reparameterizations as the causal mechanism due to an appealing  counterfactual stability property. However, the justification requires appealing to intuition. In this work, we instead argue for choosing a causal mechanism that is best under a quantitative criteria such as minimizing variance when estimating counterfactual treatment effects. We propose a parameterized family of causal mechanisms that generalize Gumbel-max. We show that they can be trained to minimize counterfactual effect variance and other losses on a distribution of queries of interest, yielding lower variance estimates of counterfactual treatment effect than fixed alternatives, also generalizing to queries not seen at training time.


Reviews and author responses:

This paper proposes a parameterized family of causal mechanisms that generalize Gumbel-max, which can be trained to minimize counterfactual-effect variance. It is motivated by the fact that the causal mechanism is not uniquely identifiable with Gumbel-max SCMs.
  Due to the nonidentifiability, the authors instead learn a parameterized family of causal mechanisms that minimizes the counterfactual-effect variance, which I think is a smart way. 

An alternative direction is to show the identifiability condition of the counterfactual effect. [1] shows that the counterfactual outcome is identifiable under very mild conditions, even if the causal mechanism is not identifiable.
1. C. Lu, B. Huang, K. Wang, K. Zhang, J. M. Hernandez-Lobato, B. Schölkopf. Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentation. arXiv:2012.09092, 2020.

However, I do not quite get the details of the proposed method. The authors may need to consider improving the readability of the paper.

Updates:
Thanks to the authors for the responses. 

After reading the paper again, I think the writing needs much improvement. The current version is hard to follow. yes

 We are glad the reviewer found our ideas interesting, and hope to shed light on some of the points of confusion. First, regarding our goal, we would like to clarify that our goal is to find a counterfactual distribution that minimizes the variance of the treatment effect. However, our method is quite general and can work with any arbitrary loss function (not only variance) for which Wasserstein distance is a single example. We note that "maximal coupling" specifically refers to a coupling that minimizes the Wasserstein distance under $d(x, y) = 1_{x=y}$, and in fact there may be no SCM that produces a maximal coupling between counterfactuals under all interventions (Proposition 2).

Responses to the reviewer's specific questions:

* Yes, our approach is applicable to any SCM defined over discrete random variables.
* The reviewer is correct that, since we have access to all interventional distributions, we focus on Markovian models. If there are unobserved confounders, those must be handled separately before applying our method.
* The assumption of maximal coupling does not necessarily imply identifiability (and vice versa). We address the non-identifiability by favoring statistical properties: we respectfully quote Reviewer XBKm saying “Due to the nonidentifiability, the authors instead learn a parameterized family of causal mechanisms that minimizes the counterfactual-effect variance, which I think is a smart way.”
* Implicit coupling allows us to learn any level-2 quantity, which in our work is exemplified by the average treatment effect. The maximal coupling is tied to the 0/1 loss instead of the treatment effect, which is not always the desired quantity. Nevertheless, we explain maximal couplings as a way to better understand the properties and limitations of implicit couplings relative to the well-studied maximal coupling. In this respect we may consider the approximation of implicit coupling to maximal coupling. Also, as shown in Proposition 2, it is not possible to construct an SCM whose implicit couplings are always maximal. 
* Regarding sampling from the posterior: By "posterior" we simply mean the distribution of a counterfactual sample $y$ given an observed sample $x^{(obs)}$ drawn with the same shared noise, in other words $\pi(y | x^{(obs)}, p, q) \propto \pi(y | \gamma, q)\pi(x^{(obs)} | \gamma, p)$. We will clarify it. 


 The reviewer is correct that our approach involves choosing a particular joint distribution over counterfactuals from within a larger set that have the same interventional distributions, and that such a choice cannot be falsified. Nevertheless, there are many situations under which counterfactuals are desired, which requires making some sort of assumption, such as monotonicity (Pearl, 2000) or counterfactual stability (Oberst and Sontag, 2019). In contrast to these other assumptions, our assumption of minimum variance is less prescriptive and can be adapted depending on how the causal model will be used. For instance, if the goal is to estimate an average treatment effect (as we do in our experiments), every counterfactual distribution produces the same average treatment effect, and so it makes sense to choose the counterfactual distribution that minimizes the error of our estimate.

In Eq. (5), we assume that the distribution D of pairs of logits is known in advance, corresponding to the set of outcome distributions and interventional distributions we wish to estimate counterfactuals for. In 7.1 and 7.2, we choose D to generate synthetic distributions of logits. In 7.3, we use a doctor-imitation policy and a RL policy to generate pairs of actions in a medical simulator, then treat the pair of distributions of next states (given the doctor and RL actions) as a sample from D.


 We are happy that the reviewer liked our approach to addressing non-identifiability by minimizing variance. The approach in [1] is indeed very relevant and adds to the work of Oberst and Sontag, 2019 on RL and causal inference. We will add it to our references. The assumption in [1] of monotonicity of the state $S_{t+1}$ with respect to the random noise $U_{t+1}$ appears to be equivalent to assuming an inverse-CDF coupling between counterfactual states. We compare our method to inverse-CDF in our experiments and achieve lower variance in the non-monotone setting. 

This paper studies the problem of estimating counterfactual distributions from a set of interventional distributions. More specifically, let X represent treatment and Y represent outcome. Our goal is to estimate the counterfactual distribution P(y_{x=0}, y_{x=1}) from the marginal interventional distribution P(y_{x=0}) and P(y_{x=1}). In the most general cases, the counterfactual distribution is not identifiable. The authors make several parametric assumptions. First, the exogenous variables U are drawn from a Gumbel distribution. Once values of U are fixed, values of Y_{x=0} and Y_{x=1} are decided by a Gumbel-max function. Finally, we would like to find a counterfactual distribution P(y_{x=0}, y_{x=1}) that minimizes the Wasserstein distance between Y_{x=0} and Y_{x=1}, called the maximal coupling. The authors describe an algorithm to approximate such maximal coupling. Simulations are performed to validate the findings.  Counterfactual queries or “what-if” questions are prevalent in various fields, including healthcare, social science and econometrics. Since counterfactual queries involve inference in alternative, hypothetical worlds, it is unclear how one could design an experimental study to answer such queries. This paper considers one possible approach to address the challenge of non-identifiability, i.e., counterfactual identification under reasonable parametric assumptions. I think the authors may have some interesting ideas. However, the writing of this paper is somewhat confusing, and many technical details deserve further explanation. I will summarize my questions/concerns as follows.

- This paper uses the semantical framework of the structural causal model (SCM). Is the proposed method applicable to any SCM with arbitrary structure? The paper seems to focus on a rather simple model which only involves treatment and an outcome.

- Is the proposed method applicable to causal models with the presence of unobserved confounders (UCs)? My understanding is that since we have access to all interventional distributions, one could focus on Markovian models without the presence of UCs. Is my reading correct?

- Does the assumption of maximal couplings imply identifiability? There could be multiple instances of counterfactual distributions \pi that achieve the minimum Wasserstein distance. In this case, how could one address the issues of non-identifiability?

- After explaining maximal couplings, the authors proceed to define the problem of finding implicit couplings. This is somewhat confusing since the authors also show that maximal couplings and implicit coupling are not equivalent in general. What are the challenges in directly finding maximal couplings? Is it due to the 0/1 loss function?

- The authors argue that implicit couplings could be a good approximation for maximal couplings. However, the approximation power in Corollary 2 appears to be rather weak. For instance, for a Gumbel-max coupling with probability alpha=0.6, then the probability under the maximal coupling is at most 1.6, which is essentially unbounded.

- The explanation for gadgets of learning implicit coupling could be improved. For instance, what is variable Z in Line 167, Page 5? It does not seem to be defined anywhere before.

- In Line 187, Page 5, the authors discuss “sampling from the posterior”. This is somewhat surprising since this paper has not mentioned anything about Bayesian learning at this point. What are the prior distributions? What are the complete conditionals given observations? What is the truncated Gumbel here? I was unable to follow after this point.
 - This work is mainly theoretical. Its long-term societal impact remains unclear at this point.

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.   ### Main comments

1. **Gumbel-max SCM**: More background should be provided on the Gumbel-max structural causal model, for the sake of readers who are interested in causal inference but not necessarily familiar with the specific model class. Further, the objective of the paper should be better motivated in the context of Gumbel-max SCM and causal inference. It is not entirely clear to me why "learning implicit coupling" in the form of Sec 3. should be of interest. 

2. **Choice of causal mechanism**: If I understand correctly, this means choosing a particular (type of) joint distribution of the counterfactuals among a whole class of alternatives that yield the same observed distribution. This paper is motivated by choosing one such that the variance of the effect is minimized. 

However, this amounts to posing assumptions that cannot be falsified even by a hypothetical randomized experiment. In my opinion, the unfalsifiable nature of these assumptions does not imply that one should make assumptions freely to suit a convenient purpose. Rather, methods should be developed that make as few such assumptions as possible. 

### Minor comments

1. Eq. (5): How is the distribution $\mathcal{D}$ over pairs of logits obtained?







 I do not see any potential negative societal impact of the work in its current form. 

Meta-review: The authors study couplings between discrete distributions that allow to infer counterfactuals. Overall, this is an important topic and in many applications, counterfactuals are desired. An inherent difficulty of inferring counterfactuals is the need to make unverifiable assumptions. The authors argue to choose a causal mechanism that is best under a quantitative criterion such as minimizing the variance when estimating causal effects. The “minimum variance” assumption seems hard to reason about in practice. In contrast, the monotonicity assumption (Pearl, 2000) seems much easier to defend. In the discussion, the authors write that the minimum variance condition is attractive for average treatment effect estimation (see the author’s answer to a comment by ybLo). This is not clear to me, since average treatment effect estimation does not require inferring counterfactuals. More generally, I am hesitant to deal with non-identifiability by “favoring statistical properties” (reponse of the authors to gci3 and comment by ybLo). Reviewer XBKm, gci3, and I found the paper hard to read.

Decision: accept