# Deep Counterfactual Estimation with Categorical Background Variables

Edward De Brouwer
ESAT-STADIUS
KU Leuven
<EMAIL_ADDRESS>

###### 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.

##  1 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
[23] and constitutes the second rung of the causation ladder [19].
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 [22]
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 [28]. The importance
of counterfactual reasoning in the human cognitive process has thus motivated
researchers to endow artificially intelligent systems with the same ability
[21]. However, counterfactual inference is not possible from observational and
interventional data alone [20].

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 [21] or a specific functional form thereof [1, 5, 21]. 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 [17, 24]. 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 can rely on
recent results in identifiability of mixture distributions [3] to 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
distributions.

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
[15].

##### 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.

##  2 Background

###  2.1 Problem Setup : Counterfactual Estimation

We consider the general causal model $M=\langle U,V,F\rangle$ depicted in
Figure 1(a) 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_{\epsilon},W\\}$ are hidden exogenous random variables
that determine the values of the observed variables $V=\\{X,T,Y\\}$.
Covariates $X\in\mathcal{X}$ represent the information available before
treatment assignment, $T\in\mathcal{T}$ is the treatment assignment and
$Y\in\mathcal{Y}$ is the observed response to the treatment. We refer to the
space of probability measures on $\mathcal{Y}$ as $\mathcal{P}(\mathcal{Y})$.
Observed variables $V$ are generated following the structural equations
$F=\\{f_{X},f_{T},f_{Y},f_{\epsilon}\\}$, such that $X=f_{X}(U_{x},W)$,
$T=f_{T}(U_{T},X)$, $U_{\epsilon}=f_{\epsilon}(W)$,
$Y=f_{Y}(X,T,U_{\epsilon})$. We further assume _strong ignorability_ (_i.e._
no hidden confounders between $T$ and $Y$).

Using notations introduced in Pearl et al. [21], 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^{\prime}$), conditioned on an observed
initial treatment response. That is, the probability of observing a different
treatment response under treatment $t^{\prime}$, after observing treatment
response $y$ for covariate $x$ and treatment $t$. The probability density
function of counterfactual $y^{\prime}$ then writes:

$\displaystyle p(Y_{t^{\prime}}=y^{\prime}\mid X=x,Y=y,T=t)$
$\displaystyle=\frac{p(Y_{t^{\prime}}=y^{\prime},X=x,T=t,Y=y)}{p(X=x,T=t,Y=y)}$
$\displaystyle=\int_{u}p(Y_{t^{\prime}}(u)=y^{\prime})p(U=u\mid X=x,T=t,Y=y),$
(1)

and we refer to the counterfactual probability measure as
$\nu_{t^{\prime}}(x,y,t)$. Equation 1 suggests a natural three step procedure
for computing the probability of counterfactual. First, the _abduction_ step
infers the density of $U$ conditioned on the observed treatment outcomes,
covariates and treatments: $p(U=u\mid X=x,T=t,Y=y)$. Second, in the _action_
step, one sets the new treatment in the causal model ($do(T=t^{\prime})$).
Lastly, in the _prediction_ step, one can propagate the values of $U=u$ and
$T=t^{\prime}$ in the causal graph, using $F$, to compute
$p(Y_{t^{\prime}}(u)=y^{\prime})$.

In practice, we only have access to a set of $N$ observations of variables
$X$, $Y$ and $T$. We refer to this dataset as
$\mathcal{D}=(\mathbf{X},\mathbf{T},\mathbf{Y})$ where
$\mathbf{X}=\\{x_{i}:i=1,...,N\\}$, $\mathbf{Y}=\\{y_{i}:i=1,...,n\\}$ and
$\mathbf{T}=\\{t_{i}:i=1,...,N\\}$. Importantly, we don’t have access to
counterfactual examples (_i.e._ a tuple $(x,y,t,t^{\prime},y^{\prime})$), such
that direclty learning a map $(x,y,t,t^{\prime})\rightarrow y_{t^{\prime}}$ 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 density
$p(X,Y,T)$ but would lead to incorrect counterfactual probabilities [19, 20].
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 $\rho$ be a metric on $\mathcal{P}(\mathcal{Y})$,
$\nu_{t^{\prime}}(X,Y,T)$ the true counterfactual probability measure and
$\hat{\nu}_{t^{\prime}}(X,Y,T)$ the estimator of the counterfactual
probability measure with $N$ data points. Counterfactuals are
$\rho$-identifiable at threshold $\delta$ if, for all
$t,t^{\prime}\in\mathcal{T},x\in\mathcal{X},y\in\mathcal{Y}$,

$\displaystyle\lim_{N\rightarrow\infty}\rho(\nu_{t^{\prime}}(x,y,y),\hat{\nu}_{t^{\prime}}(x,y,t))\leq\delta$

$Y$$X$$T$$U_{\epsilon}$$W$$U_{X}$$U_{T}$ (a) General Bayesian network for the
treatment counterfactual problem. $X$, $Y$ and $T$ are observed while
$U_{X},U_{T},W$ and $U_{\epsilon}$ are hidden background variables.

$Y$$X$$T$$U_{Z}$$U_{\eta}$$W$$U_{X}$$U_{T}$ (b) Bayesian Network embodying the
hidden categorical background variable assumption. $U_{\epsilon}$ is split in
a background categorical variable $U_{Z}$ and a continuous background variable
$U_{\eta}$.

Figure 1: Graphical model representations of the causal model $M$. We assume
strong ignorability and continuous treatments $T$, observables $X$ and
responses $Y$.

###  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
$\mathcal{X}\times\mathcal{T}$ is a connected space. Our first central
assumption posits that the hidden variable $U_{\epsilon}$ factorizes into a
categorical and a continuous variable.

######  Assumption 1 (Categorical Background Variable).

The background variable $U_{\epsilon}$ decomposes into a categorical latent
variable $U_{Z}\in[K]=\\{1,..,K\\}$ and an independent exogenous continuous
variable $U_{\eta}$.

This assumption is depicted in the graphical model of Figure 1(b) 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 density of $Y$ as a mixture model:

$\displaystyle p(Y=y\mid
X=x,T=t)=\sum_{u_{Z}\in\\{1,..,K\\}}P(U_{Z}=u_{Z})\cdot\int
p(U_{\eta}=u_{\eta})\mathbb{I}[f_{Y}(x,t,u_{Z},u_{\eta})=y]du_{\eta}$

We define $\gamma$ as the probability density function of the conditional
treatment response generated by $f_{Y}$, $U_{Z}$ and $U_{\eta}$. $\gamma$ is
thus a mixture probability density function with mixture components111The
mixture components are defined such that for any subset
$\mathcal{A}\subset\mathcal{Y}$, we have
$\int_{\mathcal{A}}{\gamma_{k}(X,T)(y)dy}=\int_{\infty}^{\infty}\mathbb{I}[f_{Y}(X,T,U_{Z}=k,U_{\eta})\in\mathcal{A}]dP(U_{\eta})$.
The mixutre weights are defined as $\alpha_{k}=P(U_{Z}=k)$.
$\gamma_{k}\in\mathcal{P}(\mathcal{Y})$ and weights $\omega_{k}$.

$\displaystyle Y\mid
X,T\sim\gamma(X,T)=\sum_{k=1}^{K}\omega_{k}\gamma_{k}(X,T)$ (2)

