theorem proposition lemma corollary definition remark

# Diffusion Based Causal Representation Learning

Amir Mohammad Karimi Mamaghan  KTH Royal Institute of Technology  Andrea
Dittadi  Helmholtz AI, Helmholtz Center Munich  Max Planck Institute for
Intelligent Systems, Tübingen, Germany
Stefan Bauer  Helmholtz AI, Helmholtz Center Munich  Technical University of
Munich  Karl Henrik Johansson  KTH Royal Institute of Technology
Francesco Quinzan  Department of Computer Science, University of Oxford

###### Abstract

Causal reasoning can be considered a cornerstone of intelligent systems.
Having access to an underlying causal graph comes with the promise of cause-
effect estimation and the identification of efficient and safe interventions.
However, learning causal representations remains a major challenge, due to the
complexity of many real-world systems. Previous works on causal representation
learning have mostly focused on Variational Auto-Encoders (VAE). These methods
only provide representations from a point estimate, and they are unsuitable to
handle high dimensions. To overcome these problems, we proposed a new
Diffusion-based Causal Representation Learning (DCRL) algorithm. This
algorithm uses diffusion-based representations for causal discovery. DCRL
offers access to infinite dimensional latent codes, which encode different
levels of information in the latent code. In a first proof of principle, we
investigate the use of DCRL for causal representation learning. We further
demonstrate experimentally that this approach performs comparably well in
identifying the causal structure and causal variables.

##  1 Introduction

Causal representation learning consists in uncovering a system’s latent causal
factors and their relationships, from observed low-level data. Causal
representation learning finds applicability in domains such as autonomous
driving [Schölkopf et al., 2021], robotics [Hellström, 2021], healthcare
[Anwar et al., 2014], climate studies [Runge et al., 2019], epidemiology
[Hernán et al., 2000, Robins et al., 2000], and finance [Hiemstra and Jones,
1994]. In these tasks, the underlying causal variables are often unknown, and
we only have access to low-level representations.

Causal representation learning is a challenging problem. In fact, identifying
latent causal factors is generally impossible from observational data. There
has been an ongoing effort to study sets of assumptions that ensure the
identifiability of causal variables and their relationships [Yang et al.,
2020, Schölkopf et al., 2021, Liu et al., 2022, Subramanian et al., 2022,
Brehmer et al., 2022]. These approaches consider the availability of
additional information or they use assumptions on the underlying causal
structure of the DGP. Interestingly, Brehmer et al. [2022] consider a weak
form of supervision in which we have access to a data pair, corresponding to
the state of the system before and after a random, unknown intervention.
Brehmer et al. [2022] prove that, in this weakly-supervised setting, the
structure and the causal variables are identifiable up to a relabeling and
element-wise reparameterization.

There has been a growing interest in leveraging generative models to learn
causal representations with specific properties. For example, disentangled and
object-centric representations have been shown to be helpful for complex
downstream tasks and generalization [Van Steenkiste et al., 2019, Dittadi et
al., 2022, Wu et al., 2022, Yoon et al., 2023, Papa et al., 2022]. Variational
Autoencoders (VAE) [Kingma and Welling, 2014] are among the most widely
studied generative models, and they have been successfully used for
disentanglement and causal representation learning [Locatello et al., 2020,
Brehmer et al., 2022]. However, the problem of learning causal representations
has not yet been approached with more powerful generative models.

Recently, diffusion models have emerged as state-of-the-art generative models,
and they have demonstrated remarkable success across several domains [Dhariwal
and Nichol, 2021a, Ramesh et al., 2022, Saharia et al., 2022, Ho et al.,
2022b, Höppe et al., 2022]. Diffusion models draw on concepts and principles
from diffusion processes to learn the data distribution [Mehrjou et al., 2017,
Ho et al., 2022a, Luhman and Luhman, 2021, Dhariwal and Nichol, 2021b, Saremi
et al., 2018, Chen et al., 2021, Cai et al., 2020, Niu et al., 2020, Sajjadi
et al., 2018, Song et al., 2021b, a, Sohl-Dickstein et al., 2015, 2015, Ho et
al., 2020, Song et al., 2021c]. These models exploit diffusion behavior to
produce diverse, high-quality, and realistic samples. Furthermore, diffusion-
based models have the appealing property of infinite-dimensional latent codes
[Abstreiter et al., 2022], which allows to efficiently learn representations
across different downstream tasks. Despite their remarkable performance and
advantages, diffusion models have not yet been employed for causal
representation learning, indicating that their potential has yet to be
explored in this context.

#### Our contribution.

In this work, we study the connection between diffusion-based models and
causal structure learning. In particular, our contributions are the following:

  * •

We propose DCRL, a diffusion-based model for causal representation learning.
We study and test the connection between the learned representations of DCRL
with causal variables. To accomplish this, we utilize both finite and
infinite-dimensional representations.

  * •

We derive the Evidence Lower Bound (ELBO) for DCRL, in the case of both finite
and infinite-dimensional representations.

  * •

We empirically illustrate that the noise and diffusion-based representations
contain equivalent information about the underlying causal variables and
causal mechanisms, and can be used interchangeably.

##  2 Related Work

#### Diffusion-based Representation Learning.

Learning representations with diffusion models remains a relatively unexplored
area. Several works try to train an external module (e.g., an encoder) along
with the score function of the diffusion model to extract representations.
Abstreiter et al. [2022] and Mittal et al. [2022] condition the score function
of a diffusion model on a time-independent and time-dependent encoder and
obtain finite and infinite-dimensional representations, respectively. Wang et
al. [2023] use the same conditioning but regularizes the objective function
with the mutual information between the input data and learned
representations. Traub [2022] does the same conditioning but they use Latent
Diffusion Models [Rombach et al., 2022] where the inputs of the diffusion
model are latent variables obtained from applying a pre-trained autoencoder on
the input. Furthermore, Kwon et al. [2022] proposes an asymmetric reverse
process that discovers the semantic latent space of a frozen diffusion model
where modification in the space synthesizes various attributes on input
images. However, in principle, diffusion models lack a semantic latent space
and it’s unclear how to efficiently learn representations using their
capabilities.

#### Causal Representation Learning.

