Diffusion Bridge AutoEncoders for Unsupervised Representation Learning

## Abstract

Abstract Diffusion-based representation learning has achieved substantial attention due to its promising capabilities in latent representation and sample generation. Recent studies have employed an auxiliary encoder to identify a corresponding representation from a sample and to adjust the dimensionality of a latent variable 𝐳 𝐳 {\mathbf{z}} bold_z . Meanwhile, this auxiliary structure invokes information split problem because the diffusion and the auxiliary encoder would divide the information from the sample into two representations for each model. Particularly, the information modeled by the diffusion becomes over-regularized because of the static prior distribution on 𝐱 T subscript 𝐱 𝑇 {\mathbf{x}}_{T} bold_x start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT . To address this problem, we introduce Diffusion Bridge AuteEncoders (DBAE), which enable 𝐳 𝐳 {\mathbf{z}} bold_z -dependent endpoint 𝐱 T subscript 𝐱 𝑇 {\mathbf{x}}_{T} bold_x start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT inference through a feed-forward architecture. This structure creates an information bottleneck at 𝐳 𝐳 {\mathbf{z}} bold_z , so 𝐱 T subscript 𝐱 𝑇 {\mathbf{x}}_{T} bold_x start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT becomes dependent on 𝐳 𝐳 {\mathbf{z}} bold_z in its generation. This results in two consequences: 1) 𝐳 𝐳 {\mathbf{z}} bold_z holds the full information of samples, and 2) 𝐱 T subscript 𝐱 𝑇 {\mathbf{x}}_{T} bold_x start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT becomes a learnable distribution, not static any further.
We propose an objective function for DBAE to enable both reconstruction and generative modeling, with their theoretical justification. Empirical evidence supports the effectiveness of the intended design in DBAE, which notably enhances downstream inference quality, reconstruction, and disentanglement. Additionally, DBAE generates high-fidelity samples in the unconditional generation.

## 1 Introduction

Unsupervised representation learning is a fundamental topic within the latent variable generative models ; ; ; ; ; . Effective representation supports better downstream inference as well as realistic data synthesis.
Variational autoencoders (VAEs)  are frequently used because they inherently include latent representations with flexible dimensionality. Whereas generative adversarial networks (GANs)  originally focused on only the generation, recent advances in identifying the latent from a sample, a.k.a. GAN inversion ; , are another method to find latent representations. Also, diffusion probabilistic models (DPMs) ; have achieved state-of-the-art performance in terms of generation quality ; , which has naturally prompted efforts to explore unsupervised representation learning within the DPM framework ; ; ; ; ; ; since researchers hypothesized that an informative latent representation has enabled such generation performance.

As surveyed above, DPMs are a type of latent variable generative model, but its inference on latent variables is not straightforward. DPMs progressively map from data ${\mathbf{x}}_{0}$ to a latent endpoint ${\mathbf{x}}_{T}$ via a predefined noise injection schedule, which does not facilitate learnable encoding, DPMs initially would be regarded as a sample generator like GAN. DDIM  introduces an invertible path between the data ${\mathbf{x}}_{0}$ and the endpoint ${\mathbf{x}}_{T}$ given a data distribution. However, the latent variable ${\mathbf{x}}_{T}$ still suffers from over-regularization problem with a Gaussian distribution, which hinders learning meaningful encodings . Moreover, the inference of ${\mathbf{x}}_{T}$ is intractable and dimension cannot be reduced, which poses disadvantages for downstream applications .

To tackle this issue, recent DPM-based representation learning studies ; ; ; ; ; ; suggest an auxiliary latent variable ${\mathbf{z}}$ with an encoder used in VAEs. This encoder-generated latent variable ${\mathbf{z}}$ is tractable without ODE solving ; , and the encoder enables dimensionality reduction, as well. This paper points out the remaining problem in DPM-based representation learning methods which we refer to as the information split problem. The information is not solely retained in the tractable latent variable ${\mathbf{z}}$; rather, a portion is also distributed into the intractable latent variable ${\mathbf{x}}_{T}$ evidenced by Figure 1(b). To address this issue, we introduce Diffusion Bridge AutoEncoders (DBAE), which incorporate a ${\mathbf{z}}$-dependent endpoint ${\mathbf{x}}_{T}$ inference using a feed-forward architecture.

The proposed model DBAE systemically resolves both over-regularization problem and information split problem. The endpoint ${\mathbf{x}}_{T}$ inference of DBAE is achieved through encoder-decoder propagation as the inferred endpoints shown in Figure 1(d). The diffusion endpoint ${\mathbf{x}}_{T}$ is no longer constrained to be Gaussian and it becomes tractable. Furthermore, unlike the two split inference paths in the previous approach (dotted line in Figure 1(a)), DBAE encourages ${\mathbf{z}}$ to become an information bottleneck in inference (dotted line in Figure 1(c)), leading ${\mathbf{z}}$ becomes more informative. In the experiments, DBAE outperforms the previous works in downstream inference quality, reconstruction, disentanglement, and unconditional generation. DBAE exhibits satisfactory results in interpolations and attribute manipulations with its qualitative benefits.

Figure: (a) Bayesian network of DiffAE
Refer to caption: extracted/5622713/Figs/Figure1_7.png

## 2 Preliminaries

### 2.1 Diffusion Probabilistic Models and Representation Learning

Diffusion probabilistic models (DPM) ; with a continuous time formulation  define a forward stochastic differential equation (SDE) in eq. 1, which starts from data distribution $q_{\textrm{data}}({\mathbf{x}}_{0})$ and gradually perturbs it to noise ${\mathbf{x}}_{T}$.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=\mathbf{f}({\mathbf{x}}_{t},t)\mathrm{ d}t+g(t)\mathrm{d}\mathbf{w}_{t},\quad{\mathbf{x}}_{0}\sim q_{\textrm{data}}({ \mathbf{x}}_{0}),$ (1) $$

where $\mathbf{w}_{t}$ denotes a standard Wiener process, $\mathbf{f}:\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{d}$ is a drift term, and $g:[0,T]\rightarrow\mathbb{R}$ is a volatility term. Let the marginal distribution of eq. 1 at time $t$ be denoted as $\tilde{q}_{t}({\mathbf{x}}_{t})$. There exists a unique reverse-time SDE  as shown in eq. 2, which traces back from noise ${\mathbf{x}}_{T}$ to data ${\mathbf{x}}_{0}$. Furthermore, the reverse-time ordinary differential equation (ODE) in eq. 3 produces an identical marginal distribution to eq. 2 for all $t$, offering an alternative generative process while conferring stochasticity of the trajectory solely to ${\mathbf{x}}_{T}$.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}( t)\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t})]\mathrm{d}t+g(t )\mathrm{d}\mathbf{\bar{w}}_{t},\quad{\mathbf{x}}_{T}\sim p_{\textrm{prior}}({ \mathbf{x}}_{T}),$ (2) $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-\frac{ 1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t})] \mathrm{d}t,\quad{\mathbf{x}}_{T}\sim p_{\textrm{prior}}({\mathbf{x}}_{T}),$ (3) $$

where $\mathbf{\bar{w}}_{t}$ denotes the reverse-time Wiener process, $\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t})$ is time-dependent score function, and $p_{\textrm{prior}}({\mathbf{x}}_{T})$ stands for the prior distribution of the reverse process, which is known to closely resembles a Gaussian distribution with the specific form of $\mathbf{f}$ and $g$ ; . To construct both reverse SDE and ODE, the diffusion model estimates a time-dependent score function $\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t})\approx\mathbf{s}_
{\boldsymbol{\theta}}({\mathbf{x}}_{t},t)$ using a neural network with score-matching objective ; ; ; .

From the perspective of representation learning, the reverse ODE in eq. 3 (a.k.a DDIM  in discrete time diffusion formulation) provides a bidirectional encoding-decoding path between data ${\mathbf{x}}_{0}$ and latent ${\mathbf{x}}_{T}$ while ${\mathbf{x}}_{T}$ is set to follow the Gaussian distribution ; . This mapping is both invertible and uniquely identifiable up to discretization error, indicating that the encoding is entirely determined by the data distribution $q_{\textrm{data}}({\mathbf{x}}_{0})$ ; . However, the latent variable ${\mathbf{x}}_{T}$ provided by ODE in eq. 3 presents some disadvantages.

Over-regularization problem: From the perspective of encoding, the foremost problem of ${\mathbf{x}}_{T}$ is the rigidity of the assumed distribution, the Gaussian distribution, which can be regarded as the over-regularization. For instance, identified that the Gaussian could be distant from the optimal prior distribution, so the paper sets aggregated posterior as optimal prior, making learnable encoding more informative. DDIM has this over-regularized prior of ${\mathbf{x}}_{T}$ that cannot be adapted via the encoding.

Some previous literature, i.e. DiffAE  and PDAE , proposes auxiliary latent space utilizing learnable encoder Enc${}_{\boldsymbol{\phi}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{l}$, which directly maps from data ${\mathbf{x}}_{0}$ to an auxiliary latent variable ${\mathbf{z}}$. Unlike DDIM, these approaches tractably obtain ${\mathbf{z}}$ from ${\mathbf{x}}_{0}$ without solving ODE, and both DiffAE and PDAE can directly learn the latent space in reduced dimensionality.
Consequently, the reverse ODE in eq. 4 becomes associated with the ${\mathbf{z}}$-conditional score function $\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}}({\mathbf{x}}_{t}|{
\mathbf{z}})$. Meanwhile, these past approaches do not solve any aspects of over-regularization in ${\mathbf{x}}_{T}$. Additionally, they introduce a new problem of information split because of the auxiliary structure.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-\frac{ 1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}}({\mathbf{ x}}_{t}|{\mathbf{z}})]\mathrm{d}t.$ (4) $$

Information split problem: The generative process in eq. 4 initiates with two latent variables ${\mathbf{z}}$ and ${\mathbf{x}}_{T}$. If the framework only relies on the tractably inferred latent variable ${\mathbf{z}}$, the reconstruction outcomes depicted in Figure 1(b) appear to fluctuate depending on the choice of ${\mathbf{x}}_{T}$. This implies that ${\mathbf{x}}_{T}$ encompasses crucial information necessary for reconstructing ${\mathbf{x}}_{0}$. To consider all information of ${\mathbf{x}}_{0}$, it is necessary to infer ${\mathbf{x}}_{T}$ by solving ODE in eq. 4 from input ${\mathbf{x}}_{0}$ to endpoint ${\mathbf{x}}_{T}$, enduring its intractability. Consequently, the persisting issue within the latent variable ${\mathbf{x}}_{T}$ remains an unresolved challenge in this framework.

Therefore, this paper intends to change this endpoint to be a learnable distribution that is generated from ${\mathbf{z}}$, so that the inferred distribution of ${\mathbf{x}}_{T}$ becomes adapted to optimize. Furthermore, we designed ${\mathbf{z}}$ to be an information bottleneck structure, encouraging all information to be compressed.

### 2.2 Diffusion Process with Fixed Endpoints

To control the information regarding the diffusion endpoint ${\mathbf{x}}_{T}$, it is imperative to specify a forward SDE that terminates at the desired endpoint. We employ Doob’s $h$-transform ; ; , which facilitates the conversion of the original forward SDE in eq. 1 into eq. 5 to terminate at a desired point (or distribution).

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}( t)\mathbf{h}({\mathbf{x}}_{t},t,\mathbf{y},T)]\mathrm{d}t+g(t)\mathrm{d} \mathbf{w}_{t},\quad{\mathbf{x}}_{0}\sim q_{\textrm{data}}({\mathbf{x}}_{0}), \quad{\mathbf{x}}_{T}=\mathbf{y},$ (5) $$

where $\mathbf{h}({\mathbf{x}}_{t},t,\mathbf{y},T):=\nabla_{{\mathbf{x}}_{t}}\log{
\tilde{q}_{t}({\mathbf{x}}_{T}|{\mathbf{x}}_{t})}|_{{\mathbf{x}}_{T}=\mathbf{y}}$ is the score function of perturbation kernel in original forward SDE, and $\mathbf{y}$ denotes the desired endpoint. Let $q_{t}({\mathbf{x}}_{t})$ denotes the marginal distribution of eq. 5 at $t$. It is noteworthy that when both ${\mathbf{x}}_{0}$ and ${\mathbf{x}}_{T}$ are given, the conditional probability of ${\mathbf{x}}_{t}$ becomes identical to the original forward SDE, i.e., $q_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})=\tilde{q}_{t}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})$ ; . In specific formulations of original forward SDE, such as VE and VP ; ; $q_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})$ becomes Gaussian distribution, enabling tractable sampling of ${\mathbf{x}}_{t}$ at any time $t$ with exact density function.

There also exists unique reverse-time SDE and ODE in eqs. 6 and 7 corresponding to $h$-transformed forward SDE of eq. 5 ; .

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}( t)\nabla_{{\mathbf{x}}_{t}}\log q_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})+g^{2} (t)\mathbf{h}({\mathbf{x}}_{t},t,\mathbf{y},T)]\mathrm{d}t+g(t)\mathrm{d} \mathbf{\bar{w}}_{t},{\mathbf{x}}_{T}=\mathbf{y},$ (6) $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-\frac{ 1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log q_{t}({\mathbf{x}}_{t}|{\mathbf{x}} _{T})+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,\mathbf{y},T)]\mathrm{d}t,\quad{ \mathbf{x}}_{T}=\mathbf{y},$ (7) $$

where $q_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ be the conditional probability defined by eq. 5. To construct reverse SDE and ODE, it is necessary to estimate $\nabla_{{\mathbf{x}}_{t}}\log q_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\approx
\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})$ through the neural network with score matching objective  in eq. 8.

$$ $\displaystyle\frac{1}{2}\int_{0}^{T}\mathbb{E}_{q_{t}({\mathbf{x}}_{t},{ \mathbf{x}}_{T})}[g^{2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t}, t,{\mathbf{x}}_{T})-\nabla_{{\mathbf{x}}_{t}}\log q_{t}({\mathbf{x}}_{t}|{ \mathbf{x}}_{T})||_{2}^{2}]\mathrm{d}t,$ (8) $$

## 3 Method

Figure: Figure 2: A schematic for Diffusion Bridge AutoEncoders. The blue line shows the latent variable inference. DBAE infers the ${\mathbf{z}}$-dependent endpoint ${\mathbf{x}}_{T}$ to make ${\mathbf{x}}_{T}$ tractable and induce ${\mathbf{z}}$ to become an information bottleneck. The paired ${\mathbf{x}}_{0}$ and ${\mathbf{x}}_{T}$ produce new forward SDE utilizing $h$-transform. The decoder and red line describe the generative process. The generation starts from the bottleneck latent variable ${\mathbf{z}}$, and decodes it to endpoint ${\mathbf{x}}_{T}$. The reverse SDE or ODE generates from ${\mathbf{x}}_{T}$ to ${\mathbf{x}}_{0}$.
Refer to caption: extracted/5622713/Figs/Figure5_11.png

