Deep Counterfactual Estimation with Categorical Background Variables
Abstract
Referred to as the third rung of the causal inference ladder, counterfactual queries typically ask the " What if ? " question retrospectively.
The standard approach to estimate counterfactuals resides in using a structural equation model that accurately reflects the underlying data generating process.
However, such models are seldom available in practice and one usually wishes to infer them from observational data alone.
Unfortunately, the correct structural equation model is in general not identifiable from the observed factual distribution.
Nevertheless, in this work, we show that under the assumption that the main latent contributors to the treatment responses 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 the background variables are categorical.
We show that our method significantly outperforms previously available deep-learning-based counterfactual methods, both theoretically and empirically on time series and image data.
Our code is available at https://github.com/edebrouwer/cfqp .
Introduction
Counterfactual queries aim at inferring the impact of a treatment conditioned on another observed treatment 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 given another treatment, everything else being equal.
An illustrative and motivating example is the case of clinical 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 given treatment B instead.
Notably, counterfactual prediction differs from interventional prediction, which is also referred to as counterfactual potential outcomes [21] and constitutes the second rung of the causation ladder [17].
Counterfactual predictions are retrospective, as they condition on an observed treatment outcome.
In contrast, interventional predictions are prospective as they only condition on observations obtained before treatment assignment.
Much more than a statistical curiosity, counterfactual reasoning reflects complex cognitive abilities that are deeply ingrained in the human brain [20] and emerges in the early stages of cognitive development [7].
The ability to reason counterfactually can indeed help to identify causes of outcomes 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 in the human cognitive process has thus motivated researchers to endow artificially intelligent systems with the same ability [19].
However, counterfactual inference is not possible from observational and interventional 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 describes the data generating process [19] or a specific functional form thereof [1, 5, 19].
Unfortunately, this assumption 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 structural equations with neural networks [15, 22].
However, despite their ability to model high-dimensional data, 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 and deep-learning-based architectures.
By assuming that the treatment and observables are continuous and that the hidden variables that contribute most to the treatment response are categorical, we show that we can approximately recover the counterfactuals using arbitrary parametric functions ( i.e. deep neural networks) to model the causal dependence between variables.
This allows us to infer counterfactuals on high-dimensional data such as time series and images.
Generally, this work explores the assumptions that can lead to approximate counterfactual reconstructions while controlling the discrepancy between the recovered and the true counterfactual distribution.
Besides the general appeal of endowing machine learning architectures with counterfactual reasoning abilities, an important motivation for our work is the counterfactual estimation of treatment effects in clinical patient trajectories.
In this motivational example, one wishes to predict the individual treatment 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.
The ability to perform counterfactual inference on patient trajectories has indeed been identified as a potential tool for improving long and costly randomized clinical trials [13].
Contributions
• We provide a new set of assumptions under which the counterfactuals are identifiable using arbitrary neural networks architecture, bridging the gap between structural equation models and deep learning architectures.
• We derive a new counterfactual identifiability result that motivates a novel counterfactual reconstruction architecture.
• We evaluate our construction on three different datasets with different high-dimensional modalities (images and time series) and demonstrate accurate counterfactual estimation.
Background 2.1 Problem Setup : Counterfactual Estimation
We consider the general causal model M = ⟨ U, V, F ⟩ depicted in Figure 1a consisting of background variables U , endogenous variables X, T and Y and the set of structural functions F .
Background variables U
= { U X , U T , U ϵ , W } are hidden exogenous random variables that determine the values of the observed variables
V
=
X, T, Y } .
Covariates X ∈ X represent the information available before treatment assignment, T ∈ T is the treatment assignment and Y ∈ Y is the observed response to the treatment.
We refer to the space of probability measures on Y as P ( Y ) .
Observed variables V are generated following the structural equations F = { f X , f T , f Y , f ϵ } , such that X =
f X
( U x , W ) ,
T = f T ( U T , X ) , U ϵ
= f
ϵ ( W ) , Y =
f Y ( X, T, U ϵ ) .
We further assume strong ignorability ( i.e. no hidden confounders between T and Y ).
Using notations introduced in Pearl et al.
[19], we define the potential response of a variable Y to an action do ( T = t ) for a particular realization of U =
u as Y t ( u ) .
Our goal is to predict the counterfactual response, for a new treatment assignment ( T = t ′ ), conditioned on an observed initial treatment response.
That is, the probability of observing a different treatment response under treatment t ′ , after observing treatment response y for covariate x and treatment t .
The probability of observing counterfactual y ′ then writes:
P
[Equation]
and we refer to the counterfactual probability measure as ν t ′ ( x, y, t ) .
Equation 1 suggests a natural three step procedure for computing the probability of counterfactual.
First, the abduction step infers the 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 the causal model ( do ( T = t ′ ) ).
Lastly, in the prediction step, one can propagate the values of U = u
and
T = t ′ in the causal graph, using F , to compute P ( Y t ′ ( u ) = y ′ ) .
In practice, we only have access to a set of N observations of variables X , Y and T .
We refer to 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 ′ , y ′ ) ), such that direclty learning a map ( x, y, t, t ′ )
→ y t ′ is excluded.
2.2 General Non-identifiability of Counterfactuals
Because the background variables U are hidden, the above three-steps procedure requires knowledge of the structural functions F = { f X , f T , f Y } .
Indeed, one can show that there exist multiple structural functions F that would lead to the same observed joint probability P ( X, Y, T ) but would lead to incorrect counterfactual probabilities [17, 18].
The correct causal model is thus in general non-identifiable, leading to non-identifiability of the counterfactual probability.
We specify what is meant by the identifiability of counterfactuals in the following definition.
Definition 1 (Identifiability of Counterfactuals) .
Let ρ be a metric on P ( Y ) , ν t ′ ( X, Y, T ) the true counterfactual probability measure and ˆ ν t ′ ( X, Y, T ) the estimator of the counterfactual probability measure with N data points.
Counterfactuals are ρ -identifiable at threshold δ if, for all t, t ′ ∈ T , x ∈
[Equation] [Equation]
[Figure]
2.3 Causal Model Assumptions for Counterfactual Idenfiability
Despite the general non-identifiability of structural equation models laid out above, we propose plausible assumptions that one can build upon to identify counterfactuals reliably.
We first 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 first central assumption posits that the hidden variable U ϵ factorizes into a categorical and a continuous variable.
Assumption 1 (Categorical Background Variable) .
The background variable U ϵ decomposes into a categorical latent variable U Z ∈ [ K ] = { 1 , .., K } and an independent exogenous continuous variable
U η .
This assumption is depicted in the graphical model of Figure 1b and embodies the intuition of different hidden groups that drive the treatment response.
For instance, a treatment could have different responses depending on the stage of the disease a patient finds themself in.
The disease stage 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 a mixture model:
[Equation]
We define γ 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 ∈ 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:
[Equation]
( µ k ( X, T ) , Σ 2 k ) , where µ k are functions mapping X and T to the mean of the mixture components and we consider different 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 are continuous functions of X and T :
[Equation]
Assumption 3 (Clusterability) .
For each ( x, t ) ∈ ( X , T ) , the density γ ( x, t ) is clusterable and the expected deviation of each γ k ( x, t ) is bounded by a constant δ ∈ R .
That is, ∀ k ∈ [ K ] , ∀ x, t ∈ ( X × T ) , with µ k ( x, t ) =
E Y ∼ γ k ( x,t )
[ Y ] :
[Equation]
In our motivating clinical example, Assumption 2 reflects that the probability of specific treatment response changes continuously over the set of observed covariates and treatments.
In particular, the expected 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 mathematical definition is given in Appendix C. Intuitively, it supposes that patients with the same observed covariates 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-group variability in the treatment response for a particular patient and treatment outcome.
Methods 3.1 Identifiability and Counterfactuals Reconstruction
For a fixed point ( X = x, T = t ) , Equation 2 is a finite mixture model, for which identifiability results are available [3].
Notably, these results guarantee identifiability 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 identifiability of the counterfactuals in the sense of definition 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 specific point ( X = x, T = t ) to predict the counterfactual at another point ( X = x, T = t ′ ) - with a different treatment assignment.
Nevertheless, using the assumptions from the previous section, we can still ensure the identifiability of the counterfactuals as the following result confirms:
Result 3.1 (Identifiabilty of Counterfactuals with Categorical Background Variables) .
Let X , T and
Y be continuous random variables generated according to the graphical model of Figure1b with the domain of X and T being connected.
Let W 1 ( · , · ) be the first Wasserstein distance on P ( Y ) , ν t ′ ( X, Y, T ) the probability distribution of Y t ′ | X, Y, T and ˆ ν Nt ′
( X, Y, T ) its estimator from N observed data points.
If Assumptions 1, 2 and 3 hold, for each ( x, t ) , the counterfactual distribution is W 1 -identifiable 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 inferred and true counterfactual distributions in the asymptotic regime.
Importantly, it does not restrict the dimension 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 number of available samples.
In the additive Gaussian case, the sample complexity for learning a K -mixture model with Y ∈ R
d within ϵ total variation distance is ˜ O
( Kd 2 /ϵ )
Fortunately, the continuity assumption (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 µ ′ r ( X, T ) .
A better sample complexity bound can then be derived with further assumptions on µ ′ r ( X, T ) .
3.2 CFQP : CounterFactual Query Prediction
Equipped with those theoretical results, we introduce CFQP, a counterfactual prediction model based on a neural Expectation-Maximization mechanism.
The basic building block of CFQP is a base-model m ( x, t ) , that predicts the treatment response y based on covariates and treatment assignment.
For each latent category k , we learn a base-model that approximates the individual treatment response in that category : m k ( x, t ) ≈ µ k ( x, t ) .
Our theoretical results require the number of classes of K to be known in advance.
Yet, this is rarely the case in practice, and we describe our architecture for an arbitrary number of classes K ′ .
The learning of the base-models follows three steps: a joint initialization, 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.
[Figure]
Initialization We first 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 the conditional 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 the data from the corresponding latent category.
We distinguish between an initial clustering and an update clustering stage.
The initial clustering assignment happens after the initialization step has converged.
We then use the clusterability assumption (Assumption 3) to drive the assignment and use K ′ -means clustering based on the residuals ∥ y − m 0 ( x, t ) ∥ 22 .
The update clustering step happens at regular intervals to improve the quality of the cluster assignment as the performance of each individual base-model increases.
A point ( x, y, t ) is assigned to cluster k = arg
[Equation]
Maximization - Learning individual base-models
Once the dataset D =
( X , T , Y ) is clustered in K ′ groups, we train the individual models m k on the corresponding cluster.
This corresponds to the maximization step of the expectation-maximization scheme.
Every ∆ epochs, we update the cluster assignment until convergence.
Counterfactual prediction For a data point ( x, t, y ) , we first infer the cluster assignment
k = arg min j ∥
[Equation]
m j
[Equation]
x, t ) ∥ 22 .
We predict the counterfactual for treatment t ′ as ˆ y ′ =
m k ( x, t ′ ) .
Related Work
Causal perspectives in machine learning have gained significant traction in the past years [24, 23].
Among them, one distinguishes between causal discovery approaches, aiming at discovering the causal 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 counterfactuals is to posit a specific structural causal model [18] or its specific functional form [5], therefore ensuring identifiability.
Synthetic controls are an example of this strategy, assuming an underlying linear structure [1].
However, these approaches are by definition restrictive in terms of the expressivity of the structural equations.
Deep learning approaches for counterfactual estimation have been recently proposed to address this issue [15, 22] but without identifiability guarantee.
For the discrete case, assumptions have been proposed to bridge this gap, such as monotonicity or generalization thereof [14, 11].
Nevertheless, because they focus on addressing the discrete case, they are not directly applicable to high-dimensional data such as time series, which is a motivation for our work.
Data: X , Y , T , the number of latent clusters K ′ , a number of epochs e max , the update-period ∆ .
Result:
A list of K ′ hidden models m k .
Initialize a single base model m 0 at random.
for epoch ← 0 to e max do
Train m minimizing end
Compute residuals r = ( m
, T
− Y
Assign the residuals of each training sample into K ′ 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 cluster.
end
Compute
Assign the residuals of each training sample into K ′ clusters (update).
endReturn trained models m k .
Experiments 5.1 Datasets
Treatment outcomes are often complex and high-dimensional.
We thus evaluate our proposed model architecture on high-dimensional data modalities: time series and images.
As counterfactual inference evaluation on real-world data is a complex and ongoing research area, we use synthetic datasets inspired 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 modification of MNIST images.
Covariates X are original MNIST images and the treatment T is a spatial rotation applied to the image.
U Z is a coloring of the image and the treatment outcome Y is the resulting colored and rotated MNIST image.
U η is either an additive Gaussian noise 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 treatment assignment.
We used two different datasets.
A Harmonic Oscillator and a Cardiovascular model.
Harmonic oscillator We simulate the angular positions of two coupled harmonic oscillators.
The treatment consists of applying different gradual offsets to the time series.
We consider three hidden groups ( K = 3 ) that modulate the treatment response differently.
Cardiovascular simulator We use a simulator of the cardiovascular system proposed in [27] to predict the impact of fluid intake on blood pressure.
Fluids are commonly administered in intensive care for treating severe hypotension.
However, the individual patient response is difficult to assess a priori, as it depends on clinically hidden variables.
We consider here two hidden groups of patients with distinct responses to fluid intake.
X and Y contain the arterial and venous blood pressure time series before and after fluid intake.
The treatment is the amount of fluid injected.
We add process noise by either considering an additive Gaussian noise on the observed treatment response or by introducing noise in the fluid 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 deep counterfactual inference architecture building on normalizing flows ; Synthetic Controls ( SC ) [1], a well known linear method for counterfactuals estimation and Deep-ITE , a deep individual treatment effect estimator such as proposed in [25, 10], using direct predictions from X and T .
[Table]
5.3 Counterfactuals Prediction
methods (best seen in color).
Here, U Z corresponds to the color of the digit and T to the rotation angle.
X is the non-rotated and non-colored image while
Y T is the colored and rotated image.
First column is the observed factual treatment outcome.
Subsequent columns show the counterfactual reconstructions for different values of T ′ .
First row is the ground truth.
We evaluate our approach in terms of the quality of the counterfactual reconstructions.
We generate counterfactuals ( t ′ , y ′ ) for each observation ( x , y , t ) in the test set, resulting in a tuple ( x , y , t , t ′ , y ′ ).
We evaluate the MSE between the true counterfactual y ′ and the estimated one ˆ y ′
as
MSE cf = ∥
y ′ − ˆ y ′
∥
We set the number of clusters K as an hyper-parameter and select the value that result in lowest validation error.
The results are presented in Table 1.
We use additive and nonadditive noise variants for each dataset.
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.
Each row corresponds to counterfactual estimations for a particular method (top is the truth), and each column represents a different treatment assignment
T ′ (the left column is the factual T ).
Because it cannot model the categorical latent variable, the
Deep-ITE model gets the expected individual treatment effect right but fails to accurately estimate the counterfactuals.
For Diff-SCM, we observe that, despite providing an accurate reconstruction for most of the treatment assignments, the lack of guarantees on the deep structural equation model leads to potential incorrect reconstructions ( e.g. at T ′ = 30 ).
To assess the quality of the counterfactual images reconstructions, we also evaluate the structural similarity index and report the results in Appendix A.1.1.
We observe that
CFQP outperforms all baselines on that metric.
5.4 Robustness Analysis
We further complement our experimental section with a comprehensive robustness analysis.
We first investigate the impact of misspecifying the number of latent classes ( i.e. when K ̸ = K ′ ).
We then investigate the impact of a correlation between X and Z with increasing strength.
Finally, we investigate the robustness of our approach to different noise responses.
5.4.1 Number of latent classes
We study the evolution of performance when the number of classes K ′ is misspecified ( K ̸ = K ′ ).
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 ′ on the image dataset with true number of latent classes K = 6 .
The validation MSE is lowest for K ′
=
6 and K ′ = 7 , which corresponds to optimal test reconstruction error.
This hyper-parameter can thus easily be tuned by monitoring the treatment prediction performance on the factual data.
In Appendix B.2, we further report the quantitative results for the other datasets in function of the number of the clusters K .
5.4.2 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, we study the impact of the correlation strength ρ between X and U Z .
More details about the definition of ρ are to be found in Appendix B.
We compare the performance of CFQP with a Deep-ITE model.
We observe that the counterfactual performance of our model remains constant, regardless of the correlation strength.
However, as the correlation becomes more important, U Z becomes a deterministic outcome of X and Deep-ITE eventually converges to the performance of CFQP.
[Figure]
5.4.3 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 models under non-additive noise responses.
We observe an overall small degradation of performance of all methods but CFQP still outperforms baselines, as predicted by our identifiability results.
5.5 Retrospective Individualized Treatment Effect Estimation
We explore the ability of our model to infer individualized treatment effects retrospectively.
Based on observed treatment outcomes, we predict the difference in treatment response between two treatment regimes for a single individual.
The metric of interest is the Precision in Estimation of Heterogeneous
Effects PEHE [9].
For two treatment regimes t ′ and t ′′ , we write
[Equation]
In Table 2, we report PEHE ( t ′ = 0 .
5 , t ′′ = 0 .
8) for test patients in the Cardiovascular dataset with different fluid 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).
5.5.1 Impact of the clustering algorithm
Our approach relies on an initial clustering step as described in Section 3.2.
Different clustering strategies can be considered for this step.
In Appendix B.5, we investigate the difference between
K-means and Gaussian mixture models.
Only minor deviations are observed in terms of performance between both approaches.
Conclusion
Estimating counterfactuals from observational studies is one of the most challenging tasks in causal inference.
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 favorable experimental 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 assumptions proposed here is not unique, and depending on the specific applications, others might be deemed more relevant.
Indeed, the space of assumptions that allow to approximately recover counterfactuals with a controlled level of error is a potentially very rich research direction.
Other model architectures compatible with our set of assumptions are also possible.
Potential negative societal impacts Counterfactual inference has been used in clinical settings and for assessing the impact of public policies [2], among others.
However, as shown in this paper, its correctness hinges on assumptions whose validity is not always met, potentially leading to misleading conclusions.
Acknowledgements Edward is funded by a FWO-SB grant.
Edward would also like to thank
Zeshan Hussain, David Sontag and Adam Arany for their comments.
A.1 Image Dataset Generation
The image dataset uses the MNIST handwritten digits images dataset.
We use the following data generating process:
• X ∈ { 0 , ..., 255 } 28 × 28 × 1 is an MNIST image sampled at random from the dataset.
We write y label the label of the digit present in the image.
• T ∈ R =
U (0 , 0 . 3) + 5 · σ
[Figure]
with σ ( · ) the sigmoid function.
This encodes confounding between the observables X and the treatment assignment T .
• p
= 1 K • p ∈ [0 , 1] K • ρ ∈ [0 , 1] is the strength of confounding between X and U z • p
[ i ] = 1 − ((1 − p 0 )
ρ + p 0 )
K −∀ i ∈
[ K ] , i ̸
=
y label mod K • p [ i ] =
((1 − p 0 )
ρ + p 0 )
∀ i ∈
[ K ] , i =
y label mod ( K ) •
U z ∼ M ultinomial ( p ) • Y ∈ { 0 , ..., 255 } 28 × 28 × 1 is the rotated image X with angle T and Gaussian blur with kernel size 5 × and standard deviation σ .
Examples of generated images Y are shown in Figure 5.
We use the original MNIST training set that we randomly divide as 70% training set, 15% validation set and 15% test set.
[Figure]
A.1.1 Evaluating the structural similarity index measure of the counterfactual images predictions
Mean square error between the ground truth and predictions might not always be the most meaningful way to assess the quality of the reconstruction.
A widely used metric in computer vision is the structural similarity index measure (SSIM).
In Table 3, we report the SSIM between the predicted and ground truth images counterfactuals.
We observe that CFQPoutperforms all baselines in that metric, producing more accurate counterfactual predictions.
[Table]
A.2 Harmonic Oscillator
We consider input times series X ∈ R d x × t x , with d x = 2 the number of temporal dimensions and t x the length of the input time series.
The temporal responses are given by Y ∈ R d y × t
y .
We use the following data generating process :
[Equation]
+ ∆ 1 ( T, t y )]
+ η y where ∆ 0 ( T ) and ∆ 1 ( T ) are defined as follows:
[Equation]
t p = 3 characterizes the time constant of the treatment response dynamics.
The latent categorical variables U Z represent hidden patient groups that drive the treatment response.
In the non-additive noise-response case, we set :
[Equation]
such that the phase ϕ is time-varying and stochastic.
Examples of the trajectories and reconstructed counterfactuals by CFQPare provided in Figure 6.
A.3 Cardiovascular Dataset
We use an ODE model of the cardiovascular system as proposed in [27, 6].
Fluid intake is commonly used for treating severe hypotension.
However, the response of patient to fluids intake is difficult to assess beforehand.
In particular, it depends on the patients cardiac contractility factor and the blood pressure at time of injection.
If blood pressure is commonly and easily measured in standard clinical practice, assessing the cardiac contractility level of a patient requires imaging techinques such as echocardiography to measure the stroke volume.
Yet, the injection of significant volume of fluids in an irreponsive patients cardiovascular system can lead to severe damage.
This lead some clinicians to advocate for fluid challenges, or limited amount of fluid injection to test the responsiveness.
This technique is still contested in the medical community and legs raising challenge, much less damaging but also less effective at assessing a patients response has been encouraged.
This lack of availability of clear guidelines for fluids intake makes it a perfect case study for counterfactual prediction.
Indeed, we’ll try to address the question of if a clinician should administer fluids to particular patient based on his clinical history and therefore help informing clinical practice.
The system of ODE used to generate the data is the following :
[Equation]
[Figure]
where
[Equation]
In the above dynamical system
[Equation]
P a
[Equation]
P v , S and SV stand for arterial blood pressure, venous blood pressure, autonomic baroreflex tone and cardiac stroke volume respectively.
I external ( t ) is the amount of fluids given the patient over time and corresponds to the exogeneous input u T ( t ) in our model.
In the data generation
[Equation]
we model it as
[Equation]
[Figure]
[Equation]
where the treatment assignment T is generated as T ∼ U (0 . 6 , 1) and the hidden group assignment
U Z ∼ B
er
and t treat is the time of treatment.
The function f introduces confounding in the treatment assignment by setting :
[Equation]
We simulate the above system of ODEs for t span = 40 seconds and sample an observation every ∆ t = 1 second.
The treatment assignment time is set to t treat = 20 .
We then set X as the first two dimensions of the dynamical system ( ( P a , P v ) for t < t treat and Y as ( ( P a , P v ) for t ≥ t treat .
An example of a generated time series is given in Figure 7.
A.3.1 Noise Responses
Additive noise In the additive noise case, we add U η ∼ N (0 , σ 2 ) t span on both X and Y
Non-additive noise In the non additive case, we modify the treatment assignment by setting:
I external ( t ) =
(1 + 2 U z + U η ( t ))
· T · 5 · f ( P a ( t = 0)) ·
e − ( t − t treat − 5)
Where U eta ( t ) is a Gaussian process with zero-mean and diagonal covariance function with variance σ 2 .
B Experiments Details B.1 Computational Resources
We ran our experiments on cluster containing two types of GPUs : NVIDIA Titan Xp and NVIDIA
Quadro GV100.
Our experiments resulted in 7 GPU-days of computation.
B.1.1 Hyper-parameters
In Table 4, we report the hyper-parameters used in our experiments.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
B.2 Number of groups
As discussed in section 5.4, the number of groups is an important hyper-parameter in our model.
In Table 5,we report the validation reconstruction loss and the corresponding test counterfactual reconstruction MSE for all datasets and different values of K .
We observe that we obtain a lower validation error for the true value of K in all datasets.
In Figures 8a and 8a, we visualize the validation and test performance graphically for the Harmonic Oscillator dataset.
We do the same on Figures 9a and 9b for the Cardiovascular dataset and on Figures 10a and 10b for the colored MNIST dataset.
B.3 Strength of Correlation Experiment
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, we show the impact of the correlation strength ρ between X and U Z .
We incorporate correlation between X and U z as follows.
The base probability p k 0 of each U z is uniform.
That is,
[Equation]
We modify this probabilty based on the label of the image (the digit label y label ∈ [9] .
For a correlation factor ρ ∈ [0 , 1] , we write
[Table]
± 0 .0 .
00198 ± 0 .
MNIST
Non
[Table]
Additive 0 .
00212 ± 0 .0 .
00217 ± 0 .
Harmonic Oscillator
Additive 0 .
1653 ± 0 .0 .
16584 ± 0 .
Harmonic Oscillator
Additive 0 .
23024 ± 0 .0 .
08058 ± 0 .
Harmonic Oscillator
Additive 9e − 05 ± 3e −0 .
01332 ± 0 .
Harmonic Oscillator
Additive 0 .
00013 ± 1 e −0 .
0199 ± 0 .
Harmonic Oscillator
Additive 0 .
00035 ± 0 .0 .
04017 ± 0 .
Harmonic Oscillator
Non
[Table]
17451 ± 0 .0 .
17459 ± 0 .
Harmonic
[Table]
Non
[Table]
15164 ± 0 .0 .
09669 ± 0 .
Harmonic
[Table]
00891 ± 0 .0 .
00956 ± 0 .
Harmonic
[Table] [Table]
± 0 .0 .
02279 ± 0 .
Harmonic
[Table]
00984 ± 0 .0 .
03464 ± 0 .
CardioVascular
Additive 1 . 01453 ± 0 .1 .
02569 ± 0 .
CardioVascular
Additive 0 .
03677 ± 0 .0 .
07701 ± 0 .
CardioVascular
Additive
[Table]
± 0 .0 .
11363 ± 0 .
CardioVascular
Non
[Table]
1 . 03865 ± 0 .1 .
09768 ± 0 .
[Table]
[Equation]
The above equation suggest that some images with specific label classes will have higher probability of sampling a latent class U Z =
u
z such that y labelmod ( K ) =
u z .
In the case when ρ = 0 , p k = p k 0 = 1 K , corresponding to a uniform probability of latent class for each class label.
When ρ = 1 , p y labelmod ( K ) =
and the other probabilities are set to 0 , corresponding to a deterministic assignment of U Z conditioned on y label .
B.4 Scalability
CFQPscales linearly with the number of groups K , though it can be easily parallelized.
Indeed, each sub-model m i can be evaluated in parallel.
Regarding scaling with respect to the dimensions
[Figure]
[Figure] [Figure]
[Figure] [Figure]
of the outcomes and treatments
[Figure]
we can use arbitrary neural networks architectures to process them, for which the complexity can vary.
In our experiments, we used CNNs for images, which scale quadratically with the size of the images and TemporalCNNs for time series which scale linearly with the length of the time series as well as number dimensions.
B.5 Clustering Algorithm
Our approach relies on an initial clustering step as described in Section 3.2.
In our main experiments, we use a k-means algorithm for this step.
However, different clustering strategies can be considered.
In Table 6, we report the performance metrics of a variant of our approach using Gaussian Mixture
Models (GMM).
For computational reasons, the GMM is fit by using a random sub-sample of the training data (1000 samples).
We observe only minor performance deviations between k-means and
GMM versions of CFQP.
[Figure]
[Table]
C Proof of Result 3.
Our proof is structured as follows.
We first show identifiability of non-parameteric mixtures at a single point ( ( X = x, T = t ) .
We then show identifiability of the permutation function over ( X ×T ) .
Finally, based on these results, we derive the bound of Result 3.1.
C.1 Identifiability of Non-Parametric Mixture Models.
To show the identifiability in the non-additive case, we re-use notation and start from the outline derived in Aragam et al.
C.1.1 Identifiability Result
We let ( X, d ) be a metric space and ( P ( X ) , ρ ) the space of regular Borel probability measures on X with finite r th moments metrized by a metric ρ .
Let’s further define P 2 ( X )
=
P ( P ( X )) , the space of mixing measures over P ( X ) .
We represent a mixture distribution over X as a new probability measure m ( · , Λ) ∈ P ( X ) such that m ( A ; Λ)
=
´ γ ( A ) d Λ( γ ) for any A ⊂ X .
Given a Borel set L ⊂ P 2 ( X ) , we define M ( L ) :
= { m (Λ) : Λ ∈ L ) as a family of mixture distributions over X .
Because we are interested in finite mixtures, we use P 2 k := { Λ ∈ P 2 ( X ) : ∥ supp (Λ) ∥ ≤ k } to denote the set of mixing measures with k components.
Finite mixtures Γ ∈
M ( P 2 k ( X )) can thus also be written in a more intuitive form,
i.e. Γ = (cid:80) Kk =1 ω k γ k where γ k ∈ P ( X ) are the mixture components and ω k are mixing weights, as in Equation 2 in the main text.
Note that we did not make any assumption regarding the distribution of the mixture components ( γ k ∈ P ( X ) ).
Now, from [3], we have the following result.
Theorem C.1.
If L is a Q L -clusterable family, then there exists a function h : M ( L )
→ L such that h ( m (Λ)) = Λ , where m : L → M ( L ) is the canonical embedding.
In particular, m is a bijection and the mixture model M ( L ) is identifiable.
We see that identifiability of the mixture therefore hinges on the notion of clusterability of the individual distributions that compose the mixture.
We now make the notion of clusterability more precise.
C.1.2 Clusterability
Projections Let’s first introduce { Q L } ∞ L =1 as an indexed collection of families of mixing measures that satisfy the following requirements:
• Q L ⊂ P 2 L ( X ) for each L ; • { Q L } is a filtration ( Q L ⊂ Q L +1 ); • The collection of mixture distributions M ( Q L ) is identifiable for each L .
Examples of such a collection is the filtration of Gaussian mixtures with L components.
We then proceed by defining the ρ -projection of a mixture distribution Γ = m (Λ) on a family of mixture distribution Q L by
[Equation]
where ρ is a metric on P ( X ) ( e.g. the Hellinger distance).
This projection intuitively maps a mixture distribution Γ to the closest element of the family Q L it is projected upon.
We futher define the map
M L : Q L → Q L that a projected distribution to its mixing measure:
[Equation]
Assignment Functions
Because it is an element of Q L , the projection T L Γ =
Q ∗ is a mixture distribution and can thus be written as T L Γ =
(cid:80) Ll =1 ω ∗ l γ ∗ l .
When projecting on mixture distributions with more mixture components than in the original distribution, i.e. L ≥ K , one can then assign mixture components in Q ∗ to mixture components in Γ .
We define the set of assignment maps α : [ L ] → [ K ] as A L → K .
Given such an assignment map, we define ψ k ( α ) :
= (cid:80) l ∈ α − 1( k ) ω
l for the weights induced by the assignment map and Q k ( α ) :
= 1 ψ k ( α ) (cid:80) l ∈ α − 1 ( k ) ω
∗ l γ ∗ l for the induced mixing components.
Regularity of a mixing measure We are now ready for introducing the definition of regularity of a mixing measure.
Definition C.1 (Regularity) .
Suppose Λ ∈ P 2 K ( X ) and Γ = m (Λ) ∈
M K ( P 2 ( X )) .
The mixing measure Λ is called Q L -regular if:
• (a) The ρ -projection Q ∗ =
T L Γ exists and is unique for each L and lim L
→∞ T L Γ = Γ ; • (b) There exists a sequence of assignment functions α = α L ∈
A L → K such that lim L →∞ Q ∗ k ( α ) =
γ k and lim L →∞ ψ ∗ k ( α ) =
ω k ∀ k
= 1 , . .
K
Any assignment function where (b) holds is called a regular assignment .
Clusterability We can now finally define clusterability of a family of mixing measures with respect to another regular family of mixing measures.
Definition C.2 (Clusterable family) .
A family of mixing measures L ⊂ P 2 ( X ) is called a Q Lclusterable family, or just a clusterable family, if
• (a) Λ is Q L -regular for all Λ ∈ L ;
• (b) For all sufficiently large L , there exists a function χ
L : M L ( L ) →
A L → K such that χ L (Ω ∗ ) is a regular assignment for every Λ ∈ L , with Ω ∗ =
M L ( T L Γ) .
The resulting mixture model M ( L ) is called a clusterable mixture model.
More intuitively, a family of mixing measures will then be Q L -clusterable if one can project every mixing measure of the family onto Q L and if, for each of these projections, there exists a function that outputs a regular assignment for each mixture distribution given as input.
This regular assignment is such that it clusters the elements of Q ∗ in a way that the distribution of each cluster converges to the distribution of a element of the original mixture distribution.
In this case, one can identify a family of mixing measures by projecting each mixing distribution onto a regular family ( e.g. Gaussian mixtures) and cluster it accordingly to χ L .
We refer the interested reader to [3] for more insight on clusterability and identifiability in non-parametric mixtures.
C.2 Proof of result 3. C.2.1 Identifiability at a Fixed Point
We start our proof by recalling the identifiability result in the case of clusterable families of mixing measures, whose proof can be found in [3].
Theorem C.2.
If L is a Q L -clusterable family, then there exists a function h : M ( L )
→ L such that h ( m (Λ)) = Λ , where m : L → M ( L ) is the canonical embedding.
In particular, m is a bijection and the mixture model M ( L ) is identifiable.
As laid out in the main text, we assume three observed random variables X ∈ X , Y ∈ Y and T ∈ T .
As suggested by Assumption 1 in the main text, the conditional distribution of Y at each point ( X = x, T = t ) is given by an unknown mixture distribution from a family M K ( L ) where each mixing measure Λ ∈ L satisfies Λ ∈ P 2 ( Y ) and | supp (Λ) | =
K .
We write the distribution of Y conditioned on X and T as:
[Equation]
Assuming that L is clusterable, we can use Theorem C.2 to deduce that the mixture distribution γ ( X, T ) is identifiable at each fixed point ( X = x, T = t ) .
Importantly, the idenfitifiablity insured by Theorem C.2 is up to a permutation of the mixture components.
We write σ x,t : [ K ] → [ K ] the permutation function at a specific point ( X = x, T = t ) .
We define the assignment function that maps the elements of a mixture distribution γ ( X, T ) to U Z as F ( x,t ) :
M ( L ) → { [ K ] , P ( Y ) } K .
We now move to showing that the assignment function is identifiable up to a constant permutation ¯ σ .
C.2.2 Identifiability of the Assignment Function
The support of random variables X and T , X × T is connected.
Assuming further (Assumption 3) that L is a clusterable family, we know that there exists a cluster function χ L, ( x,t ) : M L ( L ) →
A
L → K that maps components of the projected mixture onto the initial one.
We can then define a combined operator that take a mixture distribution as input and returns an indexed sequence of mixture components.
Let F ( x,t ) : M ( L ) → { [ K ] , P ( Y ) } K be this operator consisting of applying a ρ -projection onto a regular family, cluster the result according to χ L, ( x,t ) and return the indexed estimated distributions
Q ∗ k ( α ) .
As the per the definition of regular mixing measures, we have also have lim L →∞ Q ∗ k ( α ) =
γ k .
Note that the indexing of the elements of the mixing distribution is arbitrary at a single point ( X = x, T = t ) and thus defined up to a permutation ¯ σ .
Because X and T are continuous on a connected domain, we can evaluate the operator at ( X = x + δx, T = tδ t ) .
[Equation]
What is more, for arbitrary small ( δx, δt ) and ∀ ( x, t ) ∈
( X × T ); ∀ k, k ′ , k ′′ ∈
[ K ] , we have
[Equation]
because we require the moments of each γ k to be continuous in ( x, t ) .
It therefore implies a relation between σ x,y and σ x + δx,y + δy that satisfies:
[Equation]
The assignment function F ( x,t ) is thus identifiable up to a constant permutation ¯ σ that will determine all other permutations in X × T .
C.2.3 Counterfactual Identifiability
To recapitulate, in the previous sections, we have shown (1) identifiability of a non-parametric mixture at a point ( X = x, T = t ) based on the Assumption 3 and (2) identifiability of the permutation function over ( X × T ) .
That is, for each point in ( X = x, T, t ) , we can identify the mixture components and the mixture weights and associate it to any other point ( X = x ′ , T = ′ ) , through the assignment function F ( x,t ) .
In particular, we can identify the means µ k of each cluster component γ k .
We have
[Equation]
and
[Equation]
Because the above estimators converge, the following posterior probability converges to the true value as well:
[Equation]
for all x ∈ X , t ∈ T and y ∈ Y .
Indeed, P ( Y = y
| U Z = u Z , X = x, T = t ) is identifiable, and
P ( U Z = u Z , | X = x, T =
[Equation]
= ω u Z ( x, t
[Equation]
Let ˆ ω u z ( x, y, t ) be the estimator of P ( U Z =
u Z | Y
= y, X
= x, T = t )
We are now ready to define the estimator for the counterfactual distribution ν t ′
( x, y, t ) .
[Equation] [Equation]
This estimator is a discrete mixture distribution with mass on the means of the different mixture components, weighted by the posterior probability of the initial sample ( x, t, y ) belonging to a particular component.
C.2.4 Bounds
We can now use this estimator ν t ′ ( x, y, t ) to derive the bounds of Result 3.1.
[Equation]
Step 12 is obtained by conditional independence between y ′ and ν t ′ conditioned on X, T, Y .
Step 13 is obtained by law of total probabilities.
Step 13 is obtained by factorizing on U Z .
Note that
P ( Y ′ =
y ′
| X, T, U z
= u z ) corresponds to the mixture component γ k and the posterior assignment
P ( U Z =
u Z | X, T, Y ) corresponds to ˆ ω u Z
( X, Y, T ) , as per the definition above.
What is more, conditioned on U Z , ν t ′ ( X, T,
U z
= k ) reduces to a point estimate at ˆ µ k ( X, T ) .
We can then write
[Equation]
u
[Equation]
= δ
Where 15 holds because it 14 is a convex combination W k 1 (ˆ µ k , γ k ) .
This ends the proof of Result 3.1.
C.2.5 The additive noise case
Result 3.1 mentions the bounds reduces to zero δ = 0 when the noise is additive.
From identifiability at of a mixture at a single point, we had:
[Equation]
We also have convergence of the covariance matrix estimator
[Equation]
In the additive noise setup, the distribution of Y being normally distributed, the value of U η is actually identifiable.
Let us recall the definition of the counterfactual distribution (Eq. 1) in the main text.
[Equation]
In the additive case
[Equation]
we have Y | X
= x
[Equation]
T =
t
[Equation]
U Z =
u Z =
Σ u Z ( x, t ) U η + µ k
( x, t ) with
U η ∼ N (0 , 1 ) .
The counterfactual density can thus be written as
[Equation]
Considering the following estimator for the counterfactual distribution:
[Equation]
We can use Equations 16, 17 and 18 to show that our estimator converges to the true distribution in the limit of infinite number of samples:
[Equation]
Therefore, the distance between the estimated and true counterfactual distributions reduces to 0:
[Equation]
This concludes our proof of Result 3.1.