Given the inherent challenges of identifiability in causal representation
learning, many previous studies have tackled this issue by imposing certain
assumptions on the dataset or the causal structure. Several previous methods
rely on additional knowledge of the data generation process, such as knowledge
of the causal graph or labels for the high-level causal variables. CausalGAN
[Kocaoglu et al., 2017] requires the structure of the underlying causal graph
to be known. Yang et al. [2020] and Liu et al. [2022] assume a linear
structural equation model, and they require additional information associated
with the true causal concepts as supervising signals. Similar to Yang et al.
[2020], Komanduri et al. [2022] assume the availability of supplementary
supervision labels, but without requiring mutual independence among factors.
Von Kügelgen et al. [2021] investigate self-supervised causal representation
learning by utilizing a known, but non-trivial, causal graph between content
and style factors. Subramanian et al. [2022] applies Bayesian structure
learning in the latent space and relies on having interventional samples. For
an overview of causal representation learning we refer to Schölkopf et al.
[2021]. Other relevant work closely related to causal representation learning
includes disentangled representations and independent component analysis
[Locatello et al., 2019, Shu et al., 2019, Lachapelle et al., 2022, Hyvärinen
and Oja, 2000, Khemakhem et al., 2020, Ahuja et al., 2022].

###  2.1 Overview

The fundamental concept behind diffusion-based generative models is to learn
to generate data by inverting a diffusion process. Diffusion models comprise
two processes: a forward process and a backward process. The forward process
gradually adds noise to data and maps data to (almost) pure noise. The
backward process, on the other hand, is used to go from a noise sample back to
the original data space.

The forward process is defined by a stochastic differential equation (SDE)
across a continuous time domain $t\in[0,1]$, aiming to transform the data
distribution to a known prior distribution, typically a standard multivariate
Gaussian. Given $x_{0}$ sampled from a data distribution $p(x)$, the forward
process constructs a trajectory $(x_{t})_{t\in[0,1]}$ across the time domain.
We utilize the Variance Exploding SDE [Song et al., 2021c] for the forward
process, which is defined as:

$dx=f(x,t)+g(t)dw:=\sqrt{\frac{d[\sigma^{2}(t)]}{dt}}dw,$

where $w$ is the standard Wiener process and $\sigma^{2}(t)$ is the noise
variance of the diffusion process at time $t$. The backward process is also
formulated as an SDE in the following manner:

$dx=[f(x,t)-g^{2}(t)\nabla_{x}\log p_{t}(x)]dt+g(t)d\bar{w}\ ,$

where $\bar{w}$ is the standard Wiener process in reverse time.

#### Score matching.

To use this backward process, the score function $\nabla_{x}\log p_{t}(x)$ is
required. It is usually approximated by a neural score function
$s_{\theta}(\cdot)$ which can be trained by Explicit Score Matching [Hyvärinen
and Dayan, 2005] defined as:

$\displaystyle\mathcal{L}(\theta)=\mathbb{E}_{t}\Biggl{[}\lambda(t)\mathbb{E}_{p(x_{t})}\Big{[}||s_{\theta}(x_{t},t)-\nabla_{x_{t}}\log
p_{t}(x_{t})||^{2}\Big{]}\Biggl{]},$

However, the ground-truth score function $\nabla_{x}\log p_{t}(x)$ is
generally not known. Vincent [2011] addresses this issue by proposing
Denoising Score Matching. The approximate score function is then learned by
minimizing the loss function:

$\displaystyle\mathcal{L}(\theta)=\Biggl{[}\lambda(t)\mathbb{E}_{x_{0}}\mathbb{E}_{p(x_{t}|x_{0})}\Big{[}||s_{\theta}(x_{t},t)-\nabla_{x_{t}}\log
p_{t}(x_{t}|x_{0})||^{2}\Big{]}\Biggl{]},$

where the conditional distribution of $x_{t}$ given $x_{0}$ is
$p_{t}(x_{t}|x_{0})=\mathcal{N}(x_{t};x_{0},[\sigma^{2}(t)-\sigma^{2}(0)]\mathbf{I})$
and $\lambda(t)$ is a positive weighting function. This objective function
originates from the evidence lower bound (ELBO) of the data distribution, and
it’s been shown that with a specific weighting function, this objective
function becomes exactly a term in the ELBO Song et al. [2021c]. For more
details, see Appendix A.

#### Conditional Score Matching.

We can modify Denoising Score Matching to perform representation learning
while training the score function. Abstreiter et al. [2022] proposes
conditional denoising score matching defined as:

$\displaystyle\mathcal{L}(\theta,\phi)=\mathbb{E}_{t}\Biggl{[}\lambda(t)\mathbb{E}_{x_{0}}\mathbb{E}_{p(x_{t}|x_{0})}\Big{[}||s_{\theta}(x_{t},E_{\phi}(x_{0}),t)-\nabla_{x_{t}}\log
p_{t}(x_{t}|x_{0})||^{2}\Big{]}\Biggl{]},$ (1)

where the score function is conditioned on a module $E_{\phi}(x_{0})$ which
provides additional information about the data to the diffusion model through
a learned encoder with parameters $\phi$. In fact, the encoder learns to
extract necessary information from $x_{0}$ in a reduced-dimensional space that
helps recover $x_{0}$ by denoising $x_{t}$. Abstreiter et al. [2022] also
presents an alternative objective where the encoder is a function of time.
Formally, the new objective is

$\displaystyle\mathcal{L}(\theta,\phi)=\mathbb{E}_{t}\Biggl{[}\lambda(t)\mathbb{E}_{x_{0}}\mathbb{E}_{p(x_{t}|x_{0})}\Big{[}||s_{\theta}(x_{t},E_{\phi}(x_{0},t),t)-\nabla_{x_{t}}\log
p_{t}(x_{t}|x_{0})||^{2}\Big{]}\Biggl{]},$ (2)

With this objective, the encoder learns a representation trajectory of $x_{0}$
instead of a single representation. Training this system has the potential to
minimize the objective to zero, motivating the encoder $E_{\phi}(.)$ to learn
meaningful, distinct representations at different timesteps [Abstreiter et
al., 2022, Mittal et al., 2022].

###  2.2 Comparison with Other Generative Models