To resolve information split problem with learnable endpoint in diffusion-based representation learning ; and its subsequent applications ; ; , we introduce Diffusion Bridge AutoEncoders (DBAE) featuring ${\mathbf{z}}$-dependent endpoint ${\mathbf{x}}_{T}$ inference using a single network propagation. Figure 2 illustrates the overall schematic for DBAE. Section 3.1 explains the latent variable inference with the encoder-decoder structure and a forward SDE utilizing Doob’s $h$-transform. Section 3.2 delineates the generative process from the information bottleneck ${\mathbf{z}}$ to data ${\mathbf{x}}_{0}$. Section 3.3 elaborates on the objective function for reconstruction, unconditional generation, and its theoretical justifications.

### 3.1 Latent Variable Inference

We can access i.i.d. samples from $q_{\textrm{data}}({\mathbf{x}}_{0})$. The encoder Enc${}_{\boldsymbol{\phi}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{l}$ maps from data ${\mathbf{x}}_{0}$ to latent variable ${\mathbf{z}}$, defining the conditional probability $q_{\boldsymbol{\phi}}({\mathbf{z}}|{\mathbf{x}}_{0})$. To condense the high-level representation of ${\mathbf{x}}_{0}$, the latent dimension $l$ is set to be lower than the data dimension $d$. The decoder Dec${}_{\boldsymbol{\psi}}:\mathbb{R}^{l}\rightarrow\mathbb{R}^{d}$ maps from latent variable ${\mathbf{z}}$ to the endpoint ${\mathbf{x}}_{T}$, defining the conditional probability $q_{\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})$. Since the decoder generates endpoint ${\mathbf{x}}_{T}$ solely based on the latent variable ${\mathbf{z}}$, ${\mathbf{z}}$ becomes a bottleneck of all the information of ${\mathbf{x}}_{0}$. The encoder-decoder structure provides the endpoint distribution $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})=
\int{q_{\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})q_{\boldsymbol{\phi}}
({\mathbf{z}}|{\mathbf{x}}_{0})}\mathrm{d}{\mathbf{z}}$ for the given starting point ${\mathbf{x}}_{0}$. We now discuss a new diffusion process $\left\{{\mathbf{x}}_{t}\right\}_{t=0}^{T}$ with a given starting and endpoint pair. Here, it should be noted that DBAE has a unique feature of pairing an original data ${\mathbf{x}}_{0}$ with its reconstructed endpoint ${\mathbf{x}}_{T}$ from ${\mathbf{z}}$ unlike the other models utilizing just a noise sample from $\mathcal{N}(\mathbf{0},\mathbf{I})$. Because this pairing is solely controlled by ${\mathbf{z}}$, the endpoint has to be optimized, not static as before; so ${\mathbf{x}}_{T}$ becomes the optimal prior from the diffusion model perspective.

We define a forward SDE in eq. 9 utilizing Doob’s $h$-transform ; ; to establish the relationship between the starting point and endpoint given by the encoder-decoder. The forward SDE in eq. 9 determines the distribution of ${\mathbf{x}}_{t}$, where $t\in(0,T)$. Let the marginal distribution of eq. 9 at time $t$ as $q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t})$.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}( t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)]\mathrm{d}t+g(t)\mathrm{d} \mathbf{w}_{t},\quad{\mathbf{x}}_{0}\sim q_{\textrm{data}}({\mathbf{x}}_{0}), \,{\mathbf{x}}_{T}\sim q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T }|{\mathbf{x}}_{0}),$ (9) $$

where $\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T):=\nabla_{{\mathbf{x}}_{t}}
\log{\tilde{q}_{t}({\mathbf{x}}_{T}|{\mathbf{x}}_{t})}$ is the score function of perturbation kernel in the original forward SDE in eq. 1.

### 3.2 Generative Process

The generative process begins with the bottleneck latent variable ${\mathbf{z}}$, which can be inferred from input data ${\mathbf{x}}_{0}$ or randomly drawn from the prior distribution $p_{\text{prior}}({\mathbf{z}})$. The decoder Dec${}_{\boldsymbol{\psi}}:\mathbb{R}^{l}\rightarrow\mathbb{R}^{d}$ maps from latent variable ${\mathbf{z}}$ to the endpoint ${\mathbf{x}}_{T}$ with the probability $p_{\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})$.(^1^11The two distributions $p_{\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})$ and $q_{\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})$ are the same. However, to distinguish between inference and generation, they are respectively denoted as $p$ and $q$.) The reverse ODE corresponding to eq. 9, as described in eq. 10 regenerate the data ${\mathbf{x}}_{0}$.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-\frac{ 1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})+g^{2}(t)\mathbf{h}({ \mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)]\mathrm{d}t,$ (10) $$

where the conditional probability $q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ is defined by eq. 9. However, computing the conditional probability $q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ is intractable, so we parameterize our score model $\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T}):=\nabla_
{{\mathbf{x}}_{t}}\log p^{t}_{\boldsymbol{\theta}}({\mathbf{x}}_{t}|{\mathbf{x
}}_{T})$ to approximate $\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T})$.
Our parametrized generative process becomes eq. 11. Stochastic sampling with SDE is also naturally possible as shown in Section 2.2, but we describe only ODE for convenience.

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}=[\mathbf{f}({\mathbf{x}}_{t},t)-\frac{ 1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log p^{t}_{\boldsymbol{\theta}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T},T)]\mathrm{d}t.$ (11) $$

### 3.3 Objective Function

The objective function bifurcates depending on the tasks. For downstream inference, attribute manipulation, and interpolation, the model merely requires reconstruction capability. It is adequate to approximate $\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})\approx
\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ and unique ${\mathbf{x}}_{T}$ encodings according to each data point ${\mathbf{x}}_{0}$. For an unconditional generation, the model must possess the ability to generate random samples from the endpoint ${\mathbf{x}}_{T}$, implying that the generative endpoint distribution $p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})=\int{p_{\boldsymbol{\psi}}({\mathbf{x}
}_{T}|{\mathbf{z}})p_{\text{prior}}({\mathbf{z}})}\mathrm{d}{\mathbf{z}}$ should closely resemble the aggregated inferred distribution $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})=\int{q_{\boldsymbol{
\psi}}({\mathbf{x}}_{T}|{\mathbf{z}})q_{\boldsymbol{\phi}}({\mathbf{z}}|{
\mathbf{x}}_{0})q_{\text{data}}({\mathbf{x}}_{0})}\mathrm{d}{\mathbf{x}}_{0}
\mathrm{d}{\mathbf{z}}$.

#### 3.3.1 Reconstruction

For successful reconstruction, the model needs to fulfill two criteria: 1) encoding the latent variable ${\mathbf{x}}_{T}$ uniquely depending on the data point ${\mathbf{x}}_{0}$, and 2) regenerating from ${\mathbf{x}}_{T}$ to ${\mathbf{x}}_{0}$. The inferred latent distribution $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})$ should provide unique information for each ${\mathbf{x}}_{0}$. To achieve this, we aim to minimize the entropy $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{
x}}_{0}))$ to embed ${\mathbf{x}}_{0}$-dependent ${\mathbf{x}}_{T}$ with minimum uncertainty. On the other hands, we maximizes the entropy $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}))$ to embed different ${\mathbf{x}}_{T}$ for each ${\mathbf{x}}_{0}$. Since posterior entropy $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{
x}}_{T}))=\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}
|{\mathbf{x}}_{0}))-\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{T}))+\mathcal{H}(q_{\text{data}}({\mathbf{x}}_{0}))$ naturally includes aforementioned terms, we use this term as a regularization. Minimizing the gap between the eqs. 10 and 11 is necessary for regenerating from ${\mathbf{x}}_{T}$ to ${\mathbf{x}}_{0}$. This requires alignment between the inferred score function $\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ and the model score function $\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})$. Accordingly, we propose the score-matching objective function $\mathcal{L}_{\text{SM}}$ as described in eq. 12, similarly as eq. 8. We train DBAE with the entropy-regularized score matching objective $\mathcal{L}_{\text{AE}}$ as shown in eq. 13.

$$ $\displaystyle\mathcal{L}_{\textrm{SM}}:=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^ {t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}[ g^{2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T} )-\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t$ (12) $\displaystyle\mathcal{L}_{\textrm{AE}}:=\mathcal{L}_{\textrm{SM}}+\mathcal{H}( q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (13) $$

Theorem 1 demonstrates that the entropy-regularized score matching objective in $\mathcal{L}_{\textrm{AE}}$ becomes a tractable form of objective, and it is equivalent to the reconstruction objective. The inference distribution $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T}|{
\mathbf{x}}_{0})$ is optimized to provide the best information about ${\mathbf{x}}_{0}$ for easy reconstruction.

###### Theorem 1 .

For the objective function $\mathcal{L}_{\textrm{AE}}$, the following equality holds.

$$ $\mathcal{L}_{\textrm{AE}}=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{ \mathbf{x}}_{T})}[g^{2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t}, t,{\mathbf{x}}_{T})-\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t }|{\mathbf{x}}_{0},{\mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t$ (14) $$

Moreover, if eq. 1 is linear SDE, there exists $\alpha(t)$, $\beta(t)$, $\gamma(t)$, $\lambda(t)$, such that

$$ $\mathcal{L}_{\textrm{AE}}=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{ \mathbf{x}}_{T})}[\lambda(t)||\mathbf{x}^{0}_{\boldsymbol{\theta}}({\mathbf{x} }_{t},t,{\mathbf{x}}_{T})-{\mathbf{x}}_{0}||_{2}^{2}]}\mathrm{d}t,$ (15) $$

where $\mathbf{x}^{0}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T}):=
\alpha(t){\mathbf{x}}_{t}+\beta(t){\mathbf{x}}_{T}+\gamma(t)\mathbf{s}_{
\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})$.

The assumption and proof of Theorem 1 is in Section A.1. The $\boldsymbol{\theta}$-related term in $\mathcal{L}_{\textrm{AE}}$ is solely $\mathcal{L}_{\textrm{SM}}$, naturally leading to the optimality specified in Corollary 2.

###### Corollary 2 .

Let $\boldsymbol{\theta}^{*}=\operatorname*{arg\,min}_{\boldsymbol{\theta}}\mathcal
{L}_{\text{AE}}$. Then $\mathbf{s}_{\boldsymbol{\theta}^{*}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})=
\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ for all ${\mathbf{x}}_{t}$, $t$, ${\mathbf{x}}_{T}$.

###### Theorem 1 .

###### Corollary 2 .

#### 3.3.2 Generative Modeling

In the Section 3.3.1, the discussion focused on the objective function for reconstruction. For the generative modeling, the distribution of ${\mathbf{x}}_{T}$ should be considered. This section addresses the discrepancy between the inferred distribution $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})$ and generative prior distribution $p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})$. To address this, we propose the objective $\mathcal{L}_{\text{PR}}$ related to the generative prior.

$$ $\displaystyle\mathcal{L}_{\textrm{PR}}:=$ $\displaystyle\mathbb{E}_{q_{\text{data}}({\mathbf{x}}_{0})}[D_{\text{KL}}(q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})||p_{ \boldsymbol{\psi}}({\mathbf{x}}_{T}))]$ (16) $$

Theorem 3 demonstrates that the autoencoding objective $\mathcal{L}_{\text{AE}}$ and prior objective $\mathcal{L}_{\text{PR}}$ bound the Kullback-Leibler divergence between data distribution $q_{\text{data}}({\mathbf{x}}_{0})$ and the generative model distribution $p_{\boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{0})=\int{p_{
\boldsymbol{\theta}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T})p_{\boldsymbol{\psi}}({
\mathbf{x}}_{T}|{\mathbf{z}})p_{\text{prior}}({\mathbf{z}})}\mathrm{d}{\mathbf
{z}}\mathrm{d}{\mathbf{x}}_{T}$ up to a constant. The proof is in Section A.2

###### Theorem 3 .

$D_{\textrm{KL}}(q_{\text{data}}({\mathbf{x}}_{0})||p_{\boldsymbol{\psi},
\boldsymbol{\theta}}({\mathbf{x}}_{0}))\leq\mathcal{L}_{\textrm{AE}}+\mathcal{
L}_{\textrm{PR}}-\mathcal{H}(q_{\text{data}}({\mathbf{x}}_{0}))$.

For generative modeling, we separately optimize the terms $\mathcal{L}_{\textrm{AE}}$ and $\mathcal{L}_{\textrm{PR}}$, following ; ; . The separate training of the generative prior distribution with a powerful generative model is effective in reducing the mismatch between the prior and the aggregated posterior ; . Initially, we optimize $\mathcal{L}_{\textrm{AE}}$ with respect to the parameters of encoder ($\boldsymbol{\phi}$), decoder ($\boldsymbol{\psi}$), and score network ($\boldsymbol{\theta}$), and fix the parameters $\boldsymbol{\theta},\boldsymbol{\phi},\boldsymbol{\psi}$. Subsequently, we newly parameterize the generative prior $p_{\text{prior}}({\mathbf{z}}):=p_{\boldsymbol{\omega}}({\mathbf{z}})$ using shallow latent diffusion models, and optimize $\mathcal{L}_{\textrm{PR}}$ w.r.t $\boldsymbol{\omega}$. See Section A.3 for further details.

###### Theorem 3 .

## 4 Experiment

This section empirically validates the effectiveness of the intended design of the proposed model, DBAE. To compare DBAE with previous diffusion-based representation learning ; ; , we follow the previous experimental settings of DiffAE  as much as possible. We utilize the same encoder as well as score model architecture as described in Appendix C. We quantify both latent inference quality and generation quality through various tasks. We quantitatively evaluate the performance of downstream inference, reconstruction, disentanglement, and unconditional generation. We also show interpolation and attribute manipulation ability, qualitatively. We conducted experiments by dividing the encoder of the proposed model into two variations: 1) using a Gaussian stochastic encoder (DBAE) and 2) using a deterministic encoder (DBAE-d) for ablation studies. Task-specific detailed configurations and metric descriptions are elaborated in Appendix C.

### 4.1 Downstream Inference

