Diffusion Bridge AutoEncoders for Unsupervised Representation Learning

We appreciate for acknowledging the strengths of our paper in terms of motivation, methodology, and experimental results. We truly appreciate your thoughtful feedback on the presentation, along with detailed suggestions for improvement. Your suggestions have been incorporated into the revised manuscript and are highlighted in blue.

[Q1]. Clarifying Motivation and Significance in the Introduction

[Response to Q1]

Thank you for the suggestion. To highlight the significance of solving the information split problem, we add the following in the introduction.

• Line 41: We added a statement to highlight that diffusion models with auxiliary encoder framework are dominating generative representation learning studies.
• Lines 53-54: We added a statement that diffusion models with auxiliary encoder framework combine the strengths of diffusion models and VAEs.
• Lines 74-79: We added a fact that the primary goal of generative representation learning is reconstruction, and reconstruction capability is important to facilitate downstream inference, attribute manipulation, and interpolation. We state that the information split problem hinders the reconstruction capability.

[Q2]. Improving Precision in Terminology and Expression

[Response to Q2]

Thank you for the thoughtful suggestion. We modified the expression you mentioned for clarity. We modified the following.

• Line 45: We modified it to “to the endpoint $x_T$“ as you suggested.
• Lines 48-49: We modified it to “encoding $x_0$ into $x_T$” as you suggested.
• Lines 160–161: We modified “VP" into its full name, variance preserving.
• Lines 192-201: First, we would like to inform you that Eq.9 has been slightly revised following Reviewer EWF8's feedback Q2. The revised equation remains equivalent to the original. The previous KL term has now been replaced with a Cross-entropy term, but the explanation flow remains the same. We have modified the explanation of the reason why the Cross-entropy term is problematic. We explain: 1) how the discrepancy between $q _{\theta}^{ODE}(x_T|z, x_0)$ and $p _{prior}(x_T)$ affects the value of $CE(q _{\theta}^{ODE}(x_T|z, x_0) || p _{prior}(x_T))$, 2) how the value of $CE(q _{\theta}^{ODE}(x_T|z, x_0) || p _{prior}(x_T))$ affects the mutual information bound, 3) why the discrepancy between $q _{\theta}^{ODE}(x_T|z, x_0)$ and $p _{prior}(x_T)$ is inevitable.
• Lines 204-205: We removed the redundant citation as you suggested.
• Line 339: We removed the unnecessary “the” before “Section 4.4.1” for clarity.
• Line 350: We modified from “optimize” to “minimize” for clarity.
• Theorem 1: We added the explanation of “linear SDE” as a footnote.
• Title of Section 4.1: We have modified the title to “Encoding from $x_0$ to $x_T$ conditioned on $z$,” specifying which latent variables are inferred and clarifying their relationships.
• Equations: We modified every “eq” to “Eq” for readability. We modify the equations to appear before the referencing.

We truly appreciate the reviewer taking the time to review our response. We are well aware of the reviewer's positive points in the original review. In our previous response, we addressed the importance of the information split problem in the application cases presented in Section 5.

As you mentioned, our motivation is based on the premise that having latent representations, $z$ and $x_T$, together is beneficial in the baseline. In addition to the literature and experiments, we also aimed to explicitly explain why the use of both ($z$, $x_T$) is beneficial in the baseline in (Section 3) Motivation: Information Split Problem through Eq. 9 and further explanations in Section A.4.1. The followings elaborates the Section 3 with more details to explains why the use of both inferred ($z$, $x_T$) is beneficial in the baseline?

1. Reconstruction with random $x_T$

To compute the reconstruction quality with inferred $z$, given by $\mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[\log{p _{\theta}(x_0|z)}]$, it inherently requires drawing $x_T$. The reconstruction quality depends on how $x_T$ is drawn. Consider $p _{\theta}(x_0|z) = \int{p _{\textrm{prior}}(x_T) p _{\theta}^{ODE}(x_0|z,x_T) }dx_T$, which implies drawing $x_T \sim p _{\textrm{prior}}(x_T)$. Then, the following inequality holds (please refer to Eqs.102-109 for the derivation).

(Eq.D1). $\mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[-CE(q _{\theta}^{ODE}(x_T|z,x_0) || p _{\textrm{prior}}(x_T))] \leq \mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[\log{p _{\theta}(x_0|z)}]$

2. Reconstruction with inferred $x_T$

Instead of a random draw $x_T \sim p _{\textrm{prior}}(x_T)$, $x_T$ can also be obtained by encoding from $x_0$ to $x_T$ which implies $x_T \sim q _{\theta}^{ODE}(x_T|x_0,z)$. From this drawing, we can get the following inequality.

(Eq.D2). $\mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[-CE(q _{\theta}^{ODE}(x_T|z,x_0) || q _{\theta}^{ODE}(x_T|z,x_0))] \leq \mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[\log{p _{\theta}(x_0|z)}]$

3. Compare two reconstruction bounds