The key difference between the other generative models and diffusion-based
representations is that other generative models are only concerned with one
finite code and all the information is encoded into this single code while in
the latter, different levels of information are encoded along an infinite-
dimensional code, i.e., the encoder is conditioned on time $t$ and produces a
trajectory-based representation $(E_{\phi}(x_{0},t))_{t\in[0,1]}$. Within this
representation, various points along the trajectory contain different levels
of information, as highlighted by Mittal et al. [2022]. In this work, we first
explore a time-independent single code where we employ Eq. 1 and show that
with a certain weighting function, this objective function will become the
ELBO. Then, we apply the same experiments with infinite-dimensional latent
code (Eq. 2) and study the benefits and implications of these formulations for
causal representation learning.

##  3 Problem Description

Figure 1:  Overview of our framework. Here we have a paired image of a face
before and after an intervention (the smile). The paired image is mapped to
latent variables by a stochastic encoder. The intervention target is
determined by applying the intervention encoder to these latent variables. To
maintain the weakly supervised structure, the latent variables are projected
into a new pair and then, serve as the conditioning module for a conditional
diffusion model (The projected latent variables are diffusion-based
representations of the input pair). Finally, they are utilized in neural
solution functions together with the intervention target to obtain the latent
causal variables.

We consider a system that is described by an unknown underlying SCM on the
latent causal variable $Z$ where we have access to low-level data pairs
$(x,\tilde{x})\sim p(x,\tilde{x})$ representing the system before and after a
random, unknown, and atomic intervention. It is known that under this weakly
supervised setting, it is possible to identify the causal variables and causal
mechanisms up to a permutation and elementwise reparameterization of the
variables [Brehmer et al., 2022]. Our objective is to learn an SCM that
accurately represents the true underlying SCM associated with the given data,
up to a permutation and elementwise reparameterization of causal variables. To
this end, we train an SCM by maximizing the likelihood of data. With
sufficient data and perfect optimization, we can find the SCM that is
equivalent to the ground-truth SCM.

##  4 The DCRL Algorithm

###  4.1 Overview

Figure 1 provides a visual representation of the framework’s architecture. In
this study, we utilize a conditional diffusion model and apply it to the input
data ($x$, $\tilde{x}$) where $x,\tilde{x}\in\mathbb{R}^{3\times W\times H}$
and $W$ and $H$ are the width and height of the input, respectively. The
conditioning module is defined as the encoding module, generating high-level
diffusion-based representations $(e,\tilde{e})$ for each low-level data pair
where $e,\tilde{e}\in\mathbb{R}^{d}$ and $d$ is the number of latent causal
variables assumed to be known. We empirically show that these latent variables
contain equivalent information as in noise variables of the underlying SCM and
can be used interchangeably. Then, we infer the intervention target
$I\in\\{0,1,...,d-1\\}$ for each data pair by an intervention module and use
neural solution functions on top of the latent variables $(e,\tilde{e})$ and
the intervention target $I$ to obtain the underlying latent causal variable
$(z,\tilde{z})$. We describe each part in detail in the next paragraphs.

#### The Encoding and the Intervention Module.

The encoding module consists of two main parts: the _stochastic encoder_ and
the _projection module_. The stochastic encoder $q(e|x)$ maps data pairs ($x$,
$\tilde{x}$) to pre-projection latent variables ($e$, $\tilde{e}$). The
encoded inputs are then utilized in the intervention module $q(I|x,\tilde{x})$
to infer the intervention target $I$ for the data pair ($x$, $\tilde{x}$).
Based on our data generation process, the encoded inputs have the property
that only for the elements that are intervened upon, we have
$e_{i}\neq\tilde{e}_{i},i\in I$, and the rest will remain the same. Based on
this property, in order to infer interventions, we employ an intervention
module $q(I|e,\tilde{e})$ which is defined heuristically as

$\log q(i\in
I|x,\tilde{x})=\frac{1}{Z}(\alpha+\beta|\mu_{e}(x)_{i}-\mu_{e}(\tilde{x})_{i}|+\gamma|\mu_{e}(x)_{i}-\mu_{e}(\tilde{x})_{i}|^{2})$

Where $\mu_{e}(x)$ is the mean of the stochastic encoder $q(e|x)$, $\alpha$,
$\beta$, and $\gamma$ are learnable parameters, and $Z$ is a normalization
constant. Using this simple heuristic function, we increase the likelihood of
a component as it undergoes more significant changes in response to
interventions on the encoded input. Once the intervention is inferred from the
pre-projection latent variables, we apply the projection module. The
projection module is dependent on the inferred intervention target $I$ and
projects the encoded input $(e,\tilde{e})$ to new latent variables in a way
that for the components $e_{i}$ that are not intervened upon, $i\notin I$, the
pre-intervention and post-intervention latent components will be equal,
$e_{i}=\tilde{e}_{i}$. This prevents solution functions from deviating from
the weakly supervised structure.

We write the combination of the encoder and the projection module as
$q(e,\tilde{e}|x,\tilde{x},I)$, and refer it to as the _encoding module_. By
this definition, the encoding module $q(e,\tilde{e}|x,\tilde{x},I)$ maps the
input ($x$, $\tilde{x}$) to latent variables ($e$, $\tilde{e}$) and the
intervention module infers the intervention $I$ based on pre-projection latent
variables.

#### Prior.

Given the intervention target $I$ and latent variables ($e$, $\tilde{e}$), we
define the prior $p(e,\tilde{e},I)$ as
$p(e,\tilde{e},I)=p(I)p(e)p(\tilde{e}|e,I)$. The objective of the prior
distribution is to implicitly capture the causal structure and causal
mechanisms within the system. Specifically, $p(I)$ and $p(e)$ denote the prior
distributions over intervention targets and latent variables, respectively,
and are configured as uniform categorical and standard Gaussian distributions,
respectively. According to our data generation process, when an intervention
is applied, only the elements in the latent variables that are intervened upon
are altered; the other elements remain unchanged and independent of each
other. Consequently, we can define $p(\tilde{e}|e,I)$ as follows:

$p(\tilde{e}|e,I)=\prod_{i\notin I}\delta(\tilde{e}_{i}-e_{i})\prod_{i\in
I}p(\tilde{e}_{i}|e)$

In this equation, $\delta(.)$ is the Dirac delta function that fulfills this
property for non-intervened elements of latent variables.

#### Neural Solution Functions.