To examine the quality of latent representation in the tractable latent variable ${\mathbf{z}}$, we perform a linear-probe attribute prediction following DiTi . We train DBAE with $\mathcal{L}_{\text{AE}}$ on CelebA  training dataset and FFHQ  dataset. For evaluation, we train a linear classifier on CelebA training set with 40 binary labels and examine CelebA test set using the Average Precision (AP) metric. Additionally, we evaluate using LFW ; dataset, which is designed for attribute regression, and we measure the Pearson correlation coefficient (Pearson’s r) and MSE. Table 1 shows that DBAE outperforms other diffusion-based representation learning baselines ; ; . Since DiffAE , PDAE , and DiTi  suffer from information split problem, they induce ${\mathbf{z}}$ less informative compared to DBAE. Figure 3 shows statistics for 100 reconstructions of the same image with inferred ${\mathbf{z}}$. Since PDAE’s reconstruction varies depending on the selection of ${\mathbf{x}}_{T}$, it indicates that the intricacies of hair and facial features are contained in ${\mathbf{x}}_{T}$, which ${\mathbf{z}}$ misses. This aligns with Figure 8 in Appendix D, there is a significant performance gain of the attributes related to facial details, e.g. shadow, hair, etc. Comparing DBAE-d with DBAE reveals ablations study about the encoder’s stochasticity, and the stochastic encoder exhibits slightly better. We conjecture that the stochastic encoder leverages a broader space, which has the benefits of discriminative downstream inference.

Table: Table 1: Linear-probe attribute prediction quality comparison for models trained on CelebA and FFHQ with dim$({\mathbf{z}})=512$. ‘Gen’ indicates the generation capability. The best and second-best results are highlighted in bold and underline, respectively. We evaluate 5 times and report the average.

### 4.2 Reconstruction

Table: Table 2: Autoencoding reconstruction quality comparison. Among tractable and 512-dimensional latent variable models, the one yielding the best performance is highlighted in bold, underline for the next best performer.

We examine reconstruction quality following DiffAE  to quantify information loss on the latent variable and generation quality. We trained DBAE on FFHQ dataset and reconstructed it on CelebA-HQ  dataset. Table 2 shows the average values of SSIM , LPIPS , and MSE compared to the baselines. Tractability indicates that the inference of latent variables is performed without repeated neural network evaluations. Tractability is crucial for regularizing the latent variable for specific desires (i.e., disentanglement) during the training phase. The latent dimension indicates the dimension of the bottleneck latent variable in inference. A lower dimension is beneficial for other applications, such as downstream inference or attribute manipulation. DBAE infers ${\mathbf{x}}_{T}$ solely based on ${\mathbf{z}}$, allowing for the maintenance of a reduced bottleneck dimension even when inferring ${\mathbf{x}}_{T}$. The third block in Table 2 compares performances under the same qualitative conditions. DBAE-d exhibits performance that surpasses both DiffAE and PDAE. Naturally, DiffAE and PDAE exhibit worse performance because the information is split into ${\mathbf{x}}_{T}$ and ${\mathbf{z}}$. Unlike the downstream inference experiments in Section 4.1, the deterministic encoder performs better in this case.

### 4.3 Disentanglement

Table: Table 3: Disentanglment and sample quality comparisons on CelebA dataset. All the results are evaluated 5 times.

Unsupervised disentanglement of the latent variable ${\mathbf{z}}$ is an important application of generative representation learning, enabling controllable generation without supervision. To achieve this, we regularize to minimize total correlation (TC), i.e., $D_{\text{KL}}(q_{\boldsymbol{\phi}}({\mathbf{z}})||\Pi_{i=1}^{l}q_{\boldsymbol
{\phi}}({\mathbf{z}}_{i}))$, adopted from . TC regularization decouples the correlation between the dimensions of ${\mathbf{z}}$, allowing different information to be contained in each dimension. Disentangling regularization ; often requires obtaining the latent variable in a single feed-forward; we again highlight the importance of tractability here (i.e., ${\mathbf{x}}_{T}$ is intractable to regularize in this scheme). By following InfoDiffusion , we measure TAD and ATTRS  to quantify disentanglment on ${\mathbf{z}}$. Since sample quality and disentangling are in a trade-off , we also measure FID  between 10k samples. Table 3 shows the performance comparison, where DBAE outperforms all the baselines. Figure 4 demonstrates the effects of coefficients on TC regularization. During the exchange between FID and TAD, DBAE envelops all the baselines.

### 4.4 Unconditional Generation

Table: Table 4: Unconditional generation on FFHQ. ‘+AE’ indicates the use of the inferred distribution $q_{\boldsymbol{\phi}}({\mathbf{z}})$ instead of $p_{\boldsymbol{\omega}}({\mathbf{z}})$.

To quantify the capability of unconditional generation, we experiment on CelebA and FFHQ. For the CelebA, there is a comparison with DiffAE in Table 3, DBAE surpasses DiffAE by a large margin in FID  (22.7 vs 11.8). Table 4 shows the performance on FFHQ, which is known to be more diverse than CelebA. DBAE still performs the best among baselines in terms of Precision  and Inception Score (IS)  which is highly affected by image fidelity . DBAE shows slightly worse FID  and Recall , which is highly affected by sample diversity. To analyze it, we alter the learned generative prior $p_{\boldsymbol{\omega}}({\mathbf{z}})$ into inferred distribution $q_{\boldsymbol{\phi}}({\mathbf{z}})$ as shown in the second block of Table 4. In this autoencoding case, DBAE both captures image fidelity and diversity. We speculate that it is more sensitive to the gap between the $q_{\boldsymbol{\phi}}({\mathbf{z}})$ and $p_{\boldsymbol{\omega}}({\mathbf{z}})$ since the information is only dependent on ${\mathbf{z}}$, not depending on the joint condition of ${\mathbf{x}}_{T}$ and ${\mathbf{z}}$. A complex generative prior model $\boldsymbol{\omega}$ could potentially solve this issue ; . Figure 5 shows the generated FFHQ samples from DBAE.

### 4.5 Interpolation

For the two images ${\mathbf{x}}_{0}^{1}$ and ${\mathbf{x}}_{0}^{2}$, DBAE can mix the styles by exploring the intermediate points in the latent space. We encode images into ${\mathbf{z}}^{1}=$Enc${}_{\boldsymbol{\phi}}({\mathbf{x}}_{0}^{1})$ and ${\mathbf{z}}^{2}=$Enc${}_{\boldsymbol{\phi}}({\mathbf{x}}_{0}^{2})$. We then regenerate from ${\mathbf{z}}^{\lambda}=\lambda{\mathbf{z}}^{1}+(1-\lambda){\mathbf{z}}^{2}$ to data ${\mathbf{x}}_{0}$ using the generative process specified in eq. 11. The unique properties of DBAE offer distinct benefits here: 1) DiffAE  and PDAE  need to infer ${\mathbf{x}}_{T}^{1}$, ${\mathbf{x}}_{T}^{2}$ by solving the ODE in eq. 4 with hundreds of score function evaluations ; . They then geometrically interpolate between ${\mathbf{x}}_{T}^{1}$ and ${\mathbf{x}}_{T}^{2}$ to obtain ${\mathbf{x}}_{T}^{\lambda}$, regardless of the correspondence between ${\mathbf{z}}^{\lambda}$ and ${\mathbf{x}}_{T}^{\lambda}$. 2) DBAE directly obtains an intermediate value of ${\mathbf{x}}_{T}^{\lambda}=$Dec${}_{\boldsymbol{\psi}}({\mathbf{z}}^{\lambda})$. This does not require solving the ODE, and the correspondence between ${\mathbf{x}}_{T}^{\lambda}$ and ${\mathbf{z}}^{\lambda}$ is also naturally determined by the decoder ($\boldsymbol{\psi}$). Figure 6(a) shows the interpolation results on the LSUN Horse, Bedroom  and FFHQ datasets. The top row shows the corresponding endpoints ${\mathbf{x}}_{T}^{\lambda}$ in the interpolation, which changes smoothly between ${\mathbf{x}}_{T}^{1}$ and ${\mathbf{x}}_{T}^{2}$. The bottom row shows the interpolation results on FFHQ, which smoothly changes semantic information such as gender, glasses, and hair color.

### 4.6 Attribute Manipulation

Figure: (a) Interpolations on Horse, Bedroom with endpoints (top), and FFHQ (Bottom).
Refer to caption: extracted/5622713/Figs/Figure30_2.jpg

**Table 5: Computation costs for ${\mathbf{x}}_{T}$ inference.**
|  |  | NFE ($\downarrow$) |  | Total<br>time ($\downarrow$)<br>(ms) |
| --- | --- | --- | --- | --- |
| Method | Encϕ | Decψ | $\mathbf{s}_{\boldsymbol{\theta}}$ |  |
| PDAE | 1 | 500 | 500 | 688 |
| DiffAE | 1 | - | 250 | 265 |
| DBAE | 1 | 1 | 0 | 0.31 |

The latent space learned from unsupervised representation learning is useful for image manipulation. If we can access supervised datasets during the testing phase, it becomes possible to identify semantically meaningful directions in the learned latent space. To simulate this scenario, we experiment with DBAE trained on FFHQ and supervised dataset CelebA-HQ with 40 attributes. Following ; , we train a linear classifier and identify the direction of manipulation in the space of ${\mathbf{z}}$. DiffAE  and PDAE  additionally infer from ${\mathbf{x}}_{0}$ to ${\mathbf{x}}_{T}$ by solving eq. 4 with hundreds of score function evaluation and fix ${\mathbf{x}}_{T}$ to prevent undesirable variations in ${\mathbf{x}}_{T}$. Table 5 describes the long inference time of ${\mathbf{x}}_{T}$ in previous approaches. Moreover, if some information is split into ${\mathbf{x}}_{T}$, they cannot manage such information. On the other hand, DBAE infers ${\mathbf{x}}_{T}$ directly from manipulated ${\mathbf{z}}$, accordingly, the endpoint ${\mathbf{x}}_{T}$ is also controlled through the decoder ($\boldsymbol{\psi}$). Figure 6(b) shows the manipulation results both on CelebA-HQ images and FFHQ images with various attributes.

## 5 Conclusion

In this paper, we identify the over-regularization problem and the information split problem in diffusion-based representation learning, stemming from the fixed endpoint distribution ${\mathbf{x}}_{T}$. These issues hinder the representation capabilities of the tractable latent variable ${\mathbf{z}}$. The proposed method, Diffusion Bridge AutoEncoders, systematically addresses these challenges by constructing ${\mathbf{z}}$-dependent endpoint ${\mathbf{x}}_{T}$ inference. By transforming ${\mathbf{z}}$ into an information bottleneck, DBAE is able to extract more meaningful embeddings within the tractable latent space. The notable enhancements in the embedding quality of DBAE improve downstream inference, as well as image manipulation applications. This work lays a solid foundation for further exploration into effective representation learning in learnable diffusion inference. The usages like deep fake can potentially risk of this work.

## Appendix A Proofs and Mathematical Explanations

In this section, we follow the assumptions from Appendix A in [], and we also assume that both $\mathbf{s}_{\boldsymbol{\theta}}$ and $q_{\boldsymbol{\phi},\boldsymbol{\psi}}^{t}$ have continuous second-order derivatives and finite second moments, which are the same assumptions of Theorems 2 and 4 in [].

### A.1 Proof of Theorem 1

See 1

###### Proof.

Note that the definitions of the objective functions are

$$ $\displaystyle\mathcal{L}_{\text{SM}}:=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t }_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^ {2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})- \nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t,$ (17) $\displaystyle\mathcal{L}_{\text{AE}}:=\mathcal{L}_{\text{SM}}+\mathcal{H}(q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T})).$ (18) $$

We derive the score-matching objective $\mathcal{L}_{\text{SM}}$ with the denoising version for tractability. First, $\mathcal{L}_{\text{SM}}$ is derived as follows.

$$ $\displaystyle\mathcal{L}_{\text{SM}}=\frac{1}{2}\int_{0}^{T}\mathbb{E}_{q^{t}_ {\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}$ $\displaystyle[g^{2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})||_{2}^{2}+g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^ {2}$ $\displaystyle-2g^{2}(t)\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\mathrm{d}t.$ (19) $$

Then, the last inner product term of section A.1 can be deduced in a similar approach to []:

$$ $\displaystyle\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{ x}}_{t},{\mathbf{x}}_{T})}[\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t ,{\mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]$ (20) $\displaystyle=\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t },{\mathbf{x}}_{T})\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t }\mathrm{d}{\mathbf{x}}_{T}$ (21) $\displaystyle=\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T })q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T })\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})^{T} \nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}\mathrm{d}{\mathbf{ x}}_{T}$ (22) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{t}|{\mathbf{x}}_{T})\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t }\Big{]}$ (23) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t }\Big{]}$ (24) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\Big{\{}\nabla_{{\mathbf{x}}_{t}}\int q^{t}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T})q^{t}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0}) \mathrm{d}{\mathbf{x}}_{0}\Big{\}}\mathrm{d}{\mathbf{x}}_{t}\Big{]}$ (25) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\Big{\{}\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{0}|{\mathbf{x}}_{T})\nabla_{{\mathbf{x}}_{t}}q^{t}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0}) \mathrm{d}{\mathbf{x}}_{0}\Big{\}}\mathrm{d}{\mathbf{x}}_{t}\Big{]}$ (26) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T})^{T}\Big{\{}\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{0}|{\mathbf{x}}_{T})q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})\nabla_{{\mathbf{x}}_{t}} \log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}} _{T},{\mathbf{x}}_{0})\mathrm{d}{\mathbf{x}}_{0}\Big{\}}\mathrm{d}{\mathbf{x}} _{t}\Big{]}$ (27) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T})}\Big{[}\int\int q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{0}|{\mathbf{x}}_{T})q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})\mathbf{s}_{\boldsymbol{\theta}}({ \mathbf{x}}_{t},t,{\mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{ \mathbf{x}}_{0})\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{t}\Big{]}$ (28) $\displaystyle=\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[\mathbf{s}_{\boldsymbol{\theta}}( {\mathbf{x}}_{t},t,{\mathbf{x}}_{T})^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{ \mathbf{x}}_{0})]$ (29) $$

Next, we rewrite the second term of section A.1. To begin, we express the entropy $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{
x}}_{T}))$ with $\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{t}|{\mathbf{x}}_{T})$, which is similar to the proof of Theorem 4 in [].
Let $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{
x}}_{T})):=-\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{
\mathbf{x}}_{T})\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},
{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}\mathrm{d}{\mathbf{x}}_{T}$ be the joint entropy function of $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})$. Note that $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T},{\mathbf{
x}}_{T}))=\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}))$. Then, we have

$$ $\displaystyle\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_ {0},{\mathbf{x}}_{T}))=$ $\displaystyle\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_ {T},{\mathbf{x}}_{T}))+\int_{T}^{0}\frac{\partial\mathcal{H}_{t}({\mathbf{x}}_ {t},{\mathbf{x}}_{T})}{\partial t}\mathrm{d}t.$ (30) $$

