Diffusion Bridge AutoEncoders for Unsupervised Representation Learning
Authors: Yeongmin Kim
Kwanghyeon Lee
Minsang Park
Byeonghu Na
Il-chul Moon
Venue: ICLR 2025 Spotlight (2025)
Keywords: Diffusion Model
Represenation Learning
Autoencoders
Primary area: generative models
TL;DR: This paper introduces Diffusion Bridge Autoencoders (DBAE) to design encoder dependent endpoint inference.

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 data and to adjust the dimensionality of a latent variable $\mathbf{z}$. Meanwhile, this auxiliary structure invokes an *information split problem*; the information of each data instance $\mathbf{x}_0$ is divided into diffusion endpoint $\mathbf{x}_T$ and encoded $\mathbf{z}$ because there exist two inference paths starting from the data. The latent variable modeled by diffusion endpoint $\mathbf{x}_T$ has some disadvantages. The diffusion endpoint $\mathbf{x}_T$ is computationally expensive to obtain and inflexible in dimensionality. To address this problem, we introduce Diffusion Bridge AuteEncoders (DBAE), which enables $\mathbf{z}$-dependent endpoint $\mathbf{x}_T$ inference through a feed-forward architecture. This structure creates an information bottleneck at $\mathbf{z}$, so $\mathbf{x}_T$ becomes dependent on $\mathbf{z}$ in its generation. This results in $\mathbf{z}$ holding the full information of data. 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. Our code is
available at https://github.com/aailab-kaist/DBAE.

