Learning Generalized Gumbel-max Causal Mechanisms
Authors: Guy Lorberbom
Daniel D. Johnson
Chris J. Maddison
Daniel Tarlow
Tamir Hazan
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Gumbel max
reparameterization trick
structural causal model

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.

Learning Generalized Gumbel-max Causal Mechanisms
Guy Lorberbom"
Daniel D. JJohnson" Google Research Toronto, ON, Canada dd jaijohonnggogle... com
Technion
Haifa, Israel
guy juy_lorber@camuus hww.obbrrrrammp.tteeeiiiooooo. technion. ac. il
Chris J. Maddison
Daniel Tarlow Google Research Montreal, QC, Canada
Tamir Hazan
University of Toronto & Vector Institute Toronto, ON, Canada cmaddi toronto. edu
Technion Haifa, Israel
dtarlon@google. com
tamir www.haatecainii.ooooo ac
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 noise to samples. Unfortunately, the causal mechanism is not uniquely identified by data that can 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 reparame- terizations the causal mechanism due to intuitively appealing counterfactual stability property. In this work, we instead argue for choosing a causal mecha- nism that is best under quantitative criteria such minimizing variance when estimating counterfactual treatment effects. We propose a parameterized family causal mechanisms that generalize Gumbel-mas We show that they can be trained to minimize counterfactual effect variance and other losses on a distribution queries of interest, yielding lower variance estimates of counterfactual treatment effect than fixed alternatives, also generalizing to queries not seen at training time.
Introduction
Pearl 2009
presents "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 wish to compare treatments for patient in hospital. Level corresponds to information learnable from passive observation, e.g. correlations between treatments given in the past and their outcomes. Level 2 coresponds 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 corresponds reasoning about hypothetical interventions given that some other outcome actually occurred (called a counterfactual distribution): given that the patient recovered after receiving 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
"quual contribution
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

in the world, we cannot 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 blame, and 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 "every- thing else" fixed. Different assumptions, however, lead to different counterfactual distributions. One wewuireeruueeoeeee 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.
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 level 2 property, choosing level 3 mechanism that minimizes the variance of our estimate can lead to algorithms that converge faster. Alternatively, we can view minimizing variance kind of stability assumption on the treatment effect: 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 to a distribution of observations and
In this work, we instead view the choice of causal mechanism
flexibility to choose a causal mechanism that is specifically- interventions of interest, and in terms of a loss function that measures the quality of counterfactual sample. A key insight to lay the foundation for ppeiificall--tunee 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 specifically- tuned causal mechanisms, we prove limitations of non-tuned mechanisms (including Gumbel-max) and the power of tuned mechanisms by drawing on connections 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 level context but different when used in a level 3 context. The families contain Gumbel-man 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 variant of Gumbel-softtax relaxation Maddison et al. 2017 Jang et al. [017], and that the resulting mechanisms improve over Gumbel-mas and other fixed mechanisms. Further, we show that the learned mechanisms generalize, in the sense that we can learn a causal mechanism from training set of observed outcomes and counterfactual queries and have it generalize to 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 o; has a set of parent endogenous variables pa,, an associated exogenous variable u; and function ft. Intuitively, the function f. specifies how the value of U; depends on the other variables pa; of interest, and u; represents "everything else" that may influence the value of v:- Putting these together along with a prior distribution over u, 's defines the causal mechanism governing U; as v; = fe(pa;; u.); we emphasize that this mapping deterministic, as all of the randomness in i is captured by either pa, or u.
2

---

Marginalizing over the exogenous variable yields the interventional distribution p(u; (pss) which specifies how o; affected by modifying its parents. Counterfactual reasoning. on the other hand, amounts to holding
fixed but considering multiple values for pa; and consequently for vy. In other words, if U. (1) fi(pa; (1) uz) and ii fi u.), the counterfactual distribution is p(u (2)
(1) ppa 1 ppa;" (1) = Su, p(u, pa,
(2) (2) (u,ppu pp;; (1))
Gumbel-max SCMs The Gumbel-mar trick is a method for drawing a sample from a categorical distribution defined by logits Ex e.g. p(X k) x exp 4. It is based on the fact that, if Tkk ~ (0m for E (1,..., K), them
exp((2)
P. x argmaxy [4 + mul
(1)
Σx exp(12)
Oberst and Sontag 2019 use this as the basis for their Gumbel-man SCM, which uses vectors of Gumbel random variables y E RK the exogenous noise, and specifies the causal mechanism for each variable x with parents pa as
fpa, )) = argmax; [4 + 7x),
(2)
where la = loooog Pk = loggll p(X = klpa) is the interventional distribution of x given the choice of pa.
Oberst and Sontag 2019] motivate their Gumbel-max SCMs by appealing to property known as counterfactual stability: we observe x = under intervention p( = p. then the only way we can observe X(2) = j jf under intervention p!²) = q is the probability of j relative to i or equivalently 9: p;
has increased, i.e. if Pj P:
P/ This property generalizes the
monotonicity assumption for binary random variables [Pearl 1999 the larger class of unordered categorical random variables, and corresponds to the intuitive idea that the outcome should only change in a counterfactual scenario there is reason for it to change (i.e., some other outcome's relative probability increased more than the observed outcome's s). Oberst and Sontag (2099) 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 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 æ argmax, 14 + nk we can sample The ~ Gumbel ( and then sample the other ?/ from truncated Gumbel distributions.
We observe that sharing the exogenous noise 7 between two different logit vectors ((1) and ((2) yields joint distribution over pairs of outcomes:
Tg (x, y) = pr x = argmax, (11)) + +7k and y = argmax (10e] + 7k
(3)
We call this a Gumbel-max coupling.
Couplings. coupling between categorical distributions p(x) and ((v) is a joint distribution <(,,v) v) such that
(T,,,)) = (1)) and Σ ( = p(x).
The set all couplings between p and q is written C(p, 9). A core problem of coupling theory is find a coupling π E C(p, q) in order to estimate Few [hg (z)) By (22(w) = R2 R [bg (z) - h2(y)) for some real-valued cost functions h1, h2, with minimal variance Varz, [/y (h2))-22()) Interestingly, whenever hy (z), hy) are monotone, such a coupling can be attained by (Fx (U), F)))))))))) when (t), Fy (t) are the cumulative distribution functions (CDFs) of X, Y and U is uniform random variable over the unit interval. When h, h2 are not monotone, there are clearly cases where CDF inversion produces suboptimal couplings. Couplings are also used to define the Wasserstein distance Wi (p, q;44 d) between two distributions and (with respect to a metric d between samples):
(U))),
W1 (p. g: d) =EC(p.4)
inf Ex.y y).
(5)
When d(x, v) = Luzy, then coupling that attains this infimum is known as maximal coupling; such a coupling maximizes the probability that X and Y are the same.
Samples of Gumbel(U) random variables can be generated as log(- (oo(()))) where Uniform([0, 1])
3

