# Causal Mediation Analysis with Multi-dimensional and Indirectly Observed
Mediators

Ziyang Jiang1 Yiling Liu1 Michael H. Klein1 Ahmed Aloui1 Yiman Ren2
Keyu Li1 Vahid Tarokh1 David Carlson1
1Duke University 2University of Michigan Ross School of Business
{ziyang.jiang,yiling.liu,michael.klein413,ahmed.aloui,keyu.li,
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Causal mediation analysis (CMA) is a powerful method to dissect the total
effect of a treatment into direct and mediated effects within the potential
outcome framework. This is important in many scientific applications to
identify the underlying mechanisms of a treatment effect. However, in many
scientific applications the mediator is unobserved, but there may exist
related measurements. For example, we may want to identify how changes in
brain activity or structure mediate an antidepressant’s effect on behavior,
but we may only have access to electrophysiological or imaging brain
measurements. To date, most CMA methods assume that the mediator is one-
dimensional and observable, which oversimplifies such real-world scenarios. To
overcome this limitation, we introduce a CMA framework that can handle complex
and indirectly observed mediators based on the identifiable variational
autoencoder (iVAE) architecture. We prove that the true joint distribution
over observed and latent variables is identifiable with the proposed method.
Additionally, our framework captures a disentangled representation of the
indirectly observed mediator and yields accurate estimation of the direct and
mediated effects in synthetic and semi-synthetic experiments, providing
evidence of its potential utility in real-world applications.

##  1 Introduction

Causal inference methods are powerful tools to understand and quantify the
causal relationships between treatments and outcomes, motivating studies in
many areas [1, 2, 3, 4]. Causal inference has been combined with machine
learning in recent years to make powerful and flexible frameworks [5, 6].
While these frameworks are highly useful to estimate the total treatment
effect on an outcome, many scientific applications require understanding _how_
a treatment impacts outcomes. This knowledge can then be used to design
interventions that target the mediators to influence the outcome of interest.
For example, we may want to identify neural changes that mediate a behavioral
outcome when studying a treatment for a psychiatric disorder. Recent work has
in fact found and manipulated neural changes related to depression [7] and
social processing [8].

This need motivates the usage of causal mediation analysis (CMA), which
estimates the causal effect on an outcome of interest that is due to changes
in intermediate variables (the “mediators”) versus directly from the treatment
[9]. In specific contexts, understanding the role of the mediator is crucial
as it tells us how nature works and provides insights into the underlying
mechanisms that link variables, which enables a more accurate assessment of
the treatment’s effectiveness. In the above case, this means estimating how
much of the behavior change is explained by the treatment’s impact on the
brain, as well as how much behavioral change is unexplained by that
relationship. Early studies on mediation analysis mainly adopted linear
structural equation models (SEMs) including Wright’s method of path analysis
[10, 11] and Baron and Kenny’s method for testing mediation hypotheses [12].
In the past few decades, researchers have come up with nonparametric
generalizations for SEMs [13, 14] which do not impose any functional or
distributional forms on the causal relationships and therefore offer greater
flexibility in modeling complex dependencies between variables.

Despite these advances, a key challenge is that causal mediation analysis
typically assumes a low-dimensional, often one-dimensional, mediator, whereas
in many cases we want to identify mediation effects of complex data, such as
neuroimaging, electrophysiology, and myriad -omics studies. In this paper, we
build upon the concept of the identifiable variational autoencoder (iVAE) [15]
and introduce a novel framework for CMA that can handle _multi-dimensional_
and _indirectly observed_ mediators. We assume that there is a latent space
that generates the high-dimensional observed data (e.g., a smaller latent
space can generate the observed neural dynamics). By using an identifiable
model structure, we show that we can recover the latent space prior
conditioned on the treatment and any available covariates. In summary, our
main contributions are:

  * •

We propose a causal graph that involves both an _indirectly observed_ mediator
and observed covariates that acts as a confounder for the treatment, the
mediator, and the outcome.

  * •

We build a framework for CMA that can handle _multi-dimensional_ and
_indirectly observed_ mediators based on the proposed causal graph.

  * •

We theoretically prove that the joint distribution over observed and latent
variables in our framework is identifiable.

  * •

We show that our framework learns a disentangled representation of the
_indirectly observed_ mediator between control and treatment groups.

  * •

We empirically demonstrate the effectiveness of our framework on complex
synthetic and semi-synthetic datasets.

##  2 Related Work

#### Causal Mediation Analysis

As mentioned in the introduction, traditional mediation analysis was mainly
based on linear SEMs where the direct, mediated, and total effects are
determined by linear regression coefficients [10, 11, 12, 16, 17]. Despite its
simplicity, this approach relies on several assumptions such as normally
distributed residuals [18] and often leads to ambiguities when either the
mediator or the outcome variable is not continuous [19]. To address this
limitation, researchers formulated the causal mediation analysis (CMA)
framework based on counterfactual thinking [9, 20, 21], which can accommodate
nonlinear or nonparametric models such as targeted maximum likelihood
estimation [22], inverse propensity weighting (IPW) [23], and natural effect
models (NEMs) [24]. Within the counterfactual framework, the causal effects
are calculated as the difference between two counterfactual combinations of
mediators and outcomes, for which we will provide formal definitions in the
next section. Although causal effects are defined at the individual level, in
practice, we usually relax our estimation to their expected values over the
population as we do not generally observe both potential outcomes
simultaneously [25].

#### Causal Mediation Effect Estimation with Deep Models

Deep learning models have gained increasing attention for their capability in
estimating causal effects within the potential outcome framework [26, 27, 28,
29]. In contrast, the use of deep learning models for mediation effect
estimation has received comparatively less exploration. Xu et al. [30]
developed a semiparametric neural network-based framework to reduce the bias
in CMA. Cheng et al. [31] and Xu et al. [32] used variational autoencoders
(VAEs) to estimate the mediation effect based on a causal graph with hidden
confounders. Although these VAE-based methods share some similarities with our
proposed method, we distinguish ourselves by modeling the _mediator_ as the
latent variable rather than the covariates, resulting in a different causal
graph. Furthermore, these approaches assume that the mediator is observable
and one-dimensional, which is not necessarily the case in many scientific
applications.

#### Multi-dimensional Mediators

Compared to the many CMA methods proposed, significantly less research has
been conducted on scenarios where the mediator is multi-dimensional and not
directly observable. The majority of investigations on this subject are
situated within the domains of neuroscience [33, 34], biostatistics [35], and
bioinformatics [36, 37, 38, 39]. The approach proposed by Nath et al. [34] is
the most relevant work to our research, where the high-dimensional mediator is
first transformed into a one-dimensional variable, and the mediation effect is
estimated using an iterative maximization algorithm. Nevertheless, all these
methods primarily rely on linear SEMs and neglect the impact of any
confounding variables, thereby limiting their applicability.

##  3 Problem Setup

$T$$Z$$Y$$X$ (a)

$T$$Z$$Y$$X$$W$ (b)

Figure 1: Graphs of CMA for (a) case without observed covariates and (b) case
with observed covariates, where $T$ is the treatment assignment, $Y$ is the
outcome, $Z$ is the unobserved true mediator, $W$ is a set of observed
covariates, and $X$ is a feature caused by the unobserved true mediator $Z$
with a much higher dimension. The observed variables are colored in grey.

We assume that our causal model belongs to one of the two cases as displayed
in Figure 1. To ensure consistency with previous studies on mediation analysis
[18, 40, 41], we further assume that the treatment assignment $T$ is binary
for each observed samples, with $T=0$ indicating an assignment to the control
group and $T=1$ indicating an assignment to the treatment group. Consider the
$n^{th}$ individual in an experiment with a total of $N$ units (i.e.
$n=1,...,N$). Let
$\boldsymbol{z}_{n}(t_{n})\in\mathcal{Z}\subset\mathbb{R}^{d}$ denote the
potential value of the unobserved true mediator under the treatment assignment
$t_{n}$. Since $Y$ depends on both $T$ and $Z$, we denote
$y(t_{n},\boldsymbol{z}_{n}(t_{n}))\in\mathcal{Y}\subset\mathbb{R}$ as the
potential outcome of the $n^{th}$ individual under treatment $t_{n}$ and true
mediator $\boldsymbol{z}_{n}(t_{n})$. Following [9, 40, 42], we can define the
average causal mediation effects (ACME), the average direct effects (ADE), and
the average total effect (ATE) as follows:

$\displaystyle ACME(t)$
$\displaystyle\coloneqq\mathbb{E}\left[y(t,\boldsymbol{z}(1))-y(t,\boldsymbol{z}(0))\right],$
(1) $\displaystyle ADE(t)$
$\displaystyle\coloneqq\mathbb{E}\left[y(1,\boldsymbol{z}(t))-y(0,\boldsymbol{z}(t))\right],$
(2) $\displaystyle ATE$
$\displaystyle\coloneqq\mathbb{E}\left[y(1,\boldsymbol{z}(1))-y(0,\boldsymbol{z}(0))\right],$
(3)

where the expectations are taken over all the samples in our experiment. Our
main objective is to recover these quantities as accurately as possible. As
$\boldsymbol{z}_{n}$ is unobserved, we must infer $\boldsymbol{z}_{n}$ from
the related observed feature
$\boldsymbol{x}_{n}\in\mathcal{X}\subset\mathbb{R}^{D}$ with a much higher
dimension, i.e. $D\gg d$, as well as any other available information. In
practice, there often exists a set of observed covariates
$\boldsymbol{w}_{n}\in\mathcal{W}\subset\mathbb{R}^{m}$ that also acts as
confounders for $T$, $Z$, and $Y$ as shown in Figure 1(b). With the presence
of observed covariates, we need to make the following assumptions to make
valid inferences about the causal effects:

######  Assumption 3.1.

There exists an observed variable $X\in\mathcal{X}\subset\mathbb{R}^{D}$ that
is caused by the unobserved true mediator
$Z\in\mathcal{Z}\subset\mathbb{R}^{d}$ as shown in Figure 1.

######  Assumption 3.2.

The following two conditional independence assumptions hold sequentially.

$\displaystyle\left\\{Y(t^{\prime},z),Z(t)\right\\}$
$\displaystyle\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}T|W=w,$ (4) $\displaystyle Y(t^{\prime},z)$
$\displaystyle\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}Z(t)|T=t,W=w,$ (5)