Published as a conference paper at ICLR 2025
DIFFUSION BRIDGE AUTOENCODERS FOR UNSUPERVISED REPRESENTATION LEARNING
Yeongmin Kim', Kwanghyeon Lee', Minsang Park', Byeonghu Na', Il-Chul Moon Korea Advanced Institute of Science and Technology (KAIST), "summary.ai {alsdudr 1a10, rhkdgus042 414, pagemu, byeonghu. na, (emoooo@aaaaa ac. kr
ABSTRACT
Diffusion-ban 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 extract a corresponding representa- tion from data and adjust the dimensionality of latent variable Meanwhile, this auxiliary structure invokes an information split problem; the information of each data instance X is divided into diffusion endpoint XT and encoded z because there exist two inference paths starting from the data. The latent variable modeled by the diffusion endpoint XT has several disadvantages. The diffusion endpoint Xr is computationally expensive to obtain and inflexible in terms of dimensionality. To address this problem, we introduce Diffusion Bridge AutoEncoders (DBAE), which enables z-dependent endpoint Xy inference through feed-forward architecture. This structure creates an information bottleneck z, ensuring that XT depends z during its generation. This results in holding the full information of the data. We propose an objective function for DBAE to enable both reconstruction and generative modeling, with theoretical justification. Empirical evidence demon- strates the effectiveness the intended design in DBAE, which notably enhances downstream inference quality, reconstruction, and disentanglement. Additionally, DBAE generates high-fidelity samples in an unconditional generation. Our code is available at https: thub. com/ aa i 1ab kai St/DBAF
INTRODUCTION
Unsupervised representation learning is fundamental topic within the latent variable generative
models (Hinton et al., 2006; Kingma Welling 2014: Higgins et al. 2017 Chen et al. 2016
Jefil Alemi et al. 2018). Effective representation supports better downstream inference as well as
realistic data synthesis. Variational autoencoders (VAEs) Kingma & Welling 2014) are frequently used because they inherently include latent representations with flexible dimensionality. Generative adversarial networks (GANs) (Goodfellow et al.) 2014) with inversion (Abdal et al. 2019: (2020) another method to find latent representations. Additionally, diffusion probabilistic models (DPMs) (Ho et al., 2020 Song et al., (221c)) have achieved state-of-the-an performance in terms of generation quality (Dhariwal & Nichol 2021). naturally prompting efforts to explore unsupervised
representation learning within the DPM framework (Preechakul et al. 2022 Zhang et al. 2022 Yue
al. 2024). which have recently dominated generative representation learning studies.
DPMs are type of latent variable generative model, but inference on latent variables is not straightfor- ward. DPMs progressively map from data xy to latent endpoint Xy via a predefined noise injection schedule, which does not facilitate learnable encoding. DDIM (Song et al. (021 introduces an ODE-baaed deterministic encoding from the data Xg to the endpoint Xy- However, this encoding is determined by the choice of the forward process (Song et al. (202cc). Since the forward process with fixed noise injection is difficult to interpret as having semantic meaning, the ODE-based encoding effective semantic representation. Moreover, the encoding X into Xy is expensive because requires solving the ODE, and its inflexible dimensionality poses disadvantages for downstream applications (Sinha et al. (202).
remains challenging consider
To tackle this issue, recent DPM- based representation learning studies (Preechakul et al. 2022
Zhang et al. 2022 Kim al. 2022b Wang et al. 2023 Yang al. 2023 Yue et al. 2024 Wu &

---

Published as a conference paper at ICLR 2025
.....:: Inference
Original
Inferred z, random Xy
Generation
X
(a) Bayesian network of DiffAE
(b)
Reconstruction varying X7 ~ (O, I) in DifiAE
Original
DifiAE (265)
DBAE (0.31)
Original
DifAAE
DBAE
X
X
XT
(c) Bayesian network of DBAE
(d) Inferred X7 from DifTAE and DBAE with inference time (ms).
Figure I: Comparison between DifTAE (Preechakul et al. 2022) and DBAE. (a) depicts the simplified Bayesian network of DiffAE, illustrating two inference paths for the distinct latent variables Xy and z. (b) shows the reconstruction using the inferred z DiffAE on CelebA, where the reconstruction results perceptually vary depending the selection of Xr. (c) shows the simplified Bayesian network of DBAE with z-dependent Xy inference. (d) shows the inferred Xy from DiffAE and DBAE.
Zheng 2024) suggest an auxiliary latent variable z with an encoder used VAEs, to combine the generation performance of diffusion models and the representation learning capabilities of VAEs. The encoder- gennrereeeeertedd latent variable z is obtained without solving the ODE, and the encoder also facilitates the learning of semantic representations with dimensionality reduction. The reconstruction capability from the extracted latent representation is the primary focus of these studies, facilitating downstream inference, attribute manipulation, and interpolation.
This paper points out the remaining problem auxiliary encoder models, which we refer to as information split problem, hindering reconstruction capability. The information is not solely retained the latent variable z; rather, a portion also distributed into the latent variable xy as evidenced by Figure 1b If the auxiliary encoder models only infer and reconstruct using random Xy, the facial details of the original image are not properly reconstructed, indicating that the missing information is contained within XT. Furthermore, the inference of XT computationally expensive and inflexible in dimensionality. To address this issue, we introduce Diffusion Bridge AutoEncoders (DBAE), which incorporate z-dependent endpoint X7 inference using feed-forward architecture.
The proposed model DBAE systematically resolves the information split problem. Unlike the two split inference paths the previous approach in Figure la DBAE encourages to become an information bottleneck during inference (dotted line in Figure Ic). making z more informative. DBAE establishes this bottleneck structure by defining learnable forward process that starts from the data X and ends at the encoded endpoint Xy by utilizing Doob's 's h-transform. Moreover, DBAE does not require solving an ODE to infer endpoint X7, thereby making endpoint inference more efficient, as shown in Figure ld This efficient inference of xr benefits interpolation and attribute manipulation tasks. In experiments, DBAE outperforms the previous works in downstream inference quality, reconstruction, disentanglement, and unconditional generation. DBAE also demonstrates satisfactory results in interpolation and attribute manipulation with its qualitative advantages.
PRELIMINARIES
2.1
DIFFUSION MODELS
Diffusion probabilistic models (DPMs) (Scks-Diekstein et al. 2015 Ho et al. 2020) with a continuous time formulation (Song et al. 2021c) define forward stochastic differential equation (SDE)
dx,
f(x,, l)dt g()dd,, Xg Qáata (xm))X)),
where w, denotes standard Wiener process, f : Rd [0, + Rdd is drift term, and g |0, T] R a volatility term. Eq. (I) starts from data distribution Gdata (Xn) and gradually perturbs it into noise Xr. Let the marginal distribution of Eq. {1) time be denoted (x.). There exists unique
2

---

Published as a conference paper at ICLR 2025
reverse-time SDE (Anderson 1982)
dx, = [f(x,, t) - Vx) log qi (x:)]dt + g(t)dwe. XT Pprior (xr),
(2)
where W, denotes reverse-time Wiener process, Vx, log qr (xx) is the time-dependent score function, and (Xrior(XT) stands for the prior distribution, which closely resembles a Gaussian distribution with the specific form of f and g (Song et al. c Ho et al. 2020). Eq. (2) traces back from noise XT
to data Xy. The reverse-ting ordinary differential equation (ODE)
dx, [f(x, t) g V log (x, (x, XT Pprior (X.)
(3)
produces a marginal distribution identical Eq. 2 for all I, offering an alternative generative process while confining the stochasticity of the trajectory solely to xT. To construct both reverse SDE and ODE, the diffusion model estimates a time-dependent score function Xx, log (x(x)) R3 Sq (X,, t)
using a neural network and the score-matching objective (Vincent 2011 Song & Ermon 2019).
2.2 LATENT REPRESENTATION LEARNING WITH DIFFUSION MODELS
From the perspective of representation learning, the ODE in Eq. 33 (a.k.a DDIM (Song et al. 2021a in discrete time diffusion formulation) provides deterministic encoding from the data X, to the latent Xy. However, the latent representation Xy has some disadvantages, First, it is hard to learn its semantic meaning. This encoding determined by the forward process (f, given data distribution and assuming perfect optimization (Song et al. (201c)). The forward process (f, g) is set to fixed noise injection process, but the noise is hard to consider as a semantically meaningful encoding. Second, the dimension cannot be reduced. According to the definition of the diffusion process in Eq. ( the dimension X must the same as the data dimension. This hinders learning a compact representation, making it hard to facilitate downstream inference attribute manipulation (Sinha cct al. 2021). Finally, XT is computationally expensive to obtain. To infer xy from the data point Xp. it is necessary numerically solve the ODE in Eq. This results in high time complexity for inferring Xy, which makes it inefficient to exploit latent representations.
To resolve the problem in the latent endpoint Xy,, et al. 2022). proposes an auxiliary latent space utilizing a learnable encoder Enc, Rd - RR' which maps from data Xo an auxiliary latent variable Unlike DDIM, these approaches tractably obtain z from Xo without solving the ODE, and the encoder can directly learn the latent space in reduced dimensionality. Consequently, the generative ODE
previous literature, e.g., DiffAE (Preechakul
dx, [f(x,, t) -
g (f)sa 2, t)]dt,
becomes associated with the z-conditional score function (x z, t), which approximates Xx, log % (z) The generation starts from two distinct latent variables z and Xy, and defines the conditional probability p0 PDD xr). The ODE also provides an encoding from X. and z to xT. which defines the conditional probability d0 (XDE(( iz, xo). However, the auxiliary encoder framework encounters an information split problem which this paper raises in Section 3 This paper proposes a method to mitigate this problem.
2.3 DIFFUSION PROCESS WITH FIXED ENDPOINTS
To control the information regarding the diffusion endpoint Xy, it is imperative to specify forward SDE that terminates at the desired endpoint. We employ Doob's h-transform (Doob & Doob 1984].
which facilitates the conversion of the original forward SDE in Eq. (1) into
dx, [f(x,, () + (1)),xx, y, T))dt + g())dw;; Xo ~ (x)))
XT = y,
(5)
where h(x,, y. 7) := Vx. log q. (xr (r:))x))>>> is the score function of the perturbation kernel from the original forward SDE, and y denotes the desired endpoint. Let ge (x:) denote the marginal distribution of Eq. (5) at f. It is noteworthy that when both Xo and xT are given, the conditional probability of x, becomes identical to that of the original forward SDE, i.c., qe (x Xr X²) = qi (x, x.). If the original forward SDE in Eq. (I) is set to be specific form (c.g.. variance preserving SDE (Ho et al. 2020 then qe (x, X Xn) J follows a Gaussian distribution. This means that sampling of x, qe (x,, XT, Xo) any time i is tractable with an exact density function.
3

---

Published as a conference paper at ICLR 2025
Corresponding to the h- h-transformed forward SDE of Eq. (5). there also exist unique reverse-time SDE and ODE (Anderson, 1982 Zhou et al. 2024
dx; [f(x;,t) t) 90- (1)V (x, + (alprr)))(((()(////e)"(aa))))))/
(1))(( 7) + g(t) XT =
(6)
dx, = [(xx,))) 3 2 (1)V log((xx|xx) + (y, XT y,
where (x((xx)) is the conditional probability defined by Eq. (5) To construct the reverse SDE and ODE, is necessary to estimate Vx, log qr (x (xT) Rs So Sa L, XT) through a neural network with a score matching objective (Zhou al. 2024
2 1
E (x)
Xx, lo (x)
(8)
3 MOTIVATION: INFORMATION SPLIT PROBLEM
This paper raises problem in diffusion-based representation learning with auxiliary en- coders (Preechakul et al. 2022 Zhang et al. 2022 Wang al., 2023] Yang et al.) 2023 Yue ct al. 2024 Wu & Zheng 2024 introduced in Section 2.2 The latent variable from the encoder has benefits compared to the latent endpoint xy, but the auxiliary encoder framework encounters an information split problem: the information of the data is split into two latent variables and xy. The generative process in Eq. (4) initiates with two latent variables z and Xy- If the framework only relies on the tractably inferred latent variable the reconstruction outcomes depicted Figure [bb appear to fluctuate depending on the choice of Xy. This implies that Xy encompasses crucial information necessary for reconstructing Xg. To represent all the information of Xy, necessary to infer xy by solving the ODE in Eq. (4) from input to endpoint Xy, enduring its computational costs. Consequently, the persisting issue within the latent variable Xy remains unresolved in this framework.
To learn an informative latent representation. the mutual information between the data and the latent variable needs
maximized (Alemi et al., 2018). The information split problem hinders the maximization of the mutual information between the data x, and the latent variable z. The variational lower bound of the mutual information the auxiliary encoder framework
Exga (No 1. (eo) I E'g
ODE (xrlz, xn)
|Pprior (xprii((xx)) (xr)) + H S MM z):
(x(x0.x) where Mx0,0 2) = Egg (x0. (log gas (Xo])a represents the mutual information, H = H( (idaaa denotes the data entropy, and C'ElgD Xo) := Egg0E v2x,,) |- log Pprior (xx-)) is the cross-entropy. The cross-entropy term increases the discrepancy between 40 (xT zz Xn) and Pprior increases, resulting in a looser lower bound on the mutual information. Since ODE (xT z, Xx) inherently forms a Dirac delta distribution due to the nature of ODEs, the discrepancy between ODE (xx 2, Xo) and Pprior (XT) is inevitable in this framework. For more details, please refer to Appendix A.4.1
METHOD: DIFFUSION BRIDGE AUTOENCODERS
To resolve the information split problem in auxiliary encoder models, introduce Diffusion Bridge AutoEncoders (DBAE) featuring z-dependent endpoint Xy inference using a single network propaga- tion. The endpoint XT DBAE only depends on z, making z an information bottleneck. Figure illustrates the overall schematic for DBAE. Section 4.1 explains the latent variable inference with the encoder-decoder structure and learnable forward SDE utilizing Doob's h-transform. Section 4.2 de- lineates the generative process from the information bottleneck z to data Xg- Section 4.3 analyzes the
benefit of DBAE for mutual information maximization between Xo and z. Section 4.4 elaborates on the objective function for reconstruction, unconditional generation, and its theoretical justifications.
4.1 ENCODING FROM X TO Xy CONDITIONED ON
We can access i.idd. samples from qaaa (Xn). The encoder Enc Rd + R' maps data X, to the latent variable z, defining the conditional probability qs (2|xx)). To condense the high-level representation of
4

---

Published as a conference paper at ICLR 2025
Parametrized Forward Process (Φ) 2/)
dx, (f(x,.)) + (t) h,
xy, xy, dt + s(t) dw,
Encoder la (z Xo)
Z
Decoder 9p Pp x)
XO
XT
dx, ((x,,1) 29 (1),, log lo
(x)) + + h(x,, Xr. Ty) di
Reverse Process (0)
Figure 2: schematic for Diffusion Bridge AutoEncoders. The blue line shows the latent variable inference. DBAE infers the z-dependent endpoint XT to make XT tractable and to establish as an information bottleneck. The paired Xg and Xr define new forward SDE utilizing the h-traasformm The decoder and the red line show the generative process. The generation starts from the bottleneck latent variable z and decodes to the endpoint xy. The reverse process generates Xo from Xy-
the latent dimension I is set to be lower than the data dimension d. The decoder Dec R R maps from the latent variable z to the endpoint XT, defining the conditional probability qo (xrlz). The encoder and decoder can be deterministic (i.e., Dirac delta distribution) or stochastic (i.e., Gaussian distribution) depending on the experimental choice. Since the decoder generates the endpoint X7 solely based the latent variable encoder-decodee structure provides the endpoint distribution Lo (Xr Xx) = 9w (Xr (2) (z) x2 for given starting point Xg- We now discuss diffusion process {x, } with given starting point and endpoint pair.
becomes bottleneck for all the information in Xo- The
To establish the relationship between the starting point and endpoint given by the encoder-decoder, we utilize Doob's 's h-transform to define a new forward SDE
dx, = [f(x, t) + g (xx,, xr, T) + g(t) Xo Adata (x)), xT ~ T0. (xT (xr(xx); (xo). (10)
where l(xe, XT, 7) == Vx, log qe (xT(x.) ) is the score function of the perturbation kernel in the
XT,
original forward SDE in Eq. T The forward SDE in Eq. (1O) determines the distribution of x,,
where E (O, Let us denote the marginal distribution of Eq. (10) at time as (x)..
4.2 GENERATIVE PROCESS
The generative process begins with the bottleneck latent variable z, which can be inferred from the input data Xg or is randomly drawn from the prior distribution Ponor (z). The decoder Dec,, R Rd maps from the latent variable to the endpoint xy with the probability p& (xT (z) Corresponding to a new forward SDE in Eq. (10). there exists a reverse ODE
1 dx; f(x,, t) 9 (t) x. log g (x(xx)) + g (t) xy,
(11)
where the conditional probability o (x, (x)) is defined by Eq. 10 However, computing the con- ditional probability do (x. (xT) is intractable, Vx, log p (x:] xx) to approximate Xx, log (x, (xT). Our parametrized generative process be- comes
parameterize
model so sø(x,, x,x ==
dx; fxx, t) 3 (t) lo (x)x(( (t), T)Jdt.
(12)
Stochastic sampling with an SDE is also naturally possible as shown in Section 2.3 but we describe only the ODE for convenience.
'The two distributions (x(xrle (z) and qu (xz (z) are the same. However, to distinguish between inference and generation, they respectively denoted p and g-
5

---

Published as a conference paper at ICLR 2025
Algorithm 1: DBAE Training Algorithm for Reconstruction
Algorithm 2: Reconstruction
Input: data distribution gai (Xo). drift term f. volatility term
Input:
Frngg Desig, score network sg,
while converges
sample xo. discretized time steps
Sample time é from [o. T]
[
No nu Quax((o)) Xo).
Enc. (Xo) Dec N.... do
z Enc E (Xo) and XT (2)
N, r de Xo. xx)
àg" [x)]ssu(xx Xxr) Vxx log Xx. x
Update xx, using Eq.
12
Update db, 4b, by Liee using the gradient descent method
Output: Reconstructed sample Xo
Output: Eneg. Desy, score network
4.3 MUTUAL INFORMATION ANALYSIS
From the definition of inference and generation of DBAE in Sections 4.1 and 4.2 the variational lower bound on the mutual information between Xy and z is
Baglxo: (x)())=))) (e (xy (x) [log po (xx[xx)] D (x [z)] Pb (xr|2)) [2)] + H S MI(xu. z),
(13)
where pe xx)) is defined by the generative process in Section 4.2 Please see Appendix A..4.2
for a detailed derivation. Here, the term D z (4w (xT (z) Pa (xr 2 becomes zero because both conditional probabilities of Xy given z are the same in the inference and the generation. The Po remaining term (xn.x)) (E (x (xxl1) (2) po (Xy xx)] can be controlled by the optimization of Φ, tb.
and 0. The relation between an objective function and mutual information is declared in Theorem
4.4 OBJECTIVE FUNCTION
The objective function bifurcates depending on the specific tasks. The model requires recon- struction capability for downstream inference, attribute manipulation, and interpolation. To achieve reconstruction capability, the model needs 1) an encoding capability (Xg + z -> xT) and 2) a regeneration capability (xr + Xo). The encoding process should infer distinct latent vari- able for each data point Xg to ensure that the original information is preserved during recon- struction. The regeneration capability needs to estimate the reverse process by approximating so (xt; 1, xT) Vx, log Lo (x, XT). For an unconditional generation, the model must possess the ability to generate random samples from the endpoint X7, which implies that the generative endpoint distribution (XTT) = Pas(x (xxr///)ppinl (p) should closely match the aggregated inferred distribution Lo.s (xr) = au (xr z g) (go(z)) (xx))aas No (2) dxydz.
4.4.1 RECONSTRUCTION
For successful reconstruction, the model needs to fulfill two criteria: 1) encoding the latent variable Xy uniquely depending the data point X6, and 2) regenerating from Xy to Xp. The inferred latent distribution 9 w (xT (XT Xo) should provide unique information for each Xy. To achieve this, we aim to minimize the entropy H(gg.p (oos:<(xx/xa (xr(Xx) (xo)) embed xy-dependent Xy with minimum uncertainty. On the other hand, we maximize the entropy H(q (xr)) to embed different Xy for each Xg. Since the posterior entropy H(qs.v (Xo x)
))
H(1s, s (Xr (x)) H(g w (XT)) Higgdaa (xo)) naturally includes the aforementioned terms, we use this term as a regularization. Minimizing the gap between Eqs. (TT) and (12) necessary for regenerating from Xx to Xo. This requires alignment between the inferred score function Vx, log %o (x((x)) and the model score function so (x, xr). Similarly to Eq. (8). we propose the score-matching objective function LsM descrihed as
1 [ (x. Lsm :- -
Vx, log 4' (xx|XT) [2
(14)
We train DBAE with the entropy retroppyregullarized score matching objective CAE described
LAR = Lsm + H(gp. (Xol xr)).
(15)
The detailed training and testing procedures outlined in algorithms and respectively. Theo- rem demonstrates that the entropy-regularized score matching objective in LAE becomes a tractable form of objective, and it is equivalent to the reconstruction formulation. The inference distribution 4d.4 (xx, Xx)x)) is optimized to provide the best information about x for easy reconstruction. Theorem 1. For the objective function CAs, the following equality holds.
1 T
LAE 2 20 (xoxx (t)) t,, - Vx. log qi (x; Xx, 13]dt
(16)

---

Published as a conference paper at ICLR 2025
Moreover: if Eq. ①1 is linear SDE there exists ait), ((t), 7(t), A(t), such that
LAE {'- B. o (x0.x, [(1)]l2 (x,,
X xo11),r
(17)
where (xt, XT) = a(t)x; + 8(1)XT + (())s (x)) XT), and Lo. (x0, X XT) J "daao (dan((x))a)) (xo)g (z) Xo (o 2)): (x,/xxXXT XT; Xo Xz) following the graphical model in Fig. lc
The assumptions and proof of Theorem are in Appendix A.1 Moreover, Theorem shows the objective functions CAE is the upper bound of the negative mutual information between X and up to constant. Since the optimization direction of AA aligned with maximizing the mutual information, our objective function makes the mutual information higher, which can make z informative. The proof of Theorem is in Appendix A.5
Theorem 2. - -MI((x) z) S LAE H. where H = (X)g))aa(x))) a constant w.r.t. Φ. ab
4.4.2 GENERATIVE MODELING
In Section 4.4. the discussion focused on the objective function for reconstruction. The distribution of xT should be considered for generative modeling. This section addresses the discrepancy between the inferred distribution 90.vi v (xT) and the generative prior distribution py (xr). To address this, we propose the objective LyR related to the generative prior.
(18) Theorem 3 demonstrates that the autoencoding objective LA and prior objective LpR bound the Beulbacck-eibbler divergence between data distribution Qaata ( and the generative model distribution Pv.e (Xx)) / pe (xo XT Pw Pprior (z)dZ up to constant. The proof is in Appendix A..
LpR := [KLL(90 v (xT|x))]lpp
Theorem 3. Dxz (Admaa xo) (p (xn)) S LAF + LPR H, where HHAdaaa (xo)) is constant w.r.f. Φ, 2b. O.
For generative modeling, we separately minimize the terms CAt and LyR. following (Esser et al. 2021 Preechakul et al.) [022] Zhang et al. (002)). The separate training of the generative prior distribution with a powerful generative model effectively reduces the mismatch between the prior and the aggregated posterior (Sinha et al., 2021 Aneja al. 2021). Initially, we optimize LAA with respect to the parameters of encoder (Φ)) decoder (æb), and score network (0), and fix the parameters 0, Φ, tb. Subsequently, we newly parameterize the generative prior Pprior (z) := Pw (z) using shallow latent diffusion models, and optimize LpR w.r.t w. See Appendix A.3 for further details.
5 EXPERIMENT
This section empirically validates the effectiveness of the intended design of the proposed model, DBAE. We utilize the U-Net architecture for the score network (0), as shown in Fig. 7b Since our score network needs to account for the additional input X7, we concatenate x; and XT as the U-Net input. We employ half of the U-Net architecture the encoder (Φ) and use CNN-based upsampler as the decoder (vb), adopted from (Liu et al. 2021). The encoder and decoder architectures detailed in Fig. 7a To compare DBAE with previous srffuusion-basee representation learning approaches, we adopt the remaining experimental configurations (e.g.. batch size, learning rate) from DiffAE (Preechakul et al. 2022) closely possible. Detailed experimental configurations provided in Appendix C We evaluate both latent encoding capability and generation quality across various tasks. We quantitatively assess the performance of downstream inference, reconstruction, dis- entanglement, and unconditional generation. Additionally, we qualitatively demonstrate interpolation and attribute manipulation capabilities. Finally, we conduct experiments with two variations of the proposed model's s encoder: 1) a Gaussian stochastic encoder (DBAE) and 2) deterministic encoder (DBAE-d) for ablation studies. We use a deterministic structure for the decoder.
5.1 DOWNSTREAM INFERENCE
To examine the learned latent representation capability of Enc,, we perform a linear- probe attribute prediction following DiTi Yue et al. 2024). We train a linear classifier with parameters (w, 6) using
Fqq.
linear SDE when the drift function is linear with respect to Xx.

---

Published as a conference paper at ICLR 2025
Table 1: Linear-prohe attribute prediction quality comparison for models trained on CelebA and FFHQ with dim(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.
CelebA
FFHQ
Method
Gen
(P)
Pearson's (t)
MSE ())
(P)
Pearson's 's (1)
MSE (
SimCLR RR
Chen et al. 2020
X
0.597
0..74
0.603
0.608
0.481
0.638
8-TCVAE CCenn el al. 2018
0.450
0.378
0,573
0,432
0.335
0.608
IB-GAN (Joon et al. 2021
0.442
0.307
0,597
0,428
0.260
0.644
DiffAE PPeechakku] et al. 2022
0.603
0.598
0.421
0.605
0.606
0.410
PDAE Zhang CT al. 2022
0.602
0.596
0.410
0.597
0.603
0.416
DiTi (Yue 4024
0.623
0.617
0.392
0.614
0.622
0.384
DBAL-d
0.650
0.635
0,413
0.656
0.638
0,404
DBAE
0.655
0.643
0.369
0.664
0.675
0.332
data-attribute pairs (Xo, y). The attribute prediction y w + b based on the learned latent representation z (xcg(xg), which is fitted predict the ground-truth label y. An informative latent representation allows the linear classifier to predict the ground- label y more effectively. We evaluate (x)) trained CelebA (Liu al. 2015) and FFHQ (Karras et al., 2019). We train linear classifier on 1) CelebA with 40 binary labels, measuring accuracy as AP, and 2) LFW (Kumar al. 2009) for attribute regression, measuring accuracy using Pearson's s r and MSE. Table T shows that DBAE outperforms other diffusion- representation learning baselines. Since DiffAE, PDAE, and DiTi suffer from the information split problem, they produce a that is less informative than DBAE. Figure pressnts statistics for 100 reconstructions of the same image with inferred z. Because PDAE's s reconstruction varies depending on the selection Xy, it suggests that intricate details, such as hair and facial features, are contained in XT, which fails to capture. This observation aligns with Figure where significant performance gains are observed for attributes related to facial details, such as shadows and hair. A comparison between DBAE-d E-d and DBAE reveals that the stochastic encoder performs slightly better. We conjecture that the stochastic encoder leverages a broader latent space, which benefits discriminative downstream inference.
5.2 RECONSTRUCTION
Table 2: Autoencoding reconstruction quality comparison. Among tractable and 512 512-dimensional latent variable models, the one yielding the best perfor- mance is highlighted in bold, underline for the next best performer.
Mean
Std
Method
Tractablity
Latent dim (1)
SSIM
LPIPS (.)
MSE (1)
StyleGAN2 (W) Karns el al. 2020)
512
0.677
0.168
0.016
PDAE Statistics
SyleGAN22WW Aiklire
2019
7,168
0.827
0.114
0.006
Mean
Std
VQ-GAN Esser et di 2011
65,536
0.782
0.109
3.61e-3
VQ-VAR7
mmy
327,680
4.87-4
NVAR Valder Kinz 2020
6,005,760
0.984
4.856-5
DDIM (Infored xy) None ei al B
X
49.152
0.917
0.063
0,002
DidAE (Inferred
Prochis
X
49,664
0.991
0.011
6.07e-5
PDAE (Inferred
Zang
49,664
0.994
0.007
3.84e-5
DiTAR (Randoen Xr)
Preechakal et ml
3023
$12
0.677
0.073
0.007
DBAE-d Statistics
PDAE (Random xy)
512
5.01e-3
4.1ea33
Figure 3: Reconstruc--
DBAE-d
512
0.953
0.072
2.49e-3
tion w/ inferred z.
We examine the reconstruction quality following DiffAE (Preechakul et al.] information loss in the latent variable. For a test sample Xo, the procedure in algorithm provides reconstructed sample Sg. The reconstruction error is the distance d( d(xo, &o). where the distance function can be SSIM (Wang et al., 2003). LPIPS (Zhang et al. 2018). or MSE. Table 2 reports the averaged reconstruction eитоr over the test dataset (x.) [d] No. x,)]. We trained DBAE on FFHQ and evaluated it on CelebA-HO HQ (Karras et al., 2018). Tractability refers to the ability to perform inference on latent variables without repeated neural network evaluations. Tractability is crucial for regularizing the latent variable to achieve specific goals (e.g., disentanglement) during the training phase. The latent dimension refers to the dimension of the bottleneck latent variable during inference. A lower dimension is advantageous for applications such as downstream inference or attribute manipulation. The third block in Table compares performance 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 between Xy and z. Unlike the downstream inference experiments in Section 5.1 the deterministic encoder performs better.
2022 to quantify
8

---

Published as a conference paper at ICLR 2025
5.3 DISENTANGLEMENT
Table 3: Disentanglment and sample quality comparisons on CelebA
Method
Reg z
TAD (t)
ATTRS
FID (4)
AE
X
0.042++0.004
1,0+0.0
90,4+1.8
DifFAE Preechakul et al. 2022
x
0.155+0.0100
2.040..0
22.74:211
DBAE
x
0.124
2.2
11.8
VAE Kingma & Welling 2014
0.000:10.000
0.0::0.0
94.3::2.8
8-VAE (Higgms of al. 2011
0.088+000011
1.6+0.8
99.8:+2.44
InfoVAE Thao al. 2019
0.000++0.000
0.0+H0.0
77.8++1.6
InfoDiffusion ( Wang et al. 2023
0.299:10.000
3.0:0..0
22.3+1:.2
DisDiff Yang el
0.305+:00010
18.3++2.1
Figure 4: TAD-FID tradcoff'
DBAE+IC
0.362::0.036
3.8
13.4::0.2
compared to the baselines.
Unsupervised disentanglement the latent variable z an important application generative repre- sentation learning, as it enables controllable generation without supervision. The goal of disentangle- ment is to ensure that each dimension of the latent variable captures distinct information. To achieve minimize total correlation (TC), i.e., DRL (z)(2)) ( | (2.)). adopted from (Chen et al.) 2018). TC regularization decouples the correlation between the dimensions of z, allowing different information to be captured in each dimension. Following InfoDiffusion (Wang ict al. 2003). we measure TAD and ATTRS (Yeats et al. 20222) to quantify disentanglement in z. Since sample quality and disentanglement often involve a trade-off, we also measure FID (Heusel et al.. 2017) between 10k samples. Table 3 shows the performance comparison, where DBAE outperforms all the baselines. Figure Jemmonsratess the effects of coefficients on TC regularization, showing that DBAE envelops all the baselines. To disentangle information, well-encoded representation must first be achieved. The informative representation capability of DBAE supports this application.
this, apply regularization
5.4 UNCONDITIONAL GENERATION
Table 4: Unconditional generation on FFHQ '+AE' indicates the use of the inferred distribution (z) instead of pw (z).
Method
Prec (f)
IS (f)
FID 50k (L)
Rec (1)
DDIM
Song el al
2021a
0.697
3,14
11.27
0.451
DDPM
(Ho ci al.
2020
0.768
3.11
9.14
0.335
DiffAR
Preechaki kul et al.
2022
0.762
2.98
9.40
0,458
PDAE
Zhang et al.
2022
0.695
2.23
47.42
0.153
DBAE
0.780
3.87
11.25
0.392
Figure 5: Top two rows: un-
curated samples. Bottom two
DEEAAEAA
0.750
3.63
2.84
0.685
PDAE+AE
0.709
3.55
7.42
0.602
rows: the sampling trajectory with ODE and SDE.
DBAE+AE
0.751
3.57
1.77
0.687
To generate a sample unconditionally, the generation starts from the learned prior distribution
z ~ (z). The latent variable z is decoded into Xy = Dec and the sample xn is finally
obtained through the generative process described in Eq. 12]. For CelebA, a comparison with DiffAE
in Table 3 shows that DBAE surpasses DiffAE by large margin in FID (Heusel et al. 2017) 7) (22.7
vs. 11.8). Table shows the performance on FFHQ, which is known to be more diverse than CelebA. DBAE still performs the best among the baselines in terms of Precision (Kynkäänniemi al. 2019)
and Inception Score (IS) (Salimans et al., 2016), both of which are highly influenced by image fidelity. However, DBAE shows slightly worse FID (Heusel et al. 2017) and Recall (Kynkäänniemi jet al., 2019). which are more affected by sample diversity. To analyze this, we alter the learned generative prior pw (z) to the inferred distribution (z) as shown in the second block of Table 4 In this autoencoding case, DBAE captures both image fidelity and diversity. We speculate that it is more sensitive to the gap between 94(z) (z) and Pw Pw(z) since the information depends solely on z, not on the joint condition of xr and z. A complex generative prior model w could potentially solve this issue (Fiser et al. 2021 Vahdat et al. (02T). Figure 5 shows the randomly generated samples and sampling trajectories on FFHQ from DBAE.
5.5 INTERPOLATION
For the two images x] and x3. DBAE can mix the styles by exploring the intermediate points in the latent space. We encode images into x =Enc((xh) and x² =Finc (xf). We then regenerate
9

---

Published as a conference paper at ICLR 2025
Interpolation
Male
Young
+ Figure Interpolation (top) and attribute manipulation (bottom) with DBAE. (Red box: input image)
Smile
Wavy Hair
from x = Az' + (1 -)=xxzz to data Xy using the generative process specified in Eq. 12 The unique properties of DBAE offer distinet benefits here: 1) DiffAE (Preechakul et al. 2022 and
PDAE (Zhang et al., 2022) need to infer x). x2. by solving the ODE in Eq. with hundreds of score function evaluations (Preechakul al. 2022 Zhang al. 2022). They then geometrically interpolate between x and x to obtain x7. regardless of the correspondence between z and x². 2) DBAE directly obtains an intermediate value of x =Dec,p (z*). This does not require solving the ODE, and the correspondence between x/. and z also naturally determined by the decoder (ub). Figure [ shows the interpolation results on the LSUN Horse, Bedroom (Yu al. (015) and FFHQ datasets. The top row shows the corresponding endpoints x² in the interpolation, which changes smoothly between x and x7. The bottom row shows the interpolation results FFHQ, which smoothly changes semantic information such as gender, glasses, and hair color.
5.6 ATTRIBUTE MANIPULATION
The linear classifier used Section can can also be utilized to identify the manipulation direction z. From the prediction of a linear classifier ý = w"z + b, traversing in the direction dy decreases the logit. For image Xo, this is encoded as Enc,(xo) The encoded representation z is manipulated as now = z + Aw. The manipulated image x is obtained by decoding x" Dec (z"*), and the reverse process in Eq. (2).. DiffAE and PDAE additionally infer from X, X7 by solving Eq. (4) with hundreds of score function evaluations, fixing X7 to prevent undesirable variations in XT. T describes the long inference time for XT in previous approaches. Moreover, some information is split into XT, these methods cannot handle this information. On the other hand, DBAE infers Xy directly from manipulated z, ensuring that the endpoint X also controlled through the decoder (ạb). Figure shows the manipulation results for both CelebA-Ho images and FFHQ images with various attributes.
dz w increases
CONCLUSION
This paper identifies the information split problem diffusion- representation learning, stem- ming from separate inferences of the forward process and the auxiliary encoder. This issue hinders the representation capabilities of the tractable latent variable z. The proposed method, Diffusion Bridge AutoEncoders, systematically addresses these challenges by constructing z-dependent endpoint Xy inference. By transforming z into an information bottleneck, DBAE extracts more meaningful representations within the tractable latent space. The notable enhancements in the latent quality of DBAE improve downstream inference and image manipulation applications. This work lays a solid foundation for further exploration of effective representation in learnable diffusion inference.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by Institute of Information & Communication Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (No. RSS2024-0036113 Development technology to generate synthetic data in exceptional situations and advance artificial intelligence prediction models). This research was supported in part by the NAVER-Intel Co-Lab. The work was conducted by KAIST and reviewed by both NAVER and Intel.
REFERENCES
Rameen Abdal, Yipeng Qin, and Peter Wonka. Image2stylegan: How to embed images into the stylegan latent space? Proceedings of the IEEE/CVF international conference on computer vision, pp. 4432-4441, 2019.
Rameen Abdal, Yipeng Qin, and Peter Wonka. Image2stylegantt: How to edit the embedded images? In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 0296-8305, 2020.
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy, Fixing broken elbo. In International conference on machine learning, pp. 159-168. PMLR, 2018.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313 1982.
Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. contrastive learning approach for training variational autoencoder priors. Advances in neural information processing systems, 34:480-493, 2021.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. Advances neural information processing systems, 31, 2018.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Represen- tations, 2022. URL
net hett:/oorummii
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, llya Sutskever, and Pieter Abbeel. Info- gan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 74917799-77 2021.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8798 8794, 2021.
Joseph L Doob and JI Doob. Classical potential theory and its probabilistic counterpart, volume 262. Springer, 1984.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVE conference on computer vision and pultern recognition, pp. 12873-12883 2021.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617.
PMLR, 2018.
11

---

Published as a conference paper at ICLR 2025
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bemhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Irina Higgins. Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017. httt::::/ppeeeeee...etteeeeeeeee
Geoffrey E Hinton, Simon Osindero, and Whye Teh. A fast learning algorithm for deep belief nets. Neural computation, 18(7):1227-15555 2006.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 6851,, 2020.
Gary B Huang, Vidit Jain, and Erik Leamed-Miller. Unsupervised joint alignment of complex images. In 2007 IEEE/ I 1ith international conference on computer vision, pp. 1 1-8. IEEE, 2007.
Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James L. McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. Soda: Bottleneck diffusion models for representation learning, arXiv preprint
2023.
Donahue Jeff. Adversarial feature learning. ICLR 2017., 2017.
Insu Jeon, Wonkwang Lee, Myeongjang Pyeon, and Gunhee Kim. Ib-gan: Disentangled representa- tion learning with information bottleneck generative adversarie networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7926-7934, 2021.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. In International Conference on Learning Representations. 2018. URLpttps:: http///envvvvieieeee ew n forum?
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 4401-4411 2019.
Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings the IEEE/CVF conference computer vision and pattern recognition, pp. 8110-8119, 2020.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances Neural Information Processing Systems, 35:26565-26577, 2022.
Dongjun Kim, Byconghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, and Il-chul Moon. Maxi- mum likelihood training implicit nonlinear diffusion model. Advances in Neural Information Processing Systems, 35:32270-32284, 2022a.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International conference on machine learning, pp. 2649-26588 PMLR, 2018.
Yeongmin Kim, Dongjun Kim, HyeonMin Lee, and Il-chul Moon. Unsupervised controllable generation with score-based diffusion models: Disentangled latent code guidance. In NewrIPS 2022 Workshop on Score-Basee Methods, 2022b.
Yeongmin Kim, Byeonghu Na, Minsang Park, JoonHo Jang, Dongjun Kim, Wanmo Kang, and Il chul Moon. Training unbiased diffusion models from biased dataset. In The Twelfth International Conference on Learning Representations, 2024. URL https::////.......tt.. TRRUittss///ppeeeeeee http:///ppeeeiiiieee net /ft/rorumm
1d=39cPki jBed
12

---

Published as a conference paper at ICLR 2025
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models, Advances in neural information processing systems, 11-21696-217011 2021.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 20/4, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http
7/ farxii.orge wttpor//abaa11111 /11...111.... 6114
Neeraj Kumar, Alexander C Berg. Peter N Belhumeur, and Shree K Nayar. Attribute and simile classifiers for face verification. In 2009 IEEE 12th international conference on computer vision, pp. 365-372. IEEE, 2009.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
Bingchen Liu, Yizhe Zhu, Kunpeng Song, and Ahmed Elgammal. Towards faster and stabilized {gan} training for high-fidelity few-shot image synthesis. In International Conference Learning Representations, 2021. URL pttttp::///penneeiieee /
wwt.://wwwwwe/eo.ummte..m 2id= Fg 33
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of the IEEE international conference computer vision, pp. 3730-3738, 2015,
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
Kushagra Pandey, Avideep Mukherjee, Piyush Rai, and Abhishek Kumar. Diffuse Efficient, controllable and high-fidelity generation from low-dimensional latents. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: pttttp:://ppnneeee http////ppeeeeeeeee/eeee net. for
Id=ygoNPRiLxx
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Dif- fusion autoencoders: Toward meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10619-10629, 2022.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019.
Danilo Jimenez. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
Chris G Rogers and David Williams. Diffusiins, Markov processes and martingales: Volume 2, Itô calculus, volume 2. Cambridge university press, 2000.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-asisttd intervention-MIICAA 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234-241. Springer, 2015.
Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of Statistics, pp. 1160-1174, 1995.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-255 2015.
Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, and Romann M. Weber. CADS: Unleashing the diversity of diffusion models through condition-annealle sampling. In The Twelfth International Conference on Learning Representations, 2024. URL https: :ttppp:///peeeeee /ppnnevveww
13

---

Published as a conference paper at ICLR 2025
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Erwin Schrödinger. Sur la théorie relativiste de I' l'électron et l'interprétation de la mécanique quantique. In Annales de l'institut Henri Poincaré, volume 2, pp. 269-310, 1932.
Abhishek Sinha, Jiaming Song, Chenlin Meng, and Stefano Ermon. D2c: Diffusion-decoding models for few-shot conditional generation. Advances in Neural Information Processing Systems, 34: 12533-12548, 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International conference on machine learning. pp. 2256-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021a. /tRl/ttttsss:/opeeeee.eee net/
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1428, 2021b.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021c.
wwt.eeeww..e..ttteeeeeeee forrmm?
IddpxTIG12RRHS
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1811--826. 2016.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33: 33:19667-1967, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in neural information processing systems, 14:11287-11302, 2021.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the gan latent space. In International conference on machine learning, pp. 9786-9796. PMLR, 2020.
Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing difffuson models. In International Conference on Machine Learming, pp. 36336-36354, PMLR, 2023.
Zhou Wang, Eero Simoncelli, and Alan C Bovik. Multiscale structural similarity for image quality assessment. In The Seventh Asilomar Conference on Signals, Systems & Computers, 2003, volume 2, pp. 1398-1402. leee, 2003.
Ancong Wu and Wei-Shi Zheng. Factorized diffusion autoencoder for unsupervised disentangled representation learning. Proceedings of the AAAI Conference on Artificial Intelligence, 38(6):
5930-5939, Mar. 2024. doi: 10.1609/aaai, v38...28407 URL index. wwweexpppp/Aaaee AAII/AAAI///
/ttp///rtteeeeeeeeeeeeeeee Indees.oih//AMA//rrriceee//eeee 28407
Weihao Xia, Yulun Zhang, Yujiu Yang, Jing- Hao Xue, Bolei Zhou, and Ming -sugn Yang. Gan inversion: A survey. IEEE transactions on pattern analysis and machine intelligence, 45(3):
3121-3138, 2022.
14

---

Published as a conference paper at ICLR 2025
Tao Yang, Yuwang Wang, Yan Lu, and Nanning Zheng. Disdiff: Unsupervised disentanglement of diffusion probabilistic models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
Eric Yeats, Frank Liu, David Womble, and Hai Li. Nashae: Disentangling representations through adversarial covariance minimization. In European Conference on Computer Vision, pp. 36-51. Springer, 2022.
Fisher Yu, Yinda Zhang. Shuran Song, Ari Seff, and Jianxiong Xiao. Lsun: Construction of a large- scale image dataset using deep learning with humans in the loop. arXiv preprint 1506.033556633 2015.
Zhongqi Yue, Jiankun Wang. Qianru Sun, Lei Ji, Eric I-Chao Chang, and Hanwang Zhang. Explor- ing diffusion time-steps for unsupervised representation learning. In The Twelfth International Conference on Learning Representations, 2024. URL htthhtts::/ wwwpppewwwww.eeeeeeoeoo...eeeeee /tBp:/ttsss//ppeneviiee...e/e//e
id=bwzxht 1HP
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34:16280-16291, 2021.
Richard Zhang, Phillip Isola, Alexci A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018
Zijian Zhang, Zhou Zhao, and Zhijie Lin. Unsupervised representation learning from pre-trained diffusion probabilistic models. Advances in Neural Information Processing Systems, 35:22117- 22130, 2022.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In Proceedings of the aaai conference artificial intelligence, volume 33, pp. 5885-5892 2019.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-v3: Improved diffusion ode solver with empirical model statisties. Advances in Neural Information Processing Systems, 36, 2024.
Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models. In The Twelfth International Conference on Learning Representations, 2024. URL https
Jun-Yan Zhu, Philipp Krähenbühl, Eli Shechtman, and Alexei A Efros. Generative visual manipulation on the natural image manifold. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part V 14, 597-613. Springer, 2016.
15

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2
Preliminaries
Difiusion Models
2.2
Latent Representation Learning with Difrusion Models
2.3
Diffusion Process with Fixed Endpoints
3
Motivation: Information Split Problem
4
Method: Diffusion Bridge Encoders
Encoding from Xn to XT conditioned on
Generative Process
4.3
Mutual Information Analysis
Objective Function
4.4
Reconstruction
4.4.2
Generative Modeling
5
Experiment
Downstream Inference
Reconstruction
5.3
Disentanglement
5.4
Unconditional Generation
5.5
Interpolation
5.6
Attribute Manipulation
10
6
Conclusion
10
Proofs and Mathematical Explanations
18
Proof of Theerem
18
A.2
Proof of Theorem
22
A.3
Prior Optimization Objective
23
Mutual Information Analysis
24
A.4.1
Auxiliary encoder framework
24
A.4.2
Diffusion Bridge AutoEncoders
25
A.5
Proof of Theorem
25
B
Related Work
25
I Representation Learning in Diffusion Models
25
B.2
Parametrized Forward Diffusion
26
16

---

Published as a conference paper at ICLR 2025
Implementation details
27
Training Configuration
27
C.2
Evaluation Configuration and Metric
27
Algorithm
30
C.4
Computational Cost
30
Additional Experiments
31
D.I Downstream Inference
31
D.2
Reconstruction
33
D33
Unconditional Generation
33
Results with Intel Gaudi v2 hardware.
35
D.5
Additional Samples
35
17

---

Published as a conference paper at ICLR 2025
A PROOFS AND MATHEMATICAL EXPLANATIONS
In this section, we follow the assumptions from Appendix A in (Song et al. 2021b), and we also assume that both Se and %3.9 have continuous second-order derivatives and finite second moments, which are the same assumptions of Theorems and 4 in (Song et al. (0211).
A.I PROOF OF THEOREM 1
Theorem 1. For the funttionn CAt, the following equality holds.
1
LAE 2 (x (x.)x. (x,, - Vx, log de (xee xo, xT) l1)de
(x,,t, ① a linear SDE there exists o(l), (()), 7(1), A(t), such that
(16)
Moreover, if Eq.
/-" E "
1
2/ LAE 2
(xt,1, X 12)
(17)
where (x t, Xr) = a(t)x, + 8(t)xr + (()) (x,, t,,x,,,))) XT), and Lo (Xg, X4. x.,
(xala(Xo)gg (z/xx))//x/))))) Ir (xT (x)((x) Xo)dz, following the graphical model in Fig. 1c
Proof.
Note that the definitions of the objective functions are
LsM = 2 1. Es (g²(t)) (x,,,t,,x)) - Vx, log (x, dr
(x/
(19)
LAK = LsM + H(q v (xu x))).
(20)
We derive the score-matching objective Lsm with the denoising version for tractability. First, Lsm is derived as follows.
1 1 B 19 (x,, + (t)( log L (xxxx))) Lsm 2
2g (1)s (xe. 4, x Vx, log d w ( (d.
(21)
Then, the last inner product term of Eq. (TT) can be deduced similar approach to (Vincent Vxx
2011):
Eng (x,xr) (so (xe,x, xx)) TV.. log 2 v (x, [x))]
(22)
f vx XT)so (x, x T log d'.w )dx,dxx
(23)
I Le (XT)) (x )so (x,, (so ) Xx log L6 (xx (x)) dxx
(24)
Bogg (xx)
(x,
xx) T Vxx log Lo (x, X )dx,
(25)
-Boo (XTr)
so(xi, t, 2 (xx|xr )dx,
(26)
s0(x t, x T Xx. / (xolxxooo ) (XT, xo dxxo dxx,
-Eug y ((5r)
(27)
-Bote (xcr)
50 X1, 1, xT) (7
Low (xolx7) V
xo dxx } sax.
(28)
-Bagg (xc7)
/ so(x), (0(xx, x,) p { do (xol
xo) Vxx log d' (x, XX,, xo )xx }dx,
(29)
[
= -Bga io (x)))
do (xol X Vo Xo (o) (x)) Vxx log gó (xr|XT, xo) )xxdxx
(30)
= =Rogg (x) (so Xx, L XT) XT) Vx, log Vo ( (x|[XT,Xo)
(31)
'Eq.
linear SDE when the drift function is linear with respect to X
18

---

Published as a conference paper at ICLR 2025
Next, rewrite the second term of Eq. 21 To begin, express the entropy H(g db (x X with Vx, log & (X((XT), which is similar to the proof of Theorem 4 in (Song et al. 2021b). Let
H(q xr]] == I / (X4, Xx) log 4s s (X,, xx) xx)xx, be the joint entropy function of 2 X2, Xr). Note that Hl96. (XT, xT)) Ha9s (xr)). Then, we have
Hlqo (xo, (x0, = u xT)) + l"d
07, (x,, dt.
(32)
01
We can expand the integrand of Eq. 32 as follows.
(X,,
- / g (xt, XT) log go s (Xe, xT dxxdx
(33)
01
06
- / 10. (xT)((x))u (xx (x)) Go. (xT) + log a (x(xx))))
06
(34)
f g0. (xT)
②
Öt 190 (log 9s, (XT) log qa, (xx(x]))
(355)
=- (XT)
Io.. (xx(x)) (x + log I0. (x, (x/]x]]] (xxxx))]
(36)
Ot
We further expand the integration the last term as follows.
0t (0 [log 40 (xT) + log qo (x, (xT)] Jdx,
(37)
-/
②
2 log ao as (x (x((xT)
Ot { {sos (x, xx) ( foog (xx) log qs (x, [x)] + g (xx(xrr)
dx;
Öt
(38)
= /
0 { 100 (x, x [log a (XT) + log N (x, xx)]] O4p. (xx (x] dxc.
(39)
Ot
Ot
2 Ot (00 (x4 x ) llog as (xT) log qs. v (x x ) dxx
Ot / 9s (xec
(40)
= /
Ot { (90. (x: x [log q (XT) + log 9s v (x/ x (xx ddx
(41)
/
di (90.6 xT) log gs s (xr )xx, +
(xx (x/(xT) } log ao.s (xx xx )xx;
(42)
Ot
2 2 / 2 ux br dxe + / = log qo us (XT)
a
Hi { no, 1o (Xee
log go v (xx dxx:
(43)
0t
/
di (xx x log go w (xs (x))
(44)
Note that we use f o (x (xxxxT))xs 1 in Eqs. 41 and (4
By eq. (51) in (Zhou et al. 2024), the Fokker-Plank equation for go.w (x)) follows
ar Lo (x:(xT) = Vx, [(((xes)) (()())))))))) TT,
(1) Vx
(45)
=- (fo, (X,, (x, [x])]
(46)
where (x,, = f(x,, 1) + (1)))(() Xx, T) - 2 (1) log to (x.[xT).
19

---

Published as a conference paper at ICLR 2025
Combining Eqs.
(6)
44 and (4). we have
HH,,,,, XT)
(47)
0L
--Rgs ( (XT) // Vx. (fo (xe, () u (xx) (xx)) log gob, yb (xe) )
(48)
[ / Vxs. (Pouss (xe, E) (x, xx)) log aes vs (x /xx )dxe
- -B v (xe) (xxr)
(49)
= -Eg v (xer
For (xe,, t) s (xe xx)) log g (xT)
/ q ( fo (xe, Xx. log 9 (xx ddx;
(50)
[
/ 90. (x (x,, T Vxx x log 9 (x, |x )dxe
= -E v (xer)
(51)
/ (f(xe:1) + ( (1)h(xx, XT, T)
1 0 (1) log go. (xx(x))
(xer))
Xx, log qo. (xx) xT (6. (x
(52)
1 20 (1)Vxx log To (xxl (xx XT)}
= Eg (x. xT)
{ f(x+,1) (1)))((())))) (x,, 7)
Xx, log us (xe (x)]xx)]
(53)
-gg (X. xr)
~" 0 (D)IVxx To fo (x:
- (f((x,,) + (1)),(x,, T))"Vx log g (x, (x))
(54)
Therefore, the joint entropy function H(q0 (Xu, x) can be expressed as
H(yp.w( (xo, x))) = ((().w((T)) + [
B Bi (x))))) 1/5 (1)1V xx log dh (x(1x))11
- -f(x,, Vx, log 2 (x1(x)) (()()((),
- TT), log d (xx(x)) dt. (55)
We can re-write the above equation as follows.
((x, (t)(( log 20. (x:1xr))131
(56)
2H (70. (X0 27) x)))
(57)
+ 1/" B (x (2f(xx, t) Vxx
x, log d 2g' (t) Vx, log oo (x)(xx))dtt
22H( (Xn x))) + S/" Ey (xb. Vxx, (f(x,, + g² (1))((x, x,x,, T))
(58)
Similar to the process above, we can obtain the following results for the following joint entropy function H( (X0, X1, XT)) := u (xo, xt, xT) log go. (x0, Xt, xT) )xxx
0H X1, XT)
H(g (Xo, Xo; x) H H(q XT, XT)) +
(59)
Øt
In the following results, we utilize the Fokker-Plank equation for v (x, Xo. XT), which comes from eq. (49) in (Zou et al. 2024):
at 20.0 (x xT)
Vxi (f(x), t) + 9 (2(t))(xx,, x,x T) (x, Xo. x.)
2 (1) Vx. ao (xo: xT)
(60)
Vx,
(x,, t) 1f
(x, (Xo, xT))
(61)
20

---

Published as a conference paper at ICLR 2025
where E f(x,,t) (t))((x (t)) log 9s. 1og90.v
( (x)
Then, we have
[ B B (x0 Xx.,
1 (t)ll log Lo (x, xx. xo)1 - f(xe, Vx. log To (xx (xx(xxxxo)
0
2
- (t) t, TVV log log (x XT, xo) dt,
(62)
where the left hand side is from 0 H(gss (Xa, Xo, xr)) Hag (xx, xr, Xr)), and right hand side from lr (x,, xx) x dt. We can further derive as follows.
ot
Lo" Eg w (xo xx) 10 (t)l Vx, log I ( xT. xx)
2 !'" (x. (f(x,,t) + +g (xx),, (63)
Combining Eqs. 58 and (63) we have
L/ w (x. 19
Vx, log Cé (x. xT)I Bjdt
xx)) /'" (xo. (t) Xx, log d xo) (64)
-22
Combining all results, the score-matching objective Lsm can be expressed as
LsM = L/" Eyg (x0.
( log de (x, xT. xo)11
(x, Vx. log 9'a (x, Xx,, xo)]dt H(go.
(65)
/- 3
T Egg (xx (aal:l/s(((x...)) tx) Vx, ( xT x1 Hxgo.wxxolx)
(66)
The last equality comes from 9 d (x, x7) = Te (x, (x. xx,) which is based on the Doob's
h-transform (Doob & Doob 1984 Rogers & Williams 2000 Zhou al. 2024). Finally, we have
CAB LsM + H(40. (xulxx))
(67)
2 //"
1
(t)] (x,, x,xx) Vx, log qe (xx xT)l dt
(68)
From here, we show that the objective CAE is equivalent to the reconstruction objective. Assume that the forward SDE in Eq. (1) is linear SDE terms of x, (e.g. VP (Ho al. (2020), VE (Song et al. 2021c)). Then the transition kernel ((x: becomes Gaussian distribution. Then, we represent reparametrized form x, = a; + O,E, where a and o, are time-dependent constants determined by drift f and volatility g. and ~ (O, I). The time-dependent constant signal-to-noise
2
ratio SNR(() = = often define to discuss on diffusion process (Kingma et al. 2021 ). We define
SNR ratio, R(t)
SNN R(T) for convenient derivation. SNR(C)
Zhou et al. 2024 show the exact form of (xx x7) == N ( If1), where fa R(t) Xy a1xp(1 R()) and 6, = o 0, R(t). This Gaussian form determines the exact analytic form of the score function Vx, log q (x, X.). We plug this into objective LAE.
21

---

Published as a conference paper at ICLR 2025
T
AAB
2 20
B o xx, xx)) (g) (,x)) - Vx. log qr (Xp Xo, 11)dt
(69)
T
-Xt + (R(t) xT + 0xxo R(t)))
(Xo (x. ())
dt
of (1-)()))
(70)
/'-
)
20
( (x,, x xollidtt
(71)
where
a4 g \(t)
(72)
o6
x6 1, xT) = a(t)x; + =(()xx + (())s((())))))))))))))))))) (xe, (x,, xT),
(73)
18(t)
of.
a(l) ar - R(D))' 3(1) ar((1 - -R(I))'
7(1)
(74)
04t
A.2 PROOF OF THEOREM 3
Theorem 3. Dxi(9dam (x)) (x)))Pp,,(Xx)) S LAE + Lrn H, where H H(Adataa is constant w.E.t Φ 1b,
Proof. From the data processing inequality with our graphical model, we have the following result, similar to eq. (14) in (Song et al. 2021a).
DRt. (qsamaan (x))( S DKL (1o 2) (XO:T:2) 2))
(75)
Also, the chain rule of KL divergences, we have
DEL (yo, (Xurr z) pys (xu:r, 2))
(76)
DKL(Uo. (xT, z) IPw. (XT, z)) +
(DLL (Mo (How z)) ve (xT,, 2))],
(77)
where us.v and v are the path measures of the SDEs in Eqs. and 799 respectively:
dx, = [f(xx,)) t) + (2)))2)))))))) (t) t,y, T)jdt + (t))
Xu ~ Gduta (Xu),
Xy ~ 1 s (x2 (x))
(78)
dx, = [f(x,, 9 (V) logro (T) 4x,,,, TYJJdl g())dw) XT p (xT))
(79)
By our graphical modeling, is independent of {x,} given Xy. Therefore, we have [DxL (Ho ( [xr 2)]| [o Xy. 2))] Eg (xx) [ (( (x)])
Eg w (xr (xT
(80)
where u (. (xT) and vo xx)) are the path measures of the SDEs in Eqs.
81 and
(82 respectively:
dx, t)) (t) / log u (xx|xx) h(x,, T)]]dt g())
x(T) Xx',
(81)
dx, = [x(x,,t) g² (t) (x h(x,, y. T)]]dt g())d x(T) = X
(82)
Similar to eq. (17) (Song et al. 202 this KL divergence can be expressed using the Girsanov theorem (Oksendal 2013) and martingale property.
DKI. (H (xT) ( (-[x7)) = ///""
(x/ (x// (1) (x);;,,x)x Vx, log q (x (xx) 1
(83)
22

---

Published as a conference paper at ICLR 2025
From Eqs. 75 and (83) and Theorem have:
DKL(Adaaa (xx)) s DKL (9o. (XT, 2) o 2)) + LAE (x000x(x0lxx))
(84)
Furthermore, the first and third terms of RHS in Eq. (84) can be expressed as follows.
(1a.v (Xr, 2)llp o (XT: z)) H( (xolxT))
(85)
= | 10. (XT, 2)) log ao. (XT, z)
dxrdz+ + / yoow (xo, xT) log 9o. (xol X )xxodxx
(86)
Po. Xx,
= -f 100.0 (X0: XT, 2)
log 20. (XT: 2)
+ log Tos (xol (xolxx) dxodxrd2
(87)
Pv. (x 2)
= / 20. Xx: XT, 2)
log as (xT (aw + log ls. vs xx)) dx
(88)
Pap (XT Pap xx)
= / 10. (x: XT, 2)
log Lo (XT) + log naw (xalxT) dxodxrdz
(89)
Pepl (xT)
= / L0. (X0= XT)
log Is.e (xT )ad, t (Xo (xy)
dxodxx
(90)
Per (XT)
= f Ts. (x0 XT)
log Adata (Xo 1aa. (XT xx)) |xxdxx
(91)
Pap (XT)
= f qdatal (xo I0 vo (xT (xo)
loog" 90. (xer (xo) + log Adina (xo) dxpdxr (xT(
(92)
Eegu (xo) DKL (1o s (xer (2) H(gdau (xx))
(93)
= CpR ((ddaaa (xo))
(94)
To sum up, we have
DKL (xo) (X) S LAE + LpR H(Adata (Xo)):
(95)
n
A.3 PRIOR OPTIMIZATION OBJECTIVE
This section explains the details of the prior related objective function mentioned in Section 4.4.2 The proposed objective is LpR as shown in Eq. (96).
CpR Exg DKL. (go. (XT xo (xPx))))))
(96)
To optimize this term, we fix the parameters of the encoder (Φp - Φ²), the decoder (1/ ))",, and score network (0 + 0*), which is optimized by CAg- And we newly parameterize the generative
prior Pprior (z) - Pw (z), so the generative endpoint distribution becomes (xT) + Pvp _ (XT) We
utilize MLP-based latent diffusion models following (Preechakul et al.) 2022 Zhang et al. 2022).
The objective function Eq. (66) with respect to w is described in Eq. (97) and extends to Eq. 99) with equality. Equation (TOO) is derived from the same optimality condition. In other words,
reduces the problem of training an unconditional generative prior (z) to matching the aggregated posterior distribution 9 (z).
arg min Ex)s)(aa) [KL (40- - (xT (x0) (op) i (XT))]
(97)
+ arg min / Quana (xx) aoraoeeee .. (xe xx)) log 2 ap- (XT Xo) cxx/xx pet w (xr)
(98)
+ arg min DKL (q0- - (xT) p (xx)) + C
(99)
<> arg min (s (z) (Pww
(100)
w
23

---

Published as a conference paper at ICLR 2025
A.4 MUTUAL INFORMATION ANALYSIS
Alemi et al. (2018) shows the distortion; reconstruction enror with inferred z is the variational bound of mutual information between Xp and z in the autoencoding framework. We explain the functional form of distortion in both the auxiliary encoder framework (Appendix A.4.1) and DBAE
(Appendix A.4.2).
A.4.1 AUXILIARY ENCODER FRAMEWORK
In the auxiliary encoder framework (e.g.. DiffAE (Preechakul et al. 2022 ) the distortion
(xo) (aa(z) log po (xo|2)) and mutual information M Mx 2) Fys 2) [log 4e (Xo.2)
gave (Xo))aa
has a relation
(xon log pe (Xo 2]] + H(quarra S Mx 2),
(101)
where (xolz) = Pprior (x ( xT)dxxx when this framework reconstruct only with inferred z.
We have the followings
log (xolo())
(102)
/
log
ODF (xoll XT dxxT PPrrio( (XT
(103)
log
Ppriir (XT 1pg
ODE (xolx, XT)
40
XT Z2 X0)
dXT
(104)
40
ODE
XT Z, Xo)
XT )Pg ODE [xo] Z, XT)
90
ODL (xT (XT
12, Xo) log Pprior
40
ODE
XXT z, Xo)
dxT
(105)
E d0 (xr [xo. ox
logg
ODE (x Z. xT)
D) D 1. (90
(DE
(xr No: z)) IPprior (xT))
(106)
ODE (xT z, xx) log Pprior XxT 12
ODE (xuizz, X-)
90
dxc
(107)
9o
(x Xo)
/ 48
ODE (xT z, Xo) log prior (xT )dxT
(108)
(xr CF
ODE
x [Pprior(xr))
(109)
DDF (xF z, Xor Note that pg
Lo (X |z, X²) because the deterministic coupling of (X,, xx) is given When the coupling (Xa, X7) lies on the ODE path, both probabilities
by the ODE in Eq. 110 Po (DEE (xal 22, XT) and t0
x iz, xo) become infinite. When the coupling (Xo, XT) is outside the ODE (xg (x) and ODEE (x7lz, xx)
ODE path, both probabilities pe
xo) become zero
dx, [f(x, t) 3 (t)s (x,, 2, t)]df.
(110)
From Eq. 1001 and Eq. (109) we have the following.
E (x)) (- ODE (xT (
Xo) [ (xT))] + H H((dara(x)) S MI 2)
(111)
The discrepancy between 40 ODE (xT No; z) and Pprior (xr) makes the lower bound of mutual infor- mation between Xo and loose. This discrepancy is inevitable from the deterministic nature of ODE (xT z, Xo).
do
This discrepancy is empirically observed in Table 2 providing two cases of draw (random Xy, inferred xx.) in the auxiliary encoder models. The reconstruction gap between (random Xy- inferred xy) ) is significant in practice. However, the inference of xy is computationally expensive and inflexible in terms of dimensionality. If we only consider z inference, the information leakage is inevitable due the functional form of diffusion models with auxiliary encoder.
24

---

Published as a conference paper at ICLR 2025
A.4.2 DIFFUSION BRIDGE AUTOENCODERS
In the DBAE, the distortion gon go (x/xx)) log pa, [x0/2)] term and mutual information between Xo and z has relation in Eq.
Exssa (oa log Pe, (Xo|z)] + Higguaa (xo)) S MI(x.. z),
(112)
where pa, (xolz) - f po (xo XT (x (xx. We have followings
log p (xol iz)
(113)
log / po (x Map (xr (z)
(114)
/ (xoxx Pap xT (2) Iv X (z)
log
dxT
(115)
gup (xr (xr(z) (z)
/ ge ( (2) log po (xo XT Pub (xT
dxT
(116)
Top XT (2)
Ey (xx (x) [log po (xo [x)] DK (g (xr (z) l2)
(117)
Since DKL, (gep (0pp(x (xr [2)] lays (xT 2)) we have followings from Eq. 112 and Eq. 117
B Rasa( no (o) [a (xr (z) [log po (xo|xx))] xr) + H (Xo)) S MI ( z).
(118)
Unlike in Eq. I111, the Xy related term does not hinder maximizing mutual information between Xo and z. Moreover, the remaining term E E (x))) (u (x) (llgl/lo((ee px [xr)]] can maximized by our training, as we explain Theorem
A.5 PROOF OF THEOREM 2
Theorem 2. -MI(Xo; z) LAE H, where H -H(ysaa(x))) constant w.r.t. Φ, 1b, O.
Proof. From data processing inequality similar in Eq. 75
Eg (xT) (1o. ( (xo (xoo))))(((( Eg v (xx)
(119)
The LHS of Eq. 119 becomes followings,
Egg, (xx) [KL (1s. (x X2)) (xxlxx))))))))))) = Eg xxx. log (aaei)))-((((ooo
H(ss xr))
(120)
The RHS of Eq. 119
becomes followings from the result of Eq. 83
E ya ( (xr)
Dxi. (u (xx7)) LsM LAF H(zo (xo (xT))
(121)
From Eqs. 119
121
we have the followings
E Rs. (No xT) |- log Po (xp S CAE
(122)
We have the following to sum up Eq. 122 and Eq. (18).
- -MT(xo 2) S LAE H(ddttaa (Xddatt(xx)) (xo)
(123)
B
RELATED WORK
B. REPRESENTATION LEARNING IN DIFFUSION MODELS
Expanding the applicability of generative models to various downstream tasks depends exploring meaningful latent variables through representation learning. Methods within both variational autoen-
coders (VAEs) (Kingma & Welling 2014 Rezende al. 2014: Higgins et al. 2017 Zhao et al.
2019 Kim & Mnih| 2018) and generative adversarial networks (GANs) (Jeon et al.) 2202 Karras
al. 2020 Abdal al. 2019 2020 Chen et al. 2016) have been proposed; however, VAES suffer
25

---

Published as a conference paper at ICLR 2025
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 (Xia et al. 2022] Voynov & Karras
Bahenko, 2020 Zhuu et al. 2016 Karras et al. 2020 Abdal et al. 2019). Additionally y, 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 xy was introduced latent variable (Song et al. 202 c) with an invertible path defined by an ordinary differential equation (ODE). However, XT is difficult to consider as semantically meaningful encoding. Additionally, the dimension of XT matches that of the original data Xo, limiting the ability to learn condensed feature representation for downstream tasks (e.g., downstream inference, attribute manipulation with linear classifier). The inference of latent variables also relies on solving ODE, rendering inference intractable. This intractability not only hinders the desired regularization (e.g. disentanglment (Higgins et al. 2017: Kim & Mnih 2018: Chen et al. 2018)) of the latent variable but also slows down the downstream applications.
Diffusion AutoEncoder (DiffAE) (Preechakul et al. 2022) introduces a new framework for learning tractable latent variables in DPMs. DiffAE learns representation in the latent variable through an auxiliary encoder, with a z-conditional score network (Ronneberger et al.] 2015). The encoder- generated latent variable can learn semantic representation with a flexible dimensionality. Pre- trained DPM AutoEncoding (PDAE) (Zhang et al., 2022) proposes a method to learn unsupervised representation from pre-trained unconditional DPMs. PDAE also employs an auxiliary encoder to define z and introduces a decoder to represent Xx, logp(z))x, PDAE can parameterize the z-conditional model score combined with a pre- trained unconditional score network, utilizing the idea of classifier guidance (Dhariwal& & Nichol, 202T). 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 (Yue et al.) (004)) introduces a time-dependent latent variable on the top of PDAE to enable feature learning that depends diffusion time. InfoDiffusion (Wang et al. 2023) regularizes the latent space of DiffAE to foster an informative and disentangled representation of It should be noted that such proposed regularization in (Wang al. 2023) is also applicable with DBAE, and Section [33 demonstrates that the tradeoff between disentanglement and sample quality is better managed in DBAF than in DiffAF. FDAE (Wu & Zheng, 2024) learns disentangled
latent representation by masking image pixel content with DiffAE. DisDiff (Yang et al. (023) learns disentangled latent variable by minimizing mutual information between each latent variable from different dimensions atop PDAE. LCG-DM (Kim et al. (002b) adopts a pre-trained disentangled encoder and trains DiffAE structure with fixed encoder parameters enable unsupervised controllable generation. SODA (Hudson et al. 2023) improves the network architectures of DiffAE and training for novel image reconstruction.
All the frameworks (Preechakul et al. 2022 Zhang et al. 2022) and applications ( Yuue et al. 2024
Wang et al. 2023; Wu & Zheng 2024 Yang et al. 20233 Hudson et al. (023) utilize the encoder and do not consider the diffusion endpoint xy, leading to an information split problem. In contrast, DBAE constructs an z-dependent endpoint X7 inference with feed-forward architecture to induce as an information bottleneck. Our framework makes z more informative, which is orthogonal to
advancements in downstream applications (Kim et al. 2022b] Yue et al. 2024| Wang et al. 2023
Wu & Zheng 2024 Yang al. 2023: Hudson et al. 2023) exemplified in Section 5.3
B.2 PARAMETRIZED FORWARD DIFFUSION
The forward diffusion process with learnable parameters is a key technique in DBAE to resolve information split problem. We summarize several other methods that proposed a learnable forward process. Note that DBAE has clear technical differences from those methods.
Schödinger bridge problem (SBP) (De Bortoli et al. 2021: Chen et al.) 2022) 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-regularizeed optimal transport (Schrödinger)
1932 Genevay et al. 2018). which is often
26

---

Published as a conference paper at ICLR 2025
solved by Iterative Proportional Fitting (Ruschendorf 1995). For this optimization, samples required at any given time from the forward SDE; however, these samples are not from a Gaussian kernel like Eq. ( 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) (Zhang & Chen 2021> parameterizes the drift term in Eq. ①I using& a normalizing flow, making the endpoint of DifIf learnable. However, both training and endpoint inference are intractable because the parametrized forward SDE does not provide Gaussian kernel similar to that in SBP. Implicit nonlinear diffusion model (INDM) (Kim al. (022a) lears a diffusion model that is defined in the latent space of a normalizing flow, implicitly parameterizing both the drift and volatility terms in Eq. []- 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 inference also requires solving the ODE for encoding.
Unlike other studies, we parameterize the endpoint XT rather than the drift or volatility terms. The forward process is naturally influenced by the endpoint determined from Doob's 's h-transform. Unlike other parameterized diffusions, our approach ensures tractable learning and Xy inference, making particularly advantageous for encoding tasks.
C IMPLEMENTATION DETAILS
C.I TRAINING CONFIGURATION
Model Architecture We use the score network (0) backbone U-Net (Ronneberger et al. (013)). which are modified for diffusion models (Dhariwal & Nichol 2021 with time-embedding. Dif- fAE (Preechakul et al. 2022)) PDAE (Zhang al. 2022). and DiTi Yue et al. 2024) also utilize the same score network architecture. The only difference for DBAE is the endpoint xT conditioning, We follow DDBM (Zhou et al.) 2024) which concatenate x, and xy for the inputs as described in Figure 7b This modification only increases the input channels, so the complexity increase is marginal. While the endpoint Xy contains all the information from z, we design a score network also conditioning z for implementation to effectively utilize the latent information in the generative process. For the encoder (t)), we utilize the same structure from DiffAE (Preechakul et al., 2022). For the decoder (1²), we adopt the upsampling structure from the generator of FastGAN (Liu et al.) 202I]. while removing the intermediate stochastic element. For the generative prior (w), we utilize latent ddim from (Preechakul al. 2022). Tables 5 and 6 eexplains the network configurations for the aforementioned structures.
Optimization We follow the optimization argument from DDBM (Zhou et al. 2024 with Variance Preserving (VP) SDE. We utilize the preconditioning and time- weightigg proposed in DDBM, with the pred-x parameterization (Karras al. 2022). Table shows the remaining optimization hyperparameters. While DDBM does not include the encoder (() and the decoder (1²), we optimize jointly the parameters Φb v and 0 to minimize LAB.
C.2 EVALUATION CONHIGRRATION AND METRIC
Downstream Inference In Table I we use Average Precision (AP), Pearson Correlation Coefficient (Pearson's 's r), and Mean Squared Error (MSE) as metries for comparison. For AP measurement, we S train a linear classifier (IR² - [0, 1|10) ) to classify 40 binary attribute labels from the CelebA (Liu ct al. (2055) training dataset. The output of the encoder, Excx(x,) 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 's r and MSE, we train a linear regressor (R' - T27)) using LFW (Huang et al. 2007 Kumar et al. 2009) dataset. The regressor predicts the value of 73 attributes based on the latent variable z. Pearson's '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 ( Yuue et al. 2024) paper and adhere to the evaluation protocol found at wuthttps:: /ttp::///ttttt com yue zhonggi /dit
27

---

Published as a conference paper at ICLR 2025
X BxC
X BxCxW*H
B x dim(z)
(a) The encoder (+) structure in the left, and decoder (τb) structure in the right.
MLP()), MLP(x)
X
X X BXCXWXH
B 2W
X
(b) The score network (0) structure. While model output not directly one-step denoised sample xx, the output equivalent to X.-1 with time-dependent constant operation with accessible information.
Figure 7: The architecture overview of Diffusion Bridge AutoEncoder.
Table 5: Network architecture and training configuration of DBAE.
Parameter
CelebA 64
FFHQ 128
Horse 128
Bedroom 128
Base channels Channel multipliers
64
128
128
128
[1,2,488]
[1,1,2,3,4]
[1,1,233,4]
[1,1,233,4]
Attention resolution
[16]
[16]
[16]
[16]
Encoder base ch
64
128
128
128
Enc. atin. resolution
[16]
[16]
[16]
[16]
Encoder ch. mult. latent variable z dimension
[1,2,4.8,8]
[1,1,2,3,4.4]
[1,1,2,3,4,4]
[1,1,2,3,4,4]
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
le-4
le-4
le-4
le-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
28

---

Published as a conference paper at ICLR 2025
Table 6: Network architecture and training configuration of latent diffusion models Pa (z) for an unconditional generation, following (Preechakul et al. 2022).
Parameter
CelebA 64
FFHQ 128
Batch size
512
256
z trained
600M
600M
MLP layers (N)
10, 15
10
MLP hidden size
2048
latent variable z dimension
512
SDE
VP
3 scheduler
Constant 0.008
Learning rate
le-4
Optimizer
AdamW (weight decay 0.01)
Train Diff T Diffusion loss
1000
LI, L2
Reconstruction We quantify reconstruction error in Table though the Structural Similarity Index Measure (SSIM) (Wang et al. 2.03), Learned Perceptual Image Patch Similarity (LPIPS) ZZang al. 2018)aoo 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, like 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 (Preechakul et al. 2022) and PDAE (Zhang et al., 2022). We use Heun's 's ODE sampler (99 NFE) to evaluate SSIM and MSE and use stochastic sampler (Zhou et al. 2024) (998 NFE) to evaluate LPIPS for Table We also present performance metries for various NFE values in Tables
and 13
Disentanglment The metric Total AUROC Difference (TAD) Yeats et al. 2022 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 z. To calculate the AUROC score, first determine the dimension- wise minimum and maximum values of z. We increment the threshold from the minimum to the maximum for each dimension, converting 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 considered disentangled if at least one dimension of 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 (Yeats) al. 2022). We borrow the baseline results expect DisDiff from the InfoDiffusion (Wang al. 2023), and we follow their setting that the dim (z) = 32. DisDiff (Yang et al. 2023) utilizes the dim 192 and we borrow its performance from the original paper. We use evaluation code from www/ er cyeat / nashae - be connsntttesss
Unconditional Generation To measure unconditional generative modeling, we quantify Precision, Recall (Kynkäänniemi et al. 2019), Inception Score (IS) (Salimans et al. 2016) and the Fréchet Inception Distance (FID) (Heusel et al. 2017). Precision and Recall are measured by 10k real images and 10k generated images following (Dhariwal & Nichol| (021). Precision is the ratio of generated images belonging to real images' manifold. Recall is the ratio of real images belonging to the generated images' manifold. The manifold is constructed in 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 (Szegedy et al.] (2166) pre-trained on ImageNet (Russakovsky et al. 2015), and it computes the logits for generated samples. If an instance is predicted with high confidence for specific class,
29

---

Published as a conference paper at ICLR 2025
and predictions 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, good Inception Score (IS) can only result from high-confidence predictions based solely on sample fidelity. We measure IS for 10k generated samples. FII approximates the generated and real samples Gaussians in the feature space of an Inception Network and measures the Wasserstein distance between them. Since it measures the distance between distributions, it emphasizes the importance of sample diversity and sample fidelity. For Table
4 we measure FID between 50k random samples from the FFHQ dataset and 50k randomly generated samples. For 'AE', we measure the FID between 50k random samples from the FFHQ dataset and generate samples that reconstruct the other 50k random samples from FFHQ. In Table we measure the FID between 10k random samples from the CelebA and 10k randomly generated samples. We utilize Pittps://gt gtttps://githu.... thub. measure Precision, Recall and IS. We utilize Muuliizehhttai htt /ttp//gtttthu..... wtmc GarGaaaPaaae www.a//caaaaaaaaaaa to measure FID, In Table we loaded checkpoints for all baselines (except the generative prior of PDAE, we train to fill performance) and conducted evaluations in the same NFEs. Table 14 shows the performance under various NFEs. For CelebA training, we dim(z) 256 following (Wang et al. 2023). while FFHQ training employs a dim(z) 512 following Zhang al. 2022).
com htt.///peaaag///m ded- fus on to
fid
(Preechakul et al. 2022
C.3 ALGORITIM
This section presents the training and utilization algorithms of DBAE. Algorithm I outlines the procedure for minimizing the autoencoding objective, LAF- Algorithm explains the method for reconstruction using the trained DBAE. Algorithm describes the steps for training the generative prior, Pw- Algorithm exxlainn the procedure for unconditional generation using the trained DBAE and generative prior.
Algorithm 3: Latent DPM Training Algorithm
Input: Ence, data distribution Iama (xo), drift term f, volatility term Output: Latent DPM score network Sw
while not converges do
Sample time from '0,T]
Xu Quma (Xu)
z Enca (Xv)
Z: ~ qi (Ze 2
+ ((()) | (2,, f) Vz, (z,12)
Update by L using the gradient descent method
end
Algorithm 4: Unconditional Generation Algorithm
Input: Dec,, latent score network sw, score network latent discretized time steps (G)) discretized
time steps [t.).N
2T
1)
for No. ..., do
L Update Zı, using Eq.
XT Dec for N ..., do
L Update xt, using Eq.
Output: Unconditioned sample X
C.4 COMPUTATIONAL COST
This section presents a computational cost comparison among diffusion-baseed representation learn ing baselines. Table 7 compares DDIM (Song al. 2021a). DiffAE (Preechakul et al. 2022). PDAE (Zhang et al. 2022), and DBAE in terms of parameter size, training time, and testing time. DDIM requires only score network (99M), resulting in minimal parameter size. DiffAE involves
30

---

Published as a conference paper at ICLR 2025
Table Computational cost comparison for FFHQ128. Training time is measured in milliseconds per image per NVIDIA A A10 (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
Song et al.
2021a)
99M
9.687
0.997
DiffAE
Preechaku et al. 2022)
129M
12.088
1.059
PDAE
(Zhagg al. 2022
280M
12.163
1.375
DBAE
161M
13.190
1.024
Table Computing costs for Xy inference.
NFE (4)
Total
time (4)
Method
Enc
Decg
Sg
(ms)
PDAE
500
500
688
DiffAE
1
250
265
DBAE
0
0.31
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 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 longer sampling time.
D ADDITIONAL EXPERIMENTS
D.1 DOWNSTREAM INFERENCE
Figure shows the attribute-wise Average Precision (AP) gap between PDAE (Zhang et al. 2022) and DBAE. As discussed in Section PDAE suffers from information split problem that xy contains facial or hair details. The resulting attribute-wise gain aligns with that analysis with Figure 3 Figure [d shows the absolute attribute-wise AP of DBAE performance across the training setting varies on the encoder (déterministic/toccaatti and training datasets (CelebA training set FFHQ). The attribute-wise performance similar across the training configurations. Table 9 shows the comparsion to the other baseline Diffuse VAE (Pandey et al. (022). From the two-stage paradigm of Diffuse VAE, its latent quality is only from the latent representation capability of the VAE module.
This is an aligned result from the poor performance of f3- 3-TCVAE in Table Note that the image crop for CelebA Diffuse is not exactly the same as our setting.
Table 9: Linear-probe attribute prediction quality comparison for models trained on CelebA and CIFAR-10 with dim(z) = 512. The best and second-best results are highlighted in bold. We evaluate 5 times and report the average.
CelebA Pearson's 's (1)
CIFAR-10 AUROC (t)
Method
AP (f)
MSE (1)
DiffuseVAn VAE (Pandey et al.
2022)
0.395
0.325
0.618
0.736
DBAE
0.655
0.643
0.369
0.836
31

---

Published as a conference paper at ICLR 2025
0.01
alury
5.0 Clack Shadow
Wiaring, Hot
Gray
Receding Macoding_Haiiin Hairino
Bangs
Husihy dushy Eyebrows
dlond Heir
Wary Hair
Wearing Ceriinss
Silehurns
Eyeglessees
Rosy Aasy_Chool Chooks
Gosteo
Struight Heir
Chubby
Double Chin
Black Hair
intrective
Arched Eyebrows
Wearing Weering_eeekee Neckace
Mustache
heavy, Makeup
Young
Male
dags, Underr Eyes
Oval Oval fece
Pale Skin
Warring Lipatick
High High_Cheelae
Narrow Eyes
Smiling
No fleard
Lips
Cheekbones
Singhtly Opwn
Figure 8: Attribute-wise AP gap between PDAE and DBAE-d trained on CelebA. DBAE-d performs
better for all 40 attributes.
(a) DBAE-d trained on CelebA
(b) DBAE trained on CelebA
(c) DBAE-d trained on FFHQ
(d) DBAE trained on FFHQ
Figure 9: Attribute wise Average Precision across the training configuration of DBAF.
32

---

Published as a conference paper at ICLR 2025
Mean
(9.67)
Mean
(9.99)
Mean
Std (0.00)
Mean
Std (6.52)
DifTAE statistics (ODE/A SDE)
DBAE-d statistics (ODE SDE)
Reconstruction statistics with inferred 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 the pixel space.
Figure 10:
D.2 RECONSTRUCTION
The sampling step important for practical applications (Lu et al., 2022 Zheng et al. 2024). We compare the reconstruction results across various sampling steps among the baselines. Ta- bles 12 and 13 shows the results. The proposed model performs the best results among all NFES in (10, 20, 50, 100). We borrow the performance of DDIM, DiffAE from (Preechakul et al. 2022). We manually measure for PDAE (Zhang et al.] 2022) using an official checkpoint in https://gi www.ttttu...om com/ cwwzcj PDEE, Figure 10 shows the reconstruction statistics for single image with inferred z. Due to the information split on XT, Diff'EE shows substantial variations even when 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 stored solely at XT through the ODE path. Table TO shows that the reconstruction quality compare to Diffuse (Pandey et al. (022). Since Diffuse VAE also requires to sample random Xy for the generation, this framework also suffers from information split problem. That is the reason for poor reconstruction quality. Table [T the reconstruction quality for Horse and Bedroom datasets, which surpasses the DiffAE.
Table 10: Autoencoding reconstruction quality comparison with Diffuse VAE with 512-dimensional latent variable, the one yielding the best performance is highlighted in bold.
CelebA LPIPS (4)
Method
SSIM (1)
MSE (L)
DiffuseVAB (Pandey al. 2022
0.836
131
0.018
DBAE
0.990
0.014
4.86e-4
Table 11: More results on autoencoding reconstruction quality comparison with DiffAE with 512- dimensional latent variable, the one yielding the best performance is highlighted in bold.
Horse
Bedroom
Method
SSIM (1)
MSE (L)
SSIM (1) MSE (L)
DiffAE (Preechakul et al. 2022)
0.857
0.025
0.910
0.017
DBAE
0.902
0.012
0.948
0.007
D.3 UNCONDITIONAL GENERATION
The sampling step is also important for unconditional generation (Lu et al. 2022
Zheng al.) 2024). We reduce the NFE= in Table NFE=500 and NFE=250 Table 14 As the number of function evaluations (NFE) decreased, DDPM (Ho al. 2020 showed 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 (Kim et al.] 2024 (Sadat et al. 2024). In the of generative models, there is trade-off between fidelity and diversity (Dhariwal
33

---

Published as a conference paper at ICLR 2025
Table 12: Autoencoding reconstruction quality comparison. All the methods trained on the FFHQ dataset and evaluated on the 30K CelebA-HQ dataset. Among tractable and compact 51 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.
Method
Tractability
Latent dim (U)
NFT==1
NFB=20
NFE=50
NF 100
DDIM (Inferred Nr))
Song et al
2021a
49,152
0.600
0.760
0.878
0.917
(Inferred Xy) Pieechaku
49,66
0.927
0.991
PDAE (Inferred Xx)
Zhang et
49,664
0.822
0.901
0.966
0.987
DifFAE (Random Xx)
Preechakul et
2022
512
0.707
0.695
0.683
0.677
PIAA (Random Xy)
hang al,
0.713
DBAE
512
0,904
0.909
0.916
0.920
DBAE-d
512
0.884
0.920
0.945
0.954
Table 13: Autoencoding reconstruction quality comparison. All the methods are trained on the FFHQ dataset and evaluated the 30K CelebA-HQ dataset. Among tractable and compact 51 512-dimensional latent variable models, the one yielding the best performance was highlighted in bold, followed by an underline for the best performer. All the metric is MSE.
Tractability
Latent dim (U) NFT-
NFE=20
DDIM (Inferred Xy:)
Song et al.
2021a
X
49.152
0.019
0.008
0.003
0.002
DifFAE (Inferred Xy)
PPrechakui et al 2022
49,664
0.001
0.001
0.000
0.000
PDAE (Inferred Xr)
Zhang et
49,664
0.001
DifAA (Random X7)
Preechakul et Thang et al.
2022
512
0.006
0.007
0.007
0.007
PIDAH (Random Xy)
2022
512
0.004
0.005
0:005
0.005
DBAE
512
0.005
0.005
0.005
DBAE-d
512
0.006
0.003
0.002
0.002
& Nichol 2021 Therefore, providing balance between these virtues is important. We offer an option based on DBAE. The h-transformed forward SDE we designed in Eq. (IO) is governed by the determination of the endpoint distribution. can achieve smooth transitions between DiffAE and DBAE in terms of Xr distribution. Modeling PΦ.w [xy] (x)) Gaussian distribution (with learnable mean and covariance) with a certain variance or higher can also be considered as indirect approach.
set endpoint distribution as Eq. (124),
XT
X lo. (xT) + (1 - A) x IO,
(124)
Table 14: Unconditional generation with reduced NFE {250, 500) on FFHQ '+AE' indicates the use of the inferred distribution (2) instead of Pw (z)
NFE 500
NFE 250
Method
Prec (t)
(1)
FID 50k ())
Rec (f)
Prec
IS (?)
FID 50k (J)
Roc (1)
DDIM
Song el ul
10
0.705
3.16
11.33
0.439
0.706
3.16
11.48
0.453
DDPM
Ho el
2020
0.589
2.92
22.10
0.251
0.390
2.76
39.55
0.093
Preechakul el
2022
0.451
10.24
PDAE
Zhang
0.687
2.24
46.67
0.175
0.709
2.25
44,82
0.189
DBAE
0.774
3.91
11.71
0.391
0.758
3,90
13.88
0.381
DiTAAAAA
0.750
3.61
3.21
0.689
0.750
3.61
3.87
0.666
PDAE+AE
0.710
7.11
0.598
0.721
3.54
DBAF+AE
0.748
3.57
1.99
0.702
0.731
3.58
3.36
0.694
34

---

Published as a conference paper at ICLR 2025
D.4 RESULTS WITH INTEL GAUDI V2 HARDWARE.
We conducted evaluations across various infrastructures to assess experimental reproducibility. The performance of the trained model (DBAE-d) was evaluated both the Nvidia A100 and Intel Gaudi v2 chips. The reconstruction results for both chips are presented in Table [5] Re- construction performance on each chip was assessed using various metries, revealing negligi- ble emrors across all metrics. To facilitate reproducibility. we provide the code at https 7//ithuul Com/NAVER CNTEL-C---LbLbaaaaaaa ddbae for reproducing our experiments on Intel Gaudi v2 chips.
Table 15: Regenerated results of Table 2
across multiple hardwares.
Hardware
SSIM (t)
LPIPS (1)
MSE (L)
Nvidia A100
0.953 Intel Gaudi
0.072
2.49e-3
0.956
0.073
2.47e-3
D.5 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 X and generated image Xg each. Figure [3 compares the interpolation results with PDAE (Zhang al., (022) and DiffAE (Preechakul al. 2022) under tractable inference condition. PDAE and DiffAE result in unnatural interpolations without inferring Xy. compared DBAE.
Attribute Manipulation Figure 15 shows additional manipulation results using linear elassifier, including multiple attributes editing on single image. Figure 14 provides the variations in the manipulation method within DBAE. The top row utilizes the manipulated Xy both for the starting point of the generative process and score network condition input. The bottom row utilizes manipulated only for the score network condition input, while the starting point remains the original image's 's xT. Using manipulated Xy both for starting and conditioning results in more dramatic editing, and expect be able to adjust this according the user's 's desires.
Generation Trajectory Figure 16 shows the sampling trajectory of DBAE from Xy to X 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.
35

---

Published as a conference paper at ICLR 2025
Figure 11: FFHQ interpolations results with corresponding endpoints Xy. The leftmost and rightmost images real images.
36

---

Published as a conference paper at ICLR 2025
Figure 12: Horse and Bedroom interpolations results with corresponding endpoints Xy. The leftmost and rightmost images are real images.
37

---

Published as a conference paper at ICLR 2025
IF
IE
Figure 13: FFHQ interpolation comparison: PDAE Zhang al. 2022 (top), DiffAE
Prreehakku
et al. 2022) (middle) and DBAE (bottom).
Figure 14: Attribute manipulation on FFHQ using linear classifier and corresponding endpoints xy- The top results utilize the manipulated Xy both as the starting point of the sampling trajectory and as a condition input to the score network. The bottom results use the manipulated X7 solely as the condition input and maintain the original Xy as the starting point of the sampling trajectory. All the middle images are the original images.
38

---

Published as a conference paper at ICLR 2025
Smile
Male
Wavy Hair
Young
(a) Smooth traversals in the direction of attribute manipulation. All the middle images are the original images.
Original
Recon.
Smile
Sad
Female
Young
Old
(b)
Multiple attribute manipulation on single image.
Figure 15: Attribute manipulation using linear classifier FFHQ and CelebA-HQ.
39

---

Published as a conference paper at ICLR 2025
XT
X
(a) Sampling trajectory of DBAF trained on FFHQ.
XT
X
(b) Sampling trajectory of DBAE trained Horse.
X7
Sampling trajectory of DBAE trained Bedroom.
Figure 16: Stochastic sampling trajectory of DBAE trained on various datasets.
40

---

Published as a conference paper at ICLR 2025
3
3
5
5
(a) Generated endpoints XT
(b) Generated images Xu
Figure 17: Uncurated generated samples with corresponding endpoints from DBAE trained on FFHQ with unconditional generation.
41

---

Published as a conference paper at ICLR 2025
(a) Generated endpoints XT
(b) Generated images Xu
Figure 18: Uncurated generated samples with corresponding endpoints from DBAE trained on CelebA with unconditional generation.
42


Reviews:
[{"review_id": "eCxZ3vnY8G", "reviewer": "Reviewer_EWF8", "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\".\nThe 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.\nThus, 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.\nThis makes $z$ less useful as a representation of $x_0$ in downstream tasks (e.g., classification, clustering).\n\nThe 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$).\nThe paper discusses and empirically evaluates training objectives of their proposed architecture for representation learning, reconstruction, and unconditional generation.", "questions": "What is the significance of the information split problem in unconditional generation (see \"Weaknesses\")?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper addresses an important issue in representation learning with diffusion models.\nThe 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.\n\nThe paper is remarkably well organized.\nDespite 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.\nThe addressed problem is introduced in a didactic way (first with examples, then gradually more formal), and its proposed solution is well motivated.", "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.\n\nI 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).\n\nI did not understand the significance of the information split problem in unconditional generation (Section 4.4.2).\nI appreciate that the discussion of the objective function (Section 4.4) distinguishes between different use cases of diffusion models.\nAnd I can see that the described information split problem can be an issue for representation learning.\nBut 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$?\nWe would sample $x_T$ anyway.\nYet, empirically, it seems like the proposed DBAE+AE model outperforms the baselines even for unconditional generation.\n\nFurther, as a minor point, Eq. 9 seems suspicious to me.\nWhile 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.\nI 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\".\n\nApart 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.\nThis is a shame because it diminishes the otherwise excellent presentation.\nI 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.\nAnd 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.\nSome 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).\nI strongly recommend running the entire paper through a grammar checker, or asking an LLM to detect possible grammar issues in each paragraph.\nIt would make this otherwise great work more accessible to a broad (international) audience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JnzVMASUY0", "reviewer": "Reviewer_GqCP", "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. \n\nWith theoretical guarantees, the method outperforms the SOTA methods on various datasets and settings.", "questions": "1. Could you provide more intuition behind Eq.10?\n\n2. 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?\n\n3. 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Quality and clarity: the paper is well-written and the idea is motivated and easy to follow. \n\nOriginality: 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. \n\nThe 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.\n\nThe 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.", "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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G1cZ3vzjZq", "reviewer": "Reviewer_Ffum", "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.", "questions": "See above", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n\n2. There are thorough experiments to demonstrate the usefulness and performance of the proposed model.", "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?\n\n2. 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?\n\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eE6rK3WYWu", "reviewer": "Reviewer_rm4h", "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.", "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$?\n\nLine 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$”?\n\nLines 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?\n\nLine 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?\n\nLines 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.\n\nLines 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.\n\nLines 203–204: Since the references here were cited just two paragraphs above, the authors may want to omit them to save space.\n\nLine 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.\n\nLines 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.\n\nSection 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.\n\nSection 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?\n\nLine 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?\n\nLine 317 (Theorem 1): The term “linear SDE” appears without prior explanation. Would the authors consider providing a brief definition?\n\nLine 331: Consider removing “the” before “Section 4.4.1” for grammatical clarity.\n\nLine 342: Replacing “optimize” with “minimize” could provide greater clarity regarding the objective.\n\nLines 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?\n\nSection 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.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 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.", "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.\n\n**Clarifying Motivation and Significance in the Introduction:**\nThe 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.\n\n**Improving Precision in Terminology and Expression:**\nSeveral 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).\n\n**Providing Precise Description of Methodology/Tasks in Sections 4 & 5:**\n- *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.\n-  *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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)