We can expand the integrand of eq. 30 as follows.

$$ $\displaystyle\frac{\partial\mathcal{H}_{t}({\mathbf{x}}_{t},{\mathbf{x}}_{T})} {\partial t}$ $\displaystyle=\frac{\partial}{\partial t}\Big{[}-\int q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})\log q_{\boldsymbol{\phi} ,\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{ t}\mathrm{d}{\mathbf{x}}_{T}\Big{]}$ (31) $\displaystyle=\frac{\partial}{\partial t}\Big{[}-\int q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{T})q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})[\log q_{\boldsymbol{\phi},\boldsymbol{\psi}} ({\mathbf{x}}_{T})+\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{ t}|{\mathbf{x}}_{T})]\mathrm{d}{\mathbf{x}}_{t}\mathrm{d}{\mathbf{x}}_{T}\Big{]}$ (32) $\displaystyle=-\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}) \frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})[\log q_{\boldsymbol{\phi},\boldsymbol{\psi}} ({\mathbf{x}}_{T})+\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{ t}|{\mathbf{x}}_{T})]\big{\}}\mathrm{d}{\mathbf{x}}_{t}\mathrm{d}{\mathbf{x}}_ {T}$ (33) $\displaystyle=-\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{T})}\Big{[}\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})[\log q_{\boldsymbol{\phi },\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\big{\}}\mathrm{d}{\mathbf{x}}_{t} \Big{]}$ (34) $$

We further expand the integration in the last term as follows.

$$ $\displaystyle\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})[\log q_{\boldsymbol{\phi },\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\big{\}}\mathrm{d}{\mathbf{x}}_{t}$ (35) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}[\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]+q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\frac{\partial\log q _{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})}{ \partial t}\mathrm{d}{\mathbf{x}}_{t}$ (36) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}[\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]+\frac{\partial q_ {\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})}{ \partial t}\mathrm{d}{\mathbf{x}}_{t}$ (37) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}[\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\mathrm{d}{\mathbf {x}}_{t}+\frac{\partial}{\partial t}\int q_{\boldsymbol{\phi},\boldsymbol{\psi }}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}$ (38) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}[\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})+\log q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\mathrm{d}{\mathbf {x}}_{t}$ (39) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{ t}+\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}\log q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}$ (40) $\displaystyle=\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}) \frac{\partial}{\partial t}\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}+\int\frac{\partial }{\partial t}\big{\{}q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}| {\mathbf{x}}_{T})\big{\}}\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t}$ (41) $\displaystyle=\int\frac{\partial}{\partial t}\big{\{}q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\big{\}}\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm {d}{\mathbf{x}}_{t}$ (42) $$

Note that we use $\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T}
)\mathrm{d}{\mathbf{x}}_{t}=1$ in eqs. 39 and 42.

By eq. (51) in [], the Fokker-Plank equation for $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})$ follows

$$ $\displaystyle\frac{\partial}{\partial t}q_{\boldsymbol{\phi},\boldsymbol{\psi} }({\mathbf{x}}_{t}|{\mathbf{x}}_{T})=$ $\displaystyle-\nabla_{{\mathbf{x}}_{t}}\cdot\Big{[}(\mathbf{f}({\mathbf{x}}_{t },t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T))q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\Big{]}$ $\displaystyle+\frac{1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\cdot\nabla_{{ \mathbf{x}}_{t}}q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{ \mathbf{x}}_{T})$ (43) $\displaystyle=$ $\displaystyle-\nabla_{{\mathbf{x}}_{t}}\cdot[\tilde{\mathbf{f}}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t)q_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})],$ (44) $$

where $\tilde{\mathbf{f}}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t):=
\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{
x}}_{T},T)-\frac{1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log q_{\boldsymbol{
\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})$.

Combining eqs. 34, 42 and 44, we have

$$ $\displaystyle\frac{\partial\mathcal{H}_{t}({\mathbf{x}}_{t},{\mathbf{x}}_{T})} {\partial t}$ (45) $\displaystyle=-\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{T})}\Big{[}\int-\nabla_{{\mathbf{x}}_{t}}\cdot[\tilde{\mathbf{f}}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t)q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\log q_{\boldsymbol{\phi },\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_ {t}\Big{]}$ (46) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{T})}\Big{[}\int\nabla_{{\mathbf{x}}_{t}}\cdot[\tilde{\mathbf{f}}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t)q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\log q_{\boldsymbol{\phi },\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_ {t}\Big{]}$ (47) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{T})}\Big{[}\tilde{\mathbf{f}}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{t},t)q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x }}_{T})\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{ x}}_{T})$ $\displaystyle\quad\quad\quad\quad\quad\quad-\int q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\tilde{\mathbf{f}}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t)^{T}\nabla_{{\mathbf{x }}_{t}}\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{ x}}_{T})\mathrm{d}{\mathbf{x}}_{t}\Big{]}$ (48) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{T})}\Big{[}-\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{ \mathbf{x}}_{T})\tilde{\mathbf{f}}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t},t)^{T}\nabla_{{\mathbf{x}}_{t}}\log q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm{d}{\mathbf{x}}_{t }\Big{]}$ (49) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{T})}\Big{[}-\int\{\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf {x}}_{t},t,{\mathbf{x}}_{T},T)-\frac{1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}} \log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T} )\}^{T}$ $\displaystyle\quad\quad\quad\quad\quad\quad\nabla_{{\mathbf{x}}_{t}}\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\mathrm {d}{\mathbf{x}}_{t}\Big{]}$ (50) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{t},{\mathbf{x}}_{T})}\Big{[}\{-\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}(t) \mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)+\frac{1}{2}g^{2}(t)\nabla_{{ \mathbf{x}}_{t}}\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}| {\mathbf{x}}_{T})\}^{T}$ $\displaystyle\quad\quad\quad\quad\quad\quad\quad\nabla_{{\mathbf{x}}_{t}}\log q _{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})\Big{]}$ (51) $\displaystyle=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}} _{t},{\mathbf{x}}_{T})}\Big{[}\frac{1}{2}g^{2}(t)||\nabla_{{\mathbf{x}}_{t}} \log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T} )||_{2}^{2}$ $\displaystyle\quad\quad\quad\quad\quad\quad\quad-\{\mathbf{f}({\mathbf{x}}_{t} ,t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)\}^{T}\nabla_{{ \mathbf{x}}_{t}}\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}| {\mathbf{x}}_{T})\Big{]}.$ (52) $$

Therefore, the joint entropy function $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{
x}}_{T}))$ can be expressed as

$$ $\displaystyle\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}$ $\displaystyle({\mathbf{x}}_{0},{\mathbf{x}}_{T}))=\mathcal{H}(q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}))+\int_{T}^{0}\mathbb{E}_{q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}\Big{[ }\frac{1}{2}g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^{2}$ $\displaystyle-\mathbf{f}({\mathbf{x}}_{t},t)^{T}\nabla_{{\mathbf{x}}_{t}}\log q ^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})- g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)^{T}\nabla_{{\mathbf{ x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{ \mathbf{x}}_{T})\Big{]}\mathrm{d}t.$ (53) $$

We can re-write the above equation as follows.

$$ $\displaystyle\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi }}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q ^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})| |_{2}^{2}]\mathrm{d}t$ (54) $\displaystyle=$ $\displaystyle-2\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{0}|{\mathbf{x}}_{T}))$ (55) $\displaystyle+\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}[2\mathbf{f}({\mathbf{x}}_{t},t)^{T} \nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{t}|{\mathbf{x}}_{T})+2g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{ \mathbf{x}}_{T},T)^{T}\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})]\mathrm{d}t$ $\displaystyle=$ $\displaystyle-2\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{0}|{\mathbf{x}}_{T}))+2\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t},{\mathbf{x}}_{T})}[\nabla_{{\mathbf{x}}_{t }}\cdot\{\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t, {\mathbf{x}}_{T},T)\}]\mathrm{d}t$ (56) $$

Similar to the process above, we can obtain the following results for the following joint entropy function $\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{
x}}_{t},{\mathbf{x}}_{T})):=-\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({
\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})\log q_{\boldsymbol{\phi},
\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})\mathrm{
d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{t}\mathrm{d}{\mathbf{x}}_{T}$.

$$ $\displaystyle\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_ {0},{\mathbf{x}}_{0},{\mathbf{x}}_{T}))=\mathcal{H}(q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{T},{\mathbf{x}}_{T}))+\int_{ T}^{0}\frac{\partial\mathcal{H}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}} _{T})}{\partial t}\mathrm{d}t$ (57) $$

In the following results, we utilize the Fokker-Plank equation for $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{
\mathbf{x}}_{T})$, which comes from eq. (49) in []:

$$ $\displaystyle\frac{\partial}{\partial t}q_{\boldsymbol{\phi},\boldsymbol{\psi} }({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}_{T})=$ $\displaystyle-\nabla_{{\mathbf{x}}_{t}}\cdot\Big{[}(\mathbf{f}({\mathbf{x}}_{t },t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T))q_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}_{T}) \Big{]}$ $\displaystyle+\frac{1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\cdot\nabla_{{ \mathbf{x}}_{t}}q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{ \mathbf{x}}_{0},{\mathbf{x}}_{T})$ (58) $\displaystyle=$ $\displaystyle-\nabla_{{\mathbf{x}}_{t}}\cdot[\hat{\mathbf{f}}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t)q_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}_{T})],$ (59) $$

where $\hat{\mathbf{f}}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t},t):=
\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{
x}}_{T},T)-\frac{1}{2}g^{2}(t)\nabla_{{\mathbf{x}}_{t}}\log q_{\boldsymbol{
\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}_{T})$.

Then, we have

$$ $\displaystyle 0=\int_{T}^{0}\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}\Big{[}$ $\displaystyle\frac{1}{2}g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{ \mathbf{x}}_{0})||_{2}^{2}-\mathbf{f}({\mathbf{x}}_{t},t)\nabla_{{\mathbf{x}}_ {t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf {x}}_{T},{\mathbf{x}}_{0})$ $\displaystyle-g^{2}(t)\mathbf{h}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T},T)\nabla_ {{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})\Big{]}\mathrm{d}t,$ (60) $$

where the left hand side is from $0=\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{
\mathbf{x}}_{0},{\mathbf{x}}_{T}))-\mathcal{H}(q_{\boldsymbol{\phi},
\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{T},{\mathbf{x}}_{T}))$, and right hand side is from $\int_{T}^{0}\frac{\partial\mathcal{H}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{
\mathbf{x}}_{T})}{\partial t}\mathrm{d}t$. We can further derive as follows.

$$ $\displaystyle\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi }}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}$ $\displaystyle[g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi} ,\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})||_{2}^ {2}]\mathrm{d}t$ $\displaystyle=2\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{ \psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[\nabla_{{\mathbf{x }}_{t}}\cdot\{\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_{ t},t,{\mathbf{x}}_{T},T)\}]\mathrm{d}t$ (61) $$

Combining eqs. 56 and 61, we have

$$ $\displaystyle\int_{0}^{T}$ $\displaystyle\mathbb{E}_{q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{ x}}_{t},{\mathbf{x}}_{T})}[g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^ {2}]\mathrm{d}t$ $\displaystyle=-2\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x }}_{0}|{\mathbf{x}}_{T}))+\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}( {\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})||_{2}^{2}]\mathrm{d}t$ (62) $$

Combining all results, the score-matching objective $\mathcal{L}_{\text{SM}}$ can be expressed as

$$ $\displaystyle\mathcal{L}_{\text{SM}}$ $\displaystyle=\frac{1}{2}\int_{0}^{T}\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})||_{2 }^{2}+g^{2}(t)||\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T},{\mathbf{x}}_{0})||_{2}^{2}$ $\displaystyle\quad\quad\quad\quad\quad\quad\quad-2g^{2}(t)\mathbf{s}_{ \boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})^{T}\nabla_{{\mathbf{ x}}_{t}}\log q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{ \mathbf{x}}_{T},{\mathbf{x}}_{0})]\mathrm{d}t-\mathcal{H}(q_{\boldsymbol{\phi} ,\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (63) $\displaystyle=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})- \nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{ \mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t-\mathcal{H}(q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (64) $$

The last equality comes from $q^{t}_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},
{\mathbf{x}}_{T})=\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}
_{T})$, which is based on the Doob’s $h$-transform [, , ].
Finally, we have

$$ $\displaystyle\mathcal{L}_{\text{AE}}$ $\displaystyle=\mathcal{L}_{\text{SM}}+\mathcal{H}(q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (65) $\displaystyle=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})- \nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{ \mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t.$ (66) $$

From here, we show that the objective $\mathcal{L}_{\text{AE}}$ is equivalent to the reconstruction objective. Assume that the forward SDE in eq. 1 is linear SDE in terms of ${\mathbf{x}}_{t}$ (e.g. VP [], VE []). Then the transition kernel $\tilde{q}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ becomes Gaussian distribution. Then, we can represent reparametrized form ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}\boldsymbol{\epsilon}$, where $\alpha_{t}$ and $\sigma_{t}$ are time-dependent constants determined by drift $\mathbf{f}$ and volatility $g$, and $\boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. The time-dependent constant signal-to-noise ratio $SNR(t):=\frac{\alpha_{t}^{2}}{\sigma_{t}^{2}}$ often define to discuss on diffusion process []. We define SNR ratio, $R(t):=\frac{SNR(T)}{SNR(t)}$ for convenient derivation.

[] show the exact form of $\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{\mathbf{x}}_{T}):=\mathcal{N}
(\hat{\mu}_{t},\hat{\sigma}^{2}_{t}\mathbf{I})$, where $\hat{\mu}_{t}=R(t)\frac{\alpha_{t}}{\alpha_{T}}{\mathbf{x}}_{T}+\alpha_{t}{
\mathbf{x}}_{0}(1-R(t))$ and $\hat{\sigma_{t}}=\sigma_{t}\sqrt{1-R(t)}$. This Gaussian form determines the exact analytic form of the score function $\nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{
\mathbf{x}}_{T})$. We plug this into our objective $\mathcal{L}_{\text{AE}}$.

$$ $\displaystyle\mathcal{L}_{\textrm{AE}}$ $\displaystyle=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})- \nabla_{{\mathbf{x}}_{t}}\log\tilde{q}_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0},{ \mathbf{x}}_{T})||_{2}^{2}]}\mathrm{d}t$ (67) $\displaystyle=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[g^{2}( t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T})- \frac{-{\mathbf{x}}_{t}+(R(t)\frac{\alpha_{t}}{\alpha_{T}}{\mathbf{x}}_{T}+ \alpha_{t}{\mathbf{x}}_{0}(1-R(t)))}{\sigma^{2}_{t}(1-R(t))}||_{2}^{2}]} \mathrm{d}t$ (68) $\displaystyle=\frac{1}{2}\int_{0}^{T}{\mathbb{E}_{q^{t}_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{t},{\mathbf{x}}_{T})}[ \lambda(t)||\mathbf{x}^{0}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x }}_{T})-{\mathbf{x}}_{0}||_{2}^{2}]}\mathrm{d}t,$ (69) $$