where $0<p(T=t|W=w)<1$, $0<p(Z(t)=z|T=t,W=w)<1$, and
$t,t^{\prime}\in\\{0,1\\}$.

Assumption 3.2 is first introduced by Imai et al. [3], which is also known as
_sequential ignorability_. Note that Equation 4 is equivalent to the strong
ignorability assumption common in causal inference [43, 44]. It states that
the treatment assignment $T$ is statistically independent of potential outcome
$Y$ and potential mediators $Z$ given covariates $W$. Equation 5 states that
given the treatment and covariates, the mediator $Z$ can be viewed as if it
was randomized (in other words, there are no explained “backdoor” paths
between the mediator and outcome [18]).

##  4 Method

We leverage the model structure of identifiable variational autoencoder (iVAE)
to estimate the causal mediation effects based on the causal graphs
illustrated in Figure 1. Our primary objective is to learn a disentangled
representation of the true mediator in the latent space so that the
statistical distance between $p(\boldsymbol{z}|t=0)$ and
$p(\boldsymbol{z}|t=1)$ can be better estimated. In the following sections, we
briefly review the concepts of identifiable variational autoencoder (iVAE) in
Section 4.1, present our framework in Section 4.2, and formally state the
identifiability of our framework in Section 4.3.

Figure 2: Illustration of an iVAE where the blue nodes correspond to
probabilistic distributions.

###  4.1 Identifiable Variational Autoencoder (iVAE)

To begin with, here we provide a brief overview of iVAE [15]. We abuse the
notation slightly by redefining $\boldsymbol{x}$ and $\boldsymbol{z}$ to refer
to the observed data and the latent feature learned by a general variational
autoencoder (VAE), respectively. The primary claim made by iVAE is that a VAE
becomes identifiable up to _a linear invertible transformation_ (see Section
4.3 for full definitions) if we introduce a factorized prior distribution over
the latent variable $\boldsymbol{z}$ conditioned on an auxiliary variable
$\boldsymbol{u}$. Specifically, we have $\boldsymbol{z}$ sampled from
$p(\boldsymbol{z}|\boldsymbol{u})$ which is assumed to be conditionally
factorial with each $z_{i}\in\boldsymbol{z}$ belonging to a univariate
exponential family as specified by the following probability density function:

$p_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})=\prod_{i}\frac{Q_{i}(z_{i})}{C_{i}(\boldsymbol{u})}\exp\left[\sum_{j=1}^{k}S_{i,j}(z_{i})\lambda_{i,j}(\boldsymbol{u})\right],$
(6)

where $Q_{i}$ is the base measure, $C_{i}(\boldsymbol{u})$ is the normalizing
constant, $k$ is a pre-defined number of sufficient statistics,
$\boldsymbol{S}_{i}=(S_{i,1},...,S_{i,k})$ are the sufficient statistics, and
$\boldsymbol{\lambda}_{i}(\boldsymbol{u})=\left(\lambda_{i,1}(\boldsymbol{u}),...,\lambda_{i,k}(\boldsymbol{u})\right)$
are the natural parameters.

The architecture of the iVAE framework is displayed in Figure 2, which
consists of a variational posterior
$q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x},\boldsymbol{u})$ and a
conditional generative model
$p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z}|\boldsymbol{u})=p_{\textbf{f}}(\boldsymbol{x}|\boldsymbol{z})p_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$
where f is an injective function such that
$p_{\textbf{f}}(\boldsymbol{x}|\boldsymbol{z})=p_{\boldsymbol{\epsilon}}(\boldsymbol{x}-\textbf{f}(\boldsymbol{z}))$
and $\boldsymbol{\epsilon}$ is an independent noise variable with probability
density function $p(\boldsymbol{\epsilon})$. The parameters of the generative
model are denoted as
$\boldsymbol{\theta}=\\{\textbf{f},\boldsymbol{S},\boldsymbol{\lambda}\\}$.
When fitting iVAE on observed data, the parameter vector
$(\boldsymbol{\theta},\boldsymbol{\phi})$ is learned by maximizing the
evidence lower bound (ELBO)
$\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\boldsymbol{x},\boldsymbol{u})$:

$\log
p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{u})\geq\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\boldsymbol{x},\boldsymbol{u})\coloneqq\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x},\boldsymbol{u})}\left[\log
p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z}|\boldsymbol{u})-\log
q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x},\boldsymbol{u})\right],$
(7)

where we use the reparameterization trick to sample from
$q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x},\boldsymbol{u})$. Briefly
speaking, both the model structure and the learning process of iVAE are
similar to conventional VAEs except that the prior, the variational posterior,
and the decoder are additionally conditioned on the auxiliary variable
$\boldsymbol{u}$. However, it is important to note that $\boldsymbol{u}$ must
have some association with $\boldsymbol{x}$ and $\boldsymbol{z}$.

###  4.2 Estimating Mediation Effect with VAE

(a)

(b)

Figure 3: Illustration of the overall architecture of IMAVAE for (a) case
without observed covariates and (b) case with observed covariates. Note that
in case (b) the treatment assignment $T$ and the observed covariates $W$ are
first concatenated and then passed into the prior, encoder, and decoder.

In this section, we formally present our approach — Identifiable Mediation
Analysis with Variational Autoencoder (IMAVAE), with the overall architecture
displayed in Figure 3. The encoder, decoder, and prior components in IMAVAE
have exactly the same probabilistic form as specified in Section 4.1 and share
a similar structure with iVAE, where we take the high-dimensional feature $X$
as the input to the encoder to learn the unobserved mediator $Z$ and generate
a reconstruction $\hat{X}$ with the decoder. Importantly, we further include a
parametric model $g_{\boldsymbol{\gamma}}$ to predict the outcome $\hat{Y}$.
Figures 3(a) and 3(b) depict two variants of our framework, corresponding to
the two cases outlined in the causal graphs in Figure 1:

  * •

_Case (a)_ : Without observed covariates, the treatment assignment $T$ is
employed as the auxiliary variable and serves as input to the encoder, prior,
and predictor, as illustrated in Figure 3(a).

  * •

_Case (b)_ : With observed covariates, we first concatenate the observed
covariates $W$ and the treatment assignment $T$. The concatenated vector
$(W,T)$ is then passed into the encoder, prior, and predictor as the auxiliary
variable, as illustrated in Figure 3(b).

Similar to iVAE, we denote the parameter vector of IMAVAE as
$(\boldsymbol{\theta},\boldsymbol{\phi},\boldsymbol{\gamma})$ where
$\boldsymbol{\theta}=\\{\textbf{f},\boldsymbol{S},\boldsymbol{\lambda}\\}$.
When fitting IMAVAE to the observed data, we optimize the parameter vector by
minimizing the following objective:

$\boldsymbol{\theta}^{*},\boldsymbol{\phi}^{*},\boldsymbol{\gamma}^{*}\coloneqq\arg\min_{\boldsymbol{\theta},\boldsymbol{\phi},\boldsymbol{\gamma}}\left\\{\alpha\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\hat{\boldsymbol{x}},\boldsymbol{x})-\beta\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\boldsymbol{x},\boldsymbol{u})+\mathcal{L}_{\boldsymbol{\phi},\boldsymbol{S},\boldsymbol{\lambda},\boldsymbol{\gamma}}(\hat{y},y)\right\\},$
(8)

where $\boldsymbol{u}=t$ for case (a), $\boldsymbol{u}=(\boldsymbol{w},t)$ for
case (b),
$\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\hat{\boldsymbol{x}},\boldsymbol{x})$
is the discrepancy between the input feature $\boldsymbol{x}$ and its
reconstruction $\hat{\boldsymbol{x}}$,
$\mathcal{L}_{\boldsymbol{\phi},\boldsymbol{S},\boldsymbol{\lambda},\boldsymbol{\gamma}}(\hat{y},y)$
is the error between the predicted outcome $\hat{y}$ and the true outcome $y$.
$\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\boldsymbol{x},\boldsymbol{u})$
represents the same loss term as Equation 7. We note that this creates some
overlap as the reconstruction term on $\boldsymbol{x}$ is also in Equation 7,
but choose this form to highlight each term independently (and does not change
the overall loss with appropriately chosen weights). $\alpha$ and $\beta$ are
hyperparameters representing the importance of the reconstruction error and
the ELBO, respectively. In our experiments, we use mean squared error (MSE)
loss for both
$\mathcal{L}_{\boldsymbol{\theta},\boldsymbol{\phi}}(\hat{\boldsymbol{x}},\boldsymbol{x})$
and
$\mathcal{L}_{\boldsymbol{\phi},\boldsymbol{S},\boldsymbol{\lambda},\boldsymbol{\gamma}}(\hat{y},y)$.
The prior
$p_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$ is
set to be a multivariate normal distribution whose mean and covariance are
parameterized as a function of $\boldsymbol{u}$ using a neural network.

To give an estimation on the direct, mediated, and total effects after fitting
the parameters, we repeatedly sample $\boldsymbol{z}(t)$ from the learned
distributions (i.e.,
$p_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|t)$ for case (a) and
$p_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{w},t)$ for
case (b)). Next, we feed both $\boldsymbol{z}(t)$ and the auxiliary variables
into the predictor $g_{\boldsymbol{\gamma}}$ to obtain
$y(t,\boldsymbol{z}(t))$ for case (a) or
$y(t,\boldsymbol{w},\boldsymbol{z}(t))$ for case (b). Finally, we estimate the
ACME, ADE, and ATE according to Equations 1-3 using estimated values of $y$.

###  4.3 Identifiability of IMAVAE