---

Causal mechanisms as implicit couplings. Any causal mechanism for variable Us defines coupling between outcomes under two counterfactual interventions. In other words, for any two interventions pa; and pa; on the parent nodes pa,, sharing the exogenous noise u, yields a coupling ascM between the interventional distributions p(u, (1) [do(pa) and p(u! [do(pa) )):
U (1) 1 ( Pu, (1)) ( = =i(pa u:) and v. fi u.)
TSCM
(6)
We call this an implicit coupling because fi(-, u) is not directly defined with respect to a particular pair of marginal distributions p. 9. but instead arises from running the same causal mechanism forward with shared noise but different inputs, representing either or
This connection between SCMs and couplings enables us to translate ideas between the two domains. For instance, suppose we are interested in estimating ( Ryysu [h2(v)] between observed
outcomes from & exp (() and counterfactual outcomes from x exp (2). We might do so using counterfactual reasoning in the Gumbel-max SCM:
E2,y~~ [11 (2) h2(2)l
Σ ( (argmaa 11 - ((2) + +wll) - h2 (argmax, [x)]] 8
Here G2ly is the Gumbel distribution with location log (2) conditioned on the event that y is the maximal argument; the proof of this equality appears in Appendix However, if we are interested in minimizing the variance of a Monte Carlo estimate of the expectation, this may not be optimal.
3 Problem Statement: Building SCMs by Learning Implicit Couplings
In this section, we formalize the task of selecting causal mechanisms according to some quantitative metric. Recall our initial example of comparing two treatment policies for patients in hospital. For simplicity, we consider a single action a taken by a hypothetical treatment policy, which leads to distribution over outcomes v. (More generally, we can a be the set of parents for any variable u in SCM.) As in Oberst and Sontag 2019]. we assume we have access to all of the interventional distributions of interest without any latent confounders.
Our goal is to define a parameterized family of causal mechanisms consistent with the interventional distributions for all possible actions a. We assume that the set V of possible outcomes is finite with [V] = K, but do not restrict the space of actions a; instead, we require that our causal mechanism can produce samples from any interventional distribution p(u do(a)), expressed as vector (()) e R for which p(u k))o(a)) & exp/(").
Specifically, let IR be sample from some noise distribution (c.g., from a Gumbel(O) distribution per dimension) and let 7 E IRK be a vector of (conditional) logits defining a distribution over K categorical outcomes. We wish to learn a function fo R x RRKK {1...., K} that maps noise and logits to a sampled outcome. We require that the process produces samples from the distribution p(k) & exp4 4 when integrating over u (i.e., want reparameterization trick), and also that counterfactually sample the exogenous noise u conditioned on an observation (obee) (e.g. ~ p(ull, fe(u.0) () (o)b)))) We obtain implicit coupling by running fe with the same noise and two different logit vectors ((2) and ((2), We can think of (11) the logits governing observed outcome and ((2) their modification under an intervention.
Each setting of the parameters 0 produces different SCM. approximately minimize an objective of interest. We provide two degrees of freedom for defining this objective. First, we must choose loss function (re)),33) [[,,..., K} [l...., "lI...,K] K} - that assigns a real-valued loss to joint outcome (fo(u, (,, (()), fo(u, 1(2)), perhaps modulated by (()) and (2)_ The loss function used to specify how desirable a pair of observed and counterfactual outcomes are (e.g., we are trying to minimize variance, the squared difference (I(((()) (z))))) of scores for each outcome). Second, we must choose distribution over pairs of logits (1(1), 112).. This determines the distribution of observed outcomes and counterfactual queries of interest.
We propose to learn 0 in such a way

---

Given these choices, our main objective is as follows:
f" = argming B(PDE)DDE (ras,, (fo(u, 1 Solu, 1o(11
(7)
exple 4e exple
subject to Pu (u([o((,1) =k]
for all E RR
(8)
Relationship to 1-Wasserstein Metric. We are free to set g to be distance metric d. in which case the implicit coupling between fa(u, (0) and fo fa(u, ((2)) bears similarity to the optimal 1- Wasserstein coupling for d. However, a key difference is that fo can be used to generate samples from one side of the coupling (say p) without knowledge of what will be chosen. Thus, fo can be seen coupling p all q simultaneously, in the same way that observing particular outcome simultaneously induces counterfactual distributions for all alternative interventions. In contrast, the Wasserstein optimization requires knowledge of both p and q and then computes a coupling specific to that pair. We discuss the effect of this restriction in the next section.
Interpretation in terms of causal inference. To construct a full level 3 SCM, fo must be combined with a set of known level interventional distributions p( p(v]dd [d]))))) similar to the Gumbel- max SCM in this regard [Oberst and Sontag, 2019]. In particular fa and Gumbel-max assume there are no latent confounders, and that the set of outcomes is discrete. The objective g and distribution D serve similar role counterfactual stability or monotonicity assumptions [Oberst and Sontag 2019 Pearl 1999), in that they are a-priori choices that select the intended level 3 mechanism from the set of consistent mechanisms. The main difference is that these assumptions are made at higher level of abstraction. Instead of specifying the mechanism, we specify a family of mechanisms along with a quantitative quality measure that be optimized.
4 Properties of Implicit Couplings
Our development so far raises questions about the relationship of the proposed approach to the Gumbel- couplings underlying Gumbel-max SCMs and about the relationship of our objective to Wasserstein metrics. In this section we establish the relationships and differences. The main results are that despite being counterfactually stable, Gumbel-max couplings are not actually maximal couplings. We then go further and show that any implicit coupling limited in expressivity. This establishes the difference to Wasserstein metrics and optimal transport, which are framed in terms of minimizing over the larger space of all couplings.
4.1 Non-maximality of Gumbel-max Couplings
We know that Gumbel-max couplings are counterfactually stable. We might therefore hope that they also maximal couplings, i.e. that they assign much probability as possible to the counterfactual and observed outcomes heing the same. Unfortunately, it turns out this is not the case. Proposition 1. The probability that
i in Gumbel-max coupling is
1+E max
212
Q(1)
The full proof is in Appendix C The main idea is to express the event x = i conjunction of inequalities defining the argmax, then simplifying using properties of Gumbel distributions.
Corollary 1. Gumbel-max couplings are not maximal couplings. In particular, they are suboptimal under the az metric iff there is an such that
max ESA
P(A) P())
Σ
jjt
2(12 q(2)
<
jti max
pLD 9(1) P(T)
9(3
The proof appears in Appendix D It follows from Prop. and the fact that the probability of J = y = i maximal coupling is min (p(1), q()). On the positive side, it is straightforward to show that Gumbel-max couplings are optimal when p g and when there are only two possible outcomes (in which case they also satisfy the monotonicity assumption of Pearl (199)). We also show that Gumbel-max couplings are within constant-factor of optimal as maximal couplings.
Corollary 2. If the Gumbel-max coupling assigns probability a to the event that x = y, then the probability that x y under the maximal coupling is at most 2a.
The proof appears in Appendix
It comes from bounding the ratio of the LHS to the RHS in Eq. 9
5

---

4.2 Impossibility of Implicit Maximal Couplings
Since Gumbel-max does not always induce the maximal coupling, we might wonder if some other implicit coupling mechanism could. Here we show that it is impossible. In particular, we show that no fixed implicit coupling is maximal for every pair of distributions over the set {1, 2, 33. Thus, there will always be some pair of distributions for which an implicit coupling is non-maximal. A proof of the Proposition appears in Appendix F
There is algorithm fe that takes logits I and sample u of noise, and deterministi- cally transforms u into sample from expl, I, such that when given any two distributions p and q and using shared noise u, the joint distribution of samples always a maximal coupling of p and q.
Proposition 2.
Methods for Learning Implicit Couplings
Here develop methods for learning implicit couplings, the problem defined in Sec. B
We use the term gadget to refer to learnable, continuously parameterized family of fo. We present two gadgets in this section. Gadget does not fully satisfy the requirements laid out in Sec. B] but it is simpler and introduces some key ideas, so we present it as a warm-up. Gadget 2 fully satisfies the requirements.
5.1 Gadget
The main idea in Gadget is to learn mapping ne F + R from categorical distribution p E RR to an auxiliary joint distribution Ty (2, 2 p). represented as a matrix To(:, E IR Kx The architecture of the mapping is constrained so that marginalizing out the auxiliary variable 2 yields a distribution consistent with the given logits, i.e., Σ. (o 2/p) = p(æ). We then generate K² the auxiliary joint. We only
independent "z,= Gumbcl(O) samples and perform Gumbel-max care about the sample of x, so one way of doing this is to first maximize out the auxiliary dimension to get (p) max, (?r. + log (x, =|p)) and then return & (r)) Here & is distributed according to p because this is performing Gumbel-max on a joint distribution with correct marginals and then marginalizing out the auxiliary variable.
To create a coupling, we run this process separately for p and q but with shared realizations of the X² Gumbels. However, the place where Gadget I does not fully satisfy the requirements from Sec. iis that we transpose the Gumbels for one of the samples. That is, we sample a coupling as
h (p)]. 72, + log (1()))/))
[na(4)ly = max((>") y..
a + log Te (b,)))))
(10)
â (rg
9 = argmax (12)))()2(9)).
(11)
Gadget 1 can still be used to create a coupling, but it is more analogous to antithetical sampling, where the noise source is used differently based which sample is being drawn. Note that like in antithetical sampling, both processes draw samples from the correct distribution, since transposing a matrix of
independent Gumbels does change the distribution. We describe how to draw counterfactual samples from Gadget in Appendix
We note that if 0 is chosen such that (x klp) p(r k) and (T( =lp) = 0 for / the Gadget 1 SCM becomes identical to the Gumbel-man -ax SCM. Thus, Gadget is a generalization of the Gumbel-max causal mechanism.
5.2
Gadget 2
Gadget 1 not an implicit coupling as defined in Sec. 3 because requires Gumbels to be transposed when sampling versus y. In this section, we present gadget that is proper implicit coupling.
Gadget again invokes an auxiliary variable and parameterizes learned joint distribution, but the auxiliary variable z is no longer required to share the same sample space as x. Further, instead of performing Gumbel-max on the learned joint directly, we start by drawing a single independently of p. The gadget is defined as follows:
7
(umbel(0) for = I,....
7 Gumbel(() for = |,...,/X|
(12)
2 = argmax_(og (2) (z))
à = argmax, (log ( (2 + (x)).
(13)

---

To sample a coupling, we re-use all the y's and rum the same process with instead of This means that we additionally get ý = (loga,, (8 &, q)) + (-))), Intuitively, we can think of 2 as latent cluster identity and each cluster being associated with a different learned mapping iT (x 2 p). The learning can choose how assign clusters so that Gumbel-max coupling of me (xr 2, p) and Toly 2, q) produces joint outcomes that are favorable under g.
Architecture for [ p). Not all choices of (z 2, p) lead to correct samples. For correct- ness, need to enforce the analogous constraint in Gadget which is that when we inte- grate out the auxiliary variable, we get samples from the distribution provided as an input; i.e., p(x) for all p. Here we describe how to build an architecture for (0( that is guaranteed to satisfy the constraint.
Lx no (z)=)(=))zz,))
First, we use a neural function approximator hg R..K R that maps logits log to a nonnegative matrix Ao he (1) of probabilities for each pair (z,z). Next, we iteratively normalize the columns to have marginals p(x) and the rows of A to have marginals Tp(2) for steps (a modified version of the algorithm proposed by Sinkhorn and Knopp [1967]. The last iterate A == Ar always satisfies E. Ag.: TT(2) but may only approximately satisfy the constraint that Σg Mo (2) (2]22,
2, p) p(z). To deal with this, treat A as proposal and apply a final accept-
reject correction, falling back to an independent draw from if the z-dependent proposal is rejected.
The marginals of this process give our expression for To p,
expression
p(z)
An,a (2)
Cz An. + (1 C. T0(2)
d2 C+
Cu =
Ax,
dg
Cas
p(z)
(14)
where C. max, Encoding this expression in the architecture of no (x(12,, p) ensures that = p(x), and thus choices of yield a valid reparameterization trick for all p. H for proof. While could parameterize and learn (2)),, we have thus far fixed it
Eg
See Appendix to the uniform distribution.
We note that if we let I = 1 (i.e. we assign all outcomes to one cluster), å becomes deterministic, and thus (I E, p) < p) p(x). In this case, we recover Gumbel-max coupling of p(æ) and ( showing that Gadget 2 also generalizes the Gumbel-max SCM.
Sampling from counterfactual distributions. Given a particular outcome x ~ P. we can sam- ple a counterfactual under some intervention q by first computing the posterior To ( 2(1) x To(2) (<r 2,, (2) and sampling an auxiliary variable = that is consistent with the observation. Given obtain Gumbel-max coupling between (o(x)=, p) and To ( q) so the top-down algorithm from Oberst and Sontag 2019 can be used to sample Gumbels and a counterfactual outcome y.
5.3 Learning Gadgets
Recall from Sec. that our goal is to learn 0 so that the gadgets above produce favorable implicit couplings when measured against dataset 'D and cost function g. In both gadgets, the constraint in Eq. is is automatically satisfied by the architecture. Thus, we need only concern ourselves with Eq. which is minimization problem over L with the following form:
[())) D Bay (o(oolt.P),sol7...
(15)
We would like to use reparameterization gradient where we sample (p.q) and 7, and then differ- entiate the inner term with respect to 0. However, the loss is not a smooth function of 0 given a
realization of due to the argmax operations in Eqs. II 13 Thus, our learning strategy is to relax
these argmax operations into softmaxes, as in Jang et al. [217]) Maddison et al. [2017] This yields a smoothed fo E K-1 and a smoothed softmax surrogate loss:
[(0) = D.y E1oo,, p))a (Fo(%: Uh g( wn
(16)
This is differentiable, and we can optimize it using gradient based methods and standard techniques (either explicitly summing over all x, y or taking sample-based REINFORCE gradient).
7

---

Related Work
The variational approach for coupling relates the maximal coupling the total variation distance
- gllrv maxxc(1,...K] Ip(A) q(A)) since ||p 1o--4lpp g||Tv =in(reC(p.o) [x[X + Y]. The Wasserstein distance W (p, q; d) in Eq. is generalization of the variational principle. The Wassestein distance can be extended to the optimal transport setting, when d is any function, which
has been used extensively in machine learning. see (Frogner et al.) 2015 Arjovsky et al.) 2017
Alvarez-Melis et al., 2018 Benamou et al. 2015 Blondel et al. 2018 Courtyy et al.) 2016
2017
Cuturi 2013
Aude al. 2016 Kusner al. 2015 Luise et al. 2018 Peyré et et. 2019].
The maximal coupling of Bernoulli random variables enforces monotonicity; it is attained by sampling (rrm(0 11) and setting X = 1fu S p] and Y = 1|u S g) [Fréchet 1951). More generally, Strassen's 's theorem asserts that any two random variables satisfy fx((1) S (1) if and only if they monotone, i.e., there is coupling и for which P,[Y - X 1. The monotone coupling can be realized by using the same uniform random variable U, setting X FU) and Y = (U) Lindvall et al., 1999]. A monotone coupling â E C(p, q) of two marginal probabilities p and q. maximizes the covariance of (X, Y) and consequently minimize the variance of X Y. since Var=[X - Y] Varp[X] Var, [Y] 2Cov, (X, Y). This is equivalent to maximizing the correlation between X and Y. Minimizing the variance of X Y helps to stabilize their estimation in machine learning applications, e.g., Grathwohl et al. 2017].
Cambanis et al. 1976) give conditions when the coupling (Fx (U), (U) is optimal. Let d(z y) be supermodular, i.e., d(z1,y1) y1) + d(02, V2) 2 d(z V2) + d(22. v1) x I 22 and (a (U), (U)). Specifically, when d(æ: y) = (z) and (F( (w)) are non- non-incree increasing or non- con-dereeasing functions of u, the LHS is the coupling with maximum covariance and the RHS the coupling achieved by common random numbers paired with a CDF inversion mechanism. See also Glasserman and Yao 1992].
81 S y/. Then SuPreC(p-g) En.y (xx
Monotonicity assumptions have also been used in the causality literature enable identification of a unique level 3 SCM from interventional data, for both binary [Pearl] 1999] and categorical [Lu al. 2020 random variables. We note that the implicit coupling induced by these assumptions also corresponds to the Inverse-CDe coupling (Fx (U), (U)).
Li and Anantharam 2019 consider joint couplings of a collection of distributions that are within a constant factor of optimal for all pairs. Their approach uses Poisson processes, which are closely related to Gumbel-max [Maddison] 2016), and they introduce auxiliary latent variables to adapt the coupling to particular cost functions, similar to our gadgets. Unlike this work, Li and Anantharam do not impose a distribution over the collection of distributions or interpret the coupling as a SCM, and they focus on directly constructing couplings instead of learning them.
7 Experiments
We evaluate our approach in two settings First, we build understanding by exploring performance on several datasets of fixed and random logits E RKK Next, we learn low-variance counterfactual treatment effects in the sepsis management MDP from Oberst and Sontag 2019
7.1 Optimizing for maximality
Section shows that no implicit coupling is maximal for every pair of p and q. Here we show that we can learn near-maximal couplings if we limit attention narrower distribution D of interest. We compare our proposed learning method against fixed Gumbel- -mx couplings and a maximal coupling. We hold fixed single pair of test distributions prstt and grest and then vary distributions that are trained on. Specifically, we train Gadget 2 on pairs (ptest + p. Tpp sest + p na) where T)p and Th are vectors of unit variance Gaussian noise, and p a noise scale. When p small, it possible to learn an implicit coupling that is specific to the region of distributions around (prest, g geest) and we can achieve implicit couplings that are nearly maximal. When p large, we are asking fo to couple together all pairs of distributions, and we expect to run into the impossibility results in Sec. 4.2
An implementation of our approach and instructions for reproducing our experiments is available at https _____aa__aaaa caax_ccuusal__e causal_gadgeet
//gi tgitthu.coo www gow.ge- wowgl/g www.eaaccttteeeeaaeeeeeeeeeeeeee http:///......oroooooooooooooooe

---

Table 1: Comparison of gadgets in controlled setting, shown with standard error.
Reward h:
Fixed linear (h(x) = x))
Random monotonic
non-monotonic
p and q:
Independent
Mirrored
Fixed inc/dee
Fixed inc/dec
Independent
16,51 0.01
21.34 + 0.10
2.73 0.39
0.83 0.10
Gumbel-max
14.02 = 0.01
18.84 + 0.09
2.46 + 0.35
0.50 + 0.06
CDF-1
8.14 + 0.01
12.59 + 0.09
0.60 + 0.10
0.91 + 0.12
Optimal (I.P)
8.14 0.01
12.59 + 0.09
0.60 0.10
0.11 0.02
Gadget
14.05 0.01
16.67 + 0.45
1.42 + 0.22
0.26 + 0.03
Gadget
8.76 0.03
13.47 + 0.09
2.00 0.30
0.21 0.03
Results appear in Fig. (a), with additional visualizations in Appendix | Indeed, when noise scales are small, our gadget learns a near-maximal coupling. When the noise scale becomes large, the quality declines. Interestingly, when the noise scale is orders of magnitude larger than the signal, our gadget never becomes significantly worse than the quality of Gumbel-max couplings.
7.2 Minimizing variance over random logits
In this experiment, we consider the ability of our learned couplings reduce variance in a controlled setting. Specifically, we define a scalar "reward" ((x)) for each outcome æ, and attempt to minimize the variance of [==pp[((=)] p and q and randomness in the reward h, and show that our gadgets achieve a lower variance than both the inverse CDF and Gumbel-max methods under mon-ooootonic reward functions.
( = Bz.y-a [h(x)) - hwy)]. We explore both randomness in
In the first part, we fixed the reward to the identity function and randomized p and q in two ways: (1) independently randomly drawn p and g. (2) drawn randomly and set to have the same probabilities as p but assigned in mirrored order. This tests the abilities our gadgets to learn couple arbitrary distributions and to uncover relationships between and q.
In the second part, we test our two gadgets under varying monotonic and non-monotonic reward functions h. We fix to be linearly increasing and to be linearly decreasing (the reverse of p). examine how the gadgets perform when p is very different from q. The monotone function is constructed by taking the cumulative sum of a K- K-length vector of Uniform (0. 1). The non non-monotone function is constructed by sampling K-length vector from Gaussian distribution followed by the function R(i) sin (30 * i). At each trial, the gadgets are trained from scratch under new realiza- tion of rewards. For comparison, we also solve for the optimal coupling using linear programming Bertsimas and Titsiklis 1997] (which may not be achievable by any implicit coupling).
Table shows the results of our experiments. We find that Gadget shows strong performance across all distributions of p and g, outperforming the Gumbel-max and independent sampling, and approaching the results of the optimal coupling under non-monotone h. Gadget outperforms Gumbel-max in the mirrored setting, and also outperforms Gadget 2 in the fixed increasing decreasing q setting under monotonic h.
7.3 MDP counterfactual treatment effects
In this experiment, we use a synthetic environment sepsis management to minimize the variance of counterfactual treatment effects on an SCM for MDP. Following Oberst and Sontag (2019), we used the simulator only for obtaining the initial observed trajectorics and we do not assume access to the simulator to get counterfactual probabilities. The simulator includes four vital signs (blood pressure, oxygen concentration, and glucose levels) with discrete states (low, normal, high), as well as three treatment options (antibiotics, vasopressors, and mechanical ventilation). Our goal is to couple p(s'|s, Aductor) and p(s'|s Ointervention). i.e., the transition distributions induced by two policies: a behavior policy, which mimics the physician policy, and target policy, which is the RL policy. Our counterfactual question is what would have happened to the patient if the RL policy's 's action had been applied instead of the doctor's.
Using a trained behavior policy, we draw 20,000 patient trajectories from the simulator with maximum of 20 time steps, where the states are in space of 6 discrete variables, each with different

---

Gapperai
qptira
Coupling (Bex))
0.25
10-2
13
10a
Noise scale during training
(a)
(b)
Figure 1: (a) When the training distribution is focused, learn near- near- maximal coupling. As the distribution becomes more diffuse, the learned coupling reverts to Gumbel-max, (b) Couplings of each method for the increasing p decreasing q settings along with the counterfactual effect variance for this specific reward realization. First row: monotone reward. Second row: non non-monotone reward.
Figure 2: Variance of the treatment effect on the sepsis management simulator.
number of categories (146 states total). Based on them, we learn the parameters an MDP, and train the target policy over this MDP. Unlike Oberst and Sontag (2099), we focused on coupling single time steps. We set the total reward for each state by summing its discrete variables rewards, which were sampled from Gaussian distribution. Among all the trajectories and time steps filtered out all pairs that have less than 4 non-zero probabilities. We conducted the experiment in two settings: joint sampling, where the noise is shared between the two samples, and counterfactual sampling, where we infer the noise (enw) based on the observation (s, adoctor: s'), then sample (s'18, 8,18 dinteruention) In each setting, we tested our gadgets when (p, 4) are fixed and when (p, q) are perturbed by Gaussian sample (generalized). At testing, we compute the treatment effect variance over 2000 samples and average the result across 10 different random seeds. With each trial of the experiment we fixed pair of (p, 9) and set new random realization of reward. We repeated that process for pairs of (p, 4) and reward realizations, a total 30 trials for each setting.
The means and the standard errors of the experiments are shown in Figure Both our gadgets outperformed the fixed sampling mechanisms, and Gadget did so by significant margin under the counterfactual settings. Details on the implementation of all the experiments are in Appendix]]
8 Discussion
We have presented methods for learning level causal mechanism that is consistent with observed level 2 information. Our framework provides significant flexibility to quantitatively define what makes a causal mechanism desirable, and then uses learning to find the best mechanism. Since any such choice cannot be confirmed or rejected by running experiment, an SCM this way is unprincipled. However, principles such as counterfactual stability can still be encoded into our framework using the loss g. We thus see our gadgets powerful tools which give modelers both the freedom and also the responsibility to select appropriate criteria for causal inference tasks, instead of being restricted to assumptions that cannot be tuned for specific use case.
might argue that choosing
One limitation is that we have only considered single-step causal processes, but we believe framework can be extended to multi-step MDPs. In future work, we would also like explore other settings of D and and investigate their qualitiative properties, such as how intuitive the resulting counterfactuals are to humans.
10