where

$$ $\displaystyle\lambda(t)=\frac{\alpha_{t}}{\sigma^{2}_{t}}g^{2}(t),$ (70) $\displaystyle\mathbf{x}^{0}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{ x}}_{T}):=\alpha(t){\mathbf{x}}_{t}+\beta(t){\mathbf{x}}_{T}+\gamma(t)\mathbf{ s}_{\boldsymbol{\theta}}({\mathbf{x}}_{t},t,{\mathbf{x}}_{T}),$ (71) $\displaystyle\alpha(t)=\frac{1}{\alpha_{t}(1-R(t))},\quad\beta(t)=-\frac{R(t)} {\alpha_{T}(1-R(t))},\quad\gamma(t)=\frac{\sigma^{2}_{t}}{\alpha_{t}}.$ (72) $$

∎

###### Proof.

### A.2 Proof of Theorem 3

See 3

###### Proof.

From the data processing inequality with our graphical model, we have the following result, similar to eq. (14) in [].

$$ $\displaystyle D_{\textrm{KL}}(q_{\text{data}}({\mathbf{x}}_{0})||p_{ \boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{0}))\leq D_{\textrm{KL}}( q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0:T},{\mathbf{z}})||p_{ \boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{0:T},{\mathbf{z}}))$ (73) $$

Also, the chain rule of KL divergences, we have

$$ $\displaystyle D_{\textrm{KL}}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{0:T},{\mathbf{z}})||p_{\boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x }}_{0:T},{\mathbf{z}}))$ (74) $\displaystyle=D_{\textrm{KL}}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T},{\mathbf{z}})||p_{\boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}} _{T},{\mathbf{z}}))+\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{T},{\mathbf{z}})}[D_{\textrm{KL}}(\mu_{\boldsymbol{\phi}, \boldsymbol{\psi}}(\cdot|{\mathbf{x}}_{T},{\mathbf{z}})||\nu_{\boldsymbol{ \theta},\boldsymbol{\psi}}(\cdot|{\mathbf{x}}_{T},{\mathbf{z}}))],$ (75) $$

where $\mu_{\boldsymbol{\phi},\boldsymbol{\psi}}$ and $\nu_{\boldsymbol{\theta},\boldsymbol{\psi}}$ are the path measures of the SDEs in eqs. 76 and 77, respectively:

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}$ $\displaystyle=[\mathbf{f}({\mathbf{x}}_{t},t)+g^{2}(t)\mathbf{h}({\mathbf{x}}_ {t},t,\mathbf{y},T)]\mathrm{d}t+g(t)\mathrm{d}\mathbf{w}_{t},\quad{\mathbf{x}} _{0}\sim q_{\textrm{data}}({\mathbf{x}}_{0}),\quad{\mathbf{x}}_{T}\sim q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0}),$ (76) $\displaystyle\mathrm{d}{\mathbf{x}}_{t}$ $\displaystyle=[\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}(t)[\nabla_{{\mathbf{x}}_{t }}\log p_{\boldsymbol{\theta}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})-\mathbf{h}({ \mathbf{x}}_{t},t,\mathbf{y},T)]]\mathrm{d}t+g(t)\mathrm{d}\bar{\mathbf{w}}_{t },\quad{\mathbf{x}}_{T}\sim p_{\boldsymbol{\psi}}({\mathbf{x}}_{T}).$ (77) $$

By our graphical modeling, ${\mathbf{z}}$ is independent of $\{{\mathbf{x}}_{t}\}$ given ${\mathbf{x}}_{T}$. Therefore, we have

$$ $\displaystyle\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_ {T},{\mathbf{z}})}[D_{\textrm{KL}}(\mu_{\boldsymbol{\phi},\boldsymbol{\psi}}( \cdot|{\mathbf{x}}_{T},{\mathbf{z}})||\nu_{\boldsymbol{\theta}}(\cdot|{\mathbf {x}}_{T},{\mathbf{z}}))]=\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{T})}[D_{\textrm{KL}}(\mu_{\boldsymbol{\phi},\boldsymbol{\psi}}( \cdot|{\mathbf{x}}_{T})||\nu_{\boldsymbol{\theta}}(\cdot|{\mathbf{x}}_{T}))],$ (78) $$

where $\mu_{\boldsymbol{\phi},\boldsymbol{\psi}}(\cdot|{\mathbf{x}}_{T})$ and $\nu_{\boldsymbol{\theta}}(\cdot|{\mathbf{x}}_{T})$ are the path measures of the SDEs in eqs. 79 and 80, respectively:

$$ $\displaystyle\mathrm{d}{\mathbf{x}}_{t}$ $\displaystyle=[\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}(t)[\nabla_{{\mathbf{x}}_{t }}\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{ T})-\mathbf{h}({\mathbf{x}}_{t},t,\mathbf{y},T)]]\mathrm{d}t+g(t)\mathrm{d} \bar{\mathbf{w}}_{t},\quad{\mathbf{x}}(T)={\mathbf{x}}_{T},$ (79) $\displaystyle\mathrm{d}{\mathbf{x}}_{t}$ $\displaystyle=[\mathbf{f}({\mathbf{x}}_{t},t)-g^{2}(t)[\nabla_{{\mathbf{x}}_{t }}\log p_{\boldsymbol{\theta}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})-\mathbf{h}({ \mathbf{x}}_{t},t,\mathbf{y},T)]]\mathrm{d}t+g(t)\mathrm{d}\bar{\mathbf{w}}_{t },\quad{\mathbf{x}}(T)={\mathbf{x}}_{T}$ (80) $$

Similar to eq. (17) in [], this KL divergence can be expressed using the Girsanov theorem [] and martingale property.

$$ $\displaystyle D_{\textrm{KL}}(\mu_{\boldsymbol{\phi},\boldsymbol{\psi}}(\cdot| {\mathbf{x}}_{T})||\nu_{\boldsymbol{\theta}}(\cdot|{\mathbf{x}}_{T}))=\frac{1} {2}\int_{0}^{T}\mathbb{E}_{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x} }_{t},{\mathbf{x}}_{T})}[g^{2}(t)||\mathbf{s}_{\boldsymbol{\theta}}({\mathbf{x }}_{t},t,{\mathbf{x}}_{T})-\nabla_{{\mathbf{x}}_{t}}\log q^{t}_{\boldsymbol{ \phi},\boldsymbol{\psi}}({\mathbf{x}}_{t}|{\mathbf{x}}_{T})||_{2}^{2}]\mathrm{ d}t$ (81) $$

From eqs. 73, 75 and 81 and theorem 1, we have:

$$ $\displaystyle D_{\textrm{KL}}(q_{\text{data}}({\mathbf{x}}_{0})||p_{ \boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{0}))\leq D_{\textrm{KL}}( q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T},{\mathbf{z}})||p_{ \boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{T},{\mathbf{z}}))+ \mathcal{L}_{\textrm{AE}}-\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}( {\mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (82) $$

Furthermore, the first and third terms of RHS in eq. 82 can be expressed as follows.

$$ $\displaystyle D_{\textrm{KL}}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf {x}}_{T},{\mathbf{z}})||p_{\boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}} _{T},{\mathbf{z}}))-\mathcal{H}(q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{0}|{\mathbf{x}}_{T}))$ (83) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T},{ \mathbf{z}}))\log\frac{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T },{\mathbf{z}})}{p_{\boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{T},{ \mathbf{z}})}\mathrm{d}{\mathbf{x}}_{T}\mathrm{d}{\mathbf{z}}+\int q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{\mathbf{x}}_{T})\log q_ {\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T}) \mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}$ (84) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{ \mathbf{x}}_{T},{\mathbf{z}})\Big{[}\log\frac{q_{\boldsymbol{\phi},\boldsymbol {\psi}}({\mathbf{x}}_{T},{\mathbf{z}})}{p_{\boldsymbol{\psi},\boldsymbol{ \theta}}({\mathbf{x}}_{T},{\mathbf{z}})}+\log q_{\boldsymbol{\phi},\boldsymbol {\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T})\Big{]}\mathrm{d}{\mathbf{x}}_{0} \mathrm{d}{\mathbf{x}}_{T}\mathrm{d}{\mathbf{z}}$ (85) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{ \mathbf{x}}_{T},{\mathbf{z}})\Big{[}\log\frac{q_{\boldsymbol{\phi},\boldsymbol {\psi}}({\mathbf{x}}_{T})q_{\boldsymbol{\psi}}({\mathbf{z}}|{\mathbf{x}}_{T})} {p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})p_{\boldsymbol{\psi}}({\mathbf{z}}|{ \mathbf{x}}_{T})}+\log q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0 }|{\mathbf{x}}_{T})\Big{]}\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T} \mathrm{d}{\mathbf{z}}$ (86) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{ \mathbf{x}}_{T},{\mathbf{z}})\Big{[}\log\frac{q_{\boldsymbol{\phi},\boldsymbol {\psi}}({\mathbf{x}}_{T})}{p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})}+\log q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{\mathbf{x}}_{T})\Big{]} \mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}\mathrm{d}{\mathbf{z}}$ (87) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{ \mathbf{x}}_{T})\Big{[}\log\frac{q_{\boldsymbol{\phi},\boldsymbol{\psi}}({ \mathbf{x}}_{T})q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0}|{ \mathbf{x}}_{T})}{p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})}\big{]}\mathrm{d}{ \mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}$ (88) $\displaystyle=\int q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{0},{ \mathbf{x}}_{T})\Big{[}\log\frac{q_{\text{data}}({\mathbf{x}}_{0})q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})}{p_{ \boldsymbol{\psi}}({\mathbf{x}}_{T})}\big{]}\mathrm{d}{\mathbf{x}}_{0}\mathrm{ d}{\mathbf{x}}_{T}$ (89) $\displaystyle=\int q_{\text{data}}({\mathbf{x}}_{0})q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})\Big{[}\log\frac{q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})}{p_{ \boldsymbol{\psi}}({\mathbf{x}}_{T})}+\log q_{\text{data}}({\mathbf{x}}_{0}) \big{]}\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}$ (90) $\displaystyle=\mathbb{E}_{q_{\text{data}}({\mathbf{x}}_{0})}[D_{\textrm{KL}}(q _{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})||p_{ \boldsymbol{\psi}}({\mathbf{x}}_{T}))]-\mathcal{H}(q_{\text{data}}({\mathbf{x} }_{0}))$ (91) $\displaystyle=\mathcal{L}_{\textrm{PR}}-\mathcal{H}(q_{\text{data}}({\mathbf{x }}_{0}))$ (92) $$

To sum up, we have

$$ $\displaystyle D_{\textrm{KL}}(q_{\text{data}}({\mathbf{x}}_{0})||p_{ \boldsymbol{\psi},\boldsymbol{\theta}}({\mathbf{x}}_{0}))\leq\mathcal{L}_{ \textrm{AE}}+\mathcal{L}_{\textrm{PR}}-\mathcal{H}(q_{\text{data}}({\mathbf{x} }_{0})).$ (93) $$

∎

###### Proof.

### A.3 Prior Optimization Objective

This section explains the details of the prior related objective function mentioned in Section 3.3.2. The proposed objective is $\mathcal{L}_{\textrm{PR}}$ as shown in eq. 94.

$$ $\displaystyle\mathcal{L}_{\textrm{PR}}=$ $\displaystyle\mathbb{E}_{q_{\text{data}}({\mathbf{x}}_{0})}[D_{\text{KL}}(q_{ \boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})||p_{ \boldsymbol{\psi}}({\mathbf{x}}_{T}))]$ (94) $$

To optimize this term, we fix the parameters of the encoder ($\boldsymbol{\phi}\rightarrow\boldsymbol{\phi}^{*}$), the decoder ($\boldsymbol{\psi}\rightarrow\boldsymbol{\psi}^{*}$), and score network ($\boldsymbol{\theta}\rightarrow\boldsymbol{\theta}^{*}$), which is optimized by $\mathcal{L}_{AE}$. And we newly parameterize the generative prior $p_{\text{prior}}({\mathbf{z}})\rightarrow p_{\boldsymbol{\omega}}({\mathbf{z}})$, so the generative endpoint distribution becomes $p_{\boldsymbol{\psi}}({\mathbf{x}}_{T})\rightarrow p_{\boldsymbol{\psi}^{*},
\boldsymbol{\omega}}({\mathbf{x}}_{T})$. We utilize MLP-based latent diffusion models following [, ].

The objective function in eq. 94 with respect to $\boldsymbol{\omega}$ is described in eq. 95 and extends to eq. 99 with equality. Equation 100 is derived from the same optimality condition. In other words, it reduces the problem of training an unconditional generative prior $p_{\omega}({\mathbf{z}})$ to matching the aggregated posterior distribution $q_{\boldsymbol{\phi}^{*}}({\mathbf{z}})$.