In this section, we use similar definitions and assumptions stated by
Khemakhem et al. [15]. Specifically, let $\mathcal{Z}\subset\mathbb{R}^{d}$ be
the support of distribution of $\boldsymbol{z}$. The support of distribution
of $\boldsymbol{u}$ is $\mathcal{U}=\\{0,1\\}$ for case (a) and
$\mathcal{U}=\\{0,1\\}\times\mathcal{W}\subset\mathbb{R}^{m+1}$ for case (b).
We denote by
$\textbf{S}\coloneqq(\textbf{S}_{1},...,\textbf{S}_{d})=(S_{1,1},...,S_{d,k})\in\mathbb{R}^{dk}$
the vector of sufficient statistics of Equation 6 and
$\boldsymbol{\lambda}({\boldsymbol{u}})=(\boldsymbol{\lambda}_{1}(\boldsymbol{u}),...,\boldsymbol{\lambda}_{d}(\boldsymbol{u}))=(\lambda_{1,1}(\boldsymbol{u}),...,\lambda_{d,k}(\boldsymbol{u}))\in\mathbb{R}^{dk}$
the vector of its parameters. Following the same notations in [15], we define
$\mathcal{X}\subset\mathbb{R}^{D}$ as the image of f in Equation 6 and denote
by $\textbf{f}^{-1}:\mathcal{X}\rightarrow\mathcal{Z}$ the inverse of f.
Furthermore, we make the following assumption on the predictor:

######  Assumption 4.1.

The predictor $g_{\boldsymbol{\gamma}}(\boldsymbol{z},\boldsymbol{u})$ takes
the following form:

$g_{\boldsymbol{\gamma}}(\boldsymbol{z},\boldsymbol{u})\coloneqq
p_{\textbf{h}}(y|\boldsymbol{z},\boldsymbol{u})=p_{\boldsymbol{\xi}}(y-\textbf{h}(\boldsymbol{z},\boldsymbol{u})),$
(9)

where the function
$\textbf{h}:\mathcal{Z}\times\mathcal{U}\rightarrow\mathcal{Y}$ is injective,
$\mathcal{Y}\subset\mathbb{R}$ is the image of h, and $\boldsymbol{\xi}$ is an
independent noise variable with probability density function
$p_{\boldsymbol{\xi}}(\boldsymbol{\xi})$.

Similar to [15], for the sake of analysis, we treat h as a parameter of the
entire model and define
$\boldsymbol{\psi}\coloneqq(\textbf{f},\textbf{h}):\mathcal{Z}\times\mathcal{U}\rightarrow\mathcal{X}\times\mathcal{Y}$.
$\boldsymbol{\psi}$ remains injective since both f and h are injective, and we
consider the projection $\boldsymbol{\psi}^{-1}$ on $\mathcal{Z}$ to be
$\boldsymbol{\psi}_{|\boldsymbol{z}}^{-1}$. The domain of parameters is thus
$\Theta=\\{\boldsymbol{\theta}\coloneqq(\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda})\\}$.
To formally present our claim, we give the following definitions:

######  Definition 4.2.

Let $\sim$ be an equivalence relation on $\Theta$. We say that
$p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z},y|\boldsymbol{u})$ is
identifiable up to $\sim$ if
$p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z},y|\boldsymbol{u})=p_{\boldsymbol{\tilde{\theta}}}(\boldsymbol{x},\boldsymbol{z},y|\boldsymbol{u})\Longrightarrow\boldsymbol{\theta}\sim\boldsymbol{\tilde{\theta}}$.

######  Definition 4.3.

Let $\sim_{A}$ be the equivalence relation on $\Theta$ defined as follows:

$(\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda})\sim(\tilde{\textbf{f}},\tilde{\textbf{h}},\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}})\Longleftrightarrow\exists
A,\textbf{c}\,|\,\textbf{S}(\boldsymbol{\psi}^{-1}_{\boldsymbol{|z}}(\boldsymbol{x},y))=A\tilde{\textbf{S}}(\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}(\boldsymbol{x},y))+\textbf{c},\forall\boldsymbol{x}\in\mathcal{X};y\in\mathcal{Y},$
(10)

where $A$ is an invertible $dk\times dk$ matrix and c is a vector.

With all the assumptions and definitions stated above, we state our theorem
below as an extension of the results in [15]. The detailed proof will be
provided in Appendix A.

######  Theorem 4.4.

(Extension to Theorem 1 in Khemakhem et al. [15]) Assume that we observe data
sampled from the generative model
$p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z},y|\boldsymbol{u})=p_{\textbf{f}}(\boldsymbol{x}|\boldsymbol{z})p_{\textbf{h}}(y|\boldsymbol{z},\boldsymbol{u})p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$
where $p_{\textbf{f}}(\boldsymbol{x}|\boldsymbol{z})$,
$p_{\textbf{h}}(y|\boldsymbol{z},\boldsymbol{u})$ and
$p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$ follow
the distributional form defined in Section 4.1, Equation 9, and Equation 6,
respectively. Then the parameters
$(\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda})$ will be
$\sim_{A}$-identifiable if we assume the following holds:

  1. 1.

The set
$\\{(\boldsymbol{x},y)\in\mathcal{X}\times\mathcal{Y}\,|\,\varphi_{\epsilon}(\boldsymbol{x})=0,\varphi_{\xi}(y)=0\\}$
has measure zero, where $\varphi_{\boldsymbol{\epsilon}}$ and
$\varphi_{\boldsymbol{\xi}}$ are the characteristic functions of
$p_{\boldsymbol{\epsilon}}$ and $p_{\boldsymbol{\xi}}$ defined in Section 4.1
and Equation 9, respectively.

  2. 2.

The functions f and h are both injective.

  3. 3.

The sufficient statistics $S_{i,j}$ in Equation 6 are differentiable almost
everywhere, and $(S_{i,j})_{1\leq j\leq k}$ are linearly independent on any
subset of $\mathcal{Y}$ of measure greater than zero.

  4. 4.

There exists $dk+1$ distinct points
$\boldsymbol{u}_{0},...,\boldsymbol{u}_{dk}$ such that the matrix
$L=(\boldsymbol{\lambda}(\boldsymbol{u}_{1})-\boldsymbol{\lambda}(\boldsymbol{u}_{0}),...,\boldsymbol{\lambda}(\boldsymbol{u}_{dk})-\boldsymbol{\lambda}(\boldsymbol{u}_{0}))$
of size $dk\times dk$ is invertible.

With the aforementioned theorem, we state that the joint distribution learned
by the generative model
$p_{\boldsymbol{\theta}}(\boldsymbol{x},\boldsymbol{z},y|\boldsymbol{u})$ is
identifiable. Moreover, it is important to highlight that our extension of the
identifiability theorem, originally presented in [15], incorporates the
additional conditioning of $y$ on $\boldsymbol{u}$, thereby broadening the
scope of iVAE.

##  5 Experiments

We follow the approach using synthetic and semi-synthetic datasets used in
recent causal inference manuscripts to allow for benchmarking and comparison
of the results. We test IMAVAE on 3 datasets: 1 synthetic dataset and 2 semi-
synthetic datasets111Code will be released upon acceptance.. This allows us to
evaluate how well we estimate counterfactual values of the treatment
assignment and the mediator, and the direct, mediated, and total effects under
reasonable assumptions. The detailed experimental setup (e.g. training
details, computing resources, licenses, etc.) is given in Appendix B.

1.13