---

Acknowledgements
We would like thank the anonymous reviewers of our submission, whose excellent suggestions and requests for clarification were very helpful for improving the paper. We would also like to thank Michael Oberst and David Sontag for providing their implementation of the sepsis simulator and off off-policy evaluation logic, which we used for our experiments in Section 7.3 This research was supported by Grant No. 2029378 from the United States-Israel Binational Science Foundation (BSF).
References
David Alvarez-Melis, Tommi Jaakkola, and Stefanie Jegelka. Structured optimal transport. In International Conference on Artificial Intelligence and Statistics, pages 1771-1780. PMLR, 2018.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. arxiv 2017. arXiv preprint arXiv:1701.07875 30, 2017.
Genevay Aude, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. arXiv preprint 1r05vv/605..0000 2016.
Elias Bareinboim, JD Correa, Duligur lbeling, and Thomas Icard. On pearl's hierarchy and the foundations of causal inference. ACM Special Volume in Honor of Judea Pearl (provisional title), 2020.
Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):1111-AIII 2015.
Dimitris Bertsimas and John N Tsitsiklis. Introduction to linear optimization, volume 6. Athena Scientific Belmont, MA, 1997.
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport, In International Conference on Artificial Intelligence and Statisties, pages 380-889 PMLR, 2018.
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:// http:///tttbb com/ wtt.cooo/ooo.ooe gooa///oool/aaaa/oooo
Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactualll--ginnn policy search. arXiv preprint arXiv:1811.06277 2018.
S. Cambanis, G. Simons, and W. Stout. Inequalities for ek(x, y) when the marginals are fixed. Zeitschrift für Wahnscheinnicckeitstheerr und Verwandte Gebiete, 36:285-294, 1976.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for domain
adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853-1866 2016.
Nicolas Courty, Rémi Flamary, and Mélanie Ducoffe. Learning wasserstein embeddings. arXiv preprint arXiv:1710.07457, 2017.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26:2292-2300. 2013.
Maurice Fréchet. Sur les tableaux de corrélation dont les marges sont données. Ann. Univ. Lyon, 3 serie, Sciences, Sect. A, 14:53-77, 1951.
Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya-Polo, and Tomaso Poggio. Learning with wasserstein loss. arXiv preprint arXiv: 1506. 05439, 2015.
Paul Glasserman and David Yao. Some guidelines and guarantees for common random numbers. Management Science, 38(6):884-900 1992.