Finally, in order to encode the information about the intervened variables, we
incorporate a conditional normalizing flow $p(\tilde{e}_{i}|e)$ defined as

$p(\tilde{e}_{i}|e)=\tilde{p}(h_{i}(\tilde{e}_{i};e_{\\\
i}))\Big{|}\frac{\partial
h_{i}(\tilde{e}_{i};e_{i})}{\partial\tilde{e}_{i}}\Big{|}$

where $h(.)$ are the solution functions of the SCM. They are defined as
invertible affine transformations with parameters learned with neural
networks. Therefore, by learning solution functions, i.e., learning to
transform $e$ to $z$, we implicitly model the causal graph into the framework
and obtain the latent causal variables. For more details about the
implementation, see Appendix B.

###  4.2 The Evidence Lower Bound for DCRL

Putting everything together, we calculate the Evidence Lower Bound (ELBO) for
the proposed model which will be:

$\displaystyle\log p(x,\tilde{x})$
$\displaystyle\geq\mathbb{E}_{p(x,\tilde{x})}\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{t\sim
U(0,1)}\mathbb{E}_{q(u_{t}|x)}\mathbb{E}_{q(\tilde{u}_{t}|\tilde{x})}\Biggl{[}\log
p(I)+\log p(e)$ $\displaystyle+$ $\displaystyle\log p(\tilde{e}|e,I)-\log
q(I|x,\tilde{x})-\log
q(e,\tilde{e}|x,\tilde{x},I)+\lambda(t)||s_{\theta}(u_{t},e,t)-\nabla_{u_{t}}\log
p(u_{t}|x)||_{2}^{2}$ $\displaystyle+$
$\displaystyle\lambda(t)||s_{\theta}(\tilde{u}_{t},\tilde{e},t)-\nabla_{\tilde{u}_{t}}\log
p(\tilde{u}_{t}|\tilde{x})||_{2}^{2}\Biggl{]},$

where $\lambda(t)$ is a positive weighting function. We train the model by
minimizing a reweighted loss function reminiscent of $\beta$-VAEs:

$\displaystyle\mathcal{L}_{model}$
$\displaystyle=\mathbb{E}_{p(x,\tilde{x})}\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{t\sim
U(0,1)}\mathbb{E}_{q(u_{t}|x)}\mathbb{E}_{q(\tilde{u}_{t}|\tilde{x})}\Biggl{[}\lambda(t)||s_{\theta}(u_{t},e,t)$
$\displaystyle-\nabla_{u_{t}}\log
p(u_{t}|x)||_{2}^{2}+\lambda(t)||s_{\theta}(\tilde{u}_{t},\tilde{e},t)-\nabla_{\tilde{u}_{t}}\log
p(\tilde{u}_{t}|\tilde{x})||_{2}^{2}+\beta\Big{[}\log p(I)+\log p(e)$
$\displaystyle+\log p(\tilde{e}|e,I)-\log q(I|x,\tilde{x})-\log
q(e,\tilde{e}|x,\tilde{x},I)\Big{]}\Biggl{]},$

In case of using infinite-dimensional representations (Eq. 2), the objective
function becomes:

$\displaystyle\mathcal{L}_{model}$
$\displaystyle=\mathbb{E}_{p(x,\tilde{x})}\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{t\sim
U(0,1)}\mathbb{E}_{q(e_{t},\tilde{e_{t}}|x,\tilde{x},I)}\mathbb{E}_{q(u_{t}|x)}\mathbb{E}_{q(\tilde{u}_{t}|\tilde{x})}\Biggl{[}\lambda(t)||s_{\theta}(u_{t},e_{t},t)$
$\displaystyle-\nabla_{u_{t}}\log
p(u_{t}|x)||_{2}^{2}+\lambda(t)||s_{\theta}(\tilde{u}_{t},\tilde{e_{t}},t)-\nabla_{\tilde{u}_{t}}\log
p(\tilde{u}_{t}|\tilde{x})||_{2}^{2}+\beta\Big{[}\log p(I)+\log p(e_{t})$
$\displaystyle+\log p(\tilde{e}_{t}|e_{t},I)-\log q(I|x,\tilde{x})-\log
q(e_{t},\tilde{e}_{t}|x,\tilde{x},I)\Big{]}\Biggl{]},$ (3)