Figure 4: Distribution of the true and the estimated
$p(\boldsymbol{z}|\boldsymbol{u})$ in the latent space where the upper row
corresponds to case (a) without observed covariates, i.e. $\boldsymbol{u}=t$
and the bottom row corresponds for case (b) with observed covariates, i.e.
$\boldsymbol{u}=(\boldsymbol{w},t)$. From left to right, we present (left) the
true distribution of $p(\boldsymbol{z}|\boldsymbol{u})$, (middle left) the
estimated distribution
$\hat{p}_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$
by IMAVAE, (middle right) the estimated distribution
$\hat{p}_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$
without the reconstruction term, i.e. $\alpha=-1$, and (right) the estimated
distribution
$\hat{p}_{\boldsymbol{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$
without the ELBO term, i.e. $\beta=0$. The blue dots denote samples in control
group and the orange dots denote samples in treatment group.

###  5.1 Synthetic Dataset

| IMAVAE |  | IMAVAE
---
$\alpha=-1$
| IMAVAE
---
$\beta=0$
ACME ($t=1$) | 0.056 $\pm$ .007 | 0.078 $\pm$ .007 | 2.875 $\pm$ .016
ADE ($t=0$) | 0.058 $\pm$ .000 | 0.052 $\pm$ .000 | 0.043 $\pm$ .000
ATE | 0.003 $\pm$ .007 | 0.025 $\pm$ .006 | 2.917 $\pm$ .015

Table 1: Absolute error of ACME under treated, ADE under control, and ATE on
the synthetic dataset for IMAVAE in case (a) _without_ observed covariates

| IMAVAE |  | IMAVAE
---
$\alpha=-1$
| IMAVAE
---
$\beta=0$
ACME ($t=1$) | 0.214 $\pm$ .016 | 0.379 $\pm$ .014 | 5.912 $\pm$ .028
ADE ($t=0$) | 0.194 $\pm$ .000 | 0.385 $\pm$ .000 | 0.127 $\pm$ .000
ATE | 0.019 $\pm$ .015 | 0.011 $\pm$ .016 | 6.036 $\pm$ .028

Table 2: Absolute error of ACME under treated, ADE under control, and ATE on
the synthetic dataset for IMAVAE in case (b) _with_ observed covariates

We first construct a synthetic dataset following the causal graphs in Figure
1, where we give the details of data generation process in Appendix C. We set
the unobserved true mediator to be two-dimensional (i.e. $d=2$) for easier
visualization. We display the distributions of the true and estimated
unobserved mediator in Figure 4, where we note that IMAVAE effectively learns
_disentangled representations_ of $Z$ for the control and treatment groups in
the latent space, up to trivial indeterminacies such as rotations and sign
flips, for cases both with and without observed covariates. If we remove the
reconstruction term (i.e. $\alpha=-1$ due to the overlap of reconstruction
terms in Equation 8), the shape and orientation of the distributions become
slightly different but remain disentangled. However, if we discard the ELBO
term (i.e. $\beta=0$), the model fails to separate the distributions of
control and treatment groups. We also compute the absolute errors between the
estimated ACME, ADE, ATE, and their corresponding ground truths as shown in
Tables 1 and 2. It can be observed that IMAVAE yields slightly larger errors
when the reconstruction term is removed (i.e., $\alpha=-1$). However, without
the ELBO term (i.e., $\beta=0$), the model produces significantly larger
errors on ACME and ATE. From the obtained results, we conclude that the iVAE-
like structure in our framework is essential for learning a better
representation of the unobserved mediator, which, in turn, improves the
accuracy of mediation effect estimation.

###  5.2 Electrophysiological Dataset

As described in the introduction, causal mediation analysis holds significant
relevance for applications in systems neuroscience. Once such area is in the
emerging area of targeted neurostimulation (see [45, 46, 47] for a
description), where the brain is manipulated by optical, electrical, or
mechanical stimulation with the goal of manipulating behavior in many brain
conditions. However, while the mechanism of behavioral change is the brain,
identifying such changes is challenging with existing causal mediation
techniques due to the high-dimensional and complex nature of the brain data.
Accurate appraisal of neural changes causing behavioral change will provide a
deeper understanding of mechanisms driving neural activity and potentially
lead to more efficacious treatments.

We demonstrate capability of our method to this domain with a semisynthetic
dataset by post-processing real multi-site brain recordings using local field
potential (LFP) data from 26 mice [48], which is publicly available [49]. Each
mouse is recorded by a certain number of time steps, resulting in a total of
43,882 data points. We take the LFP signals as the observed feature $X$, while
the true mediator $Z$ is manually generated by applying principal component
analysis (PCA) to map $X$ into a lower-dimensional representation. The
treatment assignment $T$ indicates whether the mouse is recorded during an
open field exploration $(T=0)$ or a tail suspension test $(T=1)$. Furthermore,
we consider the genotype of the mouse, a binary variable, as an observed
covariate, denoted by $W\in{0,1}$. Lastly, we construct the outcome $Y$
manually as a function of the treatment, the mediator, and the genotype (only
for case (b) with observed covariate). The detailed procedure of dataset
generation is given in Appendix D.

We compare our method with two baseline models that are designed to handle
high-dimensional mediators: an integrated framework of shallow or deep neural
network and linear SEM (Shallow/Deep LSEM) [34] and a high-dimensional
mediation analysis (HIMA) framework [39]. Notably, HIMA considers each
component of $Z$ as an individual mediator instead of a multidimensional
mediator. As such, we report the mediation effect using the component with the
highest correlation. We compute and display the absolute errors of ACME, ADE,
and ATE in Table 3. Our results indicate that IMAVAE outperforms both
benchmarks by a very weide margin on all estimations except the ATE in case
(a) without covariates. The two benchmarks used in this experiment yield
significantly larger errors on ACME and ADE. We believe this is reasonable, as
both benchmarks are designed based on linear SEMs and are thus not able to
capture the correlation between the components of $Z$.

Table 3: Absolute error of ACME under treated, ADE under control, and ATE on
the tail suspension test dataset for IMAVAE and other benchmarks.

| Case (a) | Case (b)
---|---|---
|  | IMAVAE
---
(ours)
| Shallow
---
LSEM
| Deep
---
LSEM
HIMA |  | IMAVAE
---
(ours)
| Shallow
---
LSEM
| Deep
---
LSEM
HIMA
| ACME
---
$(t=1)$
2.348 $\pm$ 0.003 | 15.14 $\pm$ 0.03 | 15.48 $\pm$ 0.07 | 13.95 $\pm$ 0.03 | 0.559 $\pm$ 0.002 | 3.06 $\pm$ 2.16 | 4.80 $\pm$ 0.44 | 2.67 $\pm$ 0.02
| ADE
---
$(t=0)$
1.603 $\pm$ 0.000 | 14.71 $\pm$ 0.03 | 15.06 $\pm$ 0.07 | 3.82 $\pm$ 0.01 | 0.782 $\pm$ 0.000 | 5.03 $\pm$ 1.20 | 5.16 $\pm$ 0.51 | 3.21 $\pm$ 0.01
ATE | 0.744 $\pm$ 0.003 | 0.42 $\pm$ 0.06 | 0.42 $\pm$ 0.14 | 17.77 $\pm$ 0.03 | 0.223 $\pm$ 0.002 | 1.96 $\pm$ 3.36 | 0.36 $\pm$ 0.94 | 0.54 $\pm$ 0.02

###  5.3 Jobs II Dataset

To evaluate whether our method can generalize to real-world scenarios used in
recent causal mediation analysis frameworks, we test IMAVAE on the Jobs II
dataset [50], which aims to explore the impact of unemployment on workers’
stress and mental health and evaluate the potential benefits of participation
in a job-search skills seminar. The dataset includes a binary treatment
assignment $T$, which indicates whether a participant was assigned to attend a
job-search skills seminar ($T=1$) or to receive a booklet ($T=0$). The
mediator $Z$ is a continuous variable that measures the job-search efficacy.
All other attributes are treated as the observed covariates $W$. The outcome
variable $Y$ is also a continuous variable that represents the level of
depression reported by each participant during follow-up interviews. To obtain
the ground truth for direct and mediated effects, we followed a simulation
procedure similar to [51] to make ACMEs, ADEs, and ATE all equal to zero. The
detailed simulation procedure is given in Appendix E.

We compare the performance of our method with several benchmarks: nonlinear
SEM with interaction (LSEM-I) [3], imputing-based natural effect model (NEM-I)
[24], IPW [23], and Causal Mediation Analysis with Variational Autoencoder
(CMAVAE) [31]. It is worth noting that the Jobs II dataset presents an
observable mediator variable $Z$, which is _not_ the optimal scenario for our
proposed framework, as IMAVAE is specifically designed for CMA with
_implicitly_ observed mediators. Nonetheless, according to the results shown
in Tables 4 and 5 (where $N$ is the total number of simulated samples and
$\eta$ is a simulation parameter which stands for the magnitude of selection
into the mediator), our method still mostly outperforms the benchmarks in
terms of the estimation on ACME, ADE, and ATE with a reasonable level of
uncertainty.

##  6 Discussion

#### Design Choice of the Predictor

In Section 4.3, we prove that the true joint distribution over observed and
latent variables learned by IMAVAE is identifiable if we specify the predictor
$g_{\boldsymbol{\gamma}}$ to be a conditional distribution reparameterized by
function h. However, in practice, $g_{\boldsymbol{\gamma}}$ can be as simple
as a linear or logistic regression model since we believe the identifiability
on $(\textbf{f},\textbf{S},\boldsymbol{\lambda})$ is enough to disentangle the
representations between control and treatment groups and give an accurate
estimation on the mediation effects. We encourage readers to consider
designing $g_{\boldsymbol{\gamma}}$ in order to achieve optimal performance.

#### Limitations

As discussed in Section 3, our method relies on sequential ignorability, a
condition that is not directly testable using the observed data. However,
recent studies [31, 32] propose a potential solution by considering $X$ as a
proxy variable and accounting for hidden confounders. Exploring this approach
represents an intriguing direction for our future research.

#### Applications and Broader Impacts

We believe the proposed model architecture can be very useful for improving
interpretability for neuroscience applications. For instance, the disentangled
mediator representations obtained by IMAVAE can be used to investigate the
brain activities of individuals under different interventions. It can also be
combined with other interpretable methods such as linear factor models to
better illustrate the high-dimensional dynamics in brain networks as proposed
by Talbot et al [52]. We have not identified any potential negative societal
consequences specific to this manuscript.

Table 4: Absolute error of ACME under treated, ADE under control, and ATE on
simulated Jobs II data for IMAVAE and other benchmarks where 10% of the data
are mediated (i.e. $Z>3$).

| LSEM-I | NEM-I | IPW | CMAVAE | IMAVAE (ours)
---|---|---|---|---|---
$N$ | 500 | 1000 | 500 | 1000 | 500 | 1000 | 500 | 1000 | 500 | 1000
| ACME under treated $(t=1)$
$\eta=10$ | 0.9 $\pm$ .04 | 0.6 $\pm$ .02 | 0.6 $\pm$ .03 | 0.8 $\pm$ .01 | 0.6 $\pm$ .04 | 0.8 $\pm$ .02 | 0.2 $\pm$ .00 | 0.3 $\pm$ .00 | 0.1 $\pm$ .02 | 0.1 $\pm$ .01
$\eta=1$ | 0.0 $\pm$ .01 | 0.1 $\pm$ .01 | 0.0 $\pm$ .00 | 0.1 $\pm$ .01 | 0.0 $\pm$ .01 | 0.1 $\pm$ .01 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00 | 0.1 $\pm$ .01 | 0.0 $\pm$ .01
| ADE under control $(t=0)$
$\eta=10$ | 1.3 $\pm$ .07 | 1.6 $\pm$ .06 | 1.2 $\pm$ .06 | 1.8 $\pm$ .05 | 1.2 $\pm$ .06 | 0.2 $\pm$ .06 | 0.1 $\pm$ .00 | 0.0 $\pm$ .03 | 0.3 $\pm$ .00 | 0.3 $\pm$ .00
$\eta=1$ | 3.3 $\pm$ .08 | 0.0 $\pm$ .07 | 1.1 $\pm$ .03 | 0.2 $\pm$ .07 | 3.3 $\pm$ .08 | 0.3 $\pm$ .06 | 0.5 $\pm$ .02 | 0.4 $\pm$ .01 | 0.2 $\pm$ .00 | 0.3 $\pm$ .00
| ATE
$\eta=10$ | 2.2 $\pm$ .05 | 1.0 $\pm$ .06 | 1.8 $\pm$ .05 | 0.9 $\pm$ .06 | 0.5 $\pm$ .05 | 1.0 $\pm$ .06 | 0.3 $\pm$ .01 | 0.3 $\pm$ .03 | 0.2 $\pm$ .02 | 0.2 $\pm$ .01
$\eta=1$ | 3.3 $\pm$ .08 | 0.1 $\pm$ .07 | 3.4 $\pm$ .03 | 0.1 $\pm$ .06 | 3.2 $\pm$ .07 | 0.2 $\pm$ .05 | 0.4 $\pm$ .02 | 0.3 $\pm$ .01 | 0.2 $\pm$ .01 | 0.2 $\pm$ .01

Table 5: Absolute error of ACME under treated, ADE under control, and ATE on
simulated Jobs II data for IMAVAE and other benchmarks where 50% of the data
are mediated (i.e. $Z>3$)

| LSEM-I | NEM-I | IPW | CMAVAE | IMAVAE (ours)
---|---|---|---|---|---
$N$ | 500 | 1000 | 500 | 1000 | 500 | 1000 | 500 | 1000 | 500 | 1000
| ACME under treated $(t=1)$
$\eta=10$ | 0.9 $\pm$ .03 | 0.6 $\pm$ .03 | 0.2 $\pm$ .03 | 0.4 $\pm$ .03 | 0.2 $\pm$ .03 | 0.4 $\pm$ .03 | 0.0 $\pm$ .00 | 0.1 $\pm$ .00 | 0.0 $\pm$ .05 | 0.0 $\pm$ .03
$\eta=1$ | 0.1 $\pm$ .01 | 0.0 $\pm$ .01 | 0.2 $\pm$ .00 | 0.1 $\pm$ .01 | 0.1 $\pm$ .01 | 0.0 $\pm$ .01 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00 | 0.0 $\pm$ .05 | 0.0 $\pm$ .03
| ADE under control $(t=0)$
$\eta=10$ | 0.6 $\pm$ .06 | 0.1 $\pm$ .04 | 0.1 $\pm$ .06 | 0.1 $\pm$ .04 | 0.7 $\pm$ .07 | 0.2 $\pm$ .05 | 0.3 $\pm$ .01 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00
$\eta=1$ | 0.1 $\pm$ .10 | 0.3 $\pm$ .10 | 0.1 $\pm$ .10 | 0.3 $\pm$ .04 | 0.3 $\pm$ .10 | 0.2 $\pm$ .04 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00 | 0.1 $\pm$ .00
| ATE
$\eta=10$ | 0.3 $\pm$ .05 | 0.8 $\pm$ .03 | 0.1 $\pm$ .05 | 0.5 $\pm$ .03 | 0.9 $\pm$ .05 | 0.2 $\pm$ .04 | 0.3 $\pm$ .01 | 0.0 $\pm$ .01 | 0.1 $\pm$ .04 | 0.1 $\pm$ .03
$\eta=1$ | 0.1 $\pm$ .09 | 0.3 $\pm$ .04 | 0.3 $\pm$ .10 | 0.3 $\pm$ .04 | 0.2 $\pm$ .10 | 0.2 $\pm$ .04 | 0.0 $\pm$ .01 | 0.2 $\pm$ .01 | 0.1 $\pm$ .04 | 0.1 $\pm$ .03

##  7 Conclusion

This work makes a contribution to the field of causal mediation analysis (CMA)
by proposing a novel method, IMAVAE, that can handle situations where the
mediator is indirectly observed and observed covariates are likely to be
present. Our approach builds on existing CMA methods and leverages the
identifiable variational autoencoder (iVAE) model architecture to provide a
powerful tool for estimating direct and mediated effects. We have demonstrated
the effectiveness of IMAVAE in mediation effect estimation through theoretical
analysis and empirical evaluations. Specifically, we have proved the
identifiability of the joint distribution learned by IMAVAE and demonstrated
the disentanglement of mediators in control and treatment groups. Overall, our
proposed method offers a promising avenue for CMA in settings with much more
complex data, where traditional methods may struggle to provide accurate
estimates.

## References

  * Athey and Imbens [2017] Susan Athey and Guido W Imbens.  The state of applied econometrics: Causality and policy evaluation.  _Journal of Economic perspectives_ , 31(2):3–32, 2017.
  * Glass et al. [2013] Thomas A Glass, Steven N Goodman, Miguel A Hernán, and Jonathan M Samet.  Causal inference in public health.  _Annual review of public health_ , 34:61–75, 2013.
  * Imai et al. [2010] Kosuke Imai, Luke Keele, and Dustin Tingley.  A general approach to causal mediation analysis.  _Psychological methods_ , 15(4):309, 2010.
  * Rothman and Greenland [2005] Kenneth J Rothman and Sander Greenland.  Causation and causal inference in epidemiology.  _American journal of public health_ , 95(S1):S144–S150, 2005.
  * Guo et al. [2020] Ruocheng Guo, Lu Cheng, Jundong Li, P Richard Hahn, and Huan Liu.  A survey of learning causality with data: Problems and methods.  _ACM Computing Surveys (CSUR)_ , 53(4):1–37, 2020\.
  * Li and Zhu [2022] Zongyu Li and Zhenfeng Zhu.  A survey of deep causal model.  _arXiv preprint arXiv:2209.08860_ , 2022.
  * Hultman et al. [2018] Rainbo Hultman, Kyle Ulrich, Benjamin D. Sachs, Cameron Blount, David E. Carlson, Nkemdilim Ndubuizu, Rosemary C. Bagot, Eric M. Parise, Mai-Anh T. Vu, Neil M. Gallagher, Joyce Wang, Alcino J. Silva, Karl Deisseroth, Stephen D. Mague, Marc G. Caron, Eric J. Nestler, Lawrence Carin, and Kafui Dzirasa.  Brain-wide electrical spatiotemporal dynamics encode depression vulnerability.  _Cell_ , 173(1):166–180.e14, 2018.  ISSN 0092-8674.  doi: https://doi.org/10.1016/j.cell.2018.02.012.  URL https://www.sciencedirect.com/science/article/pii/S0092867418301569.
  * Mague et al. [2022] Stephen D. Mague, Austin Talbot, Cameron Blount, Kathryn K. Walder-Christensen, Lara J. Duffney, Elise Adamson, Alexandra L. Bey, Nkemdilim Ndubuizu, Gwenaëlle E. Thomas, Dalton N. Hughes, Yael Grossman, Rainbo Hultman, Saurabh Sinha, Alexandra M. Fink, Neil M. Gallagher, Rachel L. Fisher, Yong-Hui Jiang, David E. Carlson, and Kafui Dzirasa.  Brain-wide electrical dynamics encode individual appetitive social behavior.  _Neuron_ , 110(10):1728–1741.e7, 2022.  ISSN 0896-6273.  doi: https://doi.org/10.1016/j.neuron.2022.02.016.  URL https://www.sciencedirect.com/science/article/pii/S0896627322001817.
  * Pearl [2001] Judea Pearl.  Direct and indirect effects.  _Probabilistic and Causal Inference: The Works of Judea Pearl_ , page 373, 2001.
  * Wright [1923] Sewall Wright.  The theory of path coefficients a reply to niles’s criticism.  _Genetics_ , 8(3):239, 1923.
  * Wright [1934] Sewall Wright.  The method of path coefficients.  _The annals of mathematical statistics_ , 5(3):161–215, 1934.
  * Baron and Kenny [1986] Reuben M Baron and David A Kenny.  The moderator–mediator variable distinction in social psychological research: Conceptual, strategic, and statistical considerations.  _Journal of personality and social psychology_ , 51(6):1173, 1986.
  * Balke and Pearl [2013] Alexander Balke and Judea Pearl.  Counterfactuals and policy analysis in structural models.  _arXiv preprint arXiv:1302.4929_ , 2013.
  * Jöreskog et al. [1996] Karl G Jöreskog, Fan Yang, G Marcoulides, and R Schumacker.  Nonlinear structural equation models: The kenny-judd model with interaction effects.  _Advanced structural equation modeling: Issues and techniques_ , 3:57–88, 1996.
  * Khemakhem et al. [2020] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen.  Variational autoencoders and nonlinear ica: A unifying framework.  In _International Conference on Artificial Intelligence and Statistics_ , pages 2207–2217. PMLR, 2020.
  * MacKinnon [2012] David P MacKinnon.  _Introduction to statistical mediation analysis_.  Routledge, 2012.
  * MacKinnon and Dwyer [1993] David P MacKinnon and James H Dwyer.  Estimating mediated effects in prevention studies.  _Evaluation review_ , 17(2):144–158, 1993.
  * Pearl [2014] Judea Pearl.  Interpretation and identification of causal mediation.  _Psychological methods_ , 19(4):459, 2014.
  * Rijnhart et al. [2019] Judith JM Rijnhart, Jos WR Twisk, Iris Eekhout, and Martijn W Heymans.  Comparison of logistic-regression based methods for simple mediation analysis with a dichotomous outcome variable.  _BMC medical research methodology_ , 19:1–10, 2019.
  * Holland [1988] Paul W Holland.  Causal inference, path analysis and recursive structural equations models.  _ETS Research Report Series_ , 1988(1):i–50, 1988\.
  * Rubin [1974] Donald B Rubin.  Estimating causal effects of treatments in randomized and nonrandomized studies.  _Journal of educational Psychology_ , 66(5):688, 1974.
  * Zheng and van der Laan [2012] Wenjing Zheng and Mark J van der Laan.  Targeted maximum likelihood estimation of natural direct effects.  _The international journal of biostatistics_ , 8(1):1–40, 2012.
  * Huber et al. [2013] Martin Huber, Michael Lechner, and Conny Wunsch.  The performance of estimators based on the propensity score.  _Journal of Econometrics_ , 175(1):1–21, 2013\.
  * Lange et al. [2012] Theis Lange, Stijn Vansteelandt, and Maarten Bekaert.  A simple unified approach for estimating natural direct and indirect effects.  _American journal of epidemiology_ , 176(3):190–195, 2012.
  * Holland [1986] Paul W Holland.  Statistics and causal inference.  _Journal of the American statistical Association_ , 81(396):945–960, 1986.
  * Alaa and Van Der Schaar [2017] Ahmed M Alaa and Mihaela Van Der Schaar.  Bayesian inference of individualized treatment effects using multi-task gaussian processes.  _Advances in neural information processing systems_ , 30, 2017.
  * Jiang et al. [2023] Ziyang Jiang, Zhuoran Hou, Yiling Liu, Yiman Ren, Keyu Li, and David Carlson.  Estimating causal effects using a multi-task deep ensemble.  _arXiv preprint arXiv:2301.11351_ , 2023.
  * Louizos et al. [2017] Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling.  Causal effect inference with deep latent-variable models.  _Advances in neural information processing systems_ , 30, 2017.
  * Shalit et al. [2017] Uri Shalit, Fredrik D Johansson, and David Sontag.  Estimating individual treatment effect: generalization bounds and algorithms.  In _International Conference on Machine Learning_ , pages 3076–3085. PMLR, 2017.
  * Xu et al. [2022] Siqi Xu, Lin Liu, and Zhonghua Liu.  Deepmed: Semiparametric causal mediation analysis with debiased deep learning.  _arXiv preprint arXiv:2210.04389_ , 2022.
  * Cheng et al. [2022] Lu Cheng, Ruocheng Guo, and Huan Liu.  Causal mediation analysis with hidden confounders.  In _Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining_ , pages 113–122, 2022.
  * Xu et al. [2023] Ziqi Xu, Debo Cheng, Jiuyong Li, Jixue Liu, Lin Liu, and Ke Wang.  Disentangled representation for causal mediation analysis.  _arXiv preprint arXiv:2302.09694_ , 2023.
  * Chén et al. [2018] Oliver Y Chén, Ciprian Crainiceanu, Elizabeth L Ogburn, Brian S Caffo, Tor D Wager, and Martin A Lindquist.  High-dimensional multivariate mediation with application to neuroimaging data.  _Biostatistics_ , 19(2):121–136, 2018.
  * Nath et al. [2023] Tanmay Nath, Brian Caffo, Tor Wager, and Martin A Lindquist.  A machine learning based approach towards high-dimensional mediation analysis.  _NeuroImage_ , 268:119843, 2023.
  * Zhang et al. [2021a] Haixiang Zhang, Jun Chen, Yang Feng, Chan Wang, Huilin Li, and Lei Liu.  Mediation effect selection in high-dimensional and compositional microbiome data.  _Statistics in medicine_ , 40(4):885–896, 2021a.
  * Perera et al. [2022] Chamila Perera, Haixiang Zhang, Yinan Zheng, Lifang Hou, Annie Qu, Cheng Zheng, Ke Xie, and Lei Liu.  Hima2: high-dimensional mediation analysis and its application in epigenome-wide dna methylation data.  _BMC bioinformatics_ , 23(1):1–14, 2022.
  * Yang et al. [2021] Tianzhong Yang, Jingbo Niu, Han Chen, and Peng Wei.  Estimation of total mediation effect for high-dimensional omics mediators.  _BMC bioinformatics_ , 22:1–17, 2021.
  * Zhang et al. [2021b] Haixiang Zhang, Yinan Zheng, Lifang Hou, Cheng Zheng, and Lei Liu.  Mediation analysis for survival data with high-dimensional mediators.  _Bioinformatics_ , 37(21):3815–3821, 2021b.
  * Zhang et al. [2016] Haixiang Zhang, Yinan Zheng, Zhou Zhang, Tao Gao, Brian Joyce, Grace Yoon, Wei Zhang, Joel Schwartz, Allan Just, Elena Colicino, et al.  Estimating and testing high-dimensional mediation effects in epigenetic studies.  _Bioinformatics_ , 32(20):3150–3154, 2016.
  * Hicks and Tingley [2011] Raymond Hicks and Dustin Tingley.  Causal mediation analysis.  _The Stata Journal_ , 11(4):605–619, 2011.
  * MacKinnon et al. [2007] David P MacKinnon, Amanda J Fairchild, and Matthew S Fritz.  Mediation analysis.  _Annu. Rev. Psychol._ , 58:593–614, 2007.
  * Robins and Greenland [1992] James M Robins and Sander Greenland.  Identifiability and exchangeability for direct and indirect effects.  _Epidemiology_ , 3(2):143–155, 1992.
  * Rosenbaum and Rubin [1983] Paul R Rosenbaum and Donald B Rubin.  The central role of the propensity score in observational studies for causal effects.  _Biometrika_ , 70(1):41–55, 1983.
  * Rubin [2005] Donald B Rubin.  Causal inference using potential outcomes: Design, modeling, decisions.  _Journal of the American Statistical Association_ , 100(469):322–331, 2005.
  * Deisseroth [2011] Karl Deisseroth.  Optogenetics.  _Nature Methods_ , 8(1):26–29, Jan 2011.  ISSN 1548-7105.  doi: 10.1038/nmeth.f.324.  URL https://doi.org/10.1038/nmeth.f.324.
  * Limousin and Foltynie [2019] Patricia Limousin and Tom Foltynie.  Long-term outcomes of deep brain stimulation in parkinson disease.  _Nature Reviews Neurology_ , 15(4):234–242, Apr 2019.  ISSN 1759-4766.  doi: 10.1038/s41582-019-0145-9.  URL https://doi.org/10.1038/s41582-019-0145-9.
  * Tufail et al. [2011] Yusuf Tufail, Anna Yoshihiro, Sandipan Pati, Monica M. Li, and William J. Tyler.  Ultrasonic neuromodulation by brain stimulation with transcranial ultrasound.  _Nature Protocols_ , 6(9):1453–1470, Sep 2011\.  ISSN 1750-2799.  doi: 10.1038/nprot.2011.371.  URL https://doi.org/10.1038/nprot.2011.371.
  * Gallagher et al. [2017] Neil Gallagher, Kyle R Ulrich, Austin Talbot, Kafui Dzirasa, Lawrence Carin, and David E Carlson.  Cross-spectral factor analysis.  _Advances in neural information processing systems_ , 30, 2017.
  * Carlson et al. [2023] David Carlson, Sunil Kumar, and Kafui Dzirasa.  Multi-region local field potential recordings during a tail-suspension test.  _Duke Research Data Repository_ , 2023.
  * Vinokur et al. [1995] Amiram D Vinokur, Richard H Price, and Yaacov Schul.  Impact of the jobs intervention on unemployed workers varying in risk for depression.  _American journal of community psychology_ , 23(1):39–74, 1995.
  * Huber et al. [2016] Martin Huber, Michael Lechner, and Giovanni Mellace.  The finite sample performance of estimators for mediation analysis under sequential conditional independence.  _Journal of Business & Economic Statistics_, 34(1):139–160, 2016.
  * Talbot et al. [2020] Austin Talbot, David Dunson, Kafui Dzirasa, and David Carlson.  Supervised autoencoders learn robust joint factor models of neural activity.  _arXiv preprint arXiv:2004.05209_ , 2020.
  * Ben-Israel [1999] Adi Ben-Israel.  The change-of-variables formula using matrix volume.  _SIAM Journal on Matrix Analysis and Applications_ , 21(1):300–312, 1999.

##  Appendix A Detailed Proof of Identifiability

The proof of Theorem 4.4 closely resembles the proof presented in [15], which
consists of 3 main steps.

#### Step 1

The first step is to transform the equality of observed data distributions
into equality of noise-free distributions using a convolutional trick based on
the $1^{st}$ assumption in Theorem 4.4. Similar to [15], we introduce the
volume of a matrix denoted by $\text{vol}\,A$ as the product of singular
values of $A$. When $A$ is full column rank,
$\text{vol}\,A=\sqrt{\text{det}\,A^{\text{T}}A}$, and when $A$ is invertible,
$\text{vol}\,A=|\text{det}\,A|$. We use this matrix volume as a replacement
for the absolute determinant of Jacobian [53] in the change of variables
formula, which is most useful when the Jacobian is a rectangular matrix
$(d<D)$. Suppose we have two sets of parameters
$(\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda})$ and
$(\tilde{\textbf{f}},\tilde{\textbf{h}},\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}})$
such that
$p_{\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{x},y|\boldsymbol{u})=p_{\tilde{\textbf{f}},\tilde{\textbf{h}},\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}}}(\boldsymbol{x},y|\boldsymbol{u})$.
Recall that
$\boldsymbol{\psi}=(\textbf{f},\textbf{h}):\mathcal{Z}\times\mathcal{U}\rightarrow\mathcal{X}\times\mathcal{Y}$
and define the concatenated vector of $\boldsymbol{x}$ and $y$ as
$\boldsymbol{v}$. We have:

