A N EURAL M EAN E MBEDDING A PPROACH FOR B ACKDOOR AND F RONT - DOOR A DJUSTMENT
Arthur Gretton
Gatsby arthur.gretton@gmail.com
A BSTRACT
We consider the estimation of average and counterfactual treatment effects, under two settings: back-door adjustment and front-door adjustment .
The goal in both cases is to recover the treatment effect without having an access to a hidden confounder.
This objective is attained by first estimating the conditional mean of the desired outcome variable given relevant covariates (the “first stage” regression), and then taking the (conditional) expectation of this function as a “second stage” procedure.
We propose to compute these conditional expectations directly using a regression function to the learned input features of the first stage, thus avoiding the need for sampling or density estimation.
All functions and features (and in particular, the output features in the second stage) are neural networks learned adaptively from data, with the sole requirement that the final layer of the first stage should be linear.
The proposed method is shown to converge to the true causal parameter, and outperforms the recent state-of-the-art methods on challenging causal benchmarks, including settings involving high-dimensional image data.
I NTRODUCTION
The goal of causal inference from observational data is to predict the effect of our actions, or treatments , on the outcome without performing interventions.
Questions of interest can include what is the effect of smoking on life expectancy?
or counterfactual questions, such as given the observed health outcome for a smoker, how long would they have lived had they quit smoking?
Answering these questions becomes challenging when a confounder exists, which affects both treatment and the outcome, and causes bias in the estimation.
Causal estimation requires us to correct for this confounding bias.
A popular assumption in causal inference is the no unmeasured confounder requirement, which means that we observe all the confounders that cause the bias in the estimation.
Although a number of causal inference methods are proposed under this assumption (Hill, 2011; Shalit et al., 2017; Shi et al., 2019; Schwab et al., 2020), it rarely holds in practice.
In the smoking example, the confounder can be one’s genetic characteristics or social status, which are difficult to measure for both technical and ethical reasons.
To address this issue, Pearl (1995) proposed back-door adjustment and front-door adjustment , which recover the causal effect in the presence of hidden confounders using a back-door variable or frontdoor variable , respectively.
The back-door variable is a covariate that blocks all causal effects directed from the confounder to the treatment.
In health care, patients may have underlying predispositions to illness due to genetic or social factors (hidden), which cause measurable symptoms.
The symptoms can be used as the back-door variable if the treatment is chosen based on these.
By contrast, a front-door variable blocks the path from treatment to outcome.
In perhaps the bestknown example, the amount of tar in a smoker’s lungs serves as a front-door variable, since it is increased by smoking, shortens life expectancy, and has no direct link to underlying (hidden) sociological traits.
Pearl (1995) showed that causal quantities can be obtained by taking the (conditional) expectation of the conditional average outcome.
While Pearl (1995) only considered the discrete case, this framework was extended to the continuous case by Singh et al.
(2020), using two-stage regression (a review of this and other recent approaches for the continuous case is given in Section 5).
In the first stage, the approach regresses from the relevant covariates to the outcome of interest, expressing the function as a linear combination of non-linear feature maps.
Then, in the second stage, the causal parameters are estimated by learning the (conditional) expectation of the non-linear feature map used in the first stage.
Unlike competing methods (Colangelo & Lee, 2020; Kennedy et al., 2017), two-stage regression avoids fitting probability densities, which is challenging in high-dimensional settings (Wasserman, 2006, Section 6.5).
Singh et al.
(2020)’s method is shown to converge to the true causal parameters and exhibits better empirical performance than competing methods.
One limitation of the methods in Singh et al.
(2020) is that they use fixed pre-specified feature maps from reproducing kernel Hilbert spaces, which have a limited expressive capacity when data are complex (images, text, audio).
To overcome this, we propose to employ a neural mean embedding approach to learning task-specific adaptive feature dictionaries.
At a high level, we first employ a neural network with a linear final layer in the first stage.
For the second stage, we learn the (conditional) mean of the stage 1 features in the penultimate layer, again with a neural net.
The approach develops the technique of Xu et al.
(2021a;b) and enables the model to capture complex causal relationships for high-dimensional covariates and treatments.
Neural network feature means are also used to represent (conditional) probabilities in other machine learning settings, such as representation learning (Zaheer et al., 2017) and approximate Bayesian inference (Xu et al., 2022).
We derive the consistency of the method based on the Rademacher complexity, a result of which is of independent interest and may be relevant in establishing consistency for broader categories of neural mean embedding approaches, including Xu et al.
(2021a;b).
We empirically show that the proposed method performs better than other state-of-the-art neural causal inference methods, including those using kernel feature dictionaries.
This paper is structured as follows.
In Section 2, we introduce the causal parameters we are interested in and give a detailed description of the proposed method in Section 3.
The theoretical analysis is presented in Section 4, followed by a review of related work in Section 5.
We demonstrate the empirical performance of the proposed method in Section 6, covering two settings: a classical backdoor adjustment problem with a binary treatment, and a challenging back-door and front-door setting where the treatment consists of high-dimensional image data.
P ROBLEM S ETTING
In this section, we introduce the causal parameters and methods to estimate these causal methods, namely a back-door adjustment and front-door adjustment .
Throughout the paper, we denote a random variable in a capital letter (e.g. A ), the realization of this random variable in lowercase (e.g. a ), and the set where a random variable takes values in a calligraphic letter (e.g. A ).
We assume data is generated from a distribution P .
Causal Parameters We introduce the target causal parameters using the potential outcome framework (Rubin, 2005).
Let the treatment and the observed outcome be A ∈ A and Y ∈ Y ⊆
[ − R, R ] .
We denote the potential outcome given treatment a as Y ( a ) ∈ Y .
Here, we assume no inference , which means that we observe Y
= Y ( a ) when A = a .
We denote the hidden confounder as U ∈ U and assume conditional exchangeability ∀ a ∈ A , Y ( a ) ⊥⊥ A | U , which means that the potential outcomes are not affected by the treatment assignment.
A typical causal graph is shown in Figure 1a.
We may additionally consider the observable confounder O ∈ O , which is discussed in Appendix C.
A first goal of causal inference is to estimate the Average Treatment Effect (ATE) 1 θ ATE ( a ) =
E (cid:2) Y ( a ) (cid:3) , which is the average potential outcome of A = a .
We also consider Average Treatment
Effect on the Treated (ATT) θ ATT ( a ; a ′ ) =
E (cid:2) Y ( a ) |
A = a ′ (cid:3) , which is the expected potential outcome of A = a for those who received the treatment A = a ′ .
Given no inference and conditional exchangeability assumptions, these causal parameters can be written in the following form.
Proposition 1 (Rosenbaum & Rubin, 1983; Robins, 1986) .
Given unobserved confounder U , which satisfies no inference and conditional exchangeability, we have
[Equation]
If we observable additional confounder O , we may also consider conditional average treatment effect (CATE) : the average potential outcome for the sub-population of O = o , which is discussed
[Figure]
U
[Figure]
U
[Figure]
in Appendix C. Note that since the confounder U is not observed, we cannot recover these causal parameters only from ( A, Y
Back-door Adjustment In back-door adjustment, we assume the access to the back-door variable
X ∈ X , which blocks all causal paths from unobserved confounder U to treatment A .
See Figure 1b for a typical causal graph.
Given the back-door variable, causal parameters can be written only from observable variables ( A, Y, X ) as follows.
Proposition 2 (Pearl, 1995, Theorem 1) .
Given the back-door variable X , we have
[Equation]
where g ( a, x ) =
E [ Y |
A
= a, X = x ] .
By comparing Proposition 2 to Proposition 1, we can see that causal parameters can be learned by treating the back-door variable X as the only “confounder”, despite the presence of the additional hidden confounder U .
Hence, we may apply any method based on the “no unobservable confounder” assumption to back-door adjustment.
Front-door Adjustment Another adjustment for causal estimation is front-door adjustment , which uses the causal mechanism to determine the causal effect.
Assume we observe the front-door variable M ∈ M , which blocks all causal paths from treatment A to outcome Y , as in Figure 1c.
Then, we can recover the causal parameters as follows.
Proposition 3 (Pearl, 1995, Theorem 2) .
Given the front-door variable M , we have
[Equation]
where g ( a, m )
=
E [ Y |
A
= a, M
= m ] and A ′ ∈
A is a random variable that follows the same distribution as treatment A .
Unlike the case of the back-door adjustment, we cannot naively apply methods based on the “no unmeasured confounder” assumption here, since Proposition 3 takes a different form to Proposition 1.
A LGORITHMS
In this section, we present our proposed methods.
We first present the case with back-door adjustment and then move to front-door adjustment.
The algorithm is summarized in Appendix A.
Back-door adjustment The algorithm consists of two stages;
In the first stage, we learn the conditional expectation g = E
[ Y |
A = a, X = x ] with a specific form.
We then compute the causal parameter by estimating the expectation of the input features to g .
The conditional expectation g ( a, x ) is learned by regressing ( A, X ) to Y .
Here, we consider a specific model g ( a, x ) =
w ⊤ ( ϕ A ( a ) ⊗ ϕ X ( x )) , where
ϕ A : A → R d 1 , ϕ X : X → R d 2 are feature maps represented by neural networks, w ∈ R d 1 d 2 is a trainable weight vector, and ⊗ denotes a tensor product a ⊗ b = vec( ab ⊤ ) .
This tensor form used for g ( a, x ) [-} ni =1 ∼ P size of n , the feature maps 
 In-]explicitly separates out the treatment of the features of X and of A ; in the event that X is much higher dimension than A , then concatenating both as a single input tends to downplay the information in A .