---

Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. arXiv preprint arXiv:1711.00/23 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical Reparameterization with International Conference on Learning Representations, 2017.
In
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International conference on machine learning, pages 957-966. PMLR, 2015.
Cheuk Ting Li and Venkat Anantharam. Pairwise multi-marginal optimal transport and embedding for earth mover's distance, 2019.
Torgny Lindvall et al. On strassen's 's theorem on stochastic domination. Electronic communications in probability, 4:51-59, 1999.
Chaochao Lu, Biwei Huang, Ke Wang. José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-base data augmentation. arXiv preprint arXiv:2012.09092,, 2020.
Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. arXiv preprint arXiv: 1897, 2018.
CA Maddison. Poisson process model for monte carlo. Perturbation, Optimization, and Statistics, pages 193-232, 2016.
Chris J. Maddison, Daniel Tarlow, and Tom Minka, A Sampling. In Advances in Neural Information Processing Systems 27, 2014.
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Continuous Re- laxation of Discrete Random Variables. In International Conference on Learning Representations, 2017.
Michael Oberst and David Sontag. Counterfactual off-policy evaluation with Gumbel-max structural causal models. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4881-4890. PMLR, 09-15 Jun 2019, URL htthttp: v97/oberst19a.1 html
mlr. .ress/
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-80888888 8035. Curran Associates. Inc., 2019. URL http::///ppe http:///pppprre...rippsppppsp.ps nww.ins..ip... www.eris.....pppr........ 9015 an -n-imperativee www.iniaaaiv-styll---nteee perple--eee----foror-----e----e pdf
Judea Pearl. Probabilities of causation: three counterfactual interpretations and their identification. Synthese, 121(1):93-149 1999.
Judea Pearl. Causality. Cambridge university press, 2009.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11ee,,11155 116-66)33556077 2019.
Richard Sinkhom and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-348, 1967.
Richard Sutton, Andrew G Barto, al. Reinforcement learning: An introduction (in progress). London, England, 2017.
12