$\displaystyle\int_{\mathcal{Z}}p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})p_{\boldsymbol{\psi}}(\boldsymbol{v}|\boldsymbol{z},\boldsymbol{u})d\boldsymbol{z}$
$\displaystyle=\int_{\mathcal{Z}}p_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}}}(\boldsymbol{z}|\boldsymbol{u})p_{\tilde{\boldsymbol{\psi}}}(\boldsymbol{v}|\boldsymbol{z},\boldsymbol{u})d\boldsymbol{z},$
(11)
$\displaystyle\int_{\mathcal{Z}}p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\boldsymbol{\psi}(\boldsymbol{z},\boldsymbol{u}))d\boldsymbol{z}$
$\displaystyle=\int_{\mathcal{Z}}p_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}}}(\boldsymbol{z}|\boldsymbol{u})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\tilde{\boldsymbol{\psi}}(\boldsymbol{z},\boldsymbol{u}))d\boldsymbol{z}.$
(12)

Next, we apply change of variables
$\bar{\boldsymbol{v}}=\boldsymbol{\psi}(\boldsymbol{z},\boldsymbol{u})$ on the
left hand side (LHS) and
$\bar{\boldsymbol{v}}=\tilde{\boldsymbol{\psi}}(\boldsymbol{z},\boldsymbol{u})$
on the right hand side (RHS):

