Diffusion Bridge AutoEncoders for Unsupervised Representation Learning

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 extract a corresponding representation from data and adjust the dimensionality of a latent variable z. Meanwhile, this auxiliary structure invokes an information split problem; the information of each data instance x 0 is divided into diffusion endpoint x T and encoded z because there exist two inference paths starting from the data. The latent variable modeled by the diffusion endpoint x T has several disadvantages. The diffusion endpoint x T 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 x T inference through a feed-forward architecture. This structure creates an information bottleneck at z, ensuring that x T depends on z during its generation. This results in z 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 demonstrates the effectiveness of 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://github.com/aailab-kaist/DBAE.