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": "vujIK7LM43m", "reviewer": "Reviewer_XBKm", "summary": "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.\n", "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": "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. \n\nAn 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.\n1. 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.\n\nHowever, I do not quite get the details of the proposed method. The authors may need to consider improving the readability of the paper.\n\nUpdates:\nThanks to the authors for the responses. \n\nAfter reading the paper again, I think the writing needs much improvement. The current version is hard to follow.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "eMObqssaBJE", "reviewer": "Reviewer_ybLo", "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. ", "questions": "", "limitations": "", "rating": 3, "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": "### Main comments\n\n1. **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. \n\n2. **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. \n\nHowever, 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. \n\n### Minor comments\n\n1. Eq. (5): How is the distribution $\\mathcal{D}$ over pairs of logits obtained?\n\n\n\n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I do not see any potential negative societal impact of the work in its current form. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "7654Ln9iw2B", "reviewer": "Reviewer_gci3", "summary": "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.", "questions": "", "limitations": "", "rating": 5, "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": "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.\n\n- 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.\n\n- 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?\n\n- 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?\n\n- 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?\n\n- 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.\n\n- 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.\n\n- 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.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- This work is mainly theoretical. Its long-term societal impact remains unclear at this point.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}]

Author rebuttal:

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 (Spotlight)