$\displaystyle\int_{\mathcal{X}\times\mathcal{Y}}p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})|\boldsymbol{u})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\bar{\boldsymbol{v}})\text{vol}\,J_{\boldsymbol{\psi}^{-1}}(\bar{\boldsymbol{v}})d\bar{\boldsymbol{v}}$
(13)
$\displaystyle=\int_{\mathcal{X}\times\mathcal{Y}}p_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}}}(\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})|\boldsymbol{u})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\bar{\boldsymbol{v}})\text{vol}\,J_{\tilde{\boldsymbol{\psi}}^{-1}}(\bar{\boldsymbol{v}})d\bar{\boldsymbol{v}},$

where $J$ denotes the Jacobian and recall that
$\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}$ is the projection of
$\boldsymbol{\psi}^{-1}$ on $\mathcal{Z}$. Next, we introduce
$\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}(\boldsymbol{v})=p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})|\boldsymbol{u})\text{vol}\,J_{\boldsymbol{\psi}^{-1}}(\boldsymbol{v})\mathbb{I}_{\mathcal{X}\times\mathcal{Y}}(\boldsymbol{v})$
on the LHS and similarly on the RHS, then, following [15], Equation 13 reduces
to:

$\displaystyle\int_{\mathcal{X}\times\mathcal{Y}}\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}(\bar{\boldsymbol{v}})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\bar{\boldsymbol{v}})d\bar{\boldsymbol{v}}$
$\displaystyle=\int_{\mathcal{X}\times\mathcal{Y}}\tilde{p}_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}},\tilde{\boldsymbol{\psi}},\boldsymbol{u}}(\bar{\boldsymbol{v}})p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\boldsymbol{v}-\bar{\boldsymbol{v}})d\bar{\boldsymbol{v}},$
(14)
$\displaystyle(\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}*p_{\boldsymbol{\epsilon},\boldsymbol{\xi}})(\boldsymbol{v})$
$\displaystyle=(\tilde{p}_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}},\tilde{\boldsymbol{\psi}},\boldsymbol{u}}*p_{\boldsymbol{\epsilon},\boldsymbol{\xi}})(\boldsymbol{v}),$
(15) $\displaystyle
F[\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}](\omega)\varphi_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\omega)$
$\displaystyle=F[\tilde{p}_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}},\tilde{\boldsymbol{\psi}},\boldsymbol{u}}](\omega)\varphi_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\omega),$
(16) $\displaystyle
F[\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}](\omega)$
$\displaystyle=F[\tilde{p}_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}},\tilde{\boldsymbol{\psi}},\boldsymbol{u}}](\omega),$
(17)
$\displaystyle\tilde{p}_{\textbf{S},\boldsymbol{\lambda},\boldsymbol{\psi},\boldsymbol{u}}(\boldsymbol{v})$
$\displaystyle=\tilde{p}_{\tilde{\textbf{S}},\tilde{\boldsymbol{\lambda}},\tilde{\boldsymbol{\psi}},\boldsymbol{u}}(\boldsymbol{v}),$
(18)