In addition, we can take advantage of linearity and focus on estimating the relevant (conditional) expectation as discussed later. Given data { ( a i , y i , x i ) } ni =1 ∼ P size of n , the feature maps ϕ
A , ϕ X and the weight w can be trained by minimizing the following empirical loss:
[Equation]
We may add any regularization term to this loss, such as weight decay λ ∥ w ∥ 2 .
Let the minimizer of the loss ˆ L X 1 be ˆ w , ˆ ϕ A , ˆ ϕ X =
arg min
ˆ L X 1 and the learned regression function be ˆ g ( a, x ) =
ˆ w ⊤ ( ˆ ϕ A ( a ) ⊗ ˆ ϕ X ( x )) .
Then, by substituting ˆ g for g in Proposition 2, we have
[Equation]
This is the advantage of assuming the specific form of g ( a, x ) =
w ⊤ ( ϕ A ( a ) ⊗ ϕ X ( x )) ; From linearity, we can recover the causal parameters by estimating E [ ˆ ϕ X ( X )] , E [ ˆ ϕ X ( X ) |
A = a ′ ] .
Such (conditional) expectations of features are called (conditional) mean embedding , and thus, we name our method “ neural (conditional) mean embedding ”.
We can estimate the marginal expectation E [ ˆ ϕ X ( X )] , as a simple empirical average E [ ˆ ϕ X ( X )] ≃ n (cid:80) ni
=1 ˆ ϕ X
( x i ) .
The conditional mean embedding E [ ˆ ϕ X ( X ) |
A = a ′ ] requires more care, however: it can be learned by a technique proposed in Xu et al.
(2021a), in which we train another regression function from treatment A to the back-door feature
ˆ ϕ X
( X ) .
Specifically, we estimate E [ ˆ ϕ X ( X ) |
A = a ′ ] by ˆ f ˆ ϕ X ( a ′ ) , where the regression function ˆ f ˆ ϕ X : A → R d 2 be given by
[Equation]
Here, ∥ · ∥ denotes the Euclidean norm.
The loss ˆ L X 2 may include the additional regularization term such as a weight decay term for parameters in f .
We have
[Equation]
as the final estimator for the back-door adjustment.
The estimator for the ATE ˆ θ ATE is reduced to the average of the predictions ˆ θ ATE = 1 n (cid:80) ni
=1 ˆ g ( a, x i ) .
This coincides with other neural network causal methods (Shalit et al., 2017; Chernozhukov et al., 2022b), which do not assume g ( a, z ) =
w ⊤ ( ϕ A ( a ) ⊗ ϕ X ( x )) .
As we have seen, however, this tensor product formulation is essential for estimating ATT by back-door adjustment.
It will also be necessary for the front-door adjustment, as we will see next.
We can obtain the estimator for front-door adjustment by following the almost same procedure as the back-door adjustment.
Given data { ( a i , y i , m i ) } ni =1 , we again fit the regression model ˆ g ( a, m ) =
ˆ w ⊤ (cid:16) ˆ ϕ A ( a ) ⊗ ˆ ϕ M ( m ) (cid:17) by minimizing
[Equation]
where ϕ M : M → R d 2 is a feature map represented as the neural network.
From Proposition 3, for f ˆ ϕ M ( a ) = E (cid:104) ˆ ϕ M ( M ) (cid:12)(cid:12)(cid:12) A = a (cid:105) , we have θ ATE ( a ) ≃ ˆ w ⊤ (cid:16) E (cid:104)
ˆ ϕ A ( A ) (cid:105) ⊗ f ˆ ϕ M ( a ) (cid:17) and θ ATT ( a ; a ′ ) ≃ ˆ w ⊤ (cid:16) ˆ ϕ A ( a ′ ) ⊗ f ˆ ϕ M ( a ) (cid:17) .
Again, we estimate feature embedding by empirical average for E [ ˆ ϕ A ( A )] or solving another regression problem for µ ˆ ϕ M ( a ) .
The final estimator for front-door adjustment is given as
[Equation]
where ˆ f ˆ ϕ M is given by minimizing loss ˆ L M= 1 n (cid:80) ni
=1 ∥ ϕ M ( m i ) − f ( a i ) ∥ 2 (with additional regularization term).
T HEORETICAL A NALYSIS
In this section, we prove the consistency of the proposed method.
We focus on the back-door adjustment case, since the consistency of front-door adjustment can be derived identically.
The proposed method consists of two successive regression problems.
In the first stage, we learn the conditional expectation g , and then in the second stage, we estimate the feature embeddings.
First, we show each stage’s consistency, then present the overall convergence rate to the causal parameter.
Consistency for the first stage: In this section, we consider the hypothesis space of g as
[Equation]
Here, we denote ℓ 1 -norm and infinity norm of vector b ∈ R d as ∥ b ∥ 1 = (cid:80) di
=1 | b i | and ∥ b ∥ ∞ = max i ∈
[ d ] b i .
Note that from inequality ∥
ϕ A ( a ) ⊗ ϕ X ( x ) ∥ ∞ ≤ ∥ ϕ A ( a ) ∥ ∞
∥ ϕ X ( x ) ∥ ∞ and H ¨ older’s inequality, we can show that h ( a, x ) ∈ [ − R, R ] for all h ∈ H g .
First, we discuss the richness of this hypothesis space by the following theorem.
Theorem 1.
Let A , X ⊂ R d be compact.
Given sufficiently large R, d 1 , d 2 , for any continuous function f : A × X → R and constant ε > 0 , there exists h ∈ H g which satisfies sup a,x | f
( a, x ) − h ( a, x ) | ≤ ε.
The proof uses the modified version of universal approximation theorem (Cybenko, 1989) for neural net, which will be given in Appendix B.1.
Theorem 1 tells that we can approximate any continuous function f with an arbitrary accuracy, which suggests the richness of our function class.
Given this hypothesis space, the following lemma bounds the deviation of estimated conditional expectation ˆ g and the true one.
Lemma 1.
Given data S = { a i , y i , x i } ni =1 , let minimizer of loss ˆ L X 1 be ˆ g = arg min
ˆ L X 1 .
If the true conditional expectation g is in the hypothesis space g ∈ H g , w.p. at least 1 − 2 δ , we have
[Equation]
where ˆ R S ( H g ) is empirical Rademacher complexity of H g given data S .
The proof is given in Appendix B.3.
Here, we present the empirical Rademacher complexity when we apply a feed-forward neural network for features.
Lemma 2.
The empirical Rademacher complexity ˆ R S
( H g ) scales as ˆ R S
( H g ) ≤
O ( C L / √ n ) for some constant C if we use a specific L -layer neural net for features ϕ A , ϕ X .
See Lemma 7 in Appendix B.3 for the detailed expression of the upper bound.
Note that this may be of independent interest since the similar hypothesis class is considered in Xu et al.
(2021a;b), and no explicit upper bound is provided on the empirical Rademacher complexity in that work.
Consistency for the second stage: Next, we consider the second stage of regression.
In back-door adjustment, we estimate the feature embedding E [ ˆ ϕ X ( X )] and the conditional feature embedding
E [ ˆ ϕ X ( X ) |
A = a ′ ] .
We first state the consistency of the estimation of marginal expectation, which can be shown by Hoeffding’s inequality.
Lemma 3.
Given data { x i } ni =1 and feature map ˆ ϕ X , w.p. at least 1 − δ , we have
[Equation]
n
[Equation]
For conditional feature embedding E [ ˆ ϕ X ( X ) | A = a ′ ] , we solve the regression problem ˆ f ˆ ϕ X
=
arg min f ˆ L X 2 ( f ; ˆ ϕ X ) , the consistency of which is stated as follows.
Lemma 4.
Let hypothesis space H
f be
[Equation]
where H f is some hypothesis space of functions of f : X → [ − 1 , 1] .
Let the true function be f ˆ ϕ X ( a ) =
E [ ˆ ϕ X
( X ) |
A = a ] , and we assume f ˆ ϕ X ∈ H f .
Let ˆ f ˆ ϕ X = arg min
f ∈H f ˆ L X 2 ( f ; ˆ ϕ X )
, given data S = { ( a i , x i ) } .
Then, we have
[Equation]
w.p. at least 1 − 2 δ , where ∥ f ( A ) ∥
P ( A ) , ∞ = max i ∥
f i ∥ P ( A ) and ˆ R S
( H f ) is the empirical
Rademacher complexity of H f given data S .
The proof is identical to Lemma 1.
We use neural network hypothesis class for H f
whose empirical
Rademacher complexity is bounded by O (1 / √ n ) as discussed in Proposition 5 in Appendix B.3.
Consistency of the causal estimator Finally, we show that if these two estimators converge uniformly, we can recover the true causal parameters.
To derive the consistency of the causal parameter, we put the following assumption on hypothesis spaces in order to guarantee that convergence in ℓ 2norm leads to uniform convergence.
Assumption 1.
For functions h 1 , h 2 ∈ H g , there exists constant c > 0 and β that
[Equation]
Intuitively, this ensures that we have a non-zero probability of observing all elements in A × X .
We can see that Assumption 1 is satisfied with β = 1 and c
= min ( a,x ) ∈A×X
P ( A = a, X = x ) when treatment and back-door variables are discrete.
A similar intuition holds for the continuous case; in
Appendix B.2, we show that Assumption 1 holds when with β = 2 d + when A , X are d -dimensional intervals if the density function of P ( A, X ) is bounded away from zero and all functions in H g are
Lipschitz continuous.
Theorem 2.
Under conditions in Lemmas 1 to 3 and Assumption 1, w.p. at least 1 − 4 δ , we have
[Equation]
If we furthermore assume that for all f , ˜ f ,
[Equation]
then, w.p. at least 1 − 4 δ , we have sup a,a ′ ∈A |
[Equation]
The proof is given in Appendix B.3.
This rate is slow compared to the existing work (Singh et al., 2020), which can be as fast as O ( n − 1 / 4 ) .
However, Singh et al.
(2020) assumes that the correct regression function g is in a certain reproducing kernel Hilbert space (RKHS) , which is a stronger assumption than ours, which only assumes a Lipschitz hypothesis space.
Deriving the matching minimax rates under the Lipschitz assumption remains a topic for future work.
R ELATED W ORK
Meanwhile learning approaches to the back-door adjustment have been extensively explored in recent work, including tree models (Hill, 2011; Athey et al., 2019), kernel models (Singh et al., 2020) and neural networks (Shi et al., 2019; Chernozhukov et al., 2022b; Shalit et al., 2017), most literature considers binary treatment cases, and few methods can be applied to continuous treatments.
Schwab et al.
(2020) proposed to discretize the continuous treatments and Kennedy et al.
Colangelo & Lee (2020) conducted density estimation of P ( X ) and P ( X | A ) .
These are simple to implement but suffer from the curse of dimensionality (Wasserman, 2006, Section 6.5).
Recently, the automatic debiased machine learner (Auto-DML) approach (Chernozhukov et al., 2022a) has gained increasing attention, and can handle continuous treatments in the back-door adjustment.
Consider a functional m that maps g to causal parameter θ = E [ m ( g, ( A, X ))] .
For the ATE case, we have m ( g, ( A, X ))
=
g ( a, X ) since θ ATE ( a ) =
E [ g ( a, X )] .
We may estimate both g and the Riesz representer α that satisfies E [ m ( g, ( A, X ))]
=
E [ α ( A, X ) g
( A, X )] by the least-square regression to get the causal estimator.
Although Auto-DML can learn a complex causal relationship with neural network model (Chernozhukov et al., 2022b), it requires a considerable amount of computation when the treatment is continuous, since we have to learn a different Riesz representer α for each treatment a .
Furthermore, as discussed in Appendix B.4, the error bound on α can grow exponentially with respect to the dimension of the probability space, which may harm performance in high-dimensional settings.
Singh et al.
(2020) proposed a feature embedding approach, in which feature maps are specified as the fixed feature maps in a reproducing kernel Hilbert space (RKHS).
Although this strategy can be applied to a number of different causal parameters, the flexibility of the model is limited since it uses pre-specified features.
Our main contribution is to generalize this feature embedding approach to adaptive features which enables us to capture more complex causal relationships.
Similar techniques are used in the additional causal inference settings, such as deep feature instrumental variable method (Xu et al., 2021a) or deep proxy causal learning (Xu et al., 2021b).
By contrast with the back-door case, there is little literature that discusses non-linear front-door adjustment.
The idea was originally introduced for the discrete treatment setting (Pearl, 1995) and was later discussed using the linear causal model (Pearl, 2009).
To the best of our knowledge, Singh et al.
(2020) is the only work that considers the nonlinear front-door adjustment, where fixed kernel feature dictionaries are used.
We generalize this approach using adaptive neural feature dictionaries and obtain promising performance.
E XPERIMENTS
In this section, we evaluate the performance of the proposed method based on two scenarios.
One considers the back-door adjustment methods with binary treatment based on IHDP dataset (Gross, 1993) and ACIC dataset
(Shimoni et al., 2018).
Another tests the performance on a high-dimensional treatment based on dSprite image dataset (Matthey et al., 2017).
We first describe the training procedure we apply for our proposed method, and then report the results of each benchmark.
The details of hyperparameters used in the experiment are summarized in Appendix D.
6. T RAINING P ROCEDURE
During the training, we use the learning procedure proposed by Xu et al.
(2021a).
Let us consider the first stage regression in a back-door adjustment, in which we consider the following loss ˆ L Xwith weight decay regularization
[Equation]
To minimize ˆ L X 1 with respect to ( w , ϕ A , ϕ X ) , we can use the closed form solution of weight w .
If we fix features ϕ A , ϕ X , the minimizer of w can be written
[Equation]
where ϕ A,X ( a, x ) =
ϕ
A ( a ) ⊗ ϕ Z ( a ) .
Then, we optimize the features as ˆ ϕ A , ˆ ϕ X = arg min ϕ A , ϕ X ˆ L X 1 ( ˆ w ( ϕ A , ϕ X ) , ϕ A , ϕ X ) using Adam (Kingma & Ba, 2015).
We empirically found that this stabilizes the learning and improves the performance of the proposed method.
6. B INARY T REATMENT S CENARIO
In this section, we report the performans on two classical causal datasets: IHDP dataset and ACIC dataset.
The IHDP dataset is widely used to evaluate the performance of the estimators for the
ATE (Shi et al., 2019; Chernozhukov et al., 2022b; Athey et al., 2019).
This is a semi-synthetic dataset based on the Infant Health and Development Program (IHDP) (Gross, 1993).
Following existing work, we generate 1000 sets of 747 observations of outcomes and binary treatments based [Table] on the 25-dimensional observable confounder in the original data.
The ACIC dataset is introduced in (Shi et al., 2019), which is based on linked birth and infant death data (LBIDD) (Mathews &
MacDorman, 2006).
This is considered a more challenging benchmark dataset than IHDP since it contains data points with extreme propensity scores (i.e. P ( A = 1 | X ) can be very close to 0 or 1).
We select 101 datasets following Shi et al.
and remove outliers in each dataset using the procedure described in Appendix D.
We compare our method to competing causal methods, DragonNet (Shi et al., 2019), RieszNet (Chernozhukov et al., 2022b), and RKHS Embedding (Singh et al., 2020).
DragonNet is a neural causal inference method specially designed for the binary treatment, which applies the targeted regularization (van der Laan & Rubin, 2006) to ATE estimation.
RieszNet implements Auto-DML with a neural network, which learns the conditional expectation g and Riesz representer α jointly while sharing the intermediate features.
Given estimated ˆ g, ˆ α , it proposes three ways to calculate the causal parameter;
[Equation]
ˆ α ( A,X ) ( Y − ˆ g ( A,X ))] , where functional m maps g to the causal parameter (See Section 5 for the example of functional m ).
We report the performance of each estimator in RieszNet.
RKHS Embedding employs the feature embedding approach with a fixed kernel feature dictionaries.
The results are summarized in Table 1.
Although RieszNet(IPW) estimator performs promisingly in
IHDP, the performance degenerates for the ACIC dataset, which suggests RieszNet(IPW) is prone to extreme propensity scores.
This is not surprising, since the true Riesz representer in this case is α ( A, X ) =
A P ( A =1 | X ) − 1 − A P ( A =0 | X ) , which can be very large if P ( A = 1 | X ) becomes close to 0 or 1.
This also harms the performance of RieszNet(DR).
We can see that the proposed method outperforms all competing methods besides RieszNet(DR) in the IHDP dataset, for which the performance is comparable (0.117 ± 0.002 v.s. 0.110 ± 0.003).
6. H IGH -D IMENSIONAL T REATMENT S CENARIO
To test the performance of our method of causal inference in a more complex setting, we used dSprite data (Matthey et al., 2017), which is also used as the benchmark for other high-dimensional causal inference methods (Xu et al., 2021a;b).
The dSprite dataset consists of images that are 64 × 64 = 4096-dimensional, described by five latent parameters ( shape , scale , rotation , posX and posY ).
Throughout this paper, we fix ( shape , scale , rotation ) and use posX ∈ [0 , 1] and posY ∈ [0 , 1] as the latent parameters.
Based on this dataset, we propose two experiments; one is ATE estimation based on the back-door adjustment, and the other is ATT estimation based on front-door adjustment.
Back-door Adjustment
In our back-door adjustment experiment, we consider the case where the image is the treatment.
Let us sample hidden confounder U ∼ Unif(0 , 1) , and consider the backdoor as ( X 1 , X 2 ) =
( U cos θ + ε 1 , U sin θ + ε 2 ) where ε 1 , ε 2 ∼ N (0 , 0 . 09) , θ ∼ Unif(0 , 2 π ) .
We define treatment A as the image, where the parameters are set as posX =
X 1 +1 .
5, posY
= X 2 +1 .
We add Gaussian noise N (0 , 0 . 01) to each pixel of images.
The outcome is given as follows,
[Equation]
where A [ ij ] denotes the value of the pixel at ( i, j ) and ε Y is the noise variable sampled from ε Y ∼
N (0 , 0 . 25) .
Each dataset consists of 5000 samples of ( Y, A, X 1 , X 2 ) and we consider the problem of estimating θ ATE ( a )
=
h 2 ( A ) / 100 .
We compare the proposed method to RieszNet and RKHS
Embedding, since DragonNet is designed for binary treatments and is not applicable here.
We
[Figure] [Figure]
generate 10 datasets and the average of squared error ( θ ATE ( a ) − ˆ θ ATE ( a )) 2 at 9 test points a is reported in Figure 2.
We can see that the proposed method performs best in the setting, which shows the power of the method for complex high-dimensional inputs.
The RKHS Embedding method suffers from the limited flexibility of the model for the case of complex high-dimensional treatment, and performs worse than all neural methods besides RieszNet(IPW).
This suggests that it is difficult to estimate
Riesz representer α in a high-dimensional scenario, which is also suggested by the exponential growth of the error bound to the dimension as discussed in Appendix B.4.
We conjecture this also harms the performance of RieszNet(Direct) and RieszNet(DR), since the models for conditional expectation ˆ g and Riesz representer ˆ α share the intermediate features in the network and are jointly trained in RieszNet.
Frontdoor Adjustment We use dSprite dataset to consider front-door adjustment.
Again, we sample hidden confounder U 1 , U 2 ∼ Unif( − 1 . 5 , 1 . 5) , and we set the image to be the treatment, where the parameters are set as posX
=
U 1 +1 .
posY =
U 2 +1 .
We add Gaussian noise N (0 , 0 . 01) to each pixel of the images.
We use M = h ( A )
+ ε M as the front-door variable M , where ε M ∼
N (0 , 0 . 04) .
The outcome is given as follows,
[Equation]
We consider the problem of estimating θ ATT ( a ; a ′ ) and obtain the average squared error onpoints of a while fixing a ′ to the image of posX = 0 .
6 , posY = 0 .
We compare against RKHS
Embedding, where the result is given in Figure 3.
Note that RieszNet has not been developed for this setting.
Again, the RKHS Embedding method suffers from the limited flexibility of the model, whereas our proposed model successfully captures the complex causal relationships.
C ONCLUSION
We have proposed a novel method for back-door and front-door adjustment, based on the neural mean embedding.
We established consistency of the proposed method based on a Rademacher complexity argument, which contains a new analysis of the hypothesis space with the tensor product features.
Our empirical evaluation shows that the proposed method outperforms existing estimators, especially when high-dimensional image observations are involved.
As future work, it would be promising to apply a similar adaptive feature embedding approach to other causal parameters, such as marginal average effect ∇ a θ ATE ( a ) (Imbens & Newey, 2009).
Furthermore, it would be interesting to consider sequential treatments, as in dynamic treatment effect estimation, in which the treatment may depend on the past covariates, treatments and outcomes.
Recently, a kernel feature embedding approach (Singh et al., 2021) has been developed to estimate the dynamic treatment effect, and we expect that applying the neural mean embedding would benefit the performance. A CKNOWLEDGEMENT
This work was supported by the Gatsby Charitable Foundation.
B T ECHNICAL D ETAILS B. U NIVERSAL A PPROXIMATION T HEORY
In this section, we provide the proof of Theorem 1.
Recall our hypothesis space is
[Equation]
Let consider the feature be
[Equation]
where σ is the sigmoid function and s i , t i ∈ R D , α i , β i ∈ R are parameters.
By considering the case of d 1 = d 2 and setting “non-diagonal” elements of w to zero, we can see that
[Equation]
is a member of of H g .
Next, we present the following lemma
[Equation]
Lemma 5.
Let µ be a finite, signed regular Borel measures on A × X .
If σ satisfies the followings:
[Equation]
then, given any continuous function f : A × A → R and ε > 0 , there is a finite sum g
[Equation]
n
[Equation]
which satisfies
[Equation]
The proof is identical to Theorem 1 in Cybenko (1989).
Now, all we have to prove is that the
Sigmoid function σ satisfies (3).
This can be shown by the similar discussion as in the Lemma 1 in (Cybenko, 1989).
Proof of Theorem 1.
Assume that
[Equation]
Then, for all γ, δ ∈ R
[Equation]
we have
[Equation]
where ξ
[Equation]
We used the Lesbegue Bounded Convergence Theorem in the second equation.
From definition, we have
[Equation]
where
[Equation]
Hence for all s, α, t, β , we have
[Equation]
Based on this, we show µ = 0 .
Fix s, t and consider functional F ( h ) defined as
[Equation]
where h is bounded measurable function h ( u, v ) :
[¯ u, u ] × [¯ v, v ] → R , where
[Equation]
Let indicator function I ( b,c ] × ( d,e ] ( u, v ) defined as
I [ b,c ) ×
[ d,e )
[Equation]
u, v
[Equation]
=
[Equation]
1 ( u ∈ [ b, c ) , v ∈
[ d, e )) 0 otherwise .
Then, we have
[Equation]
× (Π X t, − c ∪ H X t, − c ) (cid:1) = 0 .
Since
I [ b,c ) ×
[ d,e ) =
I [ b, ∞ ) ×
[ d, ∞ ) − I [ c, ∞ ) ×
[ d, ∞ ) − I [ ,b ∞ )
× [ e, ∞ )
+ I [ c, ∞ ) ×
[ e, ∞ ) , we have F ( I [ b,c ) ×
[ d,e ) ) =
for all b, c, d, e ∈ R .
For linearlity, we have
[Equation]
Note that (cid:80) Ni
=1 η i I
[ b i ,c i ) ×
[ d i ,e i ) uniformly converges to any bounded measurable function h : [¯ u, u ] ×
[¯ v, v ] → R .
Hence, F ( h )
= 0 .
In particular, h ( u, v ) =
cos( u + v ) , sin( u + v ) are bounded measurable functions, and thus, (cid:90)
[Equation]
Thus, the Fourier transform of µ is 0 and so µ must be zero as well.
From Lemma 5, we see
Theorem 1.
B. I MPLICATION OF A SSUMPTION
In this section, we discuss the implication of Assumption 1, especially when the back-door and treatment variables are continuous.
First, we show the upper bound of the sup norm of Lipschitz function.
Lemma 6.
Let Z ∈
Z be the probability variable following P ( Z ) and Z ⊂ [0 , 1] d .
Then, for all
L -Lipschitz function h bounded in h ( z ) ∈ [ − R, R ] , we have
[Equation]
if the density function f ( z ) is bounded away from zero f ( z ) ≥ ε > 0 .
Proof.
Since Z is compact, there exists z ∗ such that
[Equation]
Let M = | h ( z ∗ ) | and we consider the following rectangle
[Equation]
, where z [ i ] denotes the i -th element of z .
Then, from Lipschitz continuity, for all z ∈ B , we have
[Equation]
Now, consider the volume of B .
Since
[Equation]
the events 0 ≥ z ∗ [ i ] −
M 2 R +2 √ dL and 1 ≤ z ∗ [ i ] +
M 2 R +2 √ dL do not occur simultaneously.
Therefore, we have
[Equation]
and
[Equation]
Since M = max z ∈Z | h ( z ) | , we have
[Equation]
By this, we can give the Assumption 1 follows for the interval probability space.
Corollary 1.
If A = [0 , 1] d A , X = [0 , 1] d X , and all function h ∈ H g are L − Lipschitz continuous, we have
[Equation]
where C = (cid:0) 4 c (cid:1)
[Equation]
Note that the assumption on hypothesis space is easy to satisfy since all neural network is Lipchitz function if we use the ReLU activation and regularize the operator norm of the weight in each layer.
B. C ONSISTENCY R ESULTS
Proof of Lemma 1 We use the following Rademacher bound to prove the consistency (Mohri et al., 2012).
Proposition 4.
(Mohri et al., 2012, Theorem 11.3)
Let X be a measurable space and H be a family of functions mapping from X to Y ⊆
[ − R, R ] .
Given fixed dataset S =
(( y 1 , x 1 ) , ( y 2 , x 2 ) , . .
, ( y n , x n )) ∈ ( X × Y ) n , the empirical Rademacher complexity is given by
[Equation]
where σ = ( σ 1 , . .
, σ n ) , with σ i independent random variables taking values in {− 1 , +1 } with equal probability.
Then, for any δ > 0 , with probability at least 1 − δ over the draw of an i.i.d sample S of size n , each of following holds for all h ∈ H :
[Equation]
Given Proposition
[Equation]
we can prove the consistency of conditional expectation
[Equation]
Proof of Lemma 1.
From Proposition 4 and ˆ g
[Equation]
g ∈ H g , for the probability at least 1 − 2 δ
[Equation]
we have followings.
[Equation]
From the minimality of ˆ g = arg min ˆ L X 1 , we have
[Equation]
Taking the square root of both sides completes the proof.
Empirical Rademacher Complexity of H g
We discuss the empirical Rademacher complexity of
H g when we use feed-forward neural network for features ϕ A , ϕ X here.
The discussion is based on a “peeling” argument proposed in Neyshabur et al.
Proposition 5 ((Neyshabur et al., 2015), Theorem 1) .
Let hypothesis space of L layer neural net be
H NN that
[Equation]
where σ is ReLU function and W (1)
∈ R D
× H , W ( L ) ∈ R 1 × H , W (2) , . . .
, W ( L − 1) ∈
R H
× H are weights.
The norm ∥ · ∥ p,q is matrix L p,q -norm sup x ̸ =
∥
Wx ∥ q / ∥ x ∥ p .
Then, for any L, q ≥ 1 , any 1 ≤ p ≤ ∞ , and any set S = { s 1 , . . .
, s n } , the empirical Rademacher complexity is bounded as ˆ
[Equation]
for p ∗ = 1 / (1 − 1 /p ) and [ x ] + = max { 0 , x } .
Given this, we can bound the empirical Rademacher complexity of H g when each coordinate of features is a truncated member of H NN .
Lemma 7.
Let A , X ⊂ R D and define hypothesis set H NNFeat .
( d ) that
[Equation]
where ˜ σ is a ramp function ˜ σ ( x ) = min(1 , max(0 , x )) .
Consider H g that
[Equation]
Given data set S
= { ( a 1 , x 1 ) , . . .
( a n , x n ) } , we have
[Equation]
Note that we have
[Equation]
since we apply ˜ σ in the features.
The proof is given as follows.
Proof.
Let us define the following hypothesis spaces.
[Equation]
Then, from the definition
[Equation]
we have
[Equation]
Since the maximum of a linear function of w over the constraint ∥ w ∥ ≤ R is achieved for the values satisfying ∥ w ∥
=
R , we have
[Equation]
Let ˜ H
2NN − ˜ H
2NN be the function space defined as
[Equation]
Since ˜ H 2NN contains the zero function, the final hypothesis space is the subset the convex hull of ˜ H 2 NN − ˜ H 2NN because
[Equation]
Therefore
[Equation]
we have
[Equation]
Now
[Equation]
we can bound ˆ R S ( ˜ H 2NN ) as
[Equation]
where S A = { a i } and S X = { x i } .
Here, we used Talagrand’s contraction lemma (Mohri et al., 2012, Lemma 5.11) in the inequality.
Again, from Talagrand’s contraction lemma, we have
[Equation]
since ˜ σ is an 1-Lipchitz function.
Combining them, we have
[Equation]
This and Proposition 5 completes the proof.
Now, we derive the final theorem to show the consistency of the method.
Proof of Theorem 2.
From the triangular inequality, we have
[Equation]
we have
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]
term, we have
[Table]
[Equation]
∞
Therefore
[Equation]
we have
[Equation]
Using Lemmas 1 and 3 and Assumption
[Equation]
we have
[Equation]
n with probability at least 1 − 4 δ .
Combining them and applying Lemma 7 completes the proof for
ATE bound.
For ATT, we can derive the followings with the same discussion |
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
ˆ ϕ
[Equation]
ˆ ϕ
[Equation]
( X ) |
A
[Equation]
a ′ (cid:105)(cid:13)(cid:13)(cid:13) ∞ .
Using Lemma 4 and the assumption made in Theorem 2, we have
[Equation]
If we use neural network hypothesis space H f considered in Proposition 5, we can see that the ATT bound holds.
B. L IMITATION OF S MOOTHNESS A SSUMPTION ON R IESZ R EPRESENTER
In Chernozhukov et al.
(2022b)
, we consider a functional m such that the causal parameter θ can be written as θ =
E
[ m ( g, ( A, X ))] , where g is the conditional expectation g ( a, x ) =
E [ Y |
A
= a, X = x ] .
Then, a Riesz Representer α , which satisfies E [ m ( g, ( A, X ))]
=
E [ α ( A, X ) g
( A, X )] , exists as long as
[Equation]
for all α ∈ H α and a smoothness parameter M .
When we consider ATE θ ATE ( a ) , the corresponding functional m would be
[Equation]
Chernozhukov et al.
(2021, Theorem 1) shows that the deviation of estimated the Riesz Representer ˆ α and the true one α 0 scales as linear to the smoothness parameter M .
[Equation]
[Figure]
U
[Figure]
where δ n is the critical radius that scales
[Equation]
when we consider fully connected neural networks.
Now, we show that the smoothness parameter M can have an exponential dependency on the dimension of the space, even for simple α .
Consider A =
[ − 1 , 1] d and some compact space X .
We assume the uniform distribution for P ( A, X ) .
Consider following ˜ α
[Equation]
where a [ i ] denotes i -th element of a , and here we consider ˜ α
that does not depend on x .
Say, we are interested in estimating θ ATE ( a ) of a = 0 =
. , 0] ⊤ , for which
[Equation]
Now consider B that
[Equation]
Then, since ˜ α ( a, x ) = 0 for all a / ∈ B , we have
[Equation]
We use the assumption that P ( A, X ) is the uniform distribution to have the last equality.
Hence, if ˜ α ∈ H α , the smoothness parameter M must have the exponential dependency
M ≥ 2 d .
C O BSERVABLE C ONFOUNDER
In this section, we consider the case where we have the additional observable confounder, the causal graph of which is given in Figure 4.
Given the causal graph in Figure 4, ATE and ATT are defined as follows.
[Equation]
Furthermore, we can consider another causal parameter called conditional average treatment effect (CATE) , which is a conditional average of the potential outcome given O
=
o ;
[Equation]
Given exchangeability and no inference assumption, we have
[Equation]
These causal parameters can be recovered if the back-door or the front-door variable is provided as follows.
Back-door adjustments: First, we present the Proposition stating these causal parameters can be recovered if we are given the back-door variable X .
Proposition 6 (Pearl, 1995) .
Given the back-door adjustment X in Figure 4b, we have
[Equation]
where g ( a, o, x ) =
E [ Y |
A
= a, O
= o, X = x ] .
Now, we present the deep adaptive feature embedding approach to this.
We first learn conditional expectation ˆ g as ˆ g ( a, o, x ) =
ˆ w ⊤ ( ˆ ϕ A ( a ) ⊗ ˆ ϕ O ( o )
⊗ ˆ ϕ X ( x )) , where
[Equation]
given data ( y i , a i , o i , x i ) .
Here, w is the weight and ϕ A , ϕ O , ϕ
X are the feature maps.
From
Proposition 6, we have
[Equation]
Therefore, by estimating the feature embeddings
[Equation]
we have
[Equation]
where ˆ f ˆ ϕ O ⊗ ˆ ϕ X , ˆ f ˆ ϕ X
( o ) are learned from
[Equation]
Front-door adjustment:
Given the front-door variable M , these causal parameters can be identified as follows.
Proposition 7 (Pearl, 1995) .
Given the front-door variable M in Figure 4c, we have
[Equation]
where g ( a, o, m )
=
E [ Y |
A
= a, O
= o, M = m ] and A ′ follows the identical distribution as A .
For front-door adjustment, we learn conditional expectation ˆ g as ˆ g ( a, o, x ) =
ˆ w ⊤ ( ˆ ϕ A ( a ) ⊗ ˆ ϕ O ( o )
⊗ ˆ ϕ M ( m )) , where
[Equation]
Then, from Proposition 7, we have
[Equation]
The conditional expectation E M | O
= o,A = a (cid:104) ˆ ϕ M ( M ) (cid:105) is estimated as E M |
O =
o,A = a (cid:104)
ˆ ϕ M ( M ) (cid:105) =
ˆ f ˆ ϕ M ( o, a ) , where
[Equation]
Then, by replacing the marginal expectation with the empirical average
[Equation]
we have
[Equation]
i = (cid:16) ˆ ϕ O
( o i ) ⊗ ˆ f ˆ ϕ M
( o i , a ) (cid:17)(cid:33) , ˆ θ CATE ( a ; o )
=
ˆ w ⊤ (cid:32)(cid:32)n n (cid:88) i
= ˆ ϕ A ( a i ) (cid:33) ⊗ ˆ ϕ O
( o ) ⊗ ˆ f ˆ ϕ M ( o, a ) (cid:33) .
D E XPERIMENT D ETAILS
Here, we describe the network architecture and hyper-parameters of all experiments.
Unless otherwise specified, we used Adam with learning rate = 0.001, β 1 = 0 .
9 , β 2 = 0 .
999 and ε = 10 − 8 .
For RKHS Embedding, we used Gaussian kernel for continuous variable where the bandwidth is determined by the median trick.
D. B INARY T REATMENT S CENARIO
In this scenario, all treatments are binary A ∈ { 0 , 1 } .
In RKHS Embedding and Neural Embedding, we used the feature ϕ
A given as
[Equation]
⊤ in both IHDP setting and ACIC setting.
This is equivalent to learn two models
[Equation]
) with shared nonlinear feature ϕ X ( X ) .
IHDP Dataset
We used the 1000 data used in (Chernozhukov et al., 2022b), which is publicly available at Github page of the paper.
The network structure for back-door feature ϕ X ( X ) is shown in Table 2.
Note that is much smaller network than Dragonnet or Riesznet, but increasing network size did not affect the result much.
Back
[Table]
ACIC Dataset We used the 101 data used in (Shi et al., 2019), which satisfies overlap assumption.
(i.e. Not all data points has the extreme propensity score P ( A = 1 | X ) .)
We noticed that some data contains a outliers and we only consider the data points with the outcome Y is in the range of
[Equation]
where Q 1 ( Y ) , Q 3 ( Y ) are 25% , 75% -quantile value of outcome, respectively, and IQR =
Q 3 ( Y ) −
Q 1 ( Y ) .
We run Dragonnet and RieszNet estimators with the same network architecture as IHDP dataset.
The network structure for back-door feature ϕ X ( X ) is shown in Table 3.
Note that the same structure is used in Dragonnet and Riesznet to predict conditional expectation E
[ Y | X, A ] .
[Table]
D. H IGH - DIMENSIONAL T REATMENT S CENARIO
Here, we generate all dataset by ourselves from original dSprite dataset (Matthey et al., 2017).
Back-door ATE estimation The network features for the proposed method is summarized in Table 4.
The network structures for RieszNet is given in Table 5.
Note that they share the similar feature extractor for images.
Back-door feature ϕ X
( X )
Layer
Configuration
Input( X )
FC(2, 36), ReLU
FC(36, 5), ReLU
Treatment Feature ϕ A ( A )
Layer
Configuration
Input( A )
FC(4096, 1024), SN, ReLU
FC(1024, 512), SN, ReLU, BN
FC(512, 128), SN, ReLU
FC(128, 32), SN, BN, Tanh
Front-door ATT estimation Here, we used the same network architecture as in the back-door adjustment summarized in Table 4.
[Table]