$$ $\displaystyle\arg\min_{\boldsymbol{\omega}}{\mathbb{E}_{q_{\text{data}}({ \mathbf{x}}_{0})}[D_{\text{KL}}(q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*} }({\mathbf{x}}_{T}|{\mathbf{x}}_{0})||p_{\boldsymbol{\psi}^{*},\omega}({ \mathbf{x}}_{T}))]}$ (95) $\displaystyle\Leftrightarrow$ $\displaystyle\arg\min_{\boldsymbol{\omega}}{\int q_{\text{data}}({\mathbf{x}}_ {0})q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}({\mathbf{x}}_{T}|{\mathbf{ x}}_{0})\log{\frac{q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}({\mathbf{x} }_{T}|{\mathbf{x}}_{0})}{p_{\boldsymbol{\psi}^{*},\boldsymbol{\omega}}({ \mathbf{x}}_{T})}}}\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}$ (96) $\displaystyle\Leftrightarrow$ $\displaystyle\arg\min_{\boldsymbol{\boldsymbol{\omega}}}{\int q_{\text{data}}( {\mathbf{x}}_{0})q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}({\mathbf{x}}_ {T}|{\mathbf{x}}_{0})\log{\frac{q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*} }({\mathbf{x}}_{T})q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}({\mathbf{x} }_{0}|{\mathbf{x}}_{T})}{p_{\boldsymbol{\psi}^{*},\omega}({\mathbf{x}}_{T})q_{ \text{data}}({\mathbf{x}}_{0})}}}\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{ x}}_{T}$ (97) $\displaystyle\Leftrightarrow$ $\displaystyle\arg\min_{\boldsymbol{\boldsymbol{\omega}}}{\int q_{\text{data}}( {\mathbf{x}}_{0})q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}({\mathbf{x}}_ {T}|{\mathbf{x}}_{0})\left[\log{\frac{q_{\boldsymbol{\phi}^{*},\boldsymbol{ \psi}^{*}}({\mathbf{x}}_{T})}{p_{\boldsymbol{\psi}^{*},\boldsymbol{\omega}}({ \mathbf{x}}_{T})}}+\log{\frac{q_{\boldsymbol{\phi}^{*},\boldsymbol{\psi}^{*}}( {\mathbf{x}}_{0}|{\mathbf{x}}_{T})}{q_{\text{data}}({\mathbf{x}}_{0})}}\right] }\mathrm{d}{\mathbf{x}}_{0}\mathrm{d}{\mathbf{x}}_{T}$ (98) $\displaystyle\Leftrightarrow$ $\displaystyle\arg\min_{\boldsymbol{\omega}}{D_{\text{KL}}(q_{\boldsymbol{\phi} ^{*},\boldsymbol{\psi}^{*}}({\mathbf{x}}_{T})||p_{\boldsymbol{\psi}^{*}, \boldsymbol{\omega}}({\mathbf{x}}_{T}))}+C$ (99) $\displaystyle\Leftrightarrow$ $\displaystyle\arg\min_{\boldsymbol{\omega}}{D_{\text{KL}}(q_{\boldsymbol{\phi} ^{*}}({\mathbf{z}})||p_{\boldsymbol{\omega}}({\mathbf{z}}))}$ (100) $$

## Appendix B Related Work

### B.1 Representation Learning in Diffusion Models

Expanding the applicability of generative models to various downstream tasks depends on exploring meaningful latent variables through representation learning. Methods within both variational autoencoders (VAEs) [, , , , ] and generative adversarial networks (GANs) [, , , , ] have been proposed; however, VAEs suffer from low sample quality, limiting their practical deployment in real-world scenarios. Conversely, GANs are known for their ability to produce high-quality samples with fast sampling speeds but face challenges in accessing latent variables due to their intractable model structure. This leads to computationally expensive inference methods like GAN inversion [, , , , ]. Additionally, the adversarial training objective of GANs introduces instability during the training.

In contrast, recent research has delved into representation learning within diffusion probabilistic models (DPMs), which offer stable training and high sample quality. In early studies, the diffusion endpoint ${\mathbf{x}}_{T}$ was introduced as a semantic latent variable [, ] with an invertible path defined by an ordinary differential equation (ODE). However, ${\mathbf{x}}_{T}$ is problematic due to over-regularization by the Gaussian distribution, which impedes effective encoding. Additionally, the dimension of ${\mathbf{x}}_{T}$ matches that of the original data ${\mathbf{x}}_{0}$, limiting the ability to learn condensed feature representation for downstream tasks (e.g., downstream inference, attribute manipulation with linear classifier). The inference of latent variable also relies on solving ODE, rendering inference intractable. This intractability not only hinders the desired regularization (e.g. disentanglment [, , ]) of the latent variable but also slows down the downstream applications.

Diffusion AutoEncoder (DiffAE) [] introduces a new framework for learning tractable latent variables in DPMs. DiffAE learns representation in the latent variable ${\mathbf{z}}$ through an auxiliary encoder, with a ${\mathbf{z}}$-conditional score network []. The encoder-generated latent variable ${\mathbf{z}}$ does not suffer from over-regularization, or intractability and has flexible dimensionality. Pre-trained DPM AutoEncoding (PDAE) [] proposes a method to learn unsupervised representation from pre-trained unconditional DPMs. PDAE also employs an auxiliary encoder to define ${\mathbf{z}}$ and introduces a decoder to represent $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{z}}|{\mathbf{x}})$. PDAE can parameterize the ${\mathbf{z}}$-conditional model score combined with a pre-trained unconditional score network, utilizing the idea of classifier guidance []. PDAE can use the pre-trained checkpoint from publicly available sources, but its complex decoder architecture slows down the sampling speed.

Subsequent studies have imposed additional assumptions or constraints on the encoder based on specific objectives. DiTi [] introduces a time-dependent latent variable on the top of PDAE to enable feature learning that depends on diffusion time. InfoDiffusion [] regularizes the latent space of DiffAE to foster an informative and disentangled representation of ${\mathbf{z}}$. It should be noted that such proposed regularization in [] is also applicable with DBAE, and Section 4.3 demonstrates that the tradeoff between disentanglement and sample quality is better managed in DBAE than in DiffAE. FDAE [] learns disentangled latent representation by masking image pixel content with DiffAE. DisDiff [] learns disentangled latent variable ${\mathbf{z}}$ by minimizing mutual information between each latent variable from different dimensions atop PDAE. LCG-DM [] adopts a pre-trained disentangled encoder and trains DiffAE structure with fixed encoder parameters to enable unsupervised controllable generation. SODA [] improves the network architectures of DiffAE and training for novel image reconstruction.

All the frameworks [, ] and applications [, , , , ] utilize the encoder and do not consider the diffusion endpoint ${\mathbf{x}}_{T}$, leading to a information split problem. In contrast, DBAE constructs an ${\mathbf{z}}$-dependent endpoint ${\mathbf{x}}_{T}$ inference with feed-forward architecture to induce ${\mathbf{z}}$ as an information bottleneck. Our framework makes ${\mathbf{z}}$ more informative, which is orthogonal to advancements in downstream applications [, , , , , ], as exemplified in Section 4.3.

### B.2 Parametrized Forward Diffusion

Forward diffusion process with learnable parameters is a key technique in DBAE to resolve over-regularization problem at the diffusion endpoint ${\mathbf{x}}_{T}$. Here we briefly review the previous related work that enabled parametrized forward diffusion.

Sch$\ddot{\textrm{o}}$dinger bridge problem (SBP) [, ] learns the pair of SDEs that have forward and reverse dynamics relationships. SBP identifies the joint distribution in the form of a diffusion path between two given marginal distributions. The optimization is reduced to entropy-regularized optimal transport [, ], which is often solved by Iterative Proportional Fitting []. For this optimization, samples are required at any given time $t$ from the forward SDE; however, these samples are not from a Gaussian kernel like eq. 1 or eq. 5, resulting in longer training times needed to solve the SDE numerically with intermediate particles. The formulation is also not suitable for our case, as we learn the given joint distribution through an encoder-decoder framework.

Diffusion normalizing flow (DiffFlow) [] parameterizes the drift term in eq. 1 using a normalizing flow, making the endpoint of DiffFlow learnable. However, both training and endpoint inference are intractable because the parametrized forward SDE does not provide a Gaussian kernel similar to that in SBP. Implicit nonlinear diffusion model (INDM) [] learns a diffusion model that is defined in the latent space of a normalizing flow, implicitly parameterizing both the drift and volatility terms in eq. 1. A unique benefit is its tractable training, allowing direct sampling from any diffusion time $t$. However, INDM merely progresses the existing diffusion process in the flow latent space, making it unsuitable for encoding due to technical issues such as dimensionality. The diffusion process itself is still over-regularized as Gaussian, and inference must be conducted through ODE for invertible encoding.

Unlike other studies, we parameterize the endpoint ${\mathbf{x}}_{T}$ rather than the drift or volatility terms. The forward process is naturally influenced by the endpoint determined from Doob’s $h$-transform. Unlike other parameterized diffusions, our approach ensures tractable learning and ${\mathbf{x}}_{T}$ inference, making it particularly advantageous for encoding tasks.

## Appendix C Implementation details

### C.1 Training Configuration

Model Architecture We use the score network ($\boldsymbol{\theta}$) backbone U-Net [], which are modified for diffusion models [] with time-embedding. DiffAE [], PDAE [], and DiTi [] also utilize the same score network architecture. The only difference for DBAE is the endpoint ${\mathbf{x}}_{T}$ conditioning. We follow DDBM [] which concatenate ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{T}$ for the inputs as described in Figure 7(b). This modification only increases the input channels, so the complexity increase is marginal. While the endpoint ${\mathbf{x}}_{T}$ contains all the information from ${\mathbf{z}}$, we design a score network also conditioning on ${\mathbf{z}}$ for implementation to effectively utilize the latent information in the generative process. For the encoder ($\boldsymbol{\phi}$), we utilize exactly the same structure from DiffAE []. For the decoder ($\boldsymbol{\psi}$), we adopt the upsampling structure from the generator of FastGAN [], while removing the intermediate stochastic element. For the generative prior ($\boldsymbol{\omega}$), we utilize latent ddim from []. Tables 6 and 7 explains the network configurations for the aforementioned structures.

Optimization We follow the optimization argument from DDBM [] with Variance Preserving (VP) SDE. We utilize the preconditioning and time-weighting proposed in DDBM, with the pred-x parameterization []. Table 6 shows the remaining optimization hyperparameters. While DDBM does not include the encoder ($\boldsymbol{\phi}$) and the decoder ($\boldsymbol{\psi}$), we optimize jointly the parameters $\boldsymbol{\phi}$, $\boldsymbol{\psi}$, and $\boldsymbol{\theta}$ to minimize $\mathcal{L}_{\text{AE}}$.

Figure: (a) The encoder ($\boldsymbol{\phi}$) structure in the left, and decoder ($\boldsymbol{\psi}$) structure in the right.
Refer to caption: extracted/5622713/Figs/Figure47.jpg

**Table 6: Network architecture and training configuration of DBAE based on [, , ].**
| Parameter | CelebA 64 | FFHQ 128 | Horse 128 | Bedroom 128 |
| --- | --- | --- | --- | --- |
| Base channels | 64 | 128 | 128 | 128 |
| Channel multipliers | [1,2,4,8] | [1,1,2,3,4] | [1,1,2,3,4] | [1,1,2,3,4] |
| Attention resolution | [16] | [16] | [16] | [16] |
| Encoder base ch | 64 | 128 | 128 | 128 |
| Enc. attn. resolution | [16] | [16] | [16] | [16] |
| Encoder ch. mult. | [1,2,4,8,8] | [1,1,2,3,4,4] | [1,1,2,3,4,4] | [1,1,2,3,4,4] |
| latent variable ${\mathbf{z}}$ dimension | 32, 256, 512 | 512 | 512 | 512 |
| Vanilla forward SDE | VP | VP | VP | VP |
| Images trained | 72M, 130M | 130M | 130M | 130M |
| Batch size | 128 | 128 | 128 | 128 |
| Learning rate | 1e-4 | 1e-4 | 1e-4 | 1e-4 |
| Optimizer | RAdam | RAdam | RAdam | RAdam |
| Weight decay | 0.0 | 0.0 | 0.0 | 0.0 |
| EMA rate | 0.9999 | 0.9999 | 0.9999 | 0.9999 |

**Table 7: Network architecture and training configuration of latent diffusion models $p_{\boldsymbol{\omega}}({\mathbf{z}})$ for an unconditional generation, following [].**
| Parameter | CelebA 64 | FFHQ 128 |
| --- | --- | --- |
| Batch size | 512 | 256 |
| ${\mathbf{z}}$ trained | 600M | 600M |
| MLP layers (N) | 10, 15 | 10 |
| MLP hidden size | 2048 |  |
| latent variable ${\mathbf{z}}$ dimension | 512 |  |
| SDE | VP |  |
| $\beta$ scheduler | Constant 0.008 |  |
| Learning rate | 1e-4 |  |
| Optimizer | AdamW (weight decay = 0.01) |  |
| Train Diff T | 1000 |  |
| Diffusion loss | L1, L2 |  |

### C.2 Evaluation Configuration and Metric