where we use $*$ for the convolutional operator in Equation 15, we use
$F[\cdot]$ to designate the Fourier transform in Equation 16,
$\varphi_{\boldsymbol{\epsilon},\boldsymbol{\xi}}=F[p_{\boldsymbol{\epsilon},\boldsymbol{\xi}}]$
according to the definition of the characteristic function, and we drop
$\varphi_{\boldsymbol{\epsilon},\boldsymbol{\xi}}(\omega)$ from both sides in
Equation 17 as it is non-zero almost everywhere by the $1^{st}$ assumption in
Theorem 4.4.

Equation 18 is valid for all
$(\boldsymbol{v},\boldsymbol{u})\in\mathcal{X}\times\mathcal{Y}\times\mathcal{U}$.
It indicates that for the observed data distributions to be the same, the
noise-free distributions have to be the same.

#### Step 2

The second step in [15] is about removing all terms that are either a function
of observations $\boldsymbol{x}$ and $y$ or auxiliary variables
$\boldsymbol{u}$, which is done by introducing the points provided by the
$4^{th}$ assumption in Theorem 4.4, and using $\boldsymbol{u}_{0}$ as a
“pivot”. Specifically, by taking the logarithm on both sides of Equation 18
and plugging in Equation 6 for $p_{\textbf{S},\boldsymbol{\lambda}}$, we get:

$\begin{split}&\log\text{vol}\,J_{\boldsymbol{\psi}^{-1}}(\boldsymbol{v})+\sum_{i=1}^{d}\left(\log
Q_{i}(\psi^{-1}_{i\,|\boldsymbol{z}}(\boldsymbol{v}))-\log
C_{i}(\boldsymbol{u})+\sum_{j=1}^{k}S_{i,j}(\psi^{-1}_{i\,|\boldsymbol{z}}(\boldsymbol{v}))\lambda_{i,j}(\boldsymbol{u})\right)\\\
&=\log\text{vol}\,J_{\tilde{\boldsymbol{\psi}}^{-1}}(\boldsymbol{v})+\sum_{i=1}^{d}\left(\log\tilde{Q}_{i}(\tilde{\psi}^{-1}_{i\,|\boldsymbol{z}}(\boldsymbol{v}))-\log\tilde{C}_{i}(\boldsymbol{u})+\sum_{j=1}^{k}\tilde{S}_{i,j}(\tilde{\psi}^{-1}_{i\,|\boldsymbol{z}}(\boldsymbol{v}))\tilde{\lambda}_{i,j}(\boldsymbol{u})\right).\end{split}$
(19)

Let $\boldsymbol{u}_{0},...,\boldsymbol{u}_{dk}$ be the points provided by the
$4^{th}$ assumption of Theorem 4.4, and define
$\bar{\boldsymbol{\lambda}}(\boldsymbol{u})\coloneqq\boldsymbol{\lambda}(\boldsymbol{u})-\boldsymbol{\lambda}(\boldsymbol{u}_{0})$.
Then for $l=1,...,dk$, we plug each of those $\boldsymbol{u}_{l}$ in Equation
19 to obtain $dk+1$ equations and subtract the first equation for
$\boldsymbol{u}_{0}$ from the remaining $dk$ equations to get:

$\langle\textbf{S}(\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})),\bar{\boldsymbol{\lambda}}(\boldsymbol{u}_{l})\rangle+\sum_{i=1}^{d}\log\frac{C_{i}(\boldsymbol{u}_{0})}{C_{i}(\boldsymbol{u}_{l})}=\langle\tilde{\textbf{S}}(\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})),\bar{\tilde{\boldsymbol{\lambda}}}(\boldsymbol{u}_{l})\rangle+\sum_{i=1}^{d}\log\frac{\tilde{C}_{i}(\boldsymbol{u}_{0})}{\tilde{C}_{i}(\boldsymbol{u}_{l})},$
(20)

where $\langle\cdot,\cdot\rangle$ denotes inner product. Let $L$ be the matrix
defined in the $4^{th}$ assumption of Theorem 4.4, and $\tilde{L}$ similarly
for $\tilde{\boldsymbol{\lambda}}$ ($\tilde{L}$ is not necessarily
invertible). Define
$b_{l}=\sum_{i=1}^{d}\log\frac{\tilde{C}_{i}(\boldsymbol{u}_{0})C_{i}(\boldsymbol{u}_{l})}{C_{i}(\boldsymbol{u}_{0})\tilde{C}_{i}(\boldsymbol{u}_{l})}$
and b the vector of all $b_{l}$ for $l=1,...,dk$. Expressing Equation 20 in
matrix form, we get:

$L^{\text{T}}\textbf{S}(\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v}))=\tilde{L}^{\text{T}}\tilde{\textbf{S}}(\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v}))+\textbf{b}.$
(21)

Multiplying both sides of Equation 21 by the transpose of the inverse of
$L^{\text{T}}$, we get:

$\textbf{S}(\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v}))=A\tilde{\textbf{S}}(\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v}))+\textbf{c},$
(22)

where $A=L^{-\text{T}}\tilde{L}$ and $\textbf{c}=L^{-\text{T}}\textbf{b}$.
Note that here $\boldsymbol{\psi}^{-1}_{|\boldsymbol{z}}(\boldsymbol{v})$ can
be referred to either the projection from f or from h as we map the same $Z$
into $\hat{X}$ and $Y$ through $p_{\textbf{f}}$ and $p_{\textbf{h}}$,
respectively.

#### Step 3

In the last step, Khemakhem et al. [15] show that the linear transformation
$A$ is invertible for both $k=1$ and $k>1$, resulting in an equivalence
relation in the iVAE framework. This is mainly based on the $3^{rd}$
assumption in Theorem 4.4 which guarantees the existence of the $dk\times d$
Jacobian matrix of S with rank $d$. This also implies that the Jacobian of
$\tilde{\textbf{S}}\circ\tilde{\boldsymbol{\psi}}^{-1}_{|\boldsymbol{z}}$
exists and is of rank $d$ and so is $A$. We argue that the proof of
invertibility in our framework follows the same line of reasoning as that of
the iVAE.

Therefore, with Equation 22 and the invertibility of $A$, we prove that the
parameters $(\textbf{f},\textbf{h},\textbf{S},\boldsymbol{\lambda})$ are
$\sim_{A}$-identifiable.

##  Appendix B Details of Experimental Setup

#### Model Configuration

In all experiments, the encoder
$q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x},\boldsymbol{u})$, decoder
$p_{\textbf{f}}(\boldsymbol{x}|\boldsymbol{z})$, and prior
$p_{\textbf{S},\boldsymbol{\lambda}}(\boldsymbol{z}|\boldsymbol{u})$ of IMAVAE
are configured as multivariate normal distributions whose mean and covariance
are parameterized as a function of the conditioned variables using a feed-
forward neural network. As for the predictor
$g_{\boldsymbol{\gamma}}(\boldsymbol{z},\boldsymbol{u})$, it is implemented as
a simple linear regression model. It is worth noting that while alternative
stochastic models can also be employed to replace $g_{\boldsymbol{\gamma}}$,
our experiments demonstrate that the linear regression model already achieves
superior performance. See the included code link for details on reproduction
of all experiments.

#### Training Details

