Deep Counterfactual Estimation with Categorical Background Variables
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Referred to as the third rung of the causal inference ladder, counterfactual queriestypically ask the " What if ? " question retrospectively.
The standard approach toestimate counterfactuals resides in using a structural equation model that accuratelyreﬂects the underlying data generating process.
However, such models are seldom available in practice and one usually wishes to infer them from observationaldata alone.
Unfortunately, the correct structural equation model is in general notidentiﬁable from the observed factual distribution.
Nevertheless, in this work, weshow that under the assumption that the main latent contributors to the treatmentresponses are categorical, the counterfactuals can be still reliably predicted.
Building upon this assumption, we introduce CounterFactual Query Prediction (CFQP),a novel method to infer counterfactuals from continuous observations when thebackground variables are categorical.
We show that our method signiﬁcantly outperforms previously available deep-learning-based counterfactual methods, boththeoretically and empirically on time series and image data.
Our code is availableat https://anonymous.4open.science/r/cfqp .
Introduction
Counterfactual queries aim at inferring the impact of a treatment conditioned on another observedtreatment outcome.
Typically, given an individual, a treatment assignment, and a treatment outcome,the counterfactual question asks what would have happened to that individual, had it been givenanother treatment, everything else being equal.
An illustrative and motivating example is the case ofclinical time series.
Based on the observation of the outcome of treatment A on a particular patient,counterfactual queries ask what would have been the outcome for this patient, had it been giventreatment B instead.
Notably, counterfactual prediction differs from interventional prediction, whichis also referred to as counterfactual potential outcomes [21] and constitutes the second rung of thecausation ladder [17].
Counterfactual predictions are retrospective, as they condition on an observedtreatment outcome.
In contrast, interventional predictions are prospective as they only condition onobservations obtained before treatment assignment.
Much more than a statistical curiosity, counterfactual reasoning reﬂects complex cognitive abilitiesthat are deeply ingrained in the human brain [20] and emerges in the early stages of cognitivedevelopment [7].
The ability to reason counterfactually can indeed help to identify causes ofoutcomes retrospectively, has been suggested to be central in the formation of rational intention,and supports key theories of human cognition [26].
The importance of counterfactual reasoning inthe human cognitive process has thus motivated researchers to endow artiﬁcially intelligent systemswith the same ability [19].
However, counterfactual inference is not possible from observational andinterventional data alone [18].
Counterfactual reasoning, therefore, requires making several assumptions to overcome this limitation.
The most popular one assumes knowledge of the underlying structural equation model that describesthe data generating process [19] or a speciﬁc functional form thereof [1, 5, 19].
Unfortunately, thisassumption is rarely met in practice, especially in high-dimensional data such as time series or images.
This led to the development of deep structural equation models that attempt to model the structuralequations with neural networks [15, 22].
However, despite their ability to model high-dimensionaldata, these approaches fail to provide theoretical guarantees for the reconstruction of counterfactuals.
Indeed, they focus on modeling the factual distribution which, without further assumption, can,unfortunately, lead to erroneous counterfactual distributions.
In this work, we bridge the gap between the classical structural equation model assumptions anddeep-learning-based architectures.
By assuming that the treatment and observables are continuousand that the hidden variables that contribute most to the treatment response are categorical, weshow that we can recover the counterfactuals using arbitrary parametric functions ( i.e. deep neuralnetworks) to model the causal dependence between variables.
This allows us to infer counterfactualson high-dimensional data such as time series and images.
Besides the general appeal of endowing machine learning architectures with counterfactual reasoningabilities, an important motivation for our work is the counterfactual estimation of treatment effectsin clinical patient trajectories.
In this motivational example, one wishes to predict the individualtreatment effect retrospectively.
Based on the observed treatment outcome of a particular patient,we want to predict what would have been the outcome under a different treatment assignment.
Theability to perform counterfactual inference on patient trajectories has indeed been identiﬁed as apotential tool for improving long and costly randomized clinical trials [13].
Contributions (1) We provide a new set of assumptions under which the counterfactuals areidentiﬁable using arbitrary neural networks architecture, bridging the gap between structural equationmodels and deep learning architectures.
(2) We derive a new counterfactual identiﬁability resultthat motivates a novel counterfactual reconstruction architecture.
(3) We evaluate our constructionon three different datasets with different high-dimensional modalities (images and time series) anddemonstrate accurate counterfactual estimation.
Background2.1 Problem Setup : Counterfactual Estimation
We consider the general causal model M = h U, V, F i depicted in Figure 1a consisting of backgroundvariables U , endogenous variables X, T and Y and the set of structural functions F .
Backgroundvariables U =
{ U X , U T , U ✏ , W } are hidden exogenous random variables that determine the valuesof the observed variables V =
{ X, T, Y } .
Covariates X 2 X represent the information available 69 before treatment assignment, T 2 T is the treatment assignment and Y 2 Y is the observed response 70 to the treatment.
We refer to the space of probability measures on
Y as P ( Y ) .
Observed variables V 71 are generated following the structural equations F = { f X , f T , f Y , f ✏ } , such that X
=
f
X ( U x , W ) , 72 T = f
T ( U T , X ) , U ✏ =
f ✏ ( W ) , Y = f Y ( X, T, U ✏ ) .
We further assume strong ignorability ( i.e. nohidden confounders between T and Y ).
Using notations introduced in Pearl et al.
[19], we deﬁne the potential response of a variable Yto an action do ( T = t ) for a particular realization of U =
u as Y t ( u ) .
Our goal is to predict thecounterfactual response, for a new treatment assignment ( T = t 0 ), conditioned on an observedinitial treatment response.
That is, the probability of observing a different treatment response undertreatment t 0 , after observing treatment response y for covariate x and treatment t .
The probability ofobserving counterfactual y 0 then writes:
[Equation]
and we refer to the counterfactual probability measure as ⌫ t 0 ( x, y, t ) .
Equation 1 suggests a naturalthree step procedure for computing the probability of counterfactual.
First, the abduction step infersthe probability of U conditioned on the observed treatment outcomes, covariates and treatments:
P ( U = u
| X = x, T = t, Y = y ) .
Second, in the action step, one sets the new treatment in thecausal model ( do ( T = t 0 ) ).
Lastly, in the prediction step, one can propagate the values of U = u
and
T = t 0 in the causal graph, using F , to compute P ( Y t 0 ( u ) =
y 0 ) .
In practice, we only have access to a set of N observations of variables X , Y and T .
We referto this dataset as D =
( X , T , Y ) where X = { x i : i = 1 , ..., N } , Y = { y i : i = 1 , ..., n } and
T = { t i : i = 1 , ..., N } .
Importantly, we don’t have access to counterfactual examples ( i.e. a tuple( x, y, t, t 0 , y 0 ) ), such that direclty learning a map ( x, y, t, t 0 ) !
y t 0 is excluded.
2.2 General Non-identiﬁability of Counterfactuals
Because the background variables U are hidden, the above three-steps procedure requires knowledgeof the structural functions F = { f X , f T , f Y } .
Indeed, one can show that there exist multiplestructural functions F that would lead to the same observed joint probability P ( X, Y, T ) but wouldlead to incorrect counterfactual probabilities [17, 18].
The correct causal model is thus in generalnon-identiﬁable, leading to non-identiﬁability of the counterfactual probability.
We specify what ismeant by the identiﬁability of counterfactuals in the following deﬁnition.
Deﬁnition 1 (Identiﬁability of Counterfactuals) .
Let ⇢ be a metric on P ( Y ) , ⌫ t 0 ( X, Y, T ) the truecounterfactual probability measure and ˆ ⌫ t 0 ( X, Y, T ) the estimator of the counterfactual probabilitymeasure with N data points.
Counterfactuals are ⇢ -identiﬁable
at threshold  if, for all t, t 0 2 T , x
X , y 2 Y ,
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
U X
U T
W
Y
X
T
[Equation]
W
U X
[Equation]
T
[Figure]
2.3 Causal Model Assumptions for Counterfactual Idenﬁability
Despite the general non-identiﬁability of structural equation models laid out above, we proposeplausible assumptions that one can build upon to identify counterfactuals reliably.
We ﬁrst assume
X and Y are continuous (potentially high dimensional) variables (such as images or time series).
The treatment assignment T is also assumed continuous, and X ⇥ T is a connected space.
Our ﬁrstcentral assumption posits that the hidden variable U ✏ factorizes into a categorical and a continuousvariable.
Assumption 1 (Categorical Background Variable) .
The background variable U ✏ decomposes into acategorical latent variable U Z 2 [ K ] = { 1 , .., K } and an independent exogenous continuous variable 110 U ⌘ .
This assumption is depicted in the graphical model of Figure 1b and embodies the intuition ofdifferent hidden groups that drive the treatment response.
For instance, a treatment could havedifferent responses depending on the stage of the disease a patient ﬁnds themself in.
The diseasestage is unobserved yet correlated with the observed covariates X (through W ).
Due to the categorical nature of variable U Z , one can write the conditional probability of Y as amixture model:
[Equation]
We deﬁne  as the probability density function of the conditional treatment response generated by f Y ,
U Z and U ⌘ .
is thus a mixture probability density function with mixture components 1  k 2 P ( Y )and weights !
k .
[Equation]
Without loss of generality, we assume that U ⌘ ⇠ N ( 0 , ⌃ 2 ) .
In the case of additive noise, the conditional distribution of Y becomes a mixture of Gaussians:  ( X,
[Equation]
=
[Equation]
where µ k are functions mapping X and T to the mean of the mixture components and we considerdifferent variances for each k .
We now proceed with the next assumptions.
Assumption 2 (Continuity) .
The moments of the probability density functions  k ( x, t ) exist and arecontinuous functions of X and T :
[Equation]
Assumption 3 (Clusterability) .
For each ( x, t ) 2 ( X , T ) , the density  ( x, t ) is clusterable and theexpected deviation of each  k ( x, t ) is bounded by a constant  2 R .
That is, 8 k 2 [ K ] , 8 x, t 2 129 (
X ⇥ T ) , with µ k ( x, t ) =
E Y ⇠  k ( x,t )
[ Y ] :
[Equation]
In our motivating clinical example, Assumption 2 reﬂects that the probability of speciﬁc treatmentresponse changes continuously over the set of observed covariates and treatments.
In particular, theexpected treatment outcome for a particular patient varies continuously with the treatment assignment,which is a common assumption, e.g. in clinical practice [12].
Assumption 3 posits clusterability of the mixture components  k for which a rigorous mathematicaldeﬁnition is given in Appendix C. Intuitively, it supposes that patients with the same observedcovariates and treatment assignment but different hidden group will show different treatment outcomes.
We also bound the expected deviation of the mixture components  k that characterize the inter-groupvariability in the treatment response for a particular patient and treatment outcome.
Methods
3.1 Identiﬁability and Counterfactuals Reconstruction
For a ﬁxed point ( X = x, T = t ) , Equation 2 is a ﬁnite mixture model, for which identiﬁability resultsare available [3].
Notably, these results guarantee identiﬁability up to a permutation of the latent class assignment  ( · ) :
[ K ] !
[ K ] .
That is, there exists some permutation  ( · ) :
[ K ] !
[ K ] such thatˆ   k ( k )( x,t ) ⇡  k ( x, t ) , ˆ !
k ( k ) ⇡ !
k , where ˆ  and ˆ ! are the estimated density functions and weights.
However, it does not entail identiﬁability of the counterfactuals in the sense of deﬁnition 1.
Indeed,the action step of the counterfactual strategy from Section 2.1 requires a consistent permutation across the whole domain ( X ⇥ T ) in order to reuse the inferred class assignments ˆ U Z at a speciﬁcpoint ( X = x, T = t ) to predict the counterfactual at another point ( X = x, T = t 0 ) - with adifferent treatment assignment.
Nevertheless, using the assumptions from the previous section, wecan still ensure the identiﬁability of the counterfactuals as the following result conﬁrms:
Result 3.1 (Identiﬁabilty of Counterfactuals with Categorical Background Variables) .
Let X , T and
Y be continuous random variables generated according to the graphical model of Figure1b withthe domain of X and T being connected.
Let W 1 ( · , · ) be the ﬁrst Wasserstein distance on P ( Y ) , 155 ⌫ t 0 ( X, Y, T ) the probability distribution of Y t 0 | X, Y, T and ˆ ⌫
Nt
( X, Y, T ) its estimator from Nobserved data points.
If Assumptions 1, 2 and 3 hold, for each ( x, t ) , the counterfactual distributionis W 1 -identiﬁable at threshold  :
[Equation]
In the special case when the noise response is additive, we have
[Equation]
The proof is given in Appendix C. This result gives us a bound on the distance between the inferredand true counterfactual distributions in the asymptotic regime.
Importantly, it does not restrict thedimension of X and Y , and is thus valid on challenging data modalities such as time series or images.
Continuity of distribution and complexity The result above holds asymptotically in the numberof available samples.
In the additive Gaussian case, the sample complexity for learning a K -mixturemodel with Y 2 R d within ✏ total variation distance is ˜ O
( Kd 2 / ✏ )
Fortunately, the continuityassumption (Assumption 2) saves us from having to learn an individual mixture at each point( X = x, T = t ) , by jointly learning the continuous moments functions
µ 0 r ( X, T ) .
A better samplecomplexity bound can then be derived with further assumptions on µ 0 r ( X, T ) .
3.2 CFQP : CounterFactual Query Prediction
Equipped with those theoretical results, we introduce CFQP, a counterfactual prediction model basedon a neural Expectation-Maximization mechanism.
The basic building block of CFQP is a base-modelm ( x, t ) , that predicts the treatment response y based on covariates and treatment assignment.
Foreach latent category k , we learn a base-model that approximates the individual treatment responsein that category : m k ( x, t ) ⇡ µ k ( x, t ) .
Our theoretical results require the number of classes of Kto be known in advance.
Yet, this is rarely the case in practice, and we describe our architecturefor an arbitrary number of classes K 0 .
The learning of the base-models follows three steps: a jointinitialization, an expectation phase, and a maximization phase.
The overall process is depicted in
Figure 2.
We also present a pseudo-code description of the procedure in Algorithm 1.
Initialization We ﬁrst train a single common base-model m 0 on all data points X , T and Y .
Because it discards the variability of the hidden variable U Z , this model will approximate theconditional average treatment response : m 0 ( x, t )
⇡ E [ Y | X = x, T = t ] .
Expectation - Clustering The clustering step ensures that each base-model m k is trained with thedata from the corresponding latent category.
We distinguish between an initial clustering and anupdate clustering stage.
The initial clustering assignment happens after the initialization step has converged.
We then use theclusterability assumption (Assumption 3) to drive the assignment and use K 0 -means clustering basedon the residuals k y  m 0 ( x, t ) k 22 .
[Figure]
= arg min j k y  m j ( x, t ) k 22 .
Maximization - Learning individual base-models
Once the dataset D =
( X , T , Y ) is clusteredin K 0 groups, we train the models m k on the corresponding cluster.
This
Counterfactual prediction For a data point ( x, t, y ) , we ﬁrst infer the cluster assignment k =arg min j k y  m j ( x, t ) k 22 .
We predict the counterfactual for treatment t 0 as ˆ y
= m k ( x, t 0 ) .
Algorithm
Data: X , Y , T , the number of latent clusters K 0 , a number of epochs e max , the update-period  .
Result:
A list of K 0 hidden models m k .
Initialize a single base model m 0 at random.
for epoch 0 to e max do
Train m 0 minimizing L 0 =
E ⇥ ( m 0 ( X, T )  Y )
⇤ end Compute residuals r = ( m 0 ( X , T )  Y ) .
Assign the residuals of each training sample into K 0 clusters (initial).
Initialize each m k with m 0 .
for epoch 0 to e max do for i 0 to  do
Train m k minimizing L k
= E ⇥ ( m k ( X k , T k )  Y k )
⇤ for each cluster.
end
Compute residuals R = { (
[Equation]
Assign the
residuals of each training sample into K 0 clusters (update)
endReturn trained models m k .
Related Work
Causal perspectives in machine learning have gained signiﬁcant traction in the past years [24, 23].
Among them, one distinguishes between causal discovery approaches, aiming at discovering thecausal relations between variables [8], and causal inference , aiming at building treatment effects estimators from data [16].
Our work belongs to the latter.
A common way to estimate counterfactualsis to posit a speciﬁc structural causal model [18] or its speciﬁc functional form [5], therefore ensuringidentiﬁability.
Synthetic controls are an example of this strategy, assuming an underlying linearstructure [1].
However, these approaches are by deﬁnition restrictive in terms of the expressivityof the structural equations.
Deep learning approaches for counterfactual estimation have beenrecently proposed to address this issue [15, 22] but without identiﬁability guarantee.
For the discretecase, assumptions have been proposed to bridge this gap, such as monotonicity or generalizationthereof [14, 11].
Nevertheless, because they focus on addressing the discrete case, they are notdirectly applicable to high-dimensional data such as time series, which is a motivation for our work.
Experiments
5.1 Datasets
Treatment outcomes are often complex and high-dimensional.
We thus evaluate our proposed modelarchitecture on high-dimensional data modalities: time series and images.
As counterfactual inferenceevaluation on real-world data is a complex and ongoing research area, we use synthetic datasetsinspired by case studies from the literature.
Details about the data generation are given in Appendix A.
5.1.1 Counterfactual Image Transformation
To explore the capacity of our model to operate on images, we generate a dataset inspired by [15],building upon modiﬁcation of MNIST images.
Covariates X are original MNIST images and thetreatment T is a spatial rotation applied to the image.
U Z is a coloring of the image and the treatmentoutcome Y is the resulting colored and rotated MNIST image.
U ⌘ is either an additive Gaussiannoise at each pixel or a Gaussian Blur with a random kernel.
5.1.2 Counterfactual Longitudinal Treatment Effect Prediction
Next, we explore the performance of our model on time series data.
As per our motivational example,our goal is here to infer counterfactuals from clinical trajectories with a hidden patient class.
X
and
Y are multi-dimensional time series, U Z is a latent patient group and T is a continuous treatmentassignment.
We used two different datasets.
A Harmonic Oscillator and a Cardiovascular model.
Harmonic oscillator We simulate the angularpositions of two coupled harmonic oscillators.
Thetreatment consists of applying different gradualoffsets to the time series.
We consider three hidden groups ( K = 3 ) that modulate the treatmentresponse differently.
Cardiovascular simulator We use a simulatorof the cardiovascular system proposed in [27] topredict the impact of ﬂuid intake on blood pressure.
Fluids are commonly administered in intensivecare for treating severe hypotension.
However, theindividual patient response is difﬁcult to assess apriori, as it depends on clinically hidden variables.
We consider here two hidden groups of patientswith distinct responses to ﬂuid intake.
X and Ycontain the arterial and venous blood pressure timeseries before and after ﬂuid intake.
The treatmentis the amount of ﬂuid injected.
We add processnoise by either considering an additive Gaussiannoise on the observed treatment response or byintroducing noise in the ﬂuid injection process,leading to a non-linear and complex perturbation.
5.2 Baselines
We compare our approach against multiple counterfactual inference models from the literature:
Diff-SCM [22], a recently introduced diffusion based counterfactual model; Deep-SCM [15], a deepcounterfactual inference architecture building on normalizing ﬂows ; Synthetic Controls ( SC ) [1], awell known linear method for counterfactuals estimation and Deep-ITE , a deep individual treatmenteffect estimator such as proposed in [25, 10], using direct predictions from X and T .
[Table]
5.3 Counterfactuals Prediction
We evaluate our approach in terms of the quality of the counterfactual reconstructions.
We generatecounterfactuals ( t 0 , y 0 ) for each observation ( x , y , t ) in the test set, resulting in a tuple ( x , y , t , t 0 , y 0 ).
Weevaluate the MSE between the true counterfactual y 0 and the estimated one ˆ y 0 as MSE cf =
k
y
ˆ y
k 22 .
The results are presented in Table 1.
We use additive and non-additive noise variants for eachdataset.
For each dataset and each variant, we see that our approach outperforms the other baselines,demonstrating its experimental effectiveness.
In Figure 3, we also report examples of counterfactual reconstructions on the image dataset.
Eachrow corresponds to counterfactual estimations for a particular method (top is the truth), and eachcolumn represents a different treatment assignment T 0 (the left column is the factual T ).
Becauseit cannot model the categorical latent variable, the Deep-ITE model gets the expected individualtreatment effect right but fails to accurately estimate the counterfactuals.
For Diff-SCM, we observethat, despite providing an accurate reconstruction for most of the treatment assignments, the lack ofguarantees on the deep structural equation model leads to potential incorrect reconstructions ( e.g. at
T 0 = 30 ).
5.4 Robustness Analysis
We further complement our experimental section with a comprehensive robustness analysis.
Weﬁrst investigate the impact of misspecifying the number of latent classes ( i.e. when K 6 = K 0 ).
Wethen investigate the impact of a correlation between X and Z with increasing strength.
Finally, weinvestigate the robustness of our approach to different noise responses.
Number of latent classes We study the evolution of performance when the number of classes K 0is misspeciﬁed ( K 6 = K 0 ).
In Figure 4a, we show the reconstruction error on the validation set ( i.e. ,on the factual data) and the counterfactual reconstruction error of CFQP for different values of K 0on the image dataset with true number of latent classes K = 6 .
The validation MSE is lowest for
K 0 = 6 and K 0 = 7 , which corresponds to optimal test reconstruction error.
This hyper-parametercan thus easily be tuned by monitoring the treatment prediction performance on the factual data.
5.4.1 Correlation between X and U Z
Our approach holds even when the latent background variable U Z and the covariates X are correlated.
When this is the case, information from X can be used to infer the value of U Z .
In Figure 4b, westudy the impact of the correlation strength ⇢ between X and U Z .
More details about the deﬁnitionof ⇢ are to be found in Appendix B.
We compare the performance of CFQP with a Deep-ITEmodel.
We observe that the counterfactual performance of our model remains constant, regardlessof the correlation strength.
However, as the correlation becomes more important, U Z becomes adeterministic outcome of X and Deep-ITE eventually converges to the performance of CFQP.
[Figure] [Figure]
[Figure] [Figure]
5.4.2 Non-additive Noise Responses
[Table]
Our model architecture is motivated by the additive noise case.
Yet, our theoretical results hold for arbitrary distributions.
In
Table 1, we reports the performance of the different modelsunder non-additive noise responses.
We observe an overallsmall degradation of performance of all methods but CFQP stilloutperforms baselines, as predicted by our identiﬁability results.
5.5 Retrospective Individualized Treatment Effect
Estimation
We explore the ability of our model to infer individualized treatment effects retrospectively.
Based onobserved treatment outcomes, we predict the difference in treatment response between two treatmentregimes for a single individual.
The metric of interest is the Precision in Estimation of Heterogeneous
Effects PEHE [9].
For two treatment regimes t 0 and t 00 , we write
[Equation]
In Table 2, we report PEHE ( t 0 = 0 .
5 , t
= 0 .
8) for test patients in the Cardiovascular dataset withdifferent ﬂuid intake.
We observe that the methods leveraging information about treatment outcomes,
CFQP, SC and Diff-SCM outperform the classical individual treatment effect estimators (Deep-ITE).
Conclusion
Estimating counterfactuals from observational studies is one of the most challenging tasks in causalinference.
In this work, we proposed a set of reasonable assumptions that allow computing counterfactuals on high-dimensional data while harnessing the power of modern machine learning architectures.
Based on these assumptions, we derived a new counterfactual model and demonstrated favorableexperimental performance.
In particular, we showed that this approach could be used to infer individual treatment effects a posteriori in clinical patient trajectories.
Nevertheless, the set of assumptionsproposed here is not unique, and depending on the speciﬁc applications, others might be deemedmore relevant.
Other model architectures compatible with our set of assumptions are also possible.
Potential negative societal impacts Counterfactual inference has been used in clinical settings andfor assessing the impact of public policies [2], among others.
However, as shown in this paper, itscorrectness hinges on assumptions whose validity is not always met, potentially leading to misleadingconclusions.