Downstream Inference In Table 1, we use Average Precision (AP), Pearson Correlation Coefficient (Pearson’s r), and Mean Squared Error (MSE) as metrics for comparison. For AP measurement, we train a linear classifier $(\mathbb{R}^{l}\rightarrow\mathbb{R}^{40})$ to classify 40 binary attribute labels from the CelebA [] training dataset. The output of the encoder, Enc${}_{\boldsymbol{\phi}}({\mathbf{x}}_{0})={\mathbf{z}}$, serves as the input for a linear classifier. We examine the CelebA test dataset. Precision and recall for each attribute label are calculated by computing true positives (TP), false positives (FP), and false negatives (FN) for each threshold interval divided by predicted values. The area under the precision-recall curve is obtained as AP. For Pearson’s r and MSE, we train a linear regressor $(\mathbb{R}^{l}\rightarrow\mathbb{R}^{73})$ using LFW [, ] dataset. The regressor predicts the value of 73 attributes based on the latent variable ${\mathbf{z}}$. Pearson’s r is evaluated by calculating the variance and covariance between the ground truth and predicted values for each attribute, while MSE is assessed by measuring the differences between two values. We borrow the baseline results from the DiTi [] paper and strictly adhere to the evaluation protocol found at [https://github.com/yue-zhongqi/diti](https://github.com/yue-zhongqi/diti).

Reconstruction We quantify reconstruction error in Table 2 though the Structural Similarity Index Measure (SSIM) [], Learned Perceptual Image Patch Similarity (LPIPS) [] and Mean Squared Error (MSE). This metric measures the distance between original images in CelebA-HQ and their reconstructions across all 30K samples and averages them. SSIM compares the luminance, contrast, and structure between images to measure the differences on a scale from 0 to 1, in a manner similar to human visual perception. LPIPS measures the distance in the feature space of a neural network that learns the similarity between two images. We borrow the baseline results from DiffAE [] and PDAE []. In the appendix, we also present performance metrics according to various NFE in Tables 9 and 10.

Disentanglment The metric Total AUROC Difference (TAD) [] measures how effectively the latent space is disentangled, utilizing a dataset with multiple binary ground truth labels. It calculates the correlation between attributes based on the proportion of entropy reduction given any other single attribute. Attributes that show an entropy reduction greater than 0.2 when conditioned on another attribute are considered highly correlated and therefore entangled. For each remaining attribute that is not considered entangled, we calculate the AUROC score for each dimension of the latent variable ${\mathbf{z}}$. To calculate the AUROC score, first determine the dimension-wise minimum and maximum values of ${\mathbf{z}}$. We increment the threshold from the minimum to the maximum for each dimension, converting ${\mathbf{z}}$ to a one-hot vector by comparing each dimension’s value against the threshold. This one-hot vector is then compared to the true labels to compute the AUROC score. An attribute is considered disentangled if at least one dimension of ${\mathbf{z}}$ can detect it with an AUROC score of 0.75 or higher. The sub-metric ATTRS denotes the number of such captured attributes. The TAD score is calculated as the sum of the differences between the two highest AUROC scores for each captured attribute. We randomly selected 1000 samples from the CelebA training, validation, and test sets to perform the measurement following []. We borrow the baseline results expect DisDiff from the InfoDiffusion [], and we follow their setting that the dim$({\mathbf{z}})=32$. DisDiff [] utilizes the dim$({\mathbf{z}})=192$ and we borrow its performance from the original paper. We use evaluation code from [https://github.com/ericyeats/nashae-beamsynthesis](https://github.com/ericyeats/nashae-beamsynthesis).

Unconditional Generation To measure unconditional generative modeling, we quantify Precision, Recall [], Inception Score (IS) [] and the Fr$\acute{\text{e}}$chet Inception Distance (FID) []. Precision and Recall are measured by 10k real images and 10k generated images following []. Precision is the ratio of generated images that belong to the manifold of real images. Recall is the ratio of real images that belong to the manifold of generated images. The manifold is constructed in a pre-trained feature space using the nearest neighborhoods. Precision quantifies sample fidelity, and Recall quantifies sample diversity. Both IS and FID are influenced by fidelity and diversity. IS is calculated using an Inception Network [] pre-trained on ImageNet [], and it computes the logits for generated samples. If an instance is predicted with high confidence for a specific class, and predictions are made for multiple classes across all samples, then the IS will be high. On the other hand, for samples generated from FFHQ or CelebA, predictions cannot be made for multiple classes, which does not allow for diversity to be reflected. Therefore, a good Inception Score (IS) can only result from high-confidence predictions based solely on sample fidelity. We measure IS for 10k generated samples. FID approximates the generated and real samples as Gaussians in the feature space of an Inception Network and measures the Wasserstein distance between them. Since it measures the distance between distributions, emphasizes the importance of the sample diversity as well as sample fidelity. For Table 4 we measure FID between 50k random samples from FFHQ dataset and 50k randomly generated samples. For ‘AE’, we measure the FID between 50k random samples from FFHQ dataset and generate samples that reconstruct the other 50k random samples from FFHQ. In Table 3, we measure the FID between 10k random samples from the CelebA and 10k randomly generated samples. we utilize [https://github.com/openai/guided-diffusion](https://github.com/openai/guided-diffusion) to measure Precision, Recall, and IS. We utilize [https://github.com/GaParmar/clean-fid](https://github.com/GaParmar/clean-fid) to measure FID. In Table 4, we loaded checkpoints for all baselines (except the generative prior of PDAE, we train it to fill performance) and conducted evaluations in the same NFEs. Table 11 shows the performance under various NFEs. For CelebA training, we use a dim$({\mathbf{z}})=256$ following [], while FFHQ training employs a dim$({\mathbf{z}})=512$ following  [, ].

### C.3 Algorithm

This section presents the training and utilization algorithms of DBAE. Algorithm 1 outlines the procedure for minimizing the autoencoding objective, $\mathcal{L}_{\text{AE}}$. Algorithm 2 explains the method for reconstruction using the trained DBAE. Algorithm 3 describes the steps for training the generative prior, $p_{\boldsymbol{\omega}}$. Algorithm 4 explains the procedure for unconditional generation using the trained DBAE and generative prior.

Figure: Algorithm 1 DBAE Training Algorithm for Autoencoding

Figure: Algorithm 2 Reconstruction

Figure: Algorithm 3 Latent DPM Training Algorithm

Figure: Algorithm 4 Unconditional Generation Algorithm

### C.4 Computational Cost

**Table 8: Computational cost comparison for FFHQ128. Training time is measured in milliseconds per image per NVIDIA A100 (ms/img/A100), and testing time is reported in milliseconds per one sampling step per NVIDIA A100 (ms/one sampling step/A100).**
|  | Parameter Size | Training | Testing |
| --- | --- | --- | --- |
| DDIM [] | 99M | 9.687 | 0.997 |
| DiffAE [] | 129M | 12.088 | 1.059 |
| PDAE [] | 280M | 12.163 | 1.375 |
| DBAE | 161M | 13.190 | 1.024 |

This section presents a computational cost comparison among diffusion-based representation learning baselines. Table 8 compares DDIM [], DiffAE [], PDAE [], and DBAE in terms of parameter size, training time, and testing time. DDIM requires only a score network (99M), resulting in minimal parameter size. DiffAE involves a ${\mathbf{z}}$-conditional score network (105M) and an encoder (24M), leading to an increase in parameter size. PDAE incorporates both a heavy decoder and an encoder, further increasing the parameter size. Conversely, although DBAE also includes a decoder, it is less complex (32M), resulting in a smaller relative increase in parameter size compared to PDAE. From a training time perspective, DiffAE, PDAE, and DBAE all require longer durations compared to DDIM due to their increased model sizes. DBAE’s training time is 9% longer than that of DiffAE because of the decoder module. However, the decoder does not repeatedly affect the sampling time, making it similar to DiffAE’s. In contrast, PDAE, which utilizes a decoder at every sampling step, has a longer sampling time.

## Appendix D Additional Experiments

### D.1 Downstream Inference

Figure 8 shows the attribute-wise Average Precision (AP) gap between PDAE [] and DBAE. As discussed in Section 4.1, PDAE suffers from an information split problem that ${\mathbf{x}}_{T}$ contains facial or hair details. The resulting attribute-wise gain aligns with that analysis with Figure 3. Figure 9(d) shows the absolute attribute-wise AP of DBAE performance across the training setting varies on the encoder (deterministic/stochastic) and training datasets (CelebA training set / FFHQ). The attribute-wise performance is similar across the training configurations.

Figure: Figure 8: Attribute-wise AP gap between PDAE and DBAE-d trained on CelebA. DBAE-d performs better for all 40 attributes.
Refer to caption: extracted/5622713/Figs/Figure15_11.png

Figure: (a) DBAE-d trained on CelebA
Refer to caption: extracted/5622713/Figs/AP_Plot/CELEBA_DBAE_det_AP_plot.png

### D.2 Reconstruction

The sampling step is important for practical applications [, ]. We compare the reconstruction results across various sampling steps among the baselines. Tables 9 and 10 shows the results. The proposed model performs the best results among all NFEs in $\left(10,20,50,100\right)$. We borrow the performance of DDIM, DiffAE from []. We manually measure for PDAE [] using an official checkpoint in [https://github.com/ckczzj/PDAE](https://github.com/ckczzj/PDAE). Figure 10 shows the reconstruction statistics for a single image with inferred ${\mathbf{z}}$. Due to the information split on ${\mathbf{x}}_{T}$, DiffAE shows substantial variations even utilizing ODE sampling. When DBAE also performs stochastic sampling, information is split across the sampling path, but it has less variation compared to DiffAE (9.99 vs 6.52), and DBAE induce information can be stored solely at ${\mathbf{x}}_{T}$ through the ODE path.

Figure: Figure 10: Reconstruction statistics with inferred ${\mathbf{z}}$. We quantify the mean and standard deviation of the reconstruction in the pixel space. The number in parentheses represents the dimension-wise averaged standard deviation in the pixel space.
Refer to caption: extracted/5622713/Figs/Figure51.jpg

Table: Table 9: Autoencoding reconstruction quality comparison. All the methods are trained on the FFHQ dataset and evaluated on the 30K CelebA-HQ dataset. Among tractable and compact 512-dimensional latent variable models, the one yielding the best performance was highlighted in bold, followed by an underline for the next best performer. All the metric is SSIM.

Table: Table 10: Autoencoding reconstruction quality comparison. All the methods are trained on the FFHQ dataset and evaluated on the 30K CelebA-HQ dataset. Among tractable and compact 512-dimensional latent variable models, the one yielding the best performance was highlighted in bold, followed by an underline for the next best performer. All the metric is MSE.

### D.3 Unconditional Generation

The sampling step is also important for unconditional generation [, ]. We reduce the NFE=1000 in Table 4 to NFE=500 and NFE=250 in Table 11. As the number of function evaluations (NFE) decreased, DDPM [] showed a significant drop in performance, while DBAE and the other baselines maintained a similar performance trend.

Although DBAE improves sample fidelity which is crucial for practical uses [, , , , ], sample diversity remains an important virtue depending on the specific application scenarios [, , , ]. In the area of generative models, there is a trade-off between fidelity and diversity [, , , ]. Therefore, providing a balance between these two virtues is important. We offer an option based on DBAE. The $h$-transformed forward SDE we designed in eq. 9 is governed by the determination of the endpoint distribution. If we set endpoint distribution as eq. 101, we can achieve smooth transitions between DiffAE and DBAE in terms of ${\mathbf{x}}_{T}$ distribution. Modeling $q_{\boldsymbol{\phi},\boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})$ as a Gaussian distribution (with learnable mean and covariance) with a certain variance or higher can also be considered as an indirect approach.

$$ $\displaystyle{\mathbf{x}}_{T}~{}\sim\lambda\times q_{\boldsymbol{\phi}, \boldsymbol{\psi}}({\mathbf{x}}_{T}|{\mathbf{x}}_{0})+(1-\lambda)\times \mathcal{N}(\mathbf{0},\mathbf{I})$ (101) $$

Table: Table 11: Unconditional generation with reduced NFE $\in\left\{250,500\right\}$ on FFHQ. ‘+AE’ indicates the use of the inferred distribution $q_{\boldsymbol{\phi}}({\mathbf{z}})$ instead of $p_{\boldsymbol{\omega}}({\mathbf{z}})$

### D.4 Additional Samples

Interpolation
Figures 11 and 12 shows the interpolation results of DBAE trained on FFHQ, Horse, and Bedroom. The two paired rows indicate the endpoints ${\mathbf{x}}_{T}$ and generated image ${\mathbf{x}}_{0}$ each. Figure 13 compares the interpolation results with PDAE [] and DiffAE [] under tractable inference condition. PDAE and DiffAE result in unnatural interpolations without inferring ${\mathbf{x}}_{T}$, compared to DBAE.

Attribute Manipulation
Figure 15 shows additional manipulation results using a linear classifier, including multiple attributes editing on a single image. Figure 14 provides the variations in the manipulation method within DBAE. The top row utilizes the manipulated ${\mathbf{x}}_{T}$ both for the starting point of the generative process and score network condition input. The bottom row utilizes the manipulated ${\mathbf{x}}_{T}$ only for the score network condition input, while the starting point remains the original image’s ${\mathbf{x}}_{T}$. Using manipulated ${\mathbf{x}}_{T}$ both for starting and conditioning results in more dramatic editing, and we expect to be able to adjust this according to the user’s desires.

Generation Trajectory
Figure 16 shows the sampling trajectory of DBAE from ${\mathbf{x}}_{T}$ to ${\mathbf{x}}_{0}$ with stochastic sampling for FFHQ, Horse, and Bedroom.

Unconditional Generation
Figures 17 and 18 show the randomly generated uncurated samples from DBAE for FFHQ and CelebA.

Figure: Figure 11: FFHQ interpolations results with corresponding endpoints ${\mathbf{x}}_{T}$. The leftmost and rightmost images are real images.
Refer to caption: extracted/5622713/Figs/Figure39_1.jpg

Figure: Figure 12: Horse and Bedroom interpolations results with corresponding endpoints ${\mathbf{x}}_{T}$. The leftmost and rightmost images are real images.
Refer to caption: extracted/5622713/Figs/Figure40_1.jpg

Figure: Figure 13: FFHQ interpolation comparison: PDAE [] (top), DiffAE [] (middle) and DBAE (bottom).
Refer to caption: extracted/5622713/Figs/Figure49.jpg

Figure: Figure 14: Attribute manipulation on FFHQ using a linear classifier and corresponding endpoints ${\mathbf{x}}_{T}$. The top results utilize the manipulated ${\mathbf{x}}_{T}$ both as the starting point of the sampling trajectory and as a condition input to the score network. The bottom results use the manipulated ${\mathbf{x}}_{T}$ solely as the condition input and maintain the original ${\mathbf{x}}_{T}$ as the starting point of the sampling trajectory. All the middle images are the original images.
Refer to caption: extracted/5622713/Figs/Figure50.jpg

Figure: (a) Smooth traversals in the direction of attribute manipulation. All the middle images are the original images.
Refer to caption: extracted/5622713/Figs/Figure42_3.jpg

Figure: (a) Sampling trajectory of DBAE trained on FFHQ.
Refer to caption: extracted/5622713/Figs/Figure36_2.jpg

Figure: (a) Generated endpoints ${\mathbf{x}}_{T}$
Refer to caption: extracted/5622713/Figs/Figure32.png

Figure: (a) Generated endpoints ${\mathbf{x}}_{T}$
Refer to caption: extracted/5622713/Figs/Figure34.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: What is the significance of the information split problem in unconditional generation (see "Weaknesses")?
soundness: 4
strengths: The paper addresses an important issue in representation learning with diffusion models.
The proposed solution appears theoretically sound and empirically convincing to me, but I have to admit that I am not an expert on the relevant literature regarding other solutions to the discussed issue (if any), nor on state-of-the-art empirical performance, so I will yield to the judgment of the other reviewers in this regard.

The paper is remarkably well organized.
Despite addressing a rather complicated problem in a domain that requires a lot of prior knowledge, the paper manages to be self-contained and discusses all relevant background in a brief yet (as far as I can tell) complete way.
The addressed problem is introduced in a didactic way (first with examples, then gradually more formal), and its proposed solution is well motivated.
summary: The paper identifies The paper identifies a problem in diffusion-based representation learning models with auxiliary (lower dimensional) latent variables, which the authors call the "information split problem".
The identified problem is that such models rely on both the auxiliary latents $z$ and the diffusion endpoint $x_T$ when reconstructing (or unconditionally generating) data.
Thus, even a well-trained model (i.e., a model that can reconstruct data well) may not encode all relevant features of a given data point $x_0$ in $z$ and instead rely also on $x_T$ when reconstructing it.
This makes $z$ less useful as a representation of $x_0$ in downstream tasks (e.g., classification, clustering).

The paper proposes to resolve this issue by turning $z$ into an information bottleneck, thus ensuring that $x_T$ is conditionally independent from $x_0$ given $z$ (i.e., $x_T$ cannot contain any additional information about $x_0$ that is not already encoded in $z$).
The paper discusses and empirically evaluates training objectives of their proposed architecture for representation learning, reconstruction, and unconditional generation.
weaknesses: I am not an expert on diffusion models with auxiliary encoders, so I cannot judge the novelty of the proposed method or the choice of baselines in the empirical evaluation.

I only identified minor weaknesses (although I strongly recommend addressing at least the last point below as it should be relatively easy to fix and would make the paper much more accessible).

I did not understand the significance of the information split problem in unconditional generation (Section 4.4.2).
I appreciate that the discussion of the objective function (Section 4.4) distinguishes between different use cases of diffusion models.
And I can see that the described information split problem can be an issue for representation learning.
But if all we care about are good unconditional samples, then why does it matter whether $x_T$ may contain some information that is not contained in $z$?
We would sample $x_T$ anyway.
Yet, empirically, it seems like the proposed DBAE+AE model outperforms the baselines even for unconditional generation.

Further, as a minor point, Eq. 9 seems suspicious to me.
While I understand how a bound of this form would exist in general, in the specific instance both $p_\theta^\text{ODE}(x_0 | z, x_T)$ and $p_\theta^\text{ODE}(x_T |z_0, z)$ are delta-distributions unless I am mistaken, so the log-density of the former does not exist, and I would expect the KL-divergence term that involves the latter to be infinite.
I presume that the "two infinities cancel" in some sense, but maybe the equation could be rewritten in a way that doesn't rely on "cancellation of infinities".

Apart from this, my only criticism concerns the large number of grammatical errors, which made the paper quite a bit more difficult to read for me than necessary.
This is a shame because it diminishes the otherwise excellent presentation.
I know that being overly pedantic on grammar in scientific papers can impose cultural bias, but I am fairly certain that most of the errors could be caught by automatic tools.
And while many errors are somewhat trivial (singular/plural or missing/superfluous/wrong articles), some sentences were right-out unintelligible to me and I had to pause reading and infer the intended meaning from context.
Some examples include (there are more): lines 132-133, lines 160-162 (two separate sentences that were probably meant as one, but this isn't immediately apparent due to remaining grammar issues even when joining them), and lines 287-288 (missing verb; I actually can't infer the intended meaning here).
I strongly recommend running the entire paper through a grammar checker, or asking an LLM to detect possible grammar issues in each paragraph.
It would make this otherwise great work more accessible to a broad (international) audience.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Could you provide more intuition behind Eq.10?

2. The paper proposes the AutoEncoder structure for the Diffusion-based representation learning methods. Would it also be interesting to appreciate the work from a (Variational)AE perspective? Will this work provide insights to both sides as a bridge in the middle?

3. Lastly, why and when the split of information will be problematic? For example, in the generation step, the information will be combined from both the latent variable and the endpoints. Thus, is it really necessary to have the information stored or learned in a single representation?
soundness: 3
strengths: Quality and clarity: the paper is well-written and the idea is motivated and easy to follow. 

Originality: the paper aims to solve the so-called information-split problem in the field. The effectiveness of the proposed method is well-supported by theorems and comprehensive experiments. 

The theorems indicate the proposed loss can indeed increase the mutual information between the input images and the latent variable. Moreover, the generated data distribution is guaranteed to be close to the input data distribution by the proposed loss.

The experiments show substantial improvement in the proposed DBAE method in terms of downstream tasks, reconstruction, disentanglement, and generation. Furthermore, the inference speed is also superior obviously.
summary: The paper proposes Diffusion Bridge AutoEncoders to create latent variable-dependent endpoint embeddings, resolving the information split problem due to the auxiliary encoder and fixed and inflexible dimension problems in diffusion-based models. 

With theoretical guarantees, the method outperforms the SOTA methods on various datasets and settings.
weaknesses: 1. It is unclear how the loss makes the endpoint dependent on the latent variable. The theorems only show the relationship between the input data and the latent variable instead, which is not aligned with the claim.

2. The intuition is missing in Section 4.1. For example, why the new forward SDE is defined as in Eq. 10? Providing more intuition would help readers to appreciate the forward process.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: See above
soundness: 3
strengths: 1. Combining VAE with diffusion is of practical interest. The motivation and solution seem reasonable to me. The authors also provide some theoretical justifications of the model and objective functions.

2. There are thorough experiments to demonstrate the usefulness and performance of the proposed model.
summary: This paper proposes a new Diffusion VAE model in order to combine the latent representation power of VAE and the generation power of Diffusion models. In previous Diffusion VAEs, the latent variable $z$ and $x_T$ are encoded separately from the data $x_0$. The latent variable is concatenated during the denoising process. However, the starting point $x_T$ is not dependent on $z$. The authors call this problem an information split, which can cause unfaithful reconstruction from $z$ and a large gap of mutual information. Motivated by this, the authors introduce another decoder from $z$ to $x_T$. Therefore, $x_T$ can obtain information from the original data. This dependence reduces the mutual information gap, and enhances the representation learning performance. The overall model appears to be a concatenation of VAE and the diffusion process. The authors conduct comprehensive experiments to show advantages in latent variable learning and generation.
weaknesses: 1. In the diffusion bridge Eq. (5), the input is transformed into a fixed target. However, in the proposed model, the target $x_T$ is random due to the randomness of the VAE. Does this affect the theory?

2. In Section 4.4, the authors present objective functions for reconstruction and generation tasks. In experiments, did the authors optimize different objectives and use different models for different tasks? Can we obtain both reconstruction and generation abilities using the same model?

3. I am wondering whether the training is difficult compared to without the AE part. Is there a trade-off between the expressive power of the AE and diffusion model?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Line 46: The phrase “and the endpoint $x_T$” may be unclear. Would it be more precise to revise this to “to the endpoint $x_T$” to clarify the relationship between $z$ and $x_T$?

Line 49: The term “inference of $x_T$” might be ambiguous for some readers. Could this be rephrased or annotated as “estimating $x_T$ conditioned on $x_0$” or “encoding $x_0$ into $x_T$”?

Lines 75–80: The third paragraph of the Introduction could benefit from additional clarity. The challenges described here seem to presuppose the goal of compressing $x_0$ into latent variables ($z$ and $x_T$) for reconstruction, but this goal is not explicitly stated, which could cause confusion. For example, readers focused on new sample generation might not understand why the information split issue affects reconstruction. Would the authors consider clarifying this context and explaining why reconstruction is a primary focus?

Line 99: Consider capitalizing “eq.” to “Eq.” when referencing specific equations. Also, while referencing equations before presenting them can add context, it is uncommon and might disrupt the logical flow. Would it be more effective to present equations first, then reference them?

Lines 160–161: It might be helpful to provide the full names of “VE” and “VP” (at least in a footnote), rather than only using acronyms, to support readers who may be unfamiliar with these terms.

Lines 196–197: For clarity and parallelism, consider specifying that $D_{KL}(q_{\theta}^{\textrm{ODE}}, p_{\textrm{prior}})$ represents the KL divergence between $q_{\theta}^{\textrm{ODE}}$ and $p_{\textrm{prior}}$. The phrase “$D_{KL}$ explains the information split problem” is also somewhat ambiguous. Could the authors clarify whether the KL divergence term’s presence illustrates or quantifies the problem? Additional context could improve reader understanding.

Lines 203–204: Since the references here were cited just two paragraphs above, the authors may want to omit them to save space.

Line 213: The title of Section 4.1, “latent variable inference,” may be unclear—does it refer to $z$, $x_T$, or both? Clarifying whether the authors mean both variables, and if so, how they relate, would improve clarity.

Lines 214–215: The nature of the encoder and decoder is not clearly defined. For instance, is $\text{Enc}{\phi}$ deterministic or stochastic? If it is deterministic, how does it define the conditional probability $q{\phi}$? Clarifying this would assist readers in understanding the model structure.

Section 4: A concise pseudocode format summarizing the main DBAE algorithm, including training and inference procedures, would be helpful. Consider adding boxed pseudocode for core processes, such as (1) inferring latent variables for reconstruction or sample generation, and (2) training the encoder and decoder. This structured overview could improve comprehension before diving into the details.

Section 4.1: It would be helpful if the authors could confirm the proposed procedure: (1) infer $z$ from $x_0$, (2) infer $x_T$ from $z$, and (3) use Doob’s $h$-transform via Eq. (10) to bridge $x_0$ and $x_T$ and determine the distribution of intermediate $x_t$ values. Additionally, it is unclear if the encoder $q_{\phi}$ and decoder $q_{\psi}$ require training. Could the authors clarify this?

Line 285–286: The statement “For downstream inference, attribute manipulation, and interpolation, the model requires reconstruction capability” could be introduced earlier in the Introduction to help illustrate the utility of generative modeling and contextualize the information split problem. Would the authors consider adding this as a motivation for addressing the information split issue?

Line 317 (Theorem 1): The term “linear SDE” appears without prior explanation. Would the authors consider providing a brief definition?

Line 331: Consider removing “the” before “Section 4.4.1” for grammatical clarity.

Line 342: Replacing “optimize” with “minimize” could provide greater clarity regarding the objective.

Lines 354–355: Essential experimental details should ideally appear in the main text, rather than directing readers to the Appendix. This would allow readers to understand the experimental setup, including encoder and model architecture choices, without needing to consult additional sections. Would the authors consider including these details directly?

Section 5: Could the authors provide a concise, mathematically rigorous description of each of the six experimental tasks to clearly communicate the goals, challenges, and benefits of DBAE? For example, describing the interpolation task in mathematically precise terms, as in Lines 474–476, could serve as a model for presenting all tasks consistently.
soundness: 2
strengths: This paper addresses a key challenge in latent-variable generative models, particularly in diffusion probabilistic models (DPMs) that are augmented with an auxiliary encoded latent variable $z$, which the authors refer to as the "information split problem." The authors clearly illustrate how the information split between $z$ and $x_T$ diminishes the quality of latent representations for high-fidelity tasks. To overcome this issue, they propose a principled framework that integrates autoencoders with diffusion models by (1) encoding $x_0$ to define the latent endpoint $x_T$ of the diffusion process, and (2) connecting these two points utilizing Doob’s $h$-transform. This approach combines the strengths of autoencoders—such as dimensionality reduction and efficient encoding/decoding—with the rich generative capabilities of diffusion models. Experimental results demonstrate that DBAE outperforms existing models across various benchmarks, underscoring its effectiveness in tasks like generation, reconstruction, and attribute manipulation. The paper also discusses potential extensions of DBAE to various downstream applications, highlighting its adaptability and potential for broader impact.
summary: This paper introduces the Diffusion Bridge AutoEncoders (DBAE) framework to address the "information split problem" in diffusion-based autoencoding models for unsupervised representation learning. The authors argue that in conventional latent-variable-augmented diffusion models, the information in data $x_0$ is split between two latent variables--$z$, obtained from an auxiliary encoder, and $x_T$, the endpoint of the diffusion process. This split hinders effective encoding, especially for downstream tasks such as reconstruction. DBAE proposes a solution to remedy this by creating a dependency of $x_T$ on $z$, allowing $z$ to serve as a comprehensive information bottleneck and improving representation quality. Specifically, the authors propose the process to encode $x_0$ into $z$, decode $z$ into $x_T$, and then bridge $x_0$ and $x_T$ to define a diffusion process. The authors provide empirical evidence supporting DBAE’s improvements in representation quality, reconstruction, and downstream tasks, suggesting it as a promising model for generative and reconstruction purposes.
weaknesses: Despite its strengths, I think the paper’s impact is significantly diminished by its presentation. Below, I outline major concerns with suggestions for improvement.

**Clarifying Motivation and Significance in the Introduction:**
The Introduction lacks a clear articulation of why the “information split problem” is critical to address. This issue is specific to a subset of diffusion probabilistic models (DPMs) that are augmented by an auxiliary encoded latent variable $z$. For readers to appreciate the significance of this issue, it would help to emphasize the relevance of these models, which combine strengths from both VAEs and DPMs. While the first two paragraphs provide an overview of latent-variable generative modeling, the current description introduces these models merely as proposed approaches rather than as leading techniques in the field. Highlighting their role as state-of-the-art or best-performing models would better contextualize the importance of tackling their limitations. Additionally, the Introduction could clarify why the information split problem is particularly relevant for tasks requiring accurate reconstruction of the original data, $x_0$. The authors might adapt statements like those in Lines 285–286—“For downstream inference, attribute manipulation, and interpolation, the model requires reconstruction capability”—to illustrate the utility of generative modeling and explain the advantages of integrating VAE and DPM functionalities. Establishing these goals clearly would help readers understand how the information split issue primarily affects reconstruction and how DBAE addresses this challenge.

**Improving Precision in Terminology and Expression:**
Several terms and expressions in the paper would benefit from clearer definitions and consistent use. For instance, in Line 49, “inference of $x_T$” may be ambiguous; specifying this as “estimating the latent variable $x_T$ conditioned on $x_0$” would clarify the intent. Similarly, the phrase “latent variable inference,” which appears in Section 4.1, could be made more precise by using “estimating $x_T$ conditioned on $z$” or “encoding $x_0$ into $x_T$,” depending on the context. Certain expressions, like "$D_{KL}(q_{\theta}^{\textrm{ODE}}||p_{\textrm{prior}})$ explains the information split problem” (Lines 196–197), are unclear; it would help to specify whether the KL divergence term’s presence illustrates or quantifies this problem. Ensuring precise and consistent terminology throughout would enhance readability and help avoid potential confusion (see “Questions” for further specific comments).

**Providing Precise Description of Methodology/Tasks in Sections 4 & 5:**
- *Methodology in Section 4:* The presentation of the proposed methodology in Section 4 is complex and challenging to follow. A concise pseudocode or boxed algorithm format summarizing DBAE’s key steps, including training and inference procedures, would streamline comprehension. Providing a pseudocode outline for DBAE’s main processes—inferring $z$ from $x_0$, subsequently inferring $x_T$ from $z$, and bridging $x_0$ and $x_T$ using Doob’s $h$-transform—would help organize the workflow and provide readers with a structured overview before diving into the details.
-  *Experimental Tasks in Section 5:* Section 5 could also benefit from a more systematic presentation of the six experimental tasks under investigation. Providing a precise, mathematical description of each task would make it easier for readers to interpret DBAE’s advantages across different contexts. For example, the description of Interpolation problem in Lines 474 - 476 can be a good starting point -- presenting all tasks in a similar manner would help readers understand the desiderata and challenges in each tasks as well as DBAE’s benefits across varied applications in a more structured way.

Meta-review: The paper introduces Diffusion Bridge AutoEncoders (DBAE) to address the "information split problem" in diffusion-based representation learning models. The proposed method creates a dependency between the latent variable and the diffusion endpoint through a feed-forward architecture, ensuring that the latent variable holds the full information of the data. This approach improves downstream inference, reconstruction, and disentanglement, and generates high-fidelity samples. All reviewer agreed that this paper is well-written, and the idea is motivated and easy to follow. The paper addresses an important issue in representation learning with diffusion models and provides a theoretically sound and empirically convincing solution.

Decision: Accept (Spotlight)