The cross-entropy becomes larger as the discrepancy between two distributions becomes larger. This resulted in the LHS of Eq.D2 being larger than the LHS of Eq.D1. This implies that reconstruction with inferred $x_T$ (with ODE solving) provides a tighter bound on reconstruction quality. Table 2 shows the empirical comparison that DiffAE (inferred $x_T$) performs better than DiffAE (random $x_T$), which supports inequality in Eq.D3.

(Eq.D3). $\mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[-CE(q _{\theta}^{ODE}(x_T|z,x_0) || p _{\textrm{prior}}(x_T))] \leq \mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[-CE(q _{\theta}^{ODE}(x_T|z,x_0) || q _{\theta}^{ODE}(x_T|z,x_0))]$

4. Reconstruction quality and representation quality

The reconstruction quality with inferred $z$ bounds the mutual information between $z$ and $x_0$ (please refer to Eq. 101 in the manuscript for details). Therefore, tighter bounds on reconstruction quality also result in tighter bounds on the mutual information. The higher mutual information indicates better representation quality, which can facilitate downstream inference (Sec 5.1), disentanglement (Sec 5.3), and attribute manipulation (Sec 5.6).

(Eq.D4). $\mathbb{E} _{q _{\textrm{data}}(x_0), q _{\phi} (z|x_0)}[\log{p _{\theta}(x_0|z)}] +\mathcal{H}(q _{\textrm{data}}(x_0)) \leq MI(x_0, z)$

We appreciate the reviewer for the constructive and thoughtful feedback. We answer for reviewer’s comments below.

[Q1]. 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?

[Response to Q1]

Thank you for the helpful feedback. Doob’s $h$-transform is applicable when we wish to modify the original SDE to terminate at a certain random variable (which has a distribution). Please refer to Theorem 7.11 of [C1] for the explanation of Doob’s h-transform. Its derivation does not restrict the endpoint to being a fixed point but allows it to be a random variable that has a probability distribution. Section 7.3. of [C2] also discusses doob’s h-transform with stochastic endpoints.

[Q2]. 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?

[Response to Q2]

Thank you for the good question. The objective for reconstruction is $\mathcal{L} _{AE}$ in Eq. 17, and the objective for generative modeling is $\mathcal{L} _{AE} + \mathcal{L} _{PR}$, as shown in Theorem 3. As we already explained in lines 350–356, we optimize $\mathcal{L} _{AE}$ with respect to the parameters of the encoder ($\phi$), decoder ($\psi$), and score network ($\theta$), which enables reconstruction. With the fixed parameters ($\phi, \psi, \theta$), we optimize $\mathcal{L} _{PR}$ with respect to the additional parameter $\omega$, which estimating the generative prior distribution $p _{\omega}(z)$. Consequently, the same parameters ($\phi, \psi, \theta$) are used for both experiments, while the generative prior ($\omega$) is additionally used specifically for generative modeling.

If we set the generative prior distribution $p_{prior}(z)$ to be Gaussian, it is possible to optimize $\mathcal{L} _{AE}+\mathcal{L} _{PR}$ simultaneously, performing reconstruction and generative modeling at the same time (similar to VAE). However, this approach suffers from the same issues below faced by VAEs.

• Representation Learning: The $L_{PR}$ in Eq.18 forces the matching of the encoding distribution $q_{\phi,\psi}(x_T|x_0):= \int{q_{\psi}(x_T|z)q_{\phi}(z|x_0)dz}$ and generative prior distribution $p_{\psi}(x_T):=\int p_{\psi}(x_T|z) p_{prior}(z)dz$, which occur over-regularization on the encoding $q_{\phi}(z|x_0)$ when $p_{prior}(z)$ set to be Gaussian (which hinders to learn semantic meaning). Disjoint training enables optimization on encoding parameters ($\phi,\psi$) without being regularized.

• Generative Modeling: The mismatch between the prior distribution $p_{prior}(z)$ and the aggregate posterior $q_{\phi}(z)$ is one of the reasons for the low performance of VAE in unconditional generation [C3]. Imitating learned aggregate posterior $q_{\phi}(z)$ with a powerful generative model is an effective way to mitigate the discrepancy.

Therefore, the separate training approach was used in the baselines [C4, C5] and we follow this.

[C1] Särkkä, S., & Solin, A. (2019). Applied stochastic differential equations (Vol. 10). Cambridge University Press.

[C2][ICLR 2024] Denoising Diffusion Bridge Model.

[C3][NeurIPS 2021] A contrastive learning approach for training variational autoencoder priors

[C4][CVPR 2022] Diffusion Autoencoders: Toward a Meaningful and Decodable Representation

[C5][NeurIPS 2022] Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models.

We appreciate the reviewer's sincere and helpful feedback. We answer for reviewer’s comments below.

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

[Response to Q1]

Thank you for your feedback. The conditional dependencies between all random variables are determined by the proposed graphical model design in Figure 1-(c), DBAE follows an encoding path ($x_0 \rightarrow z \rightarrow x_T$) supporting our claim that $z$ becomes an information bottleneck. This conditional dependency naturally influences the loss term in $\mathcal{L}_{AE}$.