where $(e_{t})_{t\in[0,1]}$ is the trajectory-based representation and
$e_{t}\in\mathbb{R}^{d}$ is the single point of the trajectory at time $t$.
For more details about the problem formulation, see Appendix A. To prevent a
collapse of the latent space to a lower-dimensional subspace, we add the
negative entropy of the batch-aggregate intervention posterior
($q_{I}^{batch}(I)=\mathbb{E}_{x,\tilde{x}\in batch}[q(I|x,\tilde{x}]$) as a
regularization term to the loss function:

$\mathcal{L}_{entropy}=\mathbb{E}_{batches}\Big{[}-\sum_{I}q_{I}^{batch}(I)\log
q_{I}^{batch}(I)\Big{]}$

where $\mathbb{E}_{batches}[\ \cdot\ ]$ is the expected value over all the
batches of data. After the training, the framework contains information about
the underlying causal structure and latent causal variables and it can be used
in different downstream tasks.

##  5 Experiments

###  5.1 Overview of the Experiments

Here we analyze the performance of the proposed model, DCRL, on synthetic
data. We employ DCRL for the task of causal discovery and subsequently use
ENCO [Lippe et al., 2021], a continuous optimization structure learning method
that leverages observational and interventional data, on top of DCRL to infer
the underlying causal graph. Furthermore, we evaluate the learned latent
variables with the DCI framework [Eastwood and Williams, 2018].

Figure 2: Comparison of models on different metrics when using single-point
representation. Our approach outperforms or competes favorably with the
baseline methods on all metrics. Particularly in higher dimensions, our method
excels by capturing additional information about the causal variables and the
underlying causal structure.

#### Data Generation.

In order to generate latent variables, we adopt random graphs where each edge
in a fixed topological order is sampled from a Bernoulli distribution with a
parameter that is equal to 0.5. We consider the SCM to be linear Gaussian and
we sample the weights from a multivariate Normal distribution with zero mean
and unit variance. We make sure the weights are not close to zero to avoid the
violation of the faithfulness assumption. We introduce additive Gaussian noise
with equal variances across all nodes, with its variance set to 0.1. Latent
causal variables are then sampled using ancestral sampling, and we generate
$10^{5}$ training samples, $10^{4}$ validation samples, and $10^{4}$ test
samples. Finally, to generate input data $x$, we apply a random linear
projection on the obtained latent variables. We keep the dimension of $x$
fixed to 16. We utilize an SCM with 5, 10, and 15 variables. To enhance the
robustness of the results, we generate data for 4 different seeds and repeat
our experiments for each seed.

#### Baselines.

We consider ILCM as our main baseline. To the best of our knowledge, there
aren’t any other methods that consider the same weakly-supervised assumptions.
We also evaluate the outcomes against a variation of disentanglement VAE
proposed by [Locatello et al., 2020] tailored for weakly supervised settings.
This model, referred to as d-VAE, models the weakly supervised process but
assumes unconnected variation factors instead of a causal relationship among
variables. Similarly, we apply ENCO on top of both to obtain the learned
graph.

#### Metrics.

We assess the performance of models with the following metrics:

  * •

Structural Hamming Distance (SHD) is a metric used to quantify the
dissimilarity between two directed acyclic graphs (DAGs) by measuring the
minimum number of edge additions, deletions, and reversals required to
transform one graph into another. It is calculated by summing up the absolute
differences between the entries of adjacency matrices of two graphs.

  * •

DCI Disentanglement Score is a metric used to evaluate the disentanglement
quality of a generative model and takes values between $0$ and $1$.
Disentanglement refers to the extent to which the model learns to predict the
underlying factors of variation in the data in a way that each predicted
variable captures at most one underlying factor. If a predicted factor is
important to predict a single underlying factor, the score will be $1$, and if
a predicted factor is equally important to predict all the underlying factors,
the score will be $0$ [Eastwood and Williams, 2018].

  * •

DCI Completeness Score measures how well each underlying factor of variation
is captured by a single predicted latent variable and has a value between $0$
and $1$. If a single variable contributes to one underlying factor, the score
will be 1, and if all variables equally contribute to the prediction of a
single factor, the score will be 0 [Eastwood and Williams, 2018].

###  5.2 Single-point Representations

Utilizing single-point representations where $e\in\mathbb{R}^{d}$ and is
independent of time, our method demonstrates superior or competitive
performance compared to the baselines, as indicated by the metrics shown in
Figure 2. In higher dimensions, our method excels by acquiring more
information about the causal variables and underlying causal structure.

###  5.3 Infinite-dimensional Representations

In these experiments, we utilize the infinite-dimensional representations
approach and develop trajectory-based representations for each input $x_{0}$,
denoted as $(e_{t})_{t\in[0,1]}$. In order to perform inference, we sample
points from this trajectory at intervals of $0.1$ resulting in $11$ specific
time steps. The outcomes are depicted in Figure 3 in Appendix C. Generally,
representations in the middle of the trajectory contain the most information
and are comparable to or even outperform the baselines. Going further in time,
representations appear to lose information but improve as they move towards
the end of the trajectory. This phenomenon arises because during training, as
we are further in time, the noise in the diffusion model is pretty high and
the conditioning module compensates for that by providing the necessary
information for the diffusion model to learn the score function.

##  6 Conclusion

Identifying the underlying causal variables and mechanisms of a system solely
from observational data is considered impossible without additional
assumptions. In this project, we use weak supervision as an inductive bias and
study if the information encoded in the latent code of diffusion-based
representations contains useful knowledge of causal variables and the
underlying causal graph.

## References

  * Abstreiter et al. [2022] K. Abstreiter, S. Mittal, S. Bauer, B. Schölkopf, and A. Mehrjou.  Diffusion-based representation learning.  _CoRR_ , abs/2105.14257, 2022.
  * Ahuja et al. [2022] K. Ahuja, J. S. Hartford, and Y. Bengio.  Weakly supervised representation learning with sparse perturbations.  _Proc. of NeurIPS_ , 35:15516–15528, 2022.
  * Anwar et al. [2014] A. R. Anwar, K. G. Mideska, H. Hellriegel, N. Hoogenboom, H. Krause, A. Schnitzler, G. Deuschl, J. Raethjen, U. Heute, and M. Muthuraman.  Multi-modal causality analysis of eyes-open and eyes-closed data from simultaneously recorded eeg and meg.  In _Proc. of EMBC_ , pages 2825–2828, 2014.
  * Brehmer et al. [2022] J. Brehmer, P. De Haan, P. Lippe, and T. S. Cohen.  Weakly supervised causal representation learning.  In _Proc. of NeurIPS_ , pages 2256–2265, 2022.
  * Cai et al. [2020] R. Cai, G. Yang, H. Averbuch-Elor, Z. Hao, S. J. Belongie, N. Snavely, and B. Hariharan.  Learning gradient fields for shape generation.  In _Proc. fo ECCV_ , volume 12348, pages 364–381, 2020.
  * Chen et al. [2021] N. Chen, Y. Zhang, H. Zen, R. J. Weiss, M. Norouzi, and W. Chan.  Wavegrad: Estimating gradients for waveform generation.  In _Proc. of ICLR_ , 2021.
  * Dhariwal and Nichol [2021a] P. Dhariwal and A. Nichol.  Diffusion models beat gans on image synthesis.  In _Proc. of NeurIPS_ , pages 8780–8794, 2021a.
  * Dhariwal and Nichol [2021b] P. Dhariwal and A. Q. Nichol.  Diffusion models beat gans on image synthesis.  In _Proc. of NeurIPS_ , pages 8780–8794, 2021b.
  * Dittadi et al. [2022] A. Dittadi, S. Papa, M. De Vita, B. Schölkopf, O. Winther, and F. Locatello.  Generalization and robustness implications in object-centric learning.  In _Proc. of ICML_ , pages 5221–5285, 2022.
  * Eastwood and Williams [2018] C. Eastwood and C. K. Williams.  A framework for the quantitative evaluation of disentangled representations.  In _Proc. of ICLR_ , 2018.
  * Hellström [2021] T. Hellström.  The relevance of causation in robotics: A review, categorization, and analysis.  _Paladyn, Journal of Behavioral Robotics_ , 12(1):238–255, 2021.
  * Hernán et al. [2000] M. Á. Hernán, B. Brumback, and J. M. Robins.  Marginal structural models to estimate the causal effect of zidovudine on the survival of hiv-positive men.  _Epidemiology_ , pages 561–570, 2000.
  * Hiemstra and Jones [1994] C. Hiemstra and J. D. Jones.  Testing for linear and nonlinear granger causality in the stock price-volume relation.  _The Journal of Finance_ , 49(5):1639–1664, 1994\.
  * Ho et al. [2020] J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  In _Proc. of NeurIPS_ , pages 8780–8794, 2020.
  * Ho et al. [2022a] J. Ho, C. Saharia, W. Chan, D. J. Fleet, M. Norouzi, and T. Salimans.  Cascaded diffusion models for high fidelity image generation.  _Journal of Machine Learning Research_ , 23:47:1–47:33, 2022a.
  * Ho et al. [2022b] J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet.  Video diffusion models.  _CoRR_ , abs/2204.03458, 2022b.
  * Hyvärinen and Dayan [2005] A. Hyvärinen and P. Dayan.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(4), 2005.
  * Hyvärinen and Oja [2000] A. Hyvärinen and E. Oja.  Independent component analysis: algorithms and applications.  _Neural networks_ , 13(4-5):411–430, 2000.
  * Höppe et al. [2022] T. Höppe, A. Mehrjou, S. Bauer, D. Nielsen, and A. Dittadi.  Diffusion models for video prediction and infilling.  _CoRR_ , abs/2206.07696, 2022.
  * Khemakhem et al. [2020] I. Khemakhem, D. Kingma, R. Monti, and A. Hyvarinen.  Variational autoencoders and nonlinear ica: A unifying framework.  In _Proc. of AISTATS_ , pages 2207–2217, 2020.
  * Kingma and Welling [2014] D. P. Kingma and M. Welling.  Auto-encoding variational bayes.  In Y. Bengio and Y. LeCun, editors, _Proc. of ICLR_ , 2014.
  * Kocaoglu et al. [2017] M. Kocaoglu, C. Snyder, A. G. Dimakis, and S. Vishwanath.  ausalgan: Learning causal implicit generative models with adversarial training.  _CoRR_ , abs/1709.02023, 2017.
  * Komanduri et al. [2022] A. Komanduri, Y. Wu, W. Huang, F. Chen, and X. Wu.  Scm-vae: Learning identifiable causal representations via structural knowledge.  In _IEEE Big Data_ , pages 1014–1023, 2022.
  * Kwon et al. [2022] M. Kwon, J. Jeong, and Y. Uh.  Diffusion models already have a semantic latent space.  _CoRR_ , abs/2210.10960, 2022.
  * Lachapelle et al. [2022] S. Lachapelle, P. Rodriguez, Y. Sharma, K. E. Everett, R. Le Priol, A. Lacoste, and S. Lacoste-Julien.  Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica.  In _Proc. of CLR_ , pages 428–484, 2022.
  * Lippe et al. [2021] P. Lippe, T. Cohen, and E. Gavves.  Efficient neural causal discovery without acyclicity constraints.  _CoRR_ , abs/2107.10483, 2021.
  * Liu et al. [2022] Y. Liu, Z. Zhang, D. Gong, M. Gong, B. Huang, A. v. d. Hengel, K. Zhang, and J. Q. Shi.  Identifying weight-variant latent causal models.  _CoRR_ , abs/2208.14153, 2022.
  * Locatello et al. [2019] F. Locatello, S. Bauer, M. Lucic, G. Raetsch, S. Gelly, B. Schölkopf, and O. Bachem.  Challenging common assumptions in the unsupervised learning of disentangled representations.  In _Proc. of ICML_ , pages 4114–4124, 2019.
  * Locatello et al. [2020] F. Locatello, B. Poole, G. Rätsch, B. Schölkopf, O. Bachem, and M. Tschannen.  Weakly-supervised disentanglement without compromises.  In _Proc. of ICML_ , pages 6348–6359, 2020.
  * Luhman and Luhman [2021] E. Luhman and T. Luhman.  Knowledge distillation in iterative generative models for improved sampling speed.  _CoRR_ , abs/2101.02388, 2021.
  * Luo [2022] C. Luo.  Understanding diffusion models: A unified perspective.  _CoRR_ , abs/2208.11970, 2022.
  * Mehrjou et al. [2017] A. Mehrjou, B. Schölkopf, and S. Saremi.  Annealed generative adversarial networks.  _CoRR_ , abs/1705.07505, 2017.
  * Mittal et al. [2022] S. Mittal, G. Lajoie, S. Bauer, and A. Mehrjou.  From points to functions: Infinite-dimensional representations in diffusion models, 2022.
  * Niu et al. [2020] C. Niu, Y. Song, J. Song, S. Zhao, A. Grover, and S. Ermon.  Permutation invariant graph generation via score-based generative modeling.  In _Proc. of AISTATS_ , volume 108, pages 4474–4484, 2020.
  * Papa et al. [2022] S. Papa, O. Winther, and A. Dittadi.  Inductive biases for object-centric representations in the presence of complex textures.  In _UAI 2022 Workshop on Causal Representation Learning_ , 2022.
  * Ramesh et al. [2022] A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen.  Hierarchical text-conditional image generation with clip latents.  _CoRR_ , abs/2204.06125, 2022.
  * Robins et al. [2000] J. M. Robins, M. A. Hernan, and B. Brumback.  Marginal structural models and causal inference in epidemiology.  _Epidemiology_ , pages 550–560, 2000.
  * Rombach et al. [2022] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proc. of ECCV_ , pages 10684–10695, 2022.
  * Runge et al. [2019] J. Runge, S. Bathiany, E. Bollt, G. Camps-Valls, D. Coumou, E. Deyle, C. Glymour, M. Kretschmer, M. D. Mahecha, J. Muñoz-Marí, et al.  Inferring causation from time series in earth system sciences.  _Nature Communications_ , 10(1):2553, 2019.
  * Saharia et al. [2022] C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, et al.  Photorealistic text-to-image diffusion models with deep language understanding.  In _Proc. of NeurIPS_ , 2022.
  * Sajjadi et al. [2018] M. S. M. Sajjadi, G. Parascandolo, A. Mehrjou, and B. Schölkopf.  Tempered adversarial networks.  In _Proc. of ICML_ , volume 80, pages 4448–4456, 2018.
  * Saremi et al. [2018] S. Saremi, A. Mehrjou, B. Schölkopf, and A. Hyvärinen.  Deep energy estimator networks.  _CoRR_ , abs/1805.08306, 2018.
  * 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.
  * Shu et al. [2019] R. Shu, Y. Chen, A. Kumar, S. Ermon, and B. Poole.  Weakly supervised disentanglement with guarantees.  _arXiv preprint arXiv:1910.09772_ , 2019.
  * Sohl-Dickstein et al. [2015] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _Proc. of ICML_ , pages 2256–2265, 2015.
  * Sohl-Dickstein et al. [2015] J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, and S. Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _Proc. of ICML_ , volume 37, pages 2256–2265, 2015.
  * Song et al. [2021a] J. Song, C. Meng, and S. Ermon.  Denoising diffusion implicit models.  In _Proc. of ICLR_ , 2021a.
  * Song et al. [2021b] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole.  Score-based generative modeling through stochastic differential equations.  In _Proc. of ICLR_ , 2021b.
  * Song et al. [2021c] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole.  Score-based generative modeling through stochastic differential equations.  _CoRR_ , abs/2011.13456, 2021c.
  * Subramanian et al. [2022] J. Subramanian, Y. Annadani, I. Sheth, N. R. Ke, T. Deleu, S. Bauer, D. Nowrouzezahrai, and S. E. Kahou.  Learning latent structural causal models.  _CoRR_ , abs/2210.13583, 2022.
  * Traub [2022] J. Traub.  Representation learning with diffusion models.  _arXiv preprint arXiv:2210.11058_ , 2022.
  * Van Steenkiste et al. [2019] S. Van Steenkiste, F. Locatello, J. Schmidhuber, and O. Bachem.  Are disentangled representations helpful for abstract visual reasoning?  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Vincent [2011] P. Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011\.
  * Von Kügelgen et al. [2021] J. Von Kügelgen, Y. Sharma, L. Gresele, W. Brendel, B. Schölkopf, M. Besserve, and F. Locatello.  Self-supervised learning with data augmentations provably isolates content from style.  _Proc. of NeurIPS_ , 34:16451–16467, 2021.
  * Wang et al. [2023] Y. Wang, Y. Schiff, A. Gokaslan, W. Pan, F. Wang, C. De Sa, and V. Kuleshov.  Infodiffusion: Representation learning using information maximizing diffusion models.  _arXiv preprint arXiv:2306.08757_ , 2023.
  * Wu et al. [2022] Z. Wu, N. Dvornik, K. Greff, T. Kipf, and A. Garg.  Slotformer: Unsupervised visual dynamics simulation with object-centric models.  _CoRR_ , abs/2210.05861, 2022.
  * Yang et al. [2020] M. Yang, F. Liu, Z. Chen, X. Shen, J. Hao, and J. Wang.  Causalvae: Structured causal disentanglement in variational autoencoder.  _CoRR_ , abs/2208.14153, 2020.
  * Yoon et al. [2023] J. Yoon, Y.-F. Wu, H. Bae, and S. Ahn.  An investigation into pre-training object-centric representations for reinforcement learning.  _CoRR_ , abs/2302.04419, 2023.

Appendix

##  A Problem Formulation & ELBO

The ELBO for the proposed framework will be (For simplicity, we only derive
the ELBO when using single representations independent of time, i.e.,
$e\in\mathbb{R}^{d}$. The ELBO for the infinite-dimensional case would be
similar):

$\displaystyle\log p(x,\tilde{x})$
$\displaystyle\geq\mathbb{E}_{q(e,\tilde{e},u,\tilde{u},I|x,\tilde{x})}\left[\log\frac{p(x,\tilde{x},u,\tilde{u},e,\tilde{e},I)}{q(e,\tilde{e},I,u,\tilde{u}|x,\tilde{x})}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{q(e,\tilde{e},u,\tilde{u},I|x,\tilde{x})}\left[\log\frac{p(I)}{q(I|x,\tilde{x})}+\log\frac{p(e)p(\tilde{e}|e,I)}{q(e,\tilde{e}|x,\tilde{x},I)}+\log\frac{p(x,u|e)}{q(u|x)}+\log\frac{p(\tilde{x},\tilde{u}|\tilde{e})}{q(\tilde{u}|\tilde{x})}\right]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{q(u|x)}\mathbb{E}_{q(\tilde{u}|\tilde{x})}\Bigg{[}\Big{[}\log
p(I)+\log p(e)+\log p(\tilde{e}|e,I)-\log q(I|x,\tilde{x})$ $\displaystyle-$
$\displaystyle\log
q(e,\tilde{e}|x,\tilde{x},I)\Big{]}+\left[\log\frac{p(x,u|e)}{q(u|x)}+\log\frac{p(\tilde{x},\tilde{u}|\tilde{e})}{q(\tilde{u}|\tilde{x})}\right]\Bigg{]}$

The terms in the first bracket correspond to the intervention encoder and the
noise encoding module, respectively, and the terms in the second bracket
correspond to the diffusion model conditioned on pre- and post-intervention
noise encodings.

Song et al. [2021c] shows that the discretization of SDE formulations of the
diffusion model is equivalent to discrete-time diffusion models. Therefore,
for simplicity, we derive the ELBO for discrete-time diffusion models.
Following [Luo, 2022], for a discrete-time diffusion model where $t\in[1,T]$,
we have

$\displaystyle\mathbb{E}_{q(I|x,\tilde{x})}$
$\displaystyle\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{q(u|x)}\mathbb{E}_{q(\tilde{u}|\tilde{x})}\Bigg{[}\log\frac{p(x,u|e)}{q(u|x)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{q(u|x)}\mathbb{E}_{q(\tilde{u}|\tilde{x})}\Bigg{[}\mathbb{E}_{q(u_{1}|x)}[\log
p(x|u_{1})]-D_{KL}(q(u_{T}|x)||p(u_{T}))$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(u_{t}|x)}[D_{KL}(q(u_{t-1}|u_{t},x,e)||p(u_{t-1}|u_{t},e)]\Bigg{]}$
(4)

where we have that

  * •

$\mathbb{E}_{q(u_{1}|x)}[\log p(x|u_{1})]$ is the reconstruction term and it
can be defined in a way that it is constant so it can be ignored during
training;

  * •

$D_{KL}(q(u_{T}|x)||p(u_{T}))$ is the prior matching term and can similarly be
defined in a way that it is constant;

  * •

$\mathbb{E}_{u_{t}|x}[D_{KL}(q(u_{t-1}|u_{t},x,e)||p(u_{t-1}|u_{t},e)]$ is a
denoising matching term. This term is the origin of different interpretations
of the score-based diffusion models.

For the SDE formulation of the forward process, the denoising matching term
becomes [Song et al., 2021c]

$\lambda(t)||s_{\theta}(u_{t},e,t)-\nabla_{u_{t}}\log p(u_{t}|x)||_{2}^{2}.$
(5)

The weight $\lambda(t)$ of denoising matching terms is related to the
diffusion coefficient of the forward SDE. For a Variance Exploding SDE the
weight is defined as
$\lambda(t)=2\sigma^{2}(t)\log(\sigma_{max}/\sigma_{min})$ with
$\sigma(t)=\sigma_{min}\cdot(\sigma_{max}/\sigma_{min})^{t}$.

Therefore, by combining (4) with (5), the ELBO becomes

$\displaystyle\log p(x,\tilde{x})$
$\displaystyle\geq\mathbb{E}_{p(x,\tilde{x})}\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{q(e,\tilde{e}|x,\tilde{x},I)}\mathbb{E}_{t\sim
U(0,1)}\mathbb{E}_{q(u_{t}|x)}\mathbb{E}_{q(\tilde{u}_{t}|\tilde{x})}$
$\displaystyle\Biggl{[}$ $\displaystyle\log p(I)+\log p(e)+\log
p(\tilde{e}|e,I)-\log q(I|x,\tilde{x})-\log q(e,\tilde{e}|x,\tilde{x},I)$
$\displaystyle+$
$\displaystyle\lambda(t)\Big{[}||s_{\theta}(u_{t},e,t)-\nabla_{u_{t}}\log
p(u_{t}|x)||_{2}^{2}+||s_{\theta}(\tilde{u}_{t},\tilde{e},t)-\nabla_{\tilde{u}_{t}}\log
p(\tilde{u}_{t}|\tilde{x})||_{2}^{2}\Big{]}\Biggl{]}$

For infinite-dimensional representations, we can derive the ELBO using a
similar argument. In this case, the formula for the ELBO is

$\displaystyle\log p(x,\tilde{x})$
$\displaystyle\geq\mathbb{E}_{p(x,\tilde{x})}\mathbb{E}_{q(I|x,\tilde{x})}\mathbb{E}_{t\sim
U(0,1)}\mathbb{E}_{q(e_{t},\tilde{e}_{t}|x,\tilde{x},I)}\mathbb{E}_{q(u_{t}|x)}\mathbb{E}_{q(\tilde{u}_{t}|\tilde{x})}$
$\displaystyle\Biggl{[}$ $\displaystyle\log p(I)+\log p(e_{t})+\log
p(\tilde{e}_{t}|e_{t},I)-\log q(I|x,\tilde{x})-\log
q(e_{t},\tilde{e}_{t}|x,\tilde{x},I)$ $\displaystyle+$
$\displaystyle\lambda(t)||s_{\theta}(u_{t},e_{t},t)-\nabla_{u_{t}}\log
p(u_{t}|x)||_{2}^{2}+\lambda(t)||s_{\theta}(\tilde{u}_{t},\tilde{e}_{t},t)-\nabla_{\tilde{u}_{t}}\log
p(\tilde{u}_{t}|\tilde{x})||_{2}^{2}\Biggl{]},$

##  B Implementation Details

#### Training

For the training, we follow the 4-phase training of Brehmer et al. [2022] but
consider only the first 3 phases. In summary, we consider the following steps:

  1. (1)

We begin by training the diffusion model and the encoding module together on
data pairs for $20$ epochs. This can be interpreted as a warm-up on the
diffusion model and the encoding module to extract meaningful representations
of data.

  2. (2)

We include all modules except for solution functions. We consider
$p(\tilde{e}_{i}|e)$ to be a uniform probability density. We do this phase for
$50$ epochs.

  3. (3)

We include solution functions and train the whole framework with the proposed
loss and do this for $50$ epochs.

We find out that considering our data generation process, including the fourth
training phase of Brehmer et al. [2022] has no impact on the model’s
performance. Consequently, we choose to disregard it in our analysis. We use
the loss in Eq. 3 as the objective function and consider the coefficient of
the regularization term $\mathcal{L}_{entropy}$ to be 1. Therefore, our
overall loss function is then given by
$\mathcal{L}=\mathcal{L}_{model}+\mathcal{L}_{entropy}$.

#### Architectures & Hyperparameters

We train the model for $120$ epochs and use the learning rate of 3e-4 with a
batch size of $64$. $\beta$ is initially set to $0$ and increased to $1$
during training. The noise encoder is considered Gaussian, with mean and
standard deviation parameterized as an MLP with two hidden layers and 64 units
each and ReLU activation functions. The architecture of the score function of
the diffusion model is based on NCSN++ architecture [Song et al., 2021c] with
the same set of hyperparameters. As the input $x$ is 16-dimensional and the
score model follows a convolutional architecture, we reshape the input into a
$4\times 4$ format and then feed it into the diffusion model. Furthermore, In
the forward SDE, $\sigma_{min}$ and $\sigma_{max}$ are set to 0.01 and 50,
respectively.

##  C Missing Plots

Figure 3: Comparison of models on different metrics when using infinite-
dimensional representations. From top to bottom, (a), (b), and (c) correspond
to experiments with 5, 10, and 15 causal variables, respectively. We sample
points from the trajectory at intervals of $0.1$, creating a total of $11$
specific timesteps. Typically, representations in the middle of the trajectory
carry the most information, often matching or surpassing the baseline
performance. As we move further in time, representations seem to lose some
information, but they improve as they approach the end of the trajectory.
Furthermore, the framework performs worse or on par with baselines in lower
dimensions but generally outperforms them in higher dimensions.