When training IMAVAE to minimize the objective in Equation 8, we use the Adam
optimizer and adopt parameter annealing so that the KL divergence will
gradually dominate the reconstruction error in ELBO.

#### Computing Resources

The experiments conducted on the synthetic data and the Jobs II data in
Sections 5.1 and 5.3 are of a relatively small scale and can be executed
locally. The experiment on the electrophysiological data is performed on a
computer cluster equipped with a GeForce RTX 2080 Ti GPU.

#### Data Availability

Multi-region local field potential recordings during a tail-suspension test is
an experiment comparing the electrical neural activity and behaviors of
Wildtype and Clock-$\Delta$19 genotypes of mice in the tail-suspension test.
This dataset is available to download at
https://research.repository.duke.edu/concern/datasets/zc77sr31x?locale=en for
free under a Creative Commons BY-NC Attribution-NonCommercial 4.0
International license.

Jobs II is a randomized field experiment that investigates the efficacy of a
job training intervention on unemployed workers, which can be downloaded from
the R package "mediation" or the following URL:
https://r-data.pmagunia.com/dataset/r-dataset-package-mediation-jobs.

##  Appendix C Synthetic Data Generation

The synthetic data contains $N=6000$ data points and is generated according to
the causal graphs shown in Figure 1. Specifically, for case (a) without
observed covariates, we generate $t_{i}$, $\boldsymbol{z}_{i}$,
$\boldsymbol{x}_{i}$, and $y_{i}$ for $i=1,...,N$ as follows:

$\displaystyle t_{i}$ $\displaystyle\sim\text{Bernoulli}(p),$
$\displaystyle\boldsymbol{z}_{i}$
$\displaystyle\sim\mathcal{N}(\boldsymbol{0},\sigma_{z}^{2}\textbf{I}_{d})+c\mathbb{I}(t_{i}=1)\textbf{1}_{d},$
$\displaystyle\boldsymbol{x}_{i}$
$\displaystyle=f(\boldsymbol{z}_{i})+\boldsymbol{\epsilon}_{x},$
$\displaystyle y_{i}$
$\displaystyle=\mu_{t}t_{i}+\boldsymbol{\mu}_{z}^{\text{T}}\boldsymbol{z}_{i}+\epsilon_{y},$

where $p$ is a probability parameter, $\sigma_{z}^{2}$ is a variance
parameter, $\textbf{I}_{d}$ is a $d\times d$ identity matrix,
$\mathbb{I}(\cdot)$ denotes the indicator function, $\textbf{1}_{d}$ is a
$d$-dimensional vector of all ones,
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{D}$ is a nonlinear function modeled by
an un-trained neural network, $\boldsymbol{\epsilon}_{x}\in\mathbb{R}^{D}$ and
$\epsilon_{y}\in\mathbb{R}$ are Gaussian noise terms, and $c$, $\mu_{t}$, and
$\boldsymbol{\mu}_{z}$ are coefficient constants/vector. For case (b) with
observed covariates, we generate $t_{i}$, $\boldsymbol{z}_{i}$,
$\boldsymbol{x}_{i}$, $\boldsymbol{w}_{i}$, and $y_{i}$ for $i=1,...,N$ as
follows:

$\displaystyle\boldsymbol{w}_{i}$
$\displaystyle\sim\mathcal{N}(\boldsymbol{0},\sigma_{w}^{2}\textbf{I}_{m}),$
$\displaystyle t_{i}$
$\displaystyle\sim\text{Bernoulli}(\text{sigmoid}(\boldsymbol{\mu}_{s}^{\text{T}}\boldsymbol{w}_{i})),$
$\displaystyle\boldsymbol{z}$
$\displaystyle\sim\mathcal{N}(\boldsymbol{0},\sigma_{z}^{2}\textbf{I}_{d})+c_{1}\mathbb{I}(t_{i}=1)\textbf{1}_{d}+c_{2}f_{1}(\boldsymbol{w}_{i}),$
$\displaystyle\boldsymbol{x}_{i}$
$\displaystyle=f_{2}(\boldsymbol{z}_{i})+\boldsymbol{\epsilon}_{x},$
$\displaystyle y_{i}$
$\displaystyle=\mu_{t}t_{i}+\boldsymbol{\mu}_{z}^{\text{T}}\boldsymbol{z}_{i}+\boldsymbol{\mu}_{w}^{\text{T}}\boldsymbol{w}_{i}+\epsilon_{y},$

where $\sigma_{w}^{2}$, $\sigma_{z}^{2}$ are variance parameters,
$\textbf{I}_{m},\textbf{I}_{d}$ are identity matrices with dimension $m\times
m$ and $d\times d$, respectively, $\mathbb{I}(\cdot)$ denotes the indicator
function, $\textbf{1}_{d}$ is a $d$-dimensional vector of all ones,
$f_{1}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{d},f_{2}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{D}$
are nonlinear functions modeled by un-trained neural networks,
$\boldsymbol{\epsilon}_{x}\in\mathbb{R}^{D}$ and $\epsilon_{y}\in\mathbb{R}$
are Gaussian noise terms, and
$c_{1},c_{2},\mu_{t},\boldsymbol{\mu}_{s},\boldsymbol{\mu}_{z},\boldsymbol{\mu}_{w}$
are coefficient constants/vectors.

##  Appendix D Post-processing for Tail Suspension Test Data

As elaborated in Section 5.2, we have the LFP signals from 26 mice. The full
dataset contains a total of $N=43882$ data points. The observed feature, i.e.
$\boldsymbol{x}_{i}$, represents the power spectral densities of the LFPs
recorded at 11 brain regions. These densities are evaluated from $1$ to
$56Hz$, resulting in a total of $616$ attributes, i.e.
$\boldsymbol{x}_{i}\in\mathbb{R}^{616}$. We also have treatment assignment
$t_{i}$ indicating whether the mouse corresponding to the $i^{th}$ data point
is recorded during an open field exploration $(t_{i}=0)$ or a tail suspension
test $(t_{i}=1)$.

To generate the semi-synthetic dataset, we do some post-processing on these
data. For case (a) without observed covariates, we first apply PCA to map
$\boldsymbol{x}_{i}$ into a lower-dimensional representation
$\boldsymbol{s}_{i}$. Then the true mediator $\boldsymbol{z}_{i}$ is generated
by $\boldsymbol{z}_{i}=\boldsymbol{s}_{i}+\mathbb{I}(t_{i}=1)$ where
$\mathbb{I}(\cdot)$ denotes the indicator function. The final outcome $y_{i}$
is modeled by
$y_{i}=\mu_{t}t_{i}+\boldsymbol{\mu}_{z}\boldsymbol{z}_{i}+\epsilon_{y}$ where
$\mu_{t},\boldsymbol{\mu}_{z}$ are coefficient constant/vector and
$\epsilon_{y}$ is a Gaussian noise term. For case (b), we use the genotype of
the mouse $w_{i}\in\\{0,1\\}$ as an observed covariate where $w_{i}=0$ denotes
wild type and $w_{i}=1$ denotes Clock$\Delta 19$ mutation. The true mediator
$\boldsymbol{z}_{i}$ is then generated by
$\boldsymbol{z}_{i}=\boldsymbol{s}_{i}+\mathbb{I}(t_{i}=1)+f(\boldsymbol{w}_{i})$
where $f:\mathbb{R}\rightarrow\mathbb{R}^{d}$ is a nonlinear function modeled
by a neural network. The final outcome $y_{i}$ is modeled by
$y_{i}=\mu_{t}t_{i}+\boldsymbol{\mu}_{z}\boldsymbol{z}_{i}+\mu_{w}w_{i}+\epsilon_{y}$
where $\mu_{t},\mu_{w},\boldsymbol{\mu}_{z}$ are coefficient constants/vector
and $\epsilon_{y}$ is a Gaussian noise term.

##  Appendix E Simulation Procedure for Jobs II Data

To achieve _zero_ direct, mediation, and total effects, we adopt the following
simulation procedure, similar to [31, 51], on the Jobs II dataset. Note that
all attributes other than $T,Z,Y$ are treated as observed covariates $W$ in
this analysis. $X$ in our causal graphs does not exist in this dataset.

  1. 1.

Estimate probit specifications in which we regress (a) $T$ on $W$ to get an
estimated probit coefficient $\hat{\beta}_{\text{pop}}$ and (b) $Z$ on $T$ and
$W$ to get estimated probit coefficients $\hat{\gamma}_{\text{pop}}$ and
$\hat{\delta}_{\text{pop}}$, respectively.

  2. 2.

Apply the indicator function to $Z$ so that the mediator becomes a binary
variable $Z\coloneqq\mathbb{I}(Z\geq 3)$.

  3. 3.

Discard all samples with either $T=0$ or $Z=0$, resulting in a dataset with
all non-mediated and non-treated samples.

  4. 4.

Draw independent Monte Carlo samples (500 and 1000 samples for Tables 4 and 5,
respectively) with replacement $(T^{\prime},Z^{\prime},W^{\prime},Y^{\prime})$
from the resulting dataset.

  5. 5.

Simulate the (pseudo-)treatment and (pseudo-)mediator using the following
formula:

$\displaystyle T^{\prime}$
$\displaystyle\coloneqq\mathbb{I}(W^{\prime}\hat{\beta}_{\text{pop}}+U>0),$
$\displaystyle M^{\prime}$
$\displaystyle\coloneqq\eta(T^{\prime}\hat{\gamma}_{\text{pop}}+W^{\prime}\hat{\delta}_{\text{pop}})+\alpha+V,$

where $U\sim\mathcal{N}(0,1)$, $V\sim\mathcal{N}(0,1)$, $\eta$ is a simulation
parameter which stands for the magnitude of selection into the mediator, and
we manually set $\alpha$ such that either $10\%$ or $50\%$ of the samples are
mediated (i.e. $M^{\prime}\geq 3$). Note that here the (pseudo-)mediator
$M^{\prime}$ is continuous after simulation.

With this simulation design, the ground truth of the direct, mediated, and
total effects are all _zero_.