$\mathcal{L}_{AE}=\frac{1}{2}\int_0^T \mathbb{E} _{q^{t} _{\phi,\psi}(x_0,x_t,x _T)}[\lambda (t) || x _{\theta}^0(x_t, t, x_T) - x_0||_2^2],$

where the probability under the expectation $q^t _{\phi,\psi}(x _0,x _t,x _T)$ is decomposed according to the graphical model structure in Figure 1-(c).

$q^{t} _{\phi,\psi}(x_0, x_t, x_T) =\int q _{data}(x_0) q _{\phi}(z|x_0) q _{\psi}(x_T|z) q_t(x_t|x_T, x_0) dz,$

where $q _{\phi}(z|x_0)$ is defined by the encoder, $q _{\psi}(x_T|z)$ is defined by the decoder, and $q _t(x_t|x_T,x_0)$ is derived from the forward process in Eq.10. We show the objective $\mathcal{L} _{AE}$ (which inherently incorporates our conditional dependency design) has a relationship with mutual information between $x_0$ and $z$ in Theorem 2. The original manuscript mentioned the decomposition of $q^{t} _{\phi,\psi}(x_0, x_t, x_T)$ separately in line 240, and Eq.10. We have included the entire decomposition in Theorem 1 to help readers.

[Q2]. The intuition behind Eq.10 is missing, why new SDE is required?

[Response to Q2]

Thank you for your feedback. Figure 1-(a, c) show the graphical model of baseline and DBAE. In the baseline approach, the encoding paths $(x_0 \rightarrow x_T)$ and $(x_0 \rightarrow z)$ existed independently, making it possible to reuse the vanilla diffusion process for $(x_0 \rightarrow x_T)$. However, in our new structure, the encoding path is defined as $(x_0 \rightarrow z \rightarrow x_T)$, which introduces a predefined coupling between $(x_0, x_T)$.

As a result, a new forward process starting at $x_0$ and ending at $x_T$ needed to be defined. To address this, we employed Doob's $h$-transform to define the new forward process, as shown in Eq. (10). As explained in Section 2.3, Doob's $h$-transform is a transformation that can be used to modify an existing SDE to ensure it terminates at a desired endpoint. The original manuscript explained the intuition behind the need for the newly defined forward SDE in the main text, now found in lines 239–242 of the modified manuscript.

[Q3]. 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?

[Response to Q3]

Thank you for your interesting suggestion. The encoding paths of AE and DBAE are the same, $(x_0 \rightarrow z)$. However, there is a difference in the process of generating the final sample $x_0$ starting from $z$. AE typically predicts the final sample through a single-step feed-forward architecture $(z \rightarrow x_0)$. In contrast, as explained in Section 4.2, DBAE involves a more complex generation process, which comprises a single-step feed-forward $(z \rightarrow x_T)$ and the reverse process $(x_T \rightarrow x_0)$ in Eq. 12.

If DBAE denoises Eq. 12 in one step and treats the decoder ($\psi$) and score network ($\theta$) as a single neural network, DBAE aligns exactly with AE. However, by iteratively performing Eq. 12 over multiple steps, we can generalize the generative process to be more complex, thereby creating opportunities to improve generation performance compared to AE.

We appreciate the reviewer for the constructive and thoughtful feedback. We answer for reviewer’s comments below.

[Q1]. What is the significance of the information split problem in unconditional generation?

[Response to Q1]

Thank you for your valuable feedback. As the title of the paper implies, the primary focus is on improving the representation learning capability based on reconstruction (we further highlighted in the introduction, lines 73–78, in the modified manuscript). The enhanced representation capability facilitates downstream tasks such as inference (Sec. 5.1), reconstruction (Sec. 5.2), disentanglement (Sec. 5.3), interpolation (Sec. 5.5), and attribute manipulation (Sec. 5.6). As you pointed out, the information split problem is not a significant issue in unconditional generation. Please consider the generative modeling capability as an additional feature of our model.

[Q2]. Formulation of Eq.9. without "cancellation of infinities"

[Response to Q2]

Thank you for your insightful feedback. You are correct that the phenomenon of " cancellation of infinities " occurs in the previous formulation. We have revised Eq. 9 in the modified manuscript to an equivalent formulation that avoids the issue you mentioned. The updated equation still explains the information split problem. Please refer to Eqs. 106–109 in Appendix A.4.1, where we demonstrate the equivalence between the previous and new formulations.

[Q3]. Grammar errors

[Response to Q3]

We apologize for the lack of completeness in some sentences, which may have made the paper difficult to read. We sincerely express our gratitude for thoroughly reading the paper and providing valuable feedback. We have comprehensively revised the manuscript, including the parts you mentioned. The modified parts related to your feedback are highlighted in blue. You can refer to lines 131–133, 160–162, and 292–296 in the revised manuscript.