Reviews:
[{"review_id": "l3lWoZXZ6ce", "reviewer": "Reviewer_eGhr", "summary": "The authors propose a method extending Gumbel SCM for doing counterfactual reasoning.", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper could have been written a little more clearly in places. For example, Figure 2 did not appear to be linked back to the text. It would have been helpful to have a little more explanation. “Discriminative” was misspelled line 256.\n\nFor originality, I haven't seen any other papers like it, though perhaps other reviewers have.\n\nThe quality of the theory is clear.\n\nThis is significant insofar as counterfactual reasoning is a persistent problem for causal inference. It would have been helpful if the approach had been compared to other approaches for doing counterfactual reasoning. At least a discussion could have been included as to why this covers cases different from what, say, Pearl's methods would have covered.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I don't think this work would have negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "hC7ZO7xy8DD", "reviewer": "Reviewer_4cdP", "summary": "This paper generalizes the treatment effect estimation, in terms counterfactual and observed policies, of Oberst and Sontag (2019), to couplings. They build upon the Gumbel-max SCMs from Oberst and Sontag and show that with their proposed Gumbel-max coupling, they can learns causal mechanisms that minimize the counterfactual effect variance. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this paper is quite well written for an audience that has read the paper of Oberst and Sontag (2019). Since I didn't read that paper beforehand, I had no clue what problem the authors were trying to solve. For example, the authors state in the introduction \"Casting the problem this way ...\", without explaining what the actual problem is, or describing the problem setting there in more detail. This was quite confusing, because the authors talk later about a data distribution that determines observed outcomes and counterfactuals. However, in practice, one has almost never access to those distributions IMO. Once I arrived at the Experiments section, I saw that the authors use simulations to get access to the counterfactual distributions. I think this confusion can be easily tackled by making the actual problem statement clear from the beginning.\n\nAlthough the paper is well organized and each section is written very clearly, the overall storyline of the paper can be improved IMO. I think that the different sections (and concepts) are too loosely connected. For example, clarifying how the SCM for MDPs are related to the problem that the authors are trying to solve should be very clear from the beginning onward instead of being clarified (slightly) in the Experiments section. Or, on line 12 (and also on line 45) the authors state \"We propose a parametrized family of causal mechanisms that generalize Gumbel-max\", without showing this in the paper, or at least. I couldn't find it. Are the authors not generalizing the treatment effect instead of the causal mechanisms? This is a bit confusing, after reading what the definition of a causal mechanism of an SCM is.\n\nThe actual contributions of this paper could IMO be emphasized more. For example, one of the main contributions of this paper, that is IMO, that seeing the average treatment effect for a counterfactual vs. an observed policy in terms of couplings, could be emphasized more. This contribution is, for example, a bit downplayed by stating (in line 70) that [Maddison et al., Oberst and Sontag] exploit Gumbel-max couplings, which I think they never showed, since the proof of this is actually in the Appendix of this paper and not theirs. Or, in the introduction the authors state the questions about what mechanisms leads to the lowest variance estimates of treatment effects, which they didn't put much emphasis on in the paper, except by stating it somewhere in the related work section. Why not highlight this more and state how optimality and minimizing variance of the treatment effects are connected explicitly. I think these things that motivate this work should really be emphasized and pop out. This can make this paper easier to read and more strong.\n\nAt last, I was wondering about how to interpret the main objective? By generalizing the objective from Oberst and Sontag (2019) from difference between the outcomes to some (different) coupling between the outcomes, doesn't that also change the definition of the average treatment effect? That is, can we still talk about average treatment effect in some way? What is the intuition that lays behind using couplings instead?\n\nAlthough the storyline could be improved, the rest of the paper is quite well written. Under the condition that all the proofs are correct (which I didn't check), I believe that this could potentially be a valuable contribution to NeurIPS.\n\nMinor comments:\n - line 126 \"... approach to the Gumbel-max couplings underlying Gumbel-max SCMs\": Are they really underlying Gumbel-max SCMs?\n - line 134 \"... Gumbel-max couplings are counterfactually stable ...\": Where is this shown?\n - line 75: \"casual\" -> \"causal\"?\n - Figure 3: What is CRN? Please explain this somewhere.\n\nUpdate: I updated my score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the impact of this paper can be improved by providing a better intuition about the relation between couplings and counterfactual effect estimation, especially about why couplings lead to more desirable causal mechanisms. Moreover, as said before, improving the overall storyline could make the impact of this paper stronger.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6h"}, "extra_text": {}}, {"review_id": "Etuo0ooOsHY", "reviewer": "Reviewer_cSRr", "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. \nThe 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. \nThe 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I am not familiar reparametrization or couplings, but as far as I can tell, their application to causal modeling tasks is novel. It seems that relevant work as been cited (Oberst and Sontag, 2019) and the differences have been addressed.\n\nI am unable to assess the soundness of most of the technical results as they are outside my are of expertise. However, I do have some questions regarding the specifics of the causal task that is attempted which I will discuss shortly. Overall, the paper does seem to support the claims and the experiments provided seem to support the benefits from employing the proposed approach.\n\nIn terms of clarity it is hard for me to understand what exactly is the data (observational, interventional) assumed to be known for the task. Also, what is exactly the relationship between the logits vectors and that data? Does this approach work for any type of input data or are you assuming a particular form? I believe a simple running example could help ground these aspects.\nThe introduction touches upon the different levels of the ladder of causation and the fact that it is not possible, in general, to infer level 3 data solely from level 2 data. Nevertheless, it is not clear what are the level 3 quantities that the authors aim to estimate with minimal variance. \nIt is not clear to me what kind of counterfactual is being evaluated. Also, the motivating observation \"if we only care about our models being consistent with interventional data, then we have flexibility in the choice of causal mechanism\" seems to disregard the level 3-level 2 impossibility mentioned in the introduction. I hope the authors could clarify these points in their response.\n\nIn terms of significance, I am unable to assess the applicability of the proposed method, mainly to the lack of a clear picture of the input data, the query of interest and the assumptions being made about the true underlying data generating process. It is very likely that this picture could be understood from the paper but I believe it should be more explicit and easy to grasp for the average reader trained in causal inference. On a related note and in contrast with Oberst and Sontag, 2019, the SCM formulation in the paper seems to assume the absence of latent confounding, is this the case?\n\n## Update\nAfter reading the response from the authors, I feel they have clarified my main concerns. I'm updating my score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There is discussion on technical limitations.\nThere is no negative societal impact that requires discussion at this point.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "-5neAmDa8nw", "reviewer": "Reviewer_g3MG", "summary": "This paper introduces a continuously parameterized family of causal mechanisms where all members of the family are identical when used in a level 2 context but different when used in a level 3 context. The families contain the existing Gumbel-max technique proposed by Oberst & Sontag (2019), but a wide variety of other mechanisms can be learned by using gradient-based optimization. They show that the causal mechanisms can be learnt using a variant of Gumbel-softmax relaxation and that the mechanisms improve over Gumbel-max and other fixed mechanisms. Further, they show that the learnt mechanism is generalizable from a training set of observed outcomes and counterfactual queries to a testing set of observed outcomes and counterfactual queries.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Section 1 & 2 & 3. The causal problem is not defined properly. If you are trying to solve the problem of learning causal mechanisms, please define mathematically what is a causal mechanism, what data you learn it from, what is the key challenge in general. I guess you don’t have to answer (not solve) these general questions with the Gumbel-max trick.  After presenting the general causal problem clearly, it will be better to explain the following two issues mathematically or using some toy examples of structural causal models (SCMs): (1) the non-identifiability of level 3 from level 2; (2) the requirement of counterfactuals (mentioned in lines 34-35); (3) Why we need to solve the problem using the family of causal mechanisms that generalizes the existing Gumbel-max SCMs. In the current introduction, Section 2 (Background) and Section 3 (Problem statement: Learning Implicit Coupling), I don’t see a well-defined causal inference problem. I feel like this is a coupling learning paper rather than a causal inference paper.  \n\nOverall, I suggest addressing the following three points mathematically: (1) formulate the problem of learning causal mechanisms; (2) formulate the requirement of counterfactuals; (3) highlight the advantage of learning specifically-tuned causal mechanisms.  I find the sections are well written, but the explanation for counterfactual stability in lines 88-94 is difficult to understand.\n\nSection 4. Why do we need maximal couplings? I guess the answer is on lines 64-69. It looks like the variance of the objective function you will optimize to learn the coupling. How does this variance translate into the variance of the learnt coupling or causal mechanism? For the noise variable u, distributions p and q, could you define their space clearly?\n\nSection 5. I am not very familiar with Gumbel-max and coupling theory. I must confess I don’t completely follow the discussion in this section. It is difficult for me to judge the novelty of the proposed method. Could you point out the novelties more clearly and explain why the continuously parameterized family of causal mechanisms you propose generalizes Gumbel-max?\n\nSection 6. The related works on maximal coupling look good to me. I hope I know it earlier when I am reading through the paper.\n\nSection 7. In Figure 3, you compare the variances of the estimates. Could you also compare the errors of the estimates? \n\nSummary:\nOverall, I find this paper interesting but a bit beyond my expertise. I lower my confidence level and will adjust my score appropriately after understanding the paper better by reading through your feedback and other reviews' comments.\n\nOther comments:\n\nIn lines 38-39. Is the variance of an estimator a computational behaviour or statistical property (behaviour)?\n\nLine 57:\nYou can always draw a sample from a categorical distribution.  What special about the Gumbel-max trick is that you can draw a differentiable sample from the continuous relaxation of a categorical distribution. \n\nLine 81:\n“that map states, actions, and noise” should be “that map actions, states and noise”\n\nEquation (5) under line 115, a typo: “min. \\theta”?\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they do.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper explores an understudied connection between causal inference and couplings learning, building upon very recent work from Oberst & Sontag. More specifically, it considers the general problem of constructing a SCM that allows counterfactual (level 3) reasoning, when given access to all interventional (level 2) distributions. The key challenge is that there are often many different SCMs compatible with a set of level 2 distributions. To circumvent this non-identifiability issue, they find a SCM with "good" counterfactual distributions based on interventional data, where a "good" counterfactual distribution is one with minimal loss according to some particular loss function. Here, the approach they take is to extend Gumbel-max SCMs with learned couplings.

While the reviewers initially had some concerns re. clarity/significance of the contribution, the authors' replies and subsequent discussion cleared up these concerns and a majority of reviewers agreed that it is worth publishing at neurips.
Decision: Accept (Spotlight)