Without loss of generality, we assume that
$U_{\eta}\sim\mathcal{N}(\mathbf{0},\Sigma^{2})$. In the case of additive
noise, the conditional distribution of Y becomes a mixture of Gaussians:
$\gamma(X,T)=\sum_{k=1}^{K}\omega_{k}\mathcal{N}(\mu_{k}(X,T),\Sigma_{k}^{2})$,
where $\mu_{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 $\gamma_{k}(x,t)$ exist and
are continuous functions of $X$ and $T$:

$\displaystyle\mu^{r}_{k}(x,t)=\mathbb{E}_{Y\sim\gamma_{k}(x,t)}\left[Y^{r}\right]\in
C(x,t)\quad\forall r\in\mathbb{N},k\in[K]$ (3)

######  Assumption 3 (Clusterability).

For each $(x,t)\in(\mathcal{X},\mathcal{T})$, the density $\gamma(x,t)$ is
clusterable and the expected deviation of each $\gamma_{k}(x,t)$ is bounded by
a constant $\delta\in\mathbb{R}$. That is, $\forall{k}\in[K]\text{ , }\forall
x,t\in(\mathcal{X}\times\mathcal{T})$, with
$\mu_{k}(x,t)=\mathbb{E}_{Y\sim\gamma_{k}(x,t)}[Y]$:

$\displaystyle\mathbb{E}_{Y\sim\gamma_{k}(x,t)}\big{[}\lVert
Y-\mu_{k}(x,t)\rVert_{2}\big{]}\leq\delta$ (4)

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 [14].

Assumption 3 posits _clusterability_ of the mixture components $\gamma_{k}$
for which a rigorous mathematical definition is given in Appendix C. It is
motivated by recent results on the identifiability of mixtures models[3].
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
$\gamma_{k}$ that characterize the inter-group variability in the treatment
response for a particular patient and treatment outcome.

##  3 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
$\sigma(\cdot):[K]\rightarrow[K]$. That is, there exists some permutation
$\sigma(\cdot):[K]\rightarrow[K]$ such that
$\hat{\gamma}_{\sigma_{k}(k)(x,t)}\approx\gamma_{k}(x,t)$,
$\hat{\omega}_{\sigma_{k}(k)}\approx\omega_{k}$, where $\hat{\gamma}$ and
$\hat{\omega}$ 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 $\sigma$ across the whole domain
$(\mathcal{X}\times\mathcal{T})$ in order to reuse the inferred class
assignments $\hat{U}_{Z}$ at a specific point $(X=x,T=t)$ to predict the
counterfactual at another point $(X=x,T=t^{\prime})$ — 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 Figure1(b) with the domain of $X$ and $T$ being connected.
Let $W_{1}(\cdot,\cdot)$ be the first Wasserstein distance on
$\mathcal{P}(\mathcal{Y})$, $\nu_{t^{\prime}}(X,Y,T)$ the probability
distribution of $Y_{t^{\prime}}\mid X,Y,T$ and
$\hat{\nu}_{t^{\prime}}^{N}(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 in expectation at threshold $\delta$:

$\displaystyle\lim_{N\rightarrow\infty}\mathbb{E}_{Y\sim\gamma(x,t)}\left[W_{1}(\nu_{t^{\prime}}(x,Y,t),\hat{\nu}_{t^{\prime}}^{N}(x,Y,t))\right]\leq\delta$

In the special case when the noise response is additive, we have

$\displaystyle\lim_{N\rightarrow\infty}W_{1}(\nu_{t^{\prime}}(X,Y,T),\hat{\nu}_{t^{\prime}}^{N}(X,Y,T))=0$

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 $\mathcal{X}$ and
$\mathcal{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\in\mathbb{R}^{d}$ within $\epsilon$ total variation distance is
$\tilde{O}(Kd^{2}/\epsilon)$ [4]. 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
$\mu^{\prime}_{r}(X,T)$. A better sample complexity bound can then be derived
with further assumptions on $\mu^{\prime}_{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)\approx\mu_{k}(x,t)$. Our
theoretical results require the true number of classes of $K_{0}$ 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 2: Training procedure of CFQP. We first initialize a common model
$m_{0}$. Second, we cluster the different available data points into $K$
clusters. We then use this clustering to train $K$ individual models
(Maximization step). The clustering is updated every $\Delta$ epochs
(Expectation step).

##### Initialization

We first train a single common base-model $m_{0}$ on all data points
$\mathbf{X},\mathbf{T}$ and $\mathbf{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)\approx\mathbb{E}[Y\mid 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 $\lVert
y-m_{0}(x,t)\rVert_{2}^{2}$.

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=\underset{j}{\arg\min}\lVert y-m_{j}(x,t)\rVert_{2}^{2}$.

##### Maximization - Learning individual base-models

Once the dataset $\mathcal{D}=(\mathbf{X},\mathbf{T},\mathbf{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 $\Delta$ epochs, we update the cluster assignment
until convergence.

##### Counterfactual prediction

For a data point $(x,t,y)$, we first infer the cluster assignment
$k=\underset{j}{\arg\min}\lVert y-m_{j}(x,t)\rVert_{2}^{2}$. We predict the
counterfactual for treatment $t^{\prime}$ as
$\hat{y^{\prime}}=m_{k}(x,t^{\prime})$.

Data: $\mathbf{X}$,$\mathbf{Y}$,$\mathbf{T}$, the number of latent clusters
$K$, a number of epochs $e_{max}$, the update-period $\Delta$.

Result: A list of $K$ hidden models $m_{k}$.

Initialize a single base model $m_{0}$ at random.

for  _epoch $\ \leftarrow 0$ to $e_{max}$_ do

Train $m_{0}$ minimizing
$\mathcal{L}_{0}=\mathbb{E}\left[(m_{0}(X,T)-Y)^{2}\right]$

end for

Compute residuals $\mathbf{r}=(m_{0}(\mathbf{X},\mathbf{T})-\mathbf{Y})$.

Assign the residuals of each training sample into $K^{\prime}$ clusters
(initial).

Initialize each $m_{k}$ with $m_{0}$.

for  _epoch $\ \leftarrow 0$ to $e_{max}$_ do

for  _$i\leftarrow 0$ to $\Delta$_ do

Train $m_{k}$ minimizing
$\mathcal{L}_{k}=\mathbb{E}\left[(m_{k}(X_{k},T_{k})-Y_{k})^{2}\right]$ for
each cluster.

end for

Compute residuals
$\mathcal{R}=\\{(m_{k}\left(\mathbf{X},\mathbf{T})-\mathbf{Y}\right):k=1,...,K\\}$.

Assign the residuals of each training sample into $K$ clusters (update).

end for

Return trained models $m_{k}$.

Algorithm 1 CFQP Training

##  4 Related Work

Causal perspectives in machine learning have gained significant traction in
the past years [26, 25]. Among them, one distinguishes between causal
_discovery_ approaches, aiming at discovering the causal relations between
variables [10, 9], and causal _inference_ , aiming at building treatment
effects estimators from data [18]. Our work belongs to the latter. A common
way to estimate counterfactuals is to posit a specific structural causal model
[20] 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 [17, 24] but without identifiability guarantee. For the
discrete case, assumptions have been proposed to bridge this gap, such as
monotonicity or generalization thereof [16, 13]. 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.

Our work builds upon the literature on the identifiability of finite mixture
distributions where guarantees have been achieved in the Gaussian [8] and non-
parametric cases [3], among others.

##  5 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 [17], 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_{\eta}$ 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: an 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_{0}=3$) that modulate the treatment
response differently.

##### Cardiovascular simulator.

We use a simulator of the cardiovascular system proposed in [29] 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_ [24], a recently introduced diffusion based
counterfactual model; _Deep-SCM_ [17], 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 [27, 12], using
direct predictions from $X$ and $T$.

Table 1: Test MSE of the counterfactual reconstructions for the different
datasets.

Additive Noise Non-Additive Noise Model Harmonic Oscillator colored-MNIST
Cardiovascular Harmonic Oscillator colored-MNIST Cardiovascular Deep-ITE [12]
$0.187\pm 0.006$ $0.017\pm 0.001$ $1.084\pm 0.087$ $0.174\pm 0.004$ $0.017\pm
0.001$ $1.14\pm 0.121$ SC [1] $0.177\pm 0.131$ $0.020\pm 0.001$ $1.610\pm
0.141$ $0.167\pm 0.004$ $0.020\pm 0.001$ $1.628\pm 0.144$ Deep-SCM [17]
$0.124\pm 0.005$ $0.011\pm 0.001$ $0.405\pm 0.042$ $0.123\pm 0.005$ $0.011\pm
0.001$ $0.424\pm 0.042$ Diff-SCM [24] $0.082\pm 0.023$ $0.008\pm 0.004$
$0.206\pm 0.036$ $0.106\pm 0.038$ $0.009\pm 0.002$ $0.311\pm 0.073$ CFQP
(ours) $\mathbf{0.013\pm 0.001}$ $\mathbf{0.001\pm 0.001}$ $\mathbf{0.077\pm
0.050}$ $\mathbf{0.009\pm 0.001}$ $\mathbf{0.002\pm 0.001}$ $\mathbf{0.188\pm
0.114}$

###  5.3 Counterfactuals Prediction

Figure 3: Example of reconstructed counterfactuals on the image dataset for
different 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^{\prime}$. First row
is the ground truth.

We evaluate our approach in terms of the quality of the counterfactual
reconstructions. We generate counterfactuals ($t^{\prime}$,$y^{\prime}$) for
each observation ($x$,$y$,$t$) in the test set, resulting in a tuple
($x$,$y$,$t$,$t^{\prime}$,$y^{\prime}$). We evaluate the MSE between the true
counterfactual $y^{\prime}$ and the estimated one $\hat{y}^{\prime}$ as
$\mbox{MSE}_{cf}=\lVert y^{\prime}-\hat{y}^{\prime}\rVert_{2}^{2}$. 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 non-additive 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^{\prime}$ (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^{\prime}=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_{0}\neq 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_{0}\neq K$). In Figure 4(a), 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_{0}=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 4(b), we study the impact of
the correlation strength $\rho$ between $X$ and $U_{Z}$. More details about
the definition of $\rho$ 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.

(a) Reconstruction MSE in function of the number of groups. The true number of
groups is 6.

(b) Evolution of the reconstruction MSE in function of the correlation
strength between $X$ and $U_{Z}$.

Figure 4: Robustness analysis for the number of latent classes (left) and
correlation strength (right).

####  5.4.3 Non-additive Noise Responses

Table 2: Fluids PEHE

Model Cardiovascular PEHE $\downarrow$ Deep-ITE [12] $0.594\pm 0.160$ SC [1]
$0.258\pm 0.016$ Diff-SCM [24] $0.377\pm 0.049$ CFQP (ours) $\mathbf{0.143\pm
0.108}$

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 [11]. For two treatment regimes $t^{\prime}$ and
$t^{\prime\prime}$, we write

$\displaystyle\mbox{PEHE}(t^{\prime},t^{\prime\prime})=\sqrt{\mathbb{E}_{X,T,Y}\left[\left((Y_{t^{\prime\prime}}-Y_{t^{\prime}})-(\hat{Y}_{t^{\prime\prime}}-\hat{Y}_{t^{\prime}})\right)^{2}\mid
X,Y,T\right]}.$ (5)

In Table 2, we report $\mbox{PEHE}(t^{\prime}=0.5,t^{\prime\prime}=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.

##  6 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 on preliminary
versions of this manuscript.

## References

  * Abadie [2021] A. Abadie.  Using synthetic controls: Feasibility, data requirements, and methodological aspects.  _Journal of Economic Literature_ , 59(2):391–425, 2021.
  * Abadie et al. [2010] A. Abadie, A. Diamond, and J. Hainmueller.  Synthetic control methods for comparative case studies: Estimating the effect of california’s tobacco control program.  _Journal of the American statistical Association_ , 105(490):493–505, 2010.
  * Aragam et al. [2020] B. Aragam, C. Dan, E. P. Xing, and P. Ravikumar.  Identifiability of nonparametric mixture models and bayes optimal clustering.  _The Annals of Statistics_ , 48(4):2277–2302, 2020.
  * Ashtiani et al. [2018] H. Ashtiani, S. Ben-David, N. Harvey, C. Liaw, A. Mehrabian, and Y. Plan.  Nearly tight sample complexity bounds for learning mixtures of gaussians via sample compression schemes.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Balke and Pearl [2013] A. Balke and J. Pearl.  Counterfactuals and policy analysis in structural models.  _arXiv preprint arXiv:1302.4929_ , 2013.
  * Brouwer et al. [2022] E. D. Brouwer, J. González, and S. Hyland.  Predicting the impact of treatments over time with uncertainty aware neural differential equations.  In _AISTATS 2022_ , February 2022.
  * Byrne [2016] R. M. Byrne.  Counterfactual thought.  _Annual review of psychology_ , 67:135–157, 2016.
  * Dasgupta [1999] S. Dasgupta.  Learning mixtures of gaussians.  In _40th Annual Symposium on Foundations of Computer Science (Cat. No. 99CB37039)_ , pages 634–644. IEEE, 1999.
  * De Brouwer et al. [2020] E. De Brouwer, A. Arany, J. Simm, and Y. Moreau.  Latent convergent cross mapping.  In _International Conference on Learning Representations_ , 2020.
  * Glymour et al. [2019] C. Glymour, K. Zhang, and P. Spirtes.  Review of causal discovery methods based on graphical models.  _Frontiers in genetics_ , 10:524, 2019.
  * Hill [2011] J. L. Hill.  Bayesian nonparametric modeling for causal inference.  _Journal of Computational and Graphical Statistics_ , 20(1):217–240, 2011.
  * Johansson et al. [2016] F. Johansson, U. Shalit, and D. Sontag.  Learning representations for counterfactual inference.  In _International conference on machine learning_ , pages 3020–3029. PMLR, 2016.
  * Lorberbom et al. [2021] G. Lorberbom, D. D. Johnson, C. J. Maddison, D. Tarlow, and T. Hazan.  Learning generalized gumbel-max causal mechanisms.  _Advances in Neural Information Processing Systems_ , 34:26792–26803, 2021.
  * Monnet et al. [2016] X. Monnet, P. E. Marik, and J.-L. Teboul.  Prediction of fluid responsiveness: an update.  _Annals of intensive care_ , 6(1):1–11, 2016\.
  * Nguyen et al. [2020] T.-L. Nguyen, G. S. Collins, P. Landais, and Y. Le Manach.  Counterfactual clinical prediction models could help to infer individualized treatment effects in randomized controlled trials—an illustration with the international stroke trial.  _Journal of clinical epidemiology_ , 125:47–56, 2020.
  * Oberst and Sontag [2019] M. Oberst and D. Sontag.  Counterfactual off-policy evaluation with gumbel-max structural causal models.  In _International Conference on Machine Learning_ , pages 4881–4890. PMLR, 2019.
  * Pawlowski et al. [2020] N. Pawlowski, D. Coelho de Castro, and B. Glocker.  Deep structural causal models for tractable counterfactual inference.  _Advances in Neural Information Processing Systems_ , 33:857–869, 2020.
  * Pearl [2003] J. Pearl.  Statistics and causal inference: A review.  _Test_ , 12(2):281–345, 2003.
  * Pearl [2009] J. Pearl.  Causal inference in statistics: An overview.  _Statistics surveys_ , 3:96–146, 2009.
  * Pearl [2021] J. Pearl.  7.1 causal and counterfactual inference.  _The Handbook of Rationality_ , page 427, 2021.
  * Pearl et al. [2000] J. Pearl et al.  Models, reasoning and inference.  _Cambridge, UK: CambridgeUniversityPress_ , 19:2, 2000.
  * Roese and Olson [2014] N. J. Roese and J. M. Olson.  _What might have been: The social psychology of counterfactual thinking_.  Psychology Press, 2014.
  * Rubin [1974] D. B. Rubin.  Estimating causal effects of treatments in randomized and nonrandomized studies.  _Journal of educational Psychology_ , 66(5):688, 1974.
  * Sanchez and Tsaftaris [2021] P. Sanchez and S. A. Tsaftaris.  Diffusion causal models for counterfactual estimation.  In _First Conference on Causal Learning and Reasoning_ , 2021.
  * Schölkopf [2022] B. Schölkopf.  Causality for machine learning.  In _Probabilistic and Causal Inference: The Works of Judea Pearl_ , pages 765–804. 2022.
  * Schölkopf et al. [2021] B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio.  Toward causal representation learning.  _Proceedings of the IEEE_ , 109(5):612–634, 2021\.
  * Shalit et al. [2017] U. Shalit, F. D. Johansson, and D. Sontag.  Estimating individual treatment effect: generalization bounds and algorithms.  In _International Conference on Machine Learning_ , pages 3076–3085. PMLR, 2017.
  * Starr [2021] W. Starr.  Counterfactuals.  In E. N. Zalta, editor, _The Stanford Encyclopedia of Philosophy_. Metaphysics Research Lab, Stanford University, Summer 2021 edition, 2021.
  * Zenker et al. [2007] S. Zenker, J. Rubin, and G. Clermont.  From inverse problems in mathematical physiology to quantitative differential diagnoses.  _PLoS Comput Biol_ , 3(11):e204, 2007.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes] The assumptions are
cleary defined in section 2.3. The conclusion also relfects on the limitations
of our work.

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]
Last paragraph in the conclusion discusses the potential negative societal
impacts of counterfactual inference. The main risk we identify is blind
acceptance of the assumptions required for correct inference, which can lead
to misleading conclusions in assessing treatment effects as posteriori.

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]  Our paper conforms to the ethics review guidlines.

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]
The full set of assumptions is given in Section 2.3

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] The proofs
of our results are given in Section C of the Appendix.

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]
The code is available at the following url:
https://anonymous.4open.science/r/cfqp

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] We provide the training details in Section B of
the Appendix

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes] Error bars are reported from running the
experiments 5 times (5 folds)

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] The total
amount of compute and the type of resources used is available in the Appendix
Section B

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes] The
authors of the code bases used for implementing the baselines are credited
both in the code base and cited in the text. The list of software packages
used is available at https://anonymous.4open.science/r/cfqp/licenses.txt

    2. (b)

Did you mention the license of the assets? [Yes] The licenses of all software
packages used in our experiments is available at
https://anonymous.4open.science/r/cfqp/licenses.txt.

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [Yes] We include an anonymized repository containing the code used for
the experiments. It is available at https://anonymous.4open.science/r/cfqp

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A] We are not using any personal data.

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A] None of our datasets
contains personal data.

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A] No research with human subjects

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A] No research with human
subjects

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A] No research with human
subjects

##  Appendix A Datasets

###  A.1 Image Dataset Generation

The image dataset uses the MNIST handwritten digits images dataset. We use the
following data generating process:

  * •

$X\in\\{0,...,255\\}^{28\times 28\times 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\in\mathcal{R}=\mathcal{U}(0,0.3)+5\cdot\sigma\left(\frac{\frac{1}{784}\sum_{i=1}^{28}\sum_{j=1}^{28}X_{i,j}-33}{11}\right)$
with $\sigma(\cdot)$ the sigmoid function. This encodes confounding between
the observables $X$ and the treatment assignment $T$.

  * •

$p_{0}=\frac{1}{K}$

  * •

$p\in[0,1]^{K}$

  * •

$\rho\in[0,1]$ is the strength of confounding between $X$ and $U_{Z}$

  * •

$p[i]=\frac{1-((1-p_{0})\rho+p_{0})}{K-1}\quad\forall i\in[K],i\neq
y_{label}\mbox{mod}K$

  * •

$p[i]=((1-p_{0})\rho+p_{0})\quad\forall i\in[K],i=y_{label}\mbox{mod}(K)$

  * •

$U_{Z}\sim\mathcal{M}\text{ultinomial}(p)$

  * •

$Y\in\\{0,...,255\\}^{28\times 28\times 1}$ is the rotated image $X$ with
angle $T$ and Gaussian blur with kernel size $5\times$ and standard deviation
$\sigma$.

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 5: Examples of colored and rotated MNIST images $Y$

####  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 3: Structural similarity results for the MNIST reconstructions

Additive Noise Non-Additive Noise Model colored-MNIST colored-MNIST Deep-ITE
[12] $0.490\pm 0.006$ $0.484\pm 0.004$ SC [1] $0.020\pm 0.001$ $0.020\pm
0.001$ Deep-SCM [17] $0.857\pm 0.003$ $0.856\pm 0.004$ Diff-SCM [24] $0.780\pm
0.063$ $0.750\pm 0.052$ CFQP (ours) $0.964\pm 0.003$ $0.959\pm 0.005$

###  A.2 Harmonic Oscillator

We consider input times series $X\in\mathbb{R}^{d_{x}\times 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\in\mathbb{R}^{d_{y}\times t_{y}}$. We use the following data generating
process :

$\displaystyle U_{Z}$ $\displaystyle\sim\\{0,1,2\\}$ $\displaystyle T$
$\displaystyle\sim\mathcal{U}(0.2,1)$ $\displaystyle t_{x}$
$\displaystyle=(0,...,19)$ $\displaystyle t_{y}$ $\displaystyle=(20,...,40)$
$\displaystyle\phi$ $\displaystyle\sim\mathcal{N}(0,1)$ $\displaystyle
U_{\eta_{x}}$
$\displaystyle\sim\mathcal{N}(0,\sigma^{2})\in\mathbb{R}^{t_{x},2}$
$\displaystyle U_{\eta_{y}}$
$\displaystyle\sim\mathcal{N}(0,\sigma^{2})\in\mathbb{R}^{t_{y},2}$
$\displaystyle X$
$\displaystyle=\left[\mbox{sin}(0.5t_{x}+\phi),\mbox{sin}(0.5t_{x}+2*\phi\right]+\eta_{x}$
$\displaystyle Y$
$\displaystyle=\left[\mbox{sin}(0.5t_{y}+\phi)+\Delta_{0}(T,t_{y}),sin(0.5*t_{y}+2*\phi)+\Delta_{1}(T,t_{y})\right]+\eta_{y}$

where $\Delta_{0}(T)$ and $\Delta_{1}(T)$ are defined as follows:

$\Delta_{0}(T,t)=\begin{cases}\frac{\mbox{min}(t-20,t_{p}-20)}{t_{p}-20}\cdot
T&\text{if $U_{Z}=0$ or $U_{Z}=2$}\\\ 0&\text{if $Z=1$}\\\ \end{cases}$
$\Delta_{1}(T,t)=\begin{cases}\frac{\mbox{min}(t-20,t_{p}-20)}{t_{p}-20}\cdot
T&\text{if $U_{Z}=1$ or $U_{Z}=2$}\\\ 0&\text{if $Z=0$}\\\ \end{cases}$

$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
:

$\displaystyle\phi\sim\mathcal{N}(0,\sigma^{2})^{40}$

such that the phase $\phi$ is time-varying and stochastic. Examples of the
trajectories and reconstructed counterfactuals by CFQP are provided in Figure
6.

Figure 6: Counterfactuals generated by the model for different levels of
treatment. The factual treatment is $T=0.5$.

###  A.3 Cardiovascular Dataset

We use an ODE model of the cardiovascular system as proposed in [29, 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 :

$\displaystyle\frac{dSV(t)}{dt}$ $\displaystyle=I_{\text{external }(t)}$
$\displaystyle\frac{dP_{a}(t)}{dt}$
$\displaystyle=\frac{1}{C_{a}}\left(\frac{P_{a}(t)-P_{v}(t)}{R_{TPR}(S)}-SV\cdot
f_{HR}(S)\right)$ $\displaystyle\frac{dP_{v}(t)}{dt}$
$\displaystyle=\frac{1}{C_{v}}\left(-C_{a}\frac{dP_{a}(t)}{dt}+I_{\text{external
}(t)}\right)$ $\displaystyle\frac{dS(t)}{dt}$
$\displaystyle=\frac{1}{\tau_{\text{Baro
}}}\left(1-\frac{1}{1+e^{-k_{\text{width }}\left(P_{a}(t)-P_{a_{\text{set
}}}\right)}}-S\right)$

where

$\displaystyle R_{TPR}(S(t))$
$\displaystyle=S(t)\left(R_{TPR_{Max}}-R_{TPR_{Min}}\right)+R_{TPR_{Min}}+R_{TPR_{Mod}}$
$\displaystyle f_{HR}(S(t))$
$\displaystyle=S(t)\left(f_{HR_{Max}}-f_{HR_{Min}}\right)+f_{HR_{Min}}.$

In the above dynamical system, $P_{a},P_{v},S$ and $SV$ stand for arterial
blood pressure, venous blood pressure, autonomic baroreflex tone and cardiac
stroke volume respectively. $I_{\text{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, we model it as

$\displaystyle I_{\text{external }}(t)$ $\displaystyle=(1+2U_{Z})\cdot T\cdot
5\cdot f(P_{a}(t=0))\cdot e^{-(\frac{t-t_{\text{treat}}-5}{5})^{2}}$
$\displaystyle I_{\text{external }}(t)$ $\displaystyle=0\quad\forall t\leq
t_{t}reat$

where the treatment assignment $T$ is generated as $T\sim\mathcal{U}(0.6,1)$
and the hidden group assignment $U_{Z}\sim\mathcal{B}\mbox{er}(0.5)$ and
$t_{\text{treat}}$ is the time of treatment. The function $f$ introduces
confounding in the treatment assignment by setting :

$\displaystyle f(P_{a}(t=0))$ $\displaystyle=g(0.5+(P_{a}(t=0)-0.75)/0.1)g(x)$
$\displaystyle=0.02\cdot\left(\mbox{(}cos)(5x-0.2)\cdot(5-x)^{2}\right)^{2}$

We simulate the above system of ODEs for $t_{\text{span}}=40$ seconds and
sample an observation every $\Delta_{t}=1$ second. The treatment assignment
time is set to $t_{\text{treat}}=20$.

We then set $X$ as the first two dimensions of the dynamical system
($(P_{a},P_{v}$) for $t<t_{\text{treat}}$ and $Y$ as ($(P_{a},P_{v}$) for
$t\geq t_{\text{treat}}$. An example of a generated time series is given in
Figure 7.

Figure 7: Example of time series in the cardio-vascular data set.

####  A.3.1 Noise Responses

##### Additive noise

In the additive noise case, we add
$U_{\eta}\sim\mathcal{N}(0,\sigma^{2})^{t}_{\text{span}}$ on both $X$ and $Y$

##### Non-additive noise

In the non additive case, we modify the treatment assignment by setting:

$\displaystyle I_{\text{external }}(t)$
$\displaystyle=(1+2U_{Z}+U_{\eta}(t))\cdot T\cdot 5\cdot f(P_{a}(t=0))\cdot
e^{-(\frac{t-t_{\text{treat}}-5}{5})^{2}}$

Where $U_{eta}(t)$ is a Gaussian process with zero-mean and diagonal
covariance function with variance $\sigma^{2}$.

##  Appendix 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.

Hyper-parameter Description Image Harmonic Oscillator Cardiovascular additive
non-additive additive non-additive additive non-additive $lr$ learning rate
$0.001$ $0.001$ $0.001$ $0.001$ $0.001$ $0.001$ $\Delta$ update period $10$
$10$ $20$ $20$ $20$ $20$ $\sigma$ noise variance $0.01$ $0.05$ $0.05$ $0.05$
$0.01$ $0.01$ $\mbox{epochs}_{0}$ number of epochs initialization $50$ $50$
$500$ $500$ $500$ $500$ $\mbox{epochs}_{1}$ number of epochs fine-tune $50$
$50$ $500$ $500$ $500$ $500$ $K_{0}$ number of classes $50$ $6$ $6$ $3$ $3$
$2$ bs batch size $128$ $128$ $128$ $128$ $128$ $128$ $N_{\mbox{train}}$
Number of train samples $42,000$ $42,000$ $128$ $128$ $500$ $500$
$N_{\mbox{val}}$ Number of train samples $9,000$ $9,000$ $128$ $128$ $250$
$250$ $N_{\mbox{test}}$ Number of train samples $1,000$ $1,000$ $1,000$ $128$
$1,000$ $1,000$

Table 4: Hyper-parameters used for training CFQP on the different datasets.

###  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 8(a)
and 8(a), we visualize the validation and test performance graphically for the
Harmonic Oscillator dataset. We do the same on Figures 9(a) and 9(b) for the
Cardiovascular dataset and on Figures 10(a) and 10(b) for the colored MNIST
dataset.

Data | Noise | Number of Centers | Validation MSE | Test MSE
---|---|---|---|---
MNIST | Additive | 1 | $0.0171\pm 0.00012$ | $0.01657\pm 0.00038$
MNIST | Additive | 2 | $0.02267\pm 0.00056$ | $0.01094\pm 0.00025$
MNIST | Additive | 3 | $0.02746\pm 0.00081$ | $0.00677\pm 9e-05$
MNIST | Additive | 4 | $0.02552\pm 0.0022$ | $0.00529\pm 1e-04$
MNIST | Additive | 5 | $0.02556\pm 0.01707$ | $0.00377\pm 0.00016$
MNIST | Additive | 6 | $\mathbf{0.00271\pm 0.0003}$ | $0.00145\pm 6e-05$
MNIST | Additive | 7 | $0.00272\pm 0.00028$ | $0.0016\pm 7e-05$
MNIST | Additive | 8 | $0.00278\pm 0.00027$ | $0.00167\pm 7e-05$
MNIST | Additive | 9 | $0.00289\pm 0.00031$ | $0.00179\pm 6e-05$
MNIST | Non-Additive | 1 | $0.01673\pm 9e-05$ | $0.01673\pm 0.00026$
MNIST | Non-Additive | 2 | $0.02228\pm 0.0004$ | $0.01125\pm 0.00026$
MNIST | Non-Additive | 3 | $0.02632\pm 0.00063$ | $0.00693\pm 0.00017$
MNIST | Non-Additive | 4 | $0.03182\pm 0.00653$ | $0.00552\pm 0.00017$
MNIST | Non-Additive | 5 | $0.0041\pm 0.00018$ | $0.004\pm 0.00018$
MNIST | Non-Additive | 6 | $\mathbf{0.00179\pm 0.00014}$ | $0.00171\pm 0.00012$
MNIST | Non-Additive | 7 | $0.00186\pm 0.00014$ | $0.00185\pm 0.00011$
MNIST | Non-Additive | 8 | $0.00198\pm 0.00022$ | $0.00198\pm 0.00018$
MNIST | Non-Additive | 9 | $0.00212\pm 0.00023$ | $0.00217\pm 0.00022$
Harmonic Oscillator | Additive | 1 | $0.1653\pm 0.01086$ | $0.16584\pm 0.00443$
Harmonic Oscillator | Additive | 2 | $0.23024\pm 0.04305$ | $0.08058\pm 0.01735$
Harmonic Oscillator | Additive | 3 | $\mathbf{9e-05\pm 3e-05}$ | $0.01332\pm 0.00123$
Harmonic Oscillator | Additive | 4 | $0.00013\pm 1e-04$ | $0.0199\pm 0.00604$
Harmonic Oscillator | Additive | 5 | $0.00035\pm 0.00024$ | $0.04017\pm 0.00697$
Harmonic Oscillator | Non-additive | 1 | $0.17451\pm 0.01053$ | $0.17459\pm 0.00704$
Harmonic Oscillator | Non-additive | 2 | $0.15164\pm 0.10242$ | $0.09669\pm 0.01678$
Harmonic Oscillator | Non-additive | 3 | $\mathbf{0.00891\pm 0.00035}$ | $0.00956\pm 0.00096$
Harmonic Oscillator | Non-additive | 4 | $0.00905\pm 0.00032$ | $0.02279\pm 0.00327$
Harmonic Oscillator | Non-additive | 5 | $0.00984\pm 0.00073$ | $0.03464\pm 0.01158$
CardioVascular | Additive | 1 | $1.01453\pm 0.06228$ | $1.02569\pm 0.04794$
CardioVascular | Additive | 2 | $0.03677\pm 0.01694$ | $0.07701\pm 0.05001$
CardioVascular | Additive | 3 | $0.04906\pm 0.02931$ | $0.11363\pm 0.06759$
CardioVascular | Non-additive | 1 | $1.03865\pm 0.12068$ | $1.09768\pm 0.05459$
CardioVascular | Non-additive | 2 | $\mathbf{0.1372\pm 0.05497}$ | $0.18797\pm 0.11421$
CardioVascular | Non-additive | 3 | $0.15747\pm 0.10191$ | $0.2434\pm 0.13175$
Table 5: Reconstruction validation MSE and counterfactual test MSE for all
datasets and different number of groups ($K$). Lowest validation MSE are
bolded for each dataset.

(a) Reconstruction MSE in function of the number of groups in the harmonic
oscillator dataset (additive case). The true number of groups is 3.

(b) Reconstruction MSE in function of the number of groups in the harmonic
oscillator dataset (non-additive case). The true number of groups is 3.

Figure 8: Analysis of the impact of the number of clusters $K$ on the
validation and test performance for the Harmonic Oscillator dataset.

(a) Reconstruction MSE in function of the number of groups in the
cardiovascular dataset (additive case). The true number of groups is 2.

(b) Reconstruction MSE in function of the number of groups in the
cardiovascular dataset (non-additive case). The true number of groups is 2.

Figure 9: Analysis of the impact of the number of clusters $K$ on the
validation and test performance for the Cardiovascular dataset.

(a) Reconstruction MSE in function of the number of groups in the colored
MNIST dataset (additive case). The true number of groups is 6.

(b) Reconstruction MSE in function of the number of groups in the colored
MNIST dataset (non-additive case). The true number of groups is 6.

Figure 10: Analysis of the impact of the number of clusters $K$ on the
validation and test performance 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 4(b), we show the impact of
the correlation strength $\rho$ between $X$ and $U_{Z}$.

We incorporate correlation between $X$ and $U_{Z}$ as follows. The base
probability $p_{0}^{k}$ of each $U_{Z}$ is uniform. That is,

$\displaystyle P^{0}(U_{Z}=k)=p_{0}^{k}=\frac{1}{K}$

We modify this probabilty based on the label of the image (the digit label
$y_{\mbox{label}}\in[9]$. For a correlation factor $\rho\in[0,1]$, we write

$\displaystyle
p^{k}=\begin{cases}\frac{1-((1-p_{0})\rho+p_{0})}{K-1}\quad\text{if }k\neq
y_{label}\mbox{mod}(K)\\\ ((1-p_{0})\rho+p_{0})\quad\text{if
}k=y_{label}\mbox{mod}(K)\\\ \end{cases}$

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_{\mbox{label}}\mbox{mod}(K)=u_{Z}$.

In the case when $\rho=0$, $p^{k}=p_{0}^{k}=\frac{1}{K}$, corresponding to a
uniform probability of latent class for each class label. When $\rho=1$,
$p^{y_{\mbox{label}\mbox{mod}(K)}}=1$ and the other probabilities are set to
$0$, corresponding to a deterministic assignment of $U_{Z}$ conditioned on
$y_{\mbox{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 of the outcomes and
treatments, 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.

Table 6: Test MSE of the counterfactual reconstructions for the different
datasets.

Additive Noise Non-Additive Noise Model Harmonic Oscillator colored-MNIST
Cardiovascular Harmonic Oscillator colored-MNIST Cardiovascular Deep-ITE [12]
$0.187\pm 0.006$ $0.017\pm 0.001$ $1.084\pm 0.087$ $0.174\pm 0.004$ $0.017\pm
0.001$ $1.14\pm 0.121$ SC [1] $0.177\pm 0.131$ $0.020\pm 0.001$ $1.610\pm
0.141$ $0.167\pm 0.004$ $0.020\pm 0.001$ $1.628\pm 0.144$ Deep-SCM [17]
$0.124\pm 0.005$ $0.011\pm 0.001$ $0.405\pm 0.042$ $0.123\pm 0.005$ $0.011\pm
0.001$ $0.424\pm 0.042$ Diff-SCM [24] $0.082\pm 0.023$ $0.008\pm 0.004$
$0.206\pm 0.036$ $0.106\pm 0.038$ $0.009\pm 0.002$ $0.311\pm 0.073$ CFQP
(ours) $0.013\pm 0.001$ $\mathbf{0.001\pm 0.001}$ $0.077\pm 0.050$
$\mathbf{0.009\pm 0.001}$ $\mathbf{0.002\pm 0.001}$ $\mathbf{0.188\pm 0.114}$
CFQP-GMM (ours) $\mathbf{0.001\pm 0.001}$ $0.002\pm 0.001$ $\mathbf{0.067\pm
0.042}$ $\mathbf{0.009\pm 0.001}$ $\mathbf{0.002\pm 0.001}$ $0.192\pm 0.0834$

##  Appendix C Proof of Result 3.1

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
$(\mathcal{X}\times\mathcal{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. [3].

####  C.1.1 Identifiability Result

We let $(X,d)$ be a metric space and ($\mathcal{P}(X),\rho$) the space of
regular Borel probability measures on $X$ with finite $r$th moments metrized
by a metric $\rho$. Let’s further define
$\mathcal{P}^{2}(X)=\mathcal{P}(\mathcal{P}(X))$, the space of mixing measures
over $\mathcal{P}(X)$. We represent a mixture distribution over $X$ as a new
probability measure $m(\cdot,\Lambda)\in\mathcal{P}(X)$ such that
$m(A;\Lambda)=\int\gamma(A)d\Lambda(\gamma)$ for any $A\subset X$.

Given a Borel set $\mathfrak{L}\subset\mathcal{P}^{2}(X)$, we define
$\mathcal{M}(\mathfrak{L}):=\\{m(\Lambda):\Lambda\in\mathfrak{L})$ as a family
of mixture distributions over $X$. Because we are interested in finite
mixtures, we use
$\mathcal{P}_{k}^{2}:=\\{\Lambda\in\mathcal{P}^{2}(X):\lVert\mbox{supp}(\Lambda)\rVert\leq
k\\}$ to denote the set of mixing measures with $k$ components. Finite
mixtures $\Gamma\in\mathcal{M}(\mathcal{P}_{k}^{2}(X))$ can thus also be
written in a more intuitive form, _i.e._
$\Gamma=\sum_{k=1}^{K}\omega_{k}\gamma_{k}$ where
$\gamma_{k}\in\mathcal{P}(X)$ are the _mixture components_ and $\omega_{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 ($\gamma_{k}\in\mathcal{P}(X)$). Now, from [3], we have the
following result.

######  Theorem C.1.

If $\mathfrak{L}$ is a $\mathfrak{Q}_{L}$-clusterable family, then there
exists a function $h$ : $\mathcal{M}(\mathfrak{L})\rightarrow\mathfrak{L}$
such that $h(m(\Lambda))=\Lambda$, where
$m:\mathfrak{L}\rightarrow\mathcal{M}(\mathfrak{L})$ is the canonical
embedding. In particular, $m$ is a bijection and the mixture model
$\mathcal{M}(\mathfrak{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 $\\{\mathfrak{Q}_{L}\\}_{L=1}^{\infty}$ as an indexed
collection of families of mixing measures that satisfy the following
requirements:

  * •

$\mathfrak{Q}_{L}\subset\mathcal{P}^{2}_{L}(X)$ for each $L$;

  * •

$\\{\mathfrak{Q}_{L}\\}$ is a filtration
($\mathfrak{Q}_{L}\subset\mathfrak{Q}_{L+1}$);

  * •

The collection of mixture distributions $\mathcal{M}(\mathfrak{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 $\rho$-projection of a mixture distribution
$\Gamma=m(\Lambda)$ on a family of mixture distribution $\mathfrak{Q}_{L}$ by

$\displaystyle
T_{L}{\Gamma}=\left\\{Q\in\mathfrak{Q}_{L}:\rho(Q,\Gamma)\leq\rho(P,\Gamma)\forall
P\in\mathcal{\mathfrak{Q}_{L}}\right\\}$

where $\rho$ is a metric on $\mathcal{P}(X)$ (_e.g._ the Hellinger distance).
This projection intuitively maps a mixture distribution $\Gamma$ to the
closest element of the family $\mathfrak{Q}_{L}$ it is projected upon. We
futher define the map
$M_{L}:\mathcal{\mathfrak{Q}_{L}}\rightarrow\mathfrak{Q}_{L}$ that a projected
distribution to its mixing measure:

$\displaystyle
M_{L}\Gamma=M_{L}(T_{L}\Gamma)=\left\\{\Omega\in\mathfrak{Q}_{L}:m(\Omega)\in
T_{L}\Gamma\right\\}$

##### Assignment Functions

Because it is an element of $\mathfrak{Q}_{L}$, the projection
$T_{L}\Gamma=Q^{*}$ is a mixture distribution and can thus be written as
$T_{L}\Gamma=\sum_{l=1}^{L}\omega^{*}_{l}\gamma^{*}_{l}$. When projecting on
mixture distributions with more mixture components than in the original
distribution, _i.e._ $L\geq K$, one can then assign mixture components in
$Q^{*}$ to mixture components in $\Gamma$. We define the set of assignment
maps $\alpha:[L]\rightarrow[K]$ as $\mathbb{A}_{L\rightarrow K}$. Given such
an assignment map, we define
$\psi_{k}(\alpha):=\sum_{l\in\alpha^{-}1(k)}\omega_{l}$ for the weights
induced by the assignment map and
$Q_{k}(\alpha):=\frac{1}{\psi_{k}(\alpha)}\sum_{l\in\alpha^{-1}(k)}\omega^{*}_{l}\gamma^{*}_{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 $\Lambda\in\mathcal{P}_{K}^{2}(X)$ and
$\Gamma=m(\Lambda)\in\mathcal{M}_{K}(\mathcal{P}^{2}(X))$. The mixing measure
$\Lambda$ is called $\mathfrak{Q}_{L}$-regular if:

  * •

(a) The $\rho$-projection $Q^{*}=T_{L}\Gamma$ exists and is unique for each
$L$ and $\lim_{L\rightarrow\infty}T_{L}\Gamma=$ $\Gamma$;

  * •

(b) There exists a sequence of assignment functions
$\alpha=\alpha_{L}\in\mathbb{A}_{L\rightarrow K}$ such that
$\lim_{L\rightarrow\infty}Q_{k}^{*}(\alpha)=\gamma_{k}\quad$ and
$\quad\lim_{L\rightarrow\infty}\psi_{k}^{*}(\alpha)=\omega_{k}\quad\forall
k=1,\ldots,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 $\mathfrak{L}\subset\mathcal{P}^{2}(X)$ is called
a $\mathfrak{Q}_{L}$-clusterable family, or just a clusterable family, if

  * •

(a) $\Lambda$ is $\mathfrak{Q}_{L}$-regular for all $\Lambda\in\mathfrak{L}$;

  * •

(b) For all sufficiently large $L$, there exists a function
$\chi_{L}:M_{L}(\mathfrak{L})\rightarrow\mathbb{A}_{L\rightarrow K}$ such that
$\chi_{L}\left(\Omega^{*}\right)$ is a regular assignment for every
$\Lambda\in\mathfrak{L}$, with $\Omega^{*}=M_{L}(T_{L}\Gamma)$.

The resulting mixture model $\mathcal{M}(\mathfrak{L})$ is called a
clusterable mixture model.

More intuitively, a family of mixing measures will then be
$\mathfrak{Q}_{L}$-clusterable if one can project every mixing measure of the
family onto $\mathfrak{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 $\chi_{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.1

####  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 $\mathfrak{L}$ is a $\mathfrak{Q}_{L}$-clusterable family, then there
exists a function $h$ : $\mathcal{M}(\mathfrak{L})\rightarrow\mathfrak{L}$
such that $h(m(\Lambda))=\Lambda$, where
$m:\mathfrak{L}\rightarrow\mathcal{M}(\mathfrak{L})$ is the canonical
embedding. In particular, $m$ is a bijection and the mixture model
$\mathcal{M}(\mathfrak{L})$ is identifiable.

As laid out in the main text, we assume three observed random variables
$X\in\mathcal{X}$, $Y\in\mathcal{Y}$ and $T\in\mathcal{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
$\mathcal{M}_{K}(\mathfrak{L})$ where each mixing measure
$\Lambda\in\mathfrak{L}$ satisfies $\Lambda\in\mathcal{P}^{2}(Y)$ and
$\mid\mbox{supp}(\Lambda)\mid=K$. We write the distribution of $Y$ conditioned
on $X$ and $T$ as:

$\displaystyle Y\mid
X,T\sim\gamma(X,T)=\sum_{k=1}^{K}\omega_{k}\gamma_{k}(X,T)\quad\text{with
}\gamma(X,T)\in\mathfrak{L}$ (6)

Assuming that $\mathfrak{L}$ is clusterable, we can use Theorem C.2 to deduce
that the mixture distribution $\gamma(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
$\sigma_{x,t}:[K]\rightarrow[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 $\gamma(X,T)$ to $U_{Z}$ as
$F_{(x,t)}:\mathcal{M}(\mathfrak{L})\rightarrow\\{[K],\mathcal{P}(Y)\\}^{K}$.
We now move to showing that the assignment function is identifiable up to a
constant permutation $\bar{\sigma}$.

####  C.2.2 Identifiability of the Assignment Function

The support of random variables $X$ and $T$, $\mathcal{X}\times\mathcal{T}$ is
connected. Assuming further (Assumption 3) that $\mathfrak{L}$ is a
clusterable family, we know that there exists a cluster function
$\chi_{L,(x,t)}:\mathcal{M_{L}}(\mathfrak{L})\rightarrow\mathbb{A}_{L\rightarrow
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)}:\mathcal{M}(\mathfrak{L})\rightarrow\\{[K],\mathcal{P}(Y)\\}^{K}$
be this operator consisting of applying a $\rho$-projection onto a regular
family, cluster the result according to $\chi_{L,(x,t)}$ and return the
indexed estimated distributions $Q_{k}^{*}(\alpha)$. As the per the definition
of regular mixing measures, we have also have
$\lim_{L\rightarrow\infty}Q_{k}^{*}(\alpha)=\gamma_{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 $\bar{\sigma}$.

Because $X$ and $T$ are continuous on a connected domain, we can evaluate the
operator at $(X=x+\delta x,T=t+\delta_{t})$.

$\displaystyle F_{(x,t)}$
$\displaystyle=\\{(\sigma_{x,y}(i),q_{i}^{*}(x,t)):i=1,...,K\\}$
$\displaystyle F_{(x+\delta x,t+\delta t)}$
$\displaystyle=\\{(\sigma_{x+\delta x,y+\delta y}(i),q_{i}^{*}(x+\delta
x,t+\delta t)):i=1,...,K\\}$

What is more, for arbitrary small $(\delta x,\delta t)$ and
$\forall(x,t)\in(\mathcal{X}\times\mathcal{T});\forall
k,k^{\prime},k^{\prime\prime}\in[K]$, we have

$\displaystyle\rho(\gamma_{k}(x,t),\gamma_{k^{\prime}}(x,t))>\rho(\gamma_{k^{\prime\prime}}(x,t),\gamma_{k^{\prime\prime}}(x+\delta
x,t+\delta t)$

because we require the moments of each $\gamma_{k}$ to be continuous in
$(x,t)$. It therefore implies a relation between $\sigma_{x,y}$ and
$\sigma_{x+\delta x,y+\delta y}$ that satisfies:

$\displaystyle\sigma_{x+\delta x,y+\delta
y}(i)=\mbox{argmin}_{j}\rho(q_{i}^{*}(x,t),q_{j}^{*}(x+\delta x,t+\delta t))$

The assignment function $F_{(x,t)}$ is thus identifiable up to a constant
permutation $\bar{\sigma}$ that will determine all other permutations in
$\mathcal{X}\times\mathcal{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
$(\mathcal{X}\times\mathcal{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^{\prime},T=^{\prime})$, through the assignment
function $F_{(x,t)}$.

In particular, we can identify the means $\mu_{k}$ of each cluster component
$\gamma_{k}$. We have

$\displaystyle\lim_{N\rightarrow\infty}\hat{\mu}_{k}(X,T)=\mathbb{E}_{Y\sim\gamma_{k}(X,T)}\left[Y\right]=\mu_{k}(X,T)$
(7)

and

$\displaystyle\lim_{N\rightarrow\infty}\hat{\omega}_{k}(X,T)=\omega_{k}(X,T).$
(8)

Because the above estimators converge, the following posterior probability
converges to the true value as well:

$\displaystyle P(U_{Z}=u_{Z}\mid Y=y,X=x,T=t)$ (9)
$\displaystyle=\frac{P(Y=y\mid U_{Z}=u_{Z},X=x,T=t)P(U_{Z}=u_{Z},\mid
X=x,T=t)}{\sum_{k=1}^{K}P(Y=y\mid U_{Z}=k,X=x,T=t)P(U_{Z}=k,\mid X=x,T=t)}$
(10)

for all $x\in\mathcal{X}$, $t\in\mathcal{T}$ and $y\in\mathcal{Y}$. Indeed,
$P(Y=y\mid U_{Z}=u_{Z},X=x,T=t)$ is identifiable, and $P(U_{Z}=u_{Z},\mid
X=x,T=t)=\omega_{u_{Z}}(x,t)$.

Let $\hat{\omega}_{u_{Z}}(x,y,t)$ be the estimator of $P(U_{Z}=u_{Z}\mid
Y=y,X=x,T=t)$ We are now ready to define the estimator for the counterfactual
distribution $\nu_{t^{\prime}}(x,y,t)$.

$\displaystyle\nu_{t^{\prime}}(x,y,t)=\sum_{k}\hat{\omega}_{k}(x,y,t)\delta(y=\hat{\mu}_{k}(x,t^{\prime})).$
(11)

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 $\nu_{t^{\prime}}(x,y,t)$ to derive the bounds
of Result 3.1. We first write the $W_{1}$ distance between the estimated and
true counterfactual distributions:

$\displaystyle W_{1}(\nu_{t^{\prime}}(x,t,y),Y^{\prime}(x,t,y))$

with

$\displaystyle\nu_{t}^{\prime}(x,t,y)$
$\displaystyle=\sum_{k=1}^{N}P(U_{Z}=k\mid
X=x,T=t,Y=y)\delta(\nu_{t^{\prime}}=\hat{\mu_{k}}(X=x,T=t^{\prime}))$ (13)
$\displaystyle Y^{\prime}(x,t,y)$ $\displaystyle=\sum_{k=1}^{K}P(U_{Z}=k\mid
X=x,T=t,Y=y)$ (14) $\displaystyle\quad\quad\cdot P(Y^{\prime}=y^{\prime}\mid
X=x,T=t,Y=y,U_{Z}=k)$ (15)

We thus need to compute the $W_{1}$ distance between two mixture distributions
where each component is scaled by the same factor. By restricting the
transport map to assigning each mixture component $k$ of $\nu^{\prime}_{t}$ to
the respective one in $Y^{\prime}$, we have

$\displaystyle W_{1}(\nu_{t^{\prime}}(x,t,y),Y^{\prime}(x,t,y))$ (16)
$\displaystyle\leq\sum_{k=1}^{K}P(U_{Z}=k\mid x,t,y)$ $\displaystyle
W_{1}(\delta(\nu_{t^{\prime}}=\hat{\mu_{k}}(x,T=t^{\prime})),P(Y^{\prime}=y^{\prime}\mid
x,t,y,U_{Z}=k)).$ (17)

Now, we can write the individual components of the above sum as follows

$\displaystyle P(U_{Z}=k\mid
x,t,y)W_{1}(\delta(\nu_{t^{\prime}}=\hat{\mu_{k}}(x,T=t^{\prime})),P(Y^{\prime}=y^{\prime}\mid
x,t,y,U_{Z}=k))$ $\displaystyle\quad\leq P(U_{Z}=k\mid x,t,y)$
$\displaystyle\quad\quad\cdot\iint\limits_{y^{\prime},\nu_{t^{\prime}}}\lVert
y^{\prime}-\nu_{t^{\prime}}(x,t,y)\rVert_{2}\delta(\nu_{t^{\prime}}=\hat{\mu_{k}}(x,T=t^{\prime}))P(Y^{\prime}=y^{\prime}\mid
x,t,y,U_{Z}=k))\,dy^{\prime}\,d\nu_{t^{\prime}}$
$\displaystyle\quad=P(U_{Z}=k\mid x,t,y)\int\limits_{y^{\prime}}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}P(Y^{\prime}=y^{\prime}\mid
x,t,y,U_{Z}=k))\,dy^{\prime}$
$\displaystyle\quad=\int\limits_{y^{\prime}}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}P(Y^{\prime}=y^{\prime}\mid
x,t,y,U_{Z}=k))$ $\displaystyle\quad\quad\cdot\frac{P(Y~\mid
x,t,U_{Z}=k)P(U_{Z}=k\mid x,t)}{P(Y\mid x,t)}\,dy^{\prime}$

where we bounded the $W_{1}$ distance by using the joint probability
distribution as the transport map between $\nu_{t}^{\prime}$ and $Y^{\prime}$.

Because the $W_{1}$ distances are positive, the inequality applies to the
expectation and we can write

$\displaystyle\mathbb{E}_{Y}[W_{1}(\nu_{t^{\prime}}(x,t,Y),Y^{\prime}(x,t,Y))]\leq$
$\displaystyle\sum_{k=1}^{K}\iint\limits_{y^{\prime},y}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}P(Y^{\prime}=y^{\prime}\mid
x,t,Y,U_{Z}=k))\cdot$ $\displaystyle\quad\quad\frac{P(Y~\mid
x,t,U_{Z}=k)P(U_{Z}=k\mid x,t)}{P(Y\mid x,t)}P(Y\mid x,t)\,dy^{\prime}\,dy$

We finally bound the above expectation by marginalizing $Y$:

$\displaystyle\mathbb{E}_{Y}[W_{1}(\nu_{t^{\prime}}(x,t,Y),Y^{\prime}(x,t,Y))]\leq$
(18) $\displaystyle\sum_{k=1}^{K}\iint\limits_{y^{\prime},y}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}P(Y^{\prime}=y^{\prime}\mid
x,t,Y,U_{Z}=k)\cdot$ $\displaystyle\quad\quad P(Y~\mid
x,t,U_{Z}=k)P(U_{Z}=k\mid x,t)\,dy^{\prime}\,dy~$ (19)
$\displaystyle=\sum_{k=1}^{K}P(U_{Z}=k\mid x,t)\int\limits_{y^{\prime}}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}P(Y^{\prime}=y^{\prime}\mid
x,t,U_{Z}=k)\,dy^{\prime}$ (20) $\displaystyle=\sum_{k=1}^{K}P(U_{Z}=k\mid
x,t)\int\limits_{y^{\prime}}\lVert
y^{\prime}-\hat{\mu_{k}}(x,T=t^{\prime})\rVert_{2}\gamma_{k}(y^{\prime}\mid
x,t,U_{Z}=u_{Z})\,dy^{\prime}$ (21) $\displaystyle=\sum_{k=1}^{K}P(U_{Z}=k\mid
x,t)W_{1}^{k}(\hat{\mu}_{k},\gamma_{k})$ (22)
$\displaystyle\leq\mbox{max}_{k}W_{1}^{k}(\hat{\mu}_{k},\gamma_{k})$ (23)
$\displaystyle=\delta$

Where 23 holds because it 22 is a convex combination
$W_{1}^{k}(\hat{\mu}_{k},\gamma_{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 $\delta=0$ when the noise is
additive. From identifiability at of a mixture at a single point, we had:

$\displaystyle\lim_{N\rightarrow\infty}\hat{\mu_{k}}(x,t)=\mathbb{E}_{Y\sim\gamma_{k}(x,t)}\left[Y\right]=\mu_{k}(x,t),$
(24)

$\displaystyle\lim_{N\rightarrow\infty}\hat{\omega}_{k}(x,t,y)=P(U_{Z}=u_{Z}\mid
x,t,y).$ (25)

We also have convergence of the covariance matrix estimator

$\displaystyle\lim_{N\rightarrow\infty}\hat{\Sigma_{k}}(x,t)=\mathbb{V}\mbox{ar}_{Y\sim\gamma_{k}(x,t)}\left[Y\right]=\Sigma_{k}(x,t).$
(26)

In the additive noise setup, the distribution of $Y$ being normally
distributed, the value of $U_{\eta}$ is actually identifiable. Let us recall
the definition of the counterfactual distribution (Eq. 1) in the main text.

$\displaystyle P(Y_{t^{\prime}}=y^{\prime}\mid
X=x,Y=y,T=t)=\int_{u}P(Y_{t^{\prime}}(u)=y^{\prime})P(U=u\mid X=x,T=t,Y=y)$

In the additive case, we have $Y\mid
X=x,T=t,U_{Z}=u_{Z}=\Sigma_{u_{Z}}(x,t)U_{\eta}+\mu_{k}(x,t)$ with
$U_{\eta}\sim\mathcal{N}(0,\mathbf{1})$. The counterfactual density can thus
be written as

$\displaystyle\nu_{t^{\prime}}(x,t,y)$ $\displaystyle=\sum_{k}P(U_{Z}=k\mid
x,y,t)\cdot\delta(u_{\eta,k}=\Sigma_{k}(x,t)^{-1}y-\mu_{k}(x,t))$
$\displaystyle\quad\quad\cdot\delta(y^{\prime}=\Sigma_{k}(x,t^{\prime})u_{\eta,k}+\mu_{k}(x,t^{\prime}))$
(27)

Considering the following estimator for the counterfactual distribution:

$\displaystyle\hat{\nu}_{t^{\prime}}(x,t,y)=\sum_{k}\hat{\omega}_{k}(x,t,y)\cdot\delta(u_{\eta,k}=\hat{\Sigma}_{k}(x,t)^{-1}y-\hat{\mu}_{k}(x,t))\cdot\delta(y^{\prime}=\hat{\Sigma}_{k}(x,t^{\prime})u_{\eta,k}+\hat{\mu}_{k}(x,t^{\prime}))$

We can use Equations 24, 25 and 26 to show that our estimator converges to the
true distribution in the limit of infinite number of samples:

$\displaystyle\lim_{N\rightarrow\infty}\hat{\nu}_{t^{\prime}}(x,t,y)=\nu_{t^{\prime}}(x,t,y)$
(28)

Therefore, the distance between the estimated and true counterfactual
distributions reduces to 0:

$\displaystyle\lim_{N\rightarrow\infty}W_{1}(\hat{\nu}_{t^{\prime}}(x,t,y),\nu_{t^{\prime}}(x,t,y))=0.$
(29)

This concludes our proof of Result 3.1.