Partial Identification of Treatment Effects with Implicit Generative Models
Abstract
We consider the problem of partial identification, the estimation of bounds on the treatment effects from observational data.
Although studied using discrete treatment variables or in specific causal graphs (e.g., instrumental variables), partial identification has been recently explored using tools from deep generative modeling.
We propose a new method for partial identification of average treatment effects (ATEs) in general causal graphs using implicit generative models comprising continuous and discrete random variables.
Since ATE with continuous treatment is generally non-regular, we leverage the partial derivatives of response functions to define a regular approximation of ATE, a quantity we call uniform average treatment derivative (UATD).
We prove that our algorithm converges to tight bounds on ATE in linear structural causal models (SCMs). For nonlinear SCMs, we empirically show that using UATD leads to tighter and more stable bounds than methods that directly optimize the ATE.
Introduction
Estimating average treatment effects (ATEs) is a common task that arises in fields involving decisionmaking, such as healthcare and economics.
In the presence of the gold-standard randomized controlled trial (RCT) data, one can compare the outcome variable between treated and control groups to make decisions.
But RCTs can be costly to set up and run and are, in many circumstances, infeasible.
Consequently, communities are using observational data to assist in decision-making.
Identification of treatment effects from observational data is tied to the structure of the causal graph.
For example, the treatment T and outcome Y in Figure 1b are confounded by an unobserved random variable, making it impossible to find the causal effect of T on Y only from observational data.
On the other hand, Figure 1c is identifiable, and one can adjust for confounders using the Back-door formula [Pearl, 2009].
Even in identifiable settings, non-parametric estimations such as Back-door adjustment formula can point-identify the ATE only with additional assumptions such as positivity, i.e., P ( T = t | X ) > 0 for all values of covariate X .
Observational data is finite, high-dimensional, and consequently can suffer from severe violations of such assumptions [D’Amour et al., 2021].
In lieu of the challenges of point-identification, there has been a recognition that decisions can be justified using reliable bounds on the ATE rather than its exact value.
For an oncologist treating a cancer patient, knowing that a drug has a significant, positive reduction in the patient’s risk of progression may suffice as a rationale to prescribe that drug.
This problem is known as partial identification [Manski, 2003].
Most existing methods for bounding the ATEs are only applicable in discrete/binary treatment variables [Makar et al., 2020, Zhang et al., 2021, Duarte et al., 2021,
Guo et al., 2022].
There has been recent interest in continuous treatment settings.
However, such
[Figure]
methods are applicable for special causal graphs such as the instrumental variables (IV) setting [Gunsilius, 2020, Kilbertus et al., 2020] or make parametric assumptions on the family of treatmentresponse functions [Padh et al., 2022].
An exception is the work by Hu et al.
[2021] which provides a non-parametric approach for partial identification using generative adversarial networks (GANs).
However, they only provide convergence guarantees for the special case of IV causal graphs.
Using the framework of structural causal models (SCMs) and causal graphs, one can see partial identification as a constrained optimization problem, where the objective, i.e., maximizing/minimizing the ATE, can be written as a post-intervention function of exogenous noise (a.k.a response function) and the constraint is to match the generated samples with the observational distribution.
This naturally leads to using generative neural networks such as neural causal models (NCMs) [Xia et al., 2021].
We find that directly solving the ATE optimization using flexible generative models such as GANs can lead to non-informative and degenerate solutions.
The flexibility afforded by generative models such as GANs allows them to deviate significantly from the true response curve in the neighborhood of intervention points to maximize/minimize the ATE while continuing to generate samples akin to the data distribution.
Figure 1e (left) showcases a typical solution to the ATE optimization.
Our insight is that the ATE between any two points can be approximated as an integral over the derivatives of the response function w.r.t.
the treatment variable.
Rather than directly optimizing the ATE, we optimize the partial derivatives of the response function, a quantity that we refer to as the uniform average treatment derivative (UATD).
By optimizing the UATD, the model is required to maximize/minimize the partial derivatives for all points within the treatment support, avoiding extreme local solutions as shown in Figure 1e (right).
Our contributions are as follows:
• We formally define the partial identification of average treatment effects as a distributionallyconstrained optimization problem, where we choose Wasserstein distance as our constraint metric.
• For the class of linear SCMs, we prove that the solution to our optimization problem converges to optimal bounds on the true value of ATE in infinite data for general causal graphs.
• We use the solution to partial identification of UATDs to find informative bounds on the value of ATE.
We introduce a practical algorithm to solve the distributionally-constrained optimization problem using the Lagrange multiplier formulation with alternating optimization.
We empirically
show that our algorithm results in tighter and more stable bounds than methods that directly optimize the ATE.
Problem Setup & Background
We introduce the definitions and assumptions we will use throughout the paper.
Consider the observed data as (possibly continuous) random variables V
= { X 1 , · · · , X m , T, Y } ∈ R d , where T , Y , and { X 1 , · · · , X m } denote the treatment variable, target variable, and covariates, respectively.
Data generating model.
Our approach will be based on the framework of Structural Causal Models (SCMs).
An SCM is a tuple M =
( V , U , F , P U ) , where each observed variable V i ∈ V is a deterministic function of a subset of variables
pa ( V i ) ⊆
V and latent variables
U V
i ⊆
U , i.e.,
V i =
f V i ( pa ( V i ) , U V i ) where f V i ∈ F , V i ̸∈ pa
( V i ) (1)
The only source of randomness are latent variables U with probability space (Ω , Σ , P U ) .
This induces a probability law over the observed variables P M .
We may omit the subscript M and denote the observational probability distribution by P throughout the text.
Given M , one can construct a graph with nodes V ∪ U and directed edges from nodes in pa ( V i ) ∪ U V i to V i .
We call this graph the causal graph corresponding to SCM M and denote it by G M or simply G .
We assume G is acyclic and known. Moreover, We will assume each node in the graph is 1 -dimensional.
For random variable V in the SCM M , let V M ( u ) be its deterministic value after fixing a realization u of latent variables U .
The causal effect of treatment T on target Y is:
Definition 1 (Causal Effect) .
Let Y M ( T = t ) ( u ) be the value of Y by fixing U = u
and changing the function f T to a constant function f T = t in M .
Then, we call the random variable Y M ( T = t ) the causal effect of treatment T = t on target Y .
We may simplify the notation and write it as Y t if the
SCM M and treatment variable T are known from context.
Note that Y T ( u ) ( u )
=
Y ( u ) .
When T is continuous, then we can view { Y t : t ∈ supp ( T ) } as a stochastic function defined on (Ω , Σ , P U ) .
This is referred to as the response function, partial dependence plot, and dose-response curve in the literature [Zhao and Hastie, 2021, Ritz et al., 2015, Chernozhukov et al., 2018].
Average treatment effect, average treatment derivative, and partial identification.
Our goal is to estimate bounds on the effectiveness of a treatment regime on a population from the observational distribution P and the causal graph G .
In the continuous treatment case, where there is no "on"/"off" notion of treatment, we can compare the average causal effect of an arbitrary treatment (dosage) to the average causal effect at a fixed point T = t 0 .
For example, to indicate the effect relative to not prescribing any treatment, we can choose t 0 = 0 .
This quantity is known as the average treatment effect, average level effect, or average dose effect in the literature on continuous treatment setting [Hirano and Imbens, 2004, Kennedy et al., 2017, Callaway et al., 2021].
Definition 2 (Average Treatment Effect) .
For SCM M , the average treatment effect (ATE) at T = d w.r.t.
a fixed point T = t 0 is defined as
[Equation]
Note that estimating the ATE and finding bounds on it only depends on the value of the average response function in T = d and T = t 0 .
As pointed in Gunsilius [2020], this quantity can take arbitrary values if we do not make any assumptions on the set of response functions.
Here, we assume the partial derivative of the response function w.r.t.
the treatment, i.e., ∂Y t /∂t exists and is a bounded continuous function.
We then define the average treatment derivative as the following:
Definition 3 (Average Treatment Derivative) .
For the treatment regime f T in SCM M , we define the average treatment derivative (ATD) as
[Equation]
Estimating the ATD can be seen as a proxy for the effectiveness of the prescribed treatment, where we consider the population-level average effect of an infinitesimal increase in the treatment/dosage [Rothenhäusler and Yu, 2019].
In this work, however, we leverage the regularity of this quantity to achieve smoother solutions to the ATE estimation.
We will expand on this in section 4.
Note that we cannot readily use eq. 2 (or eq. 3) to estimate the ATE (or ATD), as we only have access to the observational distribution P and the causal graph G and not the latent distribution P U .
In fact,
ATEs are generally non-identifiable, i.e., there exist multiple SCMs with the same causal graph G and generated distribution P that result in different values of ATE.
For some graphs, however, one can use non-parametric identification algorithms like do -calculus to identify the causal effect from the observational distribution [Pearl, 2009].
In practice, even for identifiable causal graphs, we cannot pinpoint the true ATE due to the uncertainty caused by sampling variation and finite sample errors.
Instead, we are interested in finding a tight set of possible solutions that will contain the true value of ATE (or ATD) with high probability.
This is known as the partial identification problem in the literature [Manski, 2003].
3 More formally, the partial identification of ATDs/ATEs is defined as:
Definition 4 (Partial Identification of ATD/ATE) .
Partial identification of ATD is the solution to the following optimization problem:
[Equation]
where M is the set of all SCMs on random variables V .
We denote the solution to the above problem as ( ATD , ATD ) .
Similarly, we can define the partial identification of ATEs by replacing ATD M ′ with
ATE M ′ ( d ) in eq.
We refer to the solution to the latter problem as ( ATE ( d ) , ATE ( d )) .
Implicit generative models.
To solve the partial identification problem, we use the expressive power of generative models to satisfy the distribution constraint in eq. 4.
Choosing distance measures such as Jensen-Shannon divergence or Wasserstein metric results in models such as GANs or Wasserstein
GANs (WGANs)
[Goodfellow et al., 2014, Arjovsky et al., 2017].
The typical way to implement these models is to solve a minimax game between the generator and a discriminator.
However, adding the ATD minimization/maximization term to the minimax loss function will result in unstable training.
Instead, in our practical algorithm, we will use Sinkhorn Generative Networks (SGNs) that use Sinkhorn divergence S ϵ , a differentiable ϵ -approximation of Wasserstein metric, as the distance measure between generated and true samples [Cuturi, 2013, Genevay et al., 2018, Feydy et al., 2019].
Due to the differentiability of Sinkhorn divergence, we will only need to train a generator network enabling us to sidestep much of the unstable minimax training in (W)GANs.
Related Work
This work builds upon partial identification and generative causal models.
Partial identification.
Finding informative bounds on treatment effects has been well-studied in the existing literature ([Robins, 1989, Manski, 1990, Evans, 2012, Ramsahai, 2012, Richardson et al., 2014, Miles et al., 2015, Finkelstein et al., 2021, Zhang and Bareinboim, 2021a,b]).
Balke and Pearl [1997] find the tightest possible bound for the discrete instrumental variable setting by converting it to a linear programming problem.
For the backdoor setting and binary treatments, Makar et al.
[2020] provide probabilistic upper/lower bounds on potential outcomes in the finite sample regime.
Recently, Zhang et al.
[2021] and Duarte et al.
[2021] independently describe a polynomial programming approach to solve the partial identification for general causal graphs.
They both use the notion of canonical SCMs to map the latent variables to the space of functions from treatment
T to outcome Y .
Though they show their polynomial programming formulation finds the optimal bound, their approach is only applicable to discrete random variables with small support.
In fact, the time complexity of their algorithm grows exponentially with the size of the support set of variables, making their algorithm intractable for continuous settings.
Gunsilius [2019] extends the commonly-used linear programming approach to partial identification of IV graphs with continuous treatments.
They use a stochastic process representation of the variables and solve the linear programming via sampling.
However, their method suffers from stability issues, as discussed in Kilbertus et al.
[2020] and is only applicable for the IV setting.
Kilbertus et al.
[2020], Padh et al.
[2022] parameterize the space of response functions by assuming them as linear combinations of a set of fixed basis functions.
Then, they match the first two moments of observed distribution while minimizing/maximizing the ATE.
However, they do not provide any theoretical guarantees on the tightness of their derived bounds.
Most similar to our work is Hu et al.
[2021] who use generative adversarial networks (GANs) to match the observed distribution and search for response functions with maximum/minimum ATEs.
They provide convergence guarantees for the instrumental variable causal graph with linear models.
Their approach is also based on the minimax game between generator and discriminator, which can result in unstable training.
Our work differs in a few important ways.
We focus on partial identification of average derivatives and use that to find bounds over the ATE.
Using this approach, we show that our derived bounds converge to the optimal bounds for linear SCM with general causal graphs, including both identifiable and non-identifiable settings.
We use Sinkhorn divergence, a differentiable approximation of Wasserstein distance, to train our implicit generative models.
Empirically, we find that this avoids the unstable training of GANs.
Guo et al.
[2022] studied the partial identification of ATE with noisy covariates.
Their work is similar to our approach in that we both use a similar robust optimization formulation.
However, they focus on identifiable causal graphs, where one can use adjustment formulas such as the Back-door formula and make parametric assumptions on the joint distribution of observed variables.
Generative causal models.
[Goudet et al., 2017, Yoon et al., 2018, Kocaoglu et al., 2018, Sauer and Geiger, 2021] use generative models to capture a causal perspective on evaluating the effect of interventions on high-dimensional data such as images.
They do not consider the problem of bounding treatment effects.
Xia et al.
[2021] introduced Neural Causal Models (NCMs) that leverages the universal approximability of neural networks to learn the SCM.
Although it is not generally possible to learn the true SCM by training on the observational data, they prove that NCMs can be used to test the identifiability of causal effects and propose an algorithm to estimate identifiable causal effects.
Their work’s theory and empirical instantiation are in the context of discrete random variable datasets.
Our work builds upon NCMs for partial identification with both continuous and discrete random variables.
Partial Identification using Implicit Generative Models
We explain our method to solve the partial identification problem in Def. 4 using implicit generative models.
In subsection 4.1, we describe partial identification of ATDs as a constrained optimization problem using G -constraint generative models [Xia et al., 2021].
Then, in subsection 4.2, we show that the solution to this constrained optimization problem converges to the optimal bounds on the
ATD in infinite data samples.
We prove our results for linear SCMs with general causal graphs, i.e., both identifiable and non-identifiable settings.
Next, we propose our approach to extend the partial identification of ATDs to ATEs.
Finally, we describe a practical algorithm to solve our distributionally-constrained optimization problem in subsection 4.3.
4.1 G -constraint generative models
To solve the partial identification problem, we need to search over the set of all possible SCMs M .
This is generally not feasible, as there is no constraint on the distribution of the latent variables P U , as well as the function family F .
Instead, we parameterize the space of all SCMs that are consistent with causal graph G using neural networks.
More specifically, we use G -constraint generative models:
Definition 5 ( G -constraint Generative Models (Def. 7 in Xia et al.
For a given causal graph G , a G -constraint generative model is a tuple M θ G =
( V , ˆ U , F θ , ˆ P ˆ U ) , where each V i ∈ V is generated from
[Equation]
where pa ( V i ) is the observed parents of node V i in G and
ˆ U C ∈
U is the latent noise corresponding to maximal C 2 -Component C ⊆ V containing node V i , i.e., each pair of variables in C have common latent parent nodes.
In addition, ˆ P ˆ U ∼ Unif (0 , 1) for each ˆ U ∈
ˆ U .
G -constraint generative models make the search over the set M feasible by limiting it to generative models with uniformly distributed latent variables that are consistent with causal graph G .
In fact, in their Theorem 3, Xia et al.
[2021] show that for any discrete SCM M ∗ with causal graph
G , there exists a G -constrained generative model M θ G that generates the same causal effect, i.e., Y
M ∗ ( T = t ) =
Y M θ G ( T = t ) a.s.
Their proof technique, however, only works for SCMs with discrete variables.
Here, we do not prove the expressiveness of G -constrained generative models for continuous SCMs.
Instead, for simplicity and completeness of our theoretical results in subsection 4.2, we assume that the true SCM is a G -constrained generative model itself.
In our experiments, we empirically show that our results hold even for SCMs with different latent distributions, such as Gaussian noise.
Assumption 1.
The true SCM M is a G -constrained generative model.
In other words, there exist θ such that M = M θ G .
Under Assumption 1, we reformulate the problem in eq.
4 using generative models, i.e.,
[Equation]
In practice, we never have access to true distribution P as we only observe a finite number of samples corresponding to the empirical distribution P n = 1 n (cid:80) ni
=1 δ v ( i ) for a given dataset { v (1) , · · · , v ( n ) } .
Also, the observed variables may be biased due to noisy measurements.
Therefore, we reformulate the problem in eq.
6 as a constrained optimization problem.
We choose the 1-Wasserstein metric as our distance measure, which naturally results in generative models such as WGANs.
We will state our theory in subsection 4.2 based on this metric.
However, in subsection 4.3, we will propose a practical algorithm that uses Sinkhorn divergence, a differentiable approximation of 1-Wasserstein distance, for more stable results.
Our constrained optimization problem is as follows:
[Equation]
where α n is a hyper-parameter that specifies the level of tightness of the bounds.
We denote the solution to eq.
7 as ( ˆ ATD , ˆ ATD ) .
In the case of noisy measurements, we need domain knowledge of how noisy the data is to determine the value of α n .
Otherwise, we can use the finite-sample convergence rate of empirical Wasserstein distance to choose an appropriate value of α n
[Weed and Bach, 2019].
As our theoretical results are concerned with the infinite-sample case, we will assume that there exist values of α n such that the true distribution lies within the Wasserstein ball.
Assumption 2.
For each n ∈ N , there exist α n >
such that W 1 ( P, P n ) ≤ α n .
4.2 Theoretical guarantees and extension to ATEs
Assumptions 1 and 2 ensure that the bound derived by eq.
7 contains the true value of ATD.
However, we do not know how informative/tight the derived bounds are.
In fact, one can always return ( −∞ , + ∞ ) as one solution to partial identification.
This part gives theoretical guarantees that our algorithm can result in tight bounds over ATD.
In particular, we focus on linear SCMs and show that, under the infinite number of samples, our algorithm converges to the optimal bound ( ATD , ATD ) for both identifiable and non-identifiable causal graphs.
See Appendix A for the proof.
Definition 6 (Linear SCMs) .
SCM M =
( V , U , F , P U ) is linear, if
[Equation]
Theorem 1 (Tight Bounds) .
Assume the dataset { v (1) , · · · , v ( n ) } is generated from a linear SCM.
Then, under assumptions 1, and 2, the solution to the constrained optimization problem in eq.converges to the optimal bound over the ATD in infinite samples, i.e., ˆ ATD → ATD and ˆ ATD → ATD.
Up until now, we have only focused on partial identification of ATDs.
Here, we discuss how to extend our results to find bounds on ATEs.
A naive solution is to replace ATD with ATE in eq.
7 and directly optimize it.
However, as demonstrated in the experiments, this approach can result in non-informative bounds.
In fact, ATE with continuous treatments is a non-regular quantity [Kennedy et al., 2017].
Instead, we claim that we can use the same generative model trained for partial identification of ATD to bound the value of ATE.
In particular, we define a new objective function by uniform intervention on the treatment, which we call uniform average treatment derivative (UATD), and show that the solution to partial identification of UATD matches the solution to partial identification of ATEs.
Definition 7 (UATD) .
For an SCM M , we define the uniform average treatment derivative (UATD) at interval [ t 0 , d ] as
[Equation]
.0 .−−
5 −  −− 5 −  −− 5 −
Treatment, T T a r g e t , Y true response upper bound treatment density ˜ µ Now, we state our result on using average derivatives to solve partial identification of ATEs:
Corollary 1.
Let θ ∗ be the solution to the partial identification of UATD at interval [ t 0 , d ] .
Then, θ ∗ is also a solution to the partial identification of ATE ( d ) .
If the true SCM M is linear, then the bound is tight, i.e., ˆ ATE ( d ) → ATE ( d ) and ˆ ATE ( d ) → ATE ( d ) as n → ∞ .
The proof is given in Appendix B.
Remark.
ATD and ATE are generally different quantities and finding bounds on one does not necessarily results in bounds on the other. In fact, ATE is not pathwise differentiable for continuous treatments [Díaz and van der Laan, 2013, Kennedy et al., 2017, Chernozhukov et al., 2018] (See Appendix C). We define UATD as a quantity to relate ATE and ATD.
The formal definition of UATD in
Def. 7 is the same as ATE up to a scale factor, and one would expect to see similar issues with ATE here as well. Therefore, we approximate UATD with a version that, instead of a uniform distribution of treatment within interval [ t 0 , d ] with zero density outside, the treatment distribution has continuous differentiable non-zero density defined over the whole support of T .
More concretely, [Equation]
(cid:34)(cid:90) supp ( T ) ∂Y M θ G ( T = t ) ∂t dµ ( t ) (cid:35) ≈ E u ∼ P ˆ U (cid:34)(cid:90) supp ( T )
∂Y M θ G ( T = t ) ∂t d ˜ µ ( t ) (cid:35) (10) where µ is a the uniform measure within interval [ t 0 , d ] and ˜ µ is its approximation with a non-zero density over the full support of T . Choosing ˜ µ trades off between regularity of the response function and the approximation error.
The more ˜ µ is close to the uniform measure µ , the more irregularity we allow in the response curve and the less informative the bound will be, while the objective function is closer to ATE. Adding more density outside of interval [ t 0 , d ] imposes regularity on the response curve (Figure 2). It is important to note that we use the formal definition of UATD (with uniform treatment distribution) in Corollary 1 since, in linear SCMs, it is not possible for the generator model to attain arbitrarily large values in points t 0 and d as the derivative of linear functions remains fixed outside of interval [ t 0 , d ] . This is why we state our theoretical results based on the uniform intervention.
4.3 Our algorithm
We describe our algorithm to solve the optimization problem in eq.
We will focus on finding ˆ ATD, a similar approach can be taken for ˆ ATD.
A general strategy is to convert the constrained problem to its unconstrained version using the method of Lagrange multiplier:
[Equation]
As the Wasserstein distance is not differentiable, we cannot directly use gradient descent to solve eq.
A common approach is to use the dual formulation of Wasserstein distance W 1 ( P M θ G , P n ) = max
|| q ϕ ||
L ≤ 1 E P n [ q ϕ ( v )] − E P M θ G [ q ϕ ( v )] and solve eq.
11 using WGANs, a similar solution used [Figure] in Hu et al.
However, this min-max-max formulation can result in unstable bounds as we show in our experiments.
Instead, we use Sinkhorn divergence, a differentiable approximation to
Wasserstein distance, as the measure of distance between distributions and solve the following:
[Equation]
To solve eq.
12, we need to evaluate ATD M θ G and calculate its gradient w.r.t.
θ .
As we are using
G -constrained generative models, we can calculate the value of Y M θ G ( T = t ) ( u ) by hard intervention
T = t , i.e., fixing the output of function f θT as t and computing Y through a topological order of calculations.
Then, we estimate ATD M θ G as follows:
[Equation]
(13) where { t ( i ) } ni =1 are samples from the treatment variable, and { u ( i ) } ni =1 are the latent variables generated from a uniform distribution.
To choose an appropriate value of α n , we first train our generator without the ATD term until the Sinkhorn loss converges to some value and use that as our choice of α n .
We then continue our training by adding the ATD term.
We note that using algorithms such as projected gradient descent to solve the constrained optimization problem requires us to project the weights of our network into the Wasserstein (Sinkhorn) ball in each step.
This can be computationally infeasible, and current methods are mainly focused on special loss functions [Mohajerin Esfahani and Kuhn, 2018, Li et al., 2019, Wong et al., 2019].
Instead, we consider an alternating optimization procedure, in which we alternate between updating the gradients for the ATD and the Sinkhorn loss.
The full details of our algorithm, its extension to ATEs, and the alternating optimization are described in Appendix D.
Experiments
We run our partial identification algorithms on a variety of simulated settings.
We mainly focus on the synthetic data generating processes as the ground truth must be known to evaluate our derived bounds properly.
Our primary goal is to show that using uniform average treatment derivatives instead of directly optimizing the average treatment effect will result in tighter and more stable bounds.
We first run our algorithm to estimate bounds on the value of ATDs for both identifiable and non-identifiable causal graphs.
We show that, as the number of samples increases, our algorithm converges to tight bounds over the true value of ATD (Figure 3).
We then focus on partial identification of ATEs and demonstrate that using partial derivatives of the response function leads to more informative bounds, while being valid (Figures 4a, 4b).
Additionally, as a sanity check, we consider two binary-treatments datasets where the optimal bounds are known and show that our approach can reach the optimal solution. We also test our method on an ACIC dataset, a case study with real-world covariates, to illustrate the performance of our algorithm on higher-dimensional datasets.
We include these additional experiments in Appendix E.
Our implementation details, as well as hyper-parameters can be found in Appendix F.
5.1 Datasets and Baseline
Datasets.
We consider various data generating processes for different causal graphs, including a linear
SCM with three-dimensional covariates and a quadratic SCM with nonlinear interaction between [Figure] [Figure] [Figure]
a [Figure] e t ,
Y E [ Y t ] [Figure] [ Y t 0 ]
Our Bounds
GAN Bounds
Real samples [Figure] the covariates and treatment for the Back-door causal graph (Figure 1c), as well as a nonlinear
SCM for the Front-door setting (Figure 1d).
For the IV graph (Figure 1b), we use four different
SCMs, two linear and two nonlinear datasets, based on the strength of the instrument variable and the confounding.
Finally, we generate a two-dimensional linear dataset for the leaky mediation causal graph (Figure 1a).
The full details of our data-generating processes can be found in Appendix G.
Baselines.
Our first baseline is the algorithm in Hu et al.
[2021] that directly optimizes the value of
ATE using GANs.
We use their default hyper-parameters with a tolerance of 0 . 0001 .
Similar to their experimental setup, we consider 50 intermediate solutions where the distance is within the tolerance and compute the bounds using the mean and one-sided confidence intervals. We also compare our algorithm with the moment-matching method in Padh et al.
They parameterize the first two moments of the generated distribution and match them with the observed samples.
They choose response functions as linear combinations of a fixed number of basis functions.
We consider the neural basis functions for our experiments with their default hyper-parameters. In particular, we train a 3-hidden layer MLP with 64 neurons in each layer to learn the target variable given the treatment. We then use the activation of the k th neuron in the last hidden layer as the k th basis function. Instead of maximizing/minimizing E [ Y T = t ] , we consider E [ Y T =
t ] − E [ Y T = t 0 ] as our goal is to bound the ATE between two points. We leave other implementation details untouched.
5.2 Results
Bounding average treatment derivatives.
We generate data with sample sizes N = { 500 , 1000 , 2000 , 5000 } from the nonlinear Front-door and linear Back-door SCMs (identifiable), as well as linear IV with strong confounding and leaky mediation settings (non-identifiable).
We run our algorithm with ten different random seeds for each setting/sample size.
Then, we choose the five runs with the lowest tolerance parameter α n and choose the upper (lower) bound as the maximum (minimum) value of the ATD within these five runs.
Figure 3 shows our derived bounds.
As expected, the algorithm is able to point-identify the value of ATD for identifiable scenarios as the number of samples increases (Figures 3 (a) and (b)).
In non-identifiable cases, our algorithm leads to tight bounds containing the true value of ATD by increasing the number of samples as depicted in
Figures 3 (c) and (d).
This is in line with our results in Theorem 1.
Bounding average treatment effects.
Here, we aim to demonstrate the effectiveness of using partial derivatives for bounding the ATE compared to the direct optimization approach.
We consider six different settings and run our algorithm for 10 different values of treatment { t i } 10 i =1 in each setting.
We compute the value of ATE w.r.t.
a fixed point t 0 .
For each value of T , we generate N = 5 ,samples and run each experiment five times.
Then, we select the maximum (minimum) value of ATE within the five runs as the upper (lower) bound.
We follow the same procedure for the baselines.
To find the bounds on ATE using our approach, we sample from a distribution with uniform density within [ t 0 , t i ] and Gaussian tails outside of [ t 0 , t i ] , and maximize/minimize the partial derivatives.
Figure 4a shows the effectiveness of this approach in comparison to the GAN baseline.
Our algorithm produces stable and tight bounds containing the true value of ATE, while the GAN baseline, which relies on the direct optimization of ATEs, results in unstable loose bounds that may not include the true value of the treatment effect.
Figure 4b compares our algorithm with the method in Padh et al.
First, we find that their method do not contain the actual response curve in the nonlinear IV setting (left).
This is not a surprising result; since they consider a fixed set of basis functions, the set of possible response curves is more restricted than ours.
If the basis functions are not carefully designed to capture the true form of the response curve, the resulting bounds can be invalid, thus not containing the actual value of ATE even if there are regimes where the approach provides a tighter fit. On the other hand, our approach considers a larger family of response functions, enabling us to both capture shape and variation in the response curve while bounding the effects.
Conclusion, Limitations and Future Work
Our work introduces a novel method to estimate bounds on average treatment effects from observational data.
Specifically, we propose optimizing the average treatment derivative, which in turn can be used to estimate the average treatment effect in treatment response curves.
Empirically we find that the use of our method recovers known bounds on treatment effects in the discrete case and outperforms other methods based on implicit models for partial identification in the continuous case.
There remain several limitations of this work.
Our work builds on the constrained optimization problem defined by Xia et al.
[2021] instantiated in the context of the ATD.
Developing new methods for function maximization/minimization approaches under distributional constraints remains an important direction for future work.
Our work primarily uses carefully designed synthetic datasets to evaluate our method under different constraints on the data distribution.
A larger-scale evaluation of our approach on real-world benchmarks will better help us assess the method’s practicality.
Moreover, the theory of our work is restricted to the linear SCM scenario. We have also made regularity assumptions throughout the paper, including Assumption 1, that the true SCM can be modeled using implicit generative models with uniform confounding distribution, as well as the approximation of
UATD with regular treatment distributions. Consequently, practitioners must exercise caution when deploying this method when there are nonlinear or irregular structures among the random variables. Finally, we have stated our results (Theorem 1 and Corollary 1) without quantifying the finite-sample estimation uncertainty. One can use the existing theory of finite-sample convergence of empirical
Wasserstein distance in Weed and Bach [2019] to extend our results with high-probability bounds. Acknowledgments We thank Tom Ginsberg and Michael Cooper for many helpful discussions. We also thank David
Alvarez Melis for his suggestion on using Sinkhorn divergence instead of Wasserstein distance. This research was supported by NSERC Discovery Award RGPIN-2022-04546 and a CIFAR AI
Chair.
Resources used in preparing this research were provided by the Vector Institute.
Paper checklist
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We have listed the key contributions of our work in the introduction.
In the experimental section we verify our theoretical contributions and perform a comparison to the closest baseline.
Our theoretical results are in idealized linear settings though we do experiment beyond that in our work.
We highlight the limitations of our theoretical results and explain the rationale for the approximations that we make in practice.
(b) Did you describe the limitations of your work?
[Yes] We explicitly highlight the limitations of our work in the context of related work, in our exposition of the theoretical results and the computation necessary for our method.
Our theoretical results make explicit the assumptions that we make and full proofs are presented in the supplementary material.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Our paper explicitly tackles the problem of bounding treatment effects from observational data.
We have highlighted that our theory is focused on the case of linear models and provided empirical results supporting the method’s utility in the general, non-linear case.
If the method is deployed in high-risk settings such as healthcare, there may be times where the method fails to estimate bounds correctly from real-data in non-linear settings.
We have encouraged practitioners to excercise caution in this regard.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] We include the complete proofs of our results in the supplementary material.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Our supplemental material makes clear the training procedure for our method, the main sets of hyperparameters and any heuristics for optimizing the models.
We will release code to reproduce our results.
Our work made no use of participants.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] Our method finds maximum/minimum values on the quantity of interest.
Instead of reporting error bars, we consider maximum/minimum values among multiple runs with different random seeds.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Experiments were run on two a single machine with 2 Nvidia RTX A5000 GPUs, 2 RTX A6000 cards GPUs, 96 CPUs and 188GB of RAM and on a larger internal cluster comprising hundreds of GPUs.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] Our evaluation is primarily on synthetic data for which we have cited the original authors who created them.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] We use publicly available datasets that do not contain unanonymized personal data. A Proof of Theorem Here, we only focus on the minimization problem, i.e., ˆ ATD → ATD. The proof of ˆ ATD → ATD will similarly follow.
Before stating our proof, we list the assumptions needed as follows.
These are mainly technical assumptions that will simplify the derivations. In this proof we assume all random variables are one-dimensional.2. Since the true SCM M is G -constrained (Assumption 1) and linear, we will consider linear G -constrained SCMs as our parameter search space.
More concretely, each random variable V i in an SCM M θ G can be written in the following form:
[Equation]
where θ V i ∈ R | pa ( V i ) | , ˆ θ V i ∈ R | ˆ U C i | , and θ ∈ Θ is the concatenation of all θ V
i , ˆ θ V i in a topological order of V i s. Note that θ ∈ R | E | , where E is the set of edges in the causal graph G , containing edges to both observed and unobserved random variables.
We assume the set of feasible parameters Θ is a bounded closed subset of R | E | .
In practice, even in non-identifiable cases with infinite bounds, we use regularization to ensure the parameters of the network are bounded.
The induced probability over observed random variables P M θ G , which we will write as P θ , and the true distribution P belong to the Wasserstein space of order p = 1 , which we refer to as P 1 .
In other words, (cid:82) R d | v 0 − v | dP θ ( v ) <
+ ∞ for any arbitrary point v 0 ∈ R d .
Again, in practice, since the support of all observed random variables is bounded, all probability distributions defined on them belong to the Wasserstein space.
Before the main proof, we will state and prove the following useful lemmas.
Lemma 1.
Let define the set of feasible parameters for n number of samples as Θ n = { θ ∈ Θ; W 1 ( P M θ G , P n ) ≤ α n } .
Then, P M θn G → P weakly for every sequence of θ n ∈ Θ n .
Proof.
First, note that the empirical distribution P n weakly converges to P as n → ∞ .
{+Since Wmetrizes P 1 , we have W 1 ( P n , P ) → 
+}
[Villani, 2009].
Hence, we can choose the sequence of α n for the distributional constraint, such that W 1 ( P, P n ) ≤ α n and α n → 0 as n → ∞ .
For any parameter θ n ∈ Θ n , we have W 1 ( P M θn G , P n ) ≤ α n } .
Therefore, W 1 ( P M θn G , P n ) → 0 .
Using the triangle inequality, we have
[Equation]
Therefore, W 1 ( P M θn G , P ) → 0 or equivalently, P M θn G → P weakly.
Lemma 2.
For any linear G -constrained SCM M θ G , we have
[Equation]
where each element in the matrix A ( θ ) ∈ R
d ×| dim (ˆ U ) | is a continuous function of θ .
Proof.
Consider a topological order of observed variables ( V 1 , · · · , V d ) .
We prove the result by induction.
For V 1 , from eq.
14, we have
[Equation]
where the elements of ϕ V 1 matches ˆ θ V 1 for ˆ U C 1 and equal to zero for ˆ U \
ˆ U C 1 .
Now, assume all variables V 1 , · · · , V d − 1 can be written as ϕ ⊤ V
i ˆ U , where ϕ V i is a continuous function of θ .
Then, from eq. 14,
[Equation]
where ϕ V d is a linear function of θ V d , ϕ V 1 , · · · , ϕ V d − 1 , and ˆ θ V d .
Defining matrix
A ( θ ) with rows ϕ
V i concludes the proof.
Now, we are ready to prove the main result.
Theorem 2 (Tight Bounds) .
Assume the dataset { v (1) , · · · , v ( n ) } is generated from a linear SCM.
Then, under assumptions 1, and 2, the solution to the constrained optimization problem in eq.converges to the optimal bound over the ATD in infinite samples, i.e., ( ˆ ATD , ˆ ATD ) → ( ATD , ATD ) .
Proof.
The goal is to show the solution to
[Equation]
converges to the solution to
[Equation]
as n → ∞ .
We first aim to re-write the value of ATD.
Note that, in a linear G -constrained
SCM, the partial derivative ∂Y M θ G ( T = t ) ( u ) ∂t is only a function of SCM parameters θ .
Let define
[Equation]
Then,
[Equation]
Therefore, we need to show the following to conclude the proof:
[Equation]
where
[Equation]
Since Θ ∞ ⊆
Θ n for each n ∈ N , we know that
[Equation]
[Equation]
It is sufficient to show that, for each ϵ > 0 , there is n 0 such that for all n > n 0 we have
[Equation]
Suppose this is not true, i.e., there exists ϵ > 0 such that for each n ∈ N we have
[Equation]
[Equation]
for all θ ∞ ∈ Θ ∞ .
Let θ n⋆ = arg min θ ∈ Θ n g ( θ ) .
The sequence ( θ n⋆ ) n ∈ N is a subset of Θ and therefore is a bounded sequence in R | E | .
Thus, by Bolzano-Weierstrass theorem, there exists a convergent sub-sequence ( θ n i ⋆ ) i ∈ N that converges to some fixed parameter θ 0 .
Also, θ 0 ∈ Θ as Θ is closed.
Now, since g is continuous, we have
[Equation]
[Equation]
for all θ ∞ ∈ Θ ∞ .
{+Hence, θ 
+}
̸∈ Θ ∞ .
On the other hand, using Lemma 2, we have V θ ni⋆ =
A ( θ n i ⋆ )
ˆ U .
Since A ( θ n i ⋆ ) is a continuous function of θ n i ⋆ , from the continuous mapping theorem, we have P θ ni⋆ → P θ 0 weakly.
Also, from
Lemma 1, we have P θ ni⋆ → P .
Therefore, P = P θ 0 .
Since θ 0 ∈ Θ , we conclude that θ 0 ∈ Θ ∞ , a contradiction.
Remark.
In this proof, we did not separate identifiable and non-identifiable cases.
In fact, the notion of identifiability can be seen as a property of the set of feasible parameters Θ n and Θ ∞ .
For example, in identifiable cases, we expect the set Θ ∞ to only contain one element while in non-identifiable cases, it consists of multiple possible solutions.
Our proof holds as long as Θ n and Θ ∞ are bounded subsets of R | E | .
B Proof of Corollary
Similar to the proof of Theorem 1, define the set of feasible parameters as Θ n = { θ ∈ Θ W 1 ( P M θ G , P n ) ≤ α n } .
Then,
[Equation]
Therefore,
[Equation]
For the second part, similar to the proof of Theorem 1, we have ∂Y M θ G ( T = t ) ∂t =
g ( θ ) for some continuous function g .
Then,
[Equation]
The rest of the proof can be directly derived from the proof of Theorem 1.
C ATE and ATD
Here, we provide more intuition on the difference between ATE and ATD.
Consider a simple setting with only two random variables T and Y , where the causal graph is T → Y .
Here, the average treatment effect is identifiable and can be computed (with infinite samples) using the following formula:
[Equation]
Now, assume we have a finite dataset D
= { ( t (1) , y (1) ) , · · · , ( t ( n ) , y ( n ) ) } ,
where t ( i ) ̸ = t 0 and t ( i )
̸
= d for all i ∈ [ n ] .
One way to find a (high probability) upper bound on the value of ATE is to solve the following direct optimization:
[Equation]
With no assumption on the regularity of the response functions Y M θ G ( T = t ) , it is possible to find a solution to eq.
32 that matches all points in D , i.e., W 1 ( P M θ G , P n ) = 0 , while getting arbitrarily values of ATE M θ G ( d ) .
See Figure 5 for a demonstration.
This shows that ATE, in the continuous treatment setting with finite number of samples, is not well-behaved and it is generally impossible to find informative non-parametric bounds on that (see also Gunsilius [2020]).
On the other hand, ATD, which is the average partial derivative of response function w.r.t.
the observed treatment distribution is defined globally over the support of T .
Therefore, it is not possible to maximize ATD arbitrarily without violating the distributional constraint in this setting.
−−  5 −
Treatment, T
T a r g e t , Y true response upper bound data
D Algorithm
Input : Dataset D =
{ X ( i ) , T ( i ) , Y ( i ) }
Ni =1 , causal graph G , Sinkhorn approximation parameter ϵ , learning rate γ , Lagrange learning rate γ L
Output : ˆ ATD: lower bounds on ATD ▷
The algorithm to find ˆ ATD is similar
Phase 1: Matching Distributions
Generate dataset
ATD M θ ( i + j ) G ←
11: # Alternate optimization 12: ¯ θ
( i j ) ← θ ( i j ) − γ ∇ θ ATD
M θ ( j ) G ▷ Update the parameters to minimize the ATD 13:
D ← { k k
Y k
} Nk =1 ∼ M θ i j )
G 14: θ j +1)
← ¯ θ ( i + j ) − γ ∇ θ S ϵ
( ˆ D , D )
▷ Update the parameters to minimize the Sinkhorn loss 15: # Lagrange multiplier update 16: λ ( j +1) ← λ ( j ) + γ L ( S ϵ ( ˆ D , D ) − α ) 17: j ← j +18: end while 19: return ATD M θ
( i + j ) G
Extension of Algorithm 1 to ATEs.
We can use a similar method to Algorithm 1 for partial identification of average treatment effects (ATEs).
The only difference is that, instead of maximizing/minimizing ATD M θ G , we optimize for the value of UATD M θ G [ t 0 , d ] .
More concretely, we use the following approximation to estimate UATD M θ G [ t 0 , d ] :
[Equation]
[Table]
where { t ( i ) } ni =1 are samples from a uniform distribution within [ t 0 , d ] with a Gaussian tail, and { u ( i ) } ni
=1 are the latent variables generated from a uniform distribution.
Note that, the only difference between eq.
33 and eq. 13 is the distribution of treatment variables, where in the former we use a uniform distribution with Gaussian tail, while the latter uses the same distribution of treatments in the observed dataset.
In our experiments, for the uniform distribution with Gaussian tail, we generate samples from N ( µ, σ ) , where µ = t 0 + dand σ = d − t 0 2 .
Then, for each sample within [ t 0 , d ] , we generate a new sample from uniform distribution Unif [ t 0 , d ] .
E Additional Experiments E.1 Discrete Setting
To showcase the generality of our framework, we study two datasets with binary treatments.
We consider the binary IV dataset described in Duarte et al.
[2021], where the true value of ATE is not identifiable, but the optimal bound is known.
We also use the Front-door binary dataset in Zhang et al.
[2021] where the causal effect is identifiable (See Appendix G).
Here, the partial derivatives do not exist, so we directly optimize the ATE.
Note that, in the discrete setting, the network cannot generate arbitrary large values in the intervention points without violating the distributional constraint.
Tableshows our derived bounds and compares them to the optimal bounds.
In the identifiable Front-door causal graph, we find a tight bound over the true ATE.
In the non-identifiable IV setting, our bound includes the optimal bound with a small gap.
E.2 ACIC Dataset
To demonstrate the applicability of our method on datasets with higher-dimensional covariates, we consider the Atlantic Causal Inference Conference (ACIC) 2019 Data Challenge [Gruber et al., 2019].
The dataset is constructed based on the spam detection data from UCI [Dua and Graff, 2017].
The outcome of interest Y is whether an email is marked as spam or not.
The treatment variable T is also a binary variable showing if the number of capital letters in an email exceeds a threshold.
There are 22 continuous covariates that correspond to certain word frequencies.
We follow a similar setup as in Guo et al.
In particular, we consider the problem of partial identification under noisy measurements.
Here, the data-generating causal graph is the Back-door setting, and the causal effect of T on Y is identifiable.
However, we assume measurement error on the covariates by imposing synthetic noise on them and aim to estimate bounds on ATE under this uncertainty.
Since our algorithm can incorporate uncertainty through the distributional constraint, it will lead to a valid and informative bound by choosing an appropriate value for α n .
We generate a dataset of 2 , 000 samples using ACIC’s data-generating process.
Similar to Guo et al.
[2022], we synthetically add five different levels of Gaussian noise with mean ∈ (0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5) and standard deviation ∈ (0 . 5 , 0 .
We then run our algorithm on these five noisy datasets and the original noiseless one, 10 different trials for each.
Figure 6 illustrates our derived bounds and the true ATE for each of the noise levels.
The results always include the actual value of ATE, while not being too conservative.
As the noise level increases, our derived bounds get naturally less informative.
F Implementation Choices
We use similar neural network architectures for each variable in the causal graph.
For hyper-parameter tuning, we search over networks with { 2 , 3 , 5 } hidden layers, { 16 , 64 , 256 } neurons in each layer,
[Figure]
learning rate within { 0 . 001 , 0 . 005 , 0 . 01 } , and Lagrange multiplier learning rate { 0 . 5 , 1 } .
The hyperparameters with the lowest Sinkhorn loss are chosen for each causal graph.
We run all experiments for 500 epochs and use all samples ( N = 5000 ) in each iteration.
To find the level of distribution constraint α n , we minimize the Sinkhorn loss and check the loss value on a validation set until it does not decrease for 30 epochs.
We then choose the minimum value of Sinkhorn loss on the training data as α n .
We use an alternate optimization approach to maximize/minimize the ADE (or ATE).
In each step, we first minimize the Sinkhorn loss and then maximize/minimize the value of ADE.
We use gradient clipping for the ADE optimizer with the value of 0 . 2 .
The learning rate of the ADE optimizer is half the learning rate for the Sinkhorn loss optimizer.
To find the upper (lower) bounds on ATD/ATE, we find their maximum (minimum) value within all steps that satisfy the distributional constraint.
All implementation is done in PyTorch Lightning using Adam optimizer.
To evaluate and calculate gradients of Sinkhorn loss, we use the "geomloss" library [Feydy et al., 2019].
G Data Generating Processes
Linear Back-door.
[Equation]
Nonlinear Back-door.
[Equation]
Front-door.
[Equation]
Linear IV (weak instrument, strong confounding).
[Equation]
[Equation]
Nonlinear IV (strong instrument, weak confounding).
[Equation]
Leaky Mediation.
[Equation]
Binary IV [Duarte et al., 2021].
We use the noncompliance IV dataset in Section D.1 from Duarte et al.
The true value of ATE is − 0 .
25 while the optimal bound is [ − 0 .
55 , − 0 .
Binary Front-door [Zhang et al., 2021].
[Equation]
Linear IV with strong confounding [Padh et al., 2022].
[Equation]
Nonlinear IV with nonlinear interaction between treatment and confounding [Padh et al., 2022].
[Equation]