Automated Learning of Semantic Embedding Representations for Diffusion Models

TL;DR

This work investigates whether diffusion models can learn semantically meaningful embeddings beyond generation. It introduces DiER, which adds a timestep-aware encoder to a Diffusion Transformer to produce embeddings $v_t^s$ across diffusion timesteps, turning the diffusion process into $T$-level DAEs and enabling self-supervised representation learning. Empirical results across six datasets show DiER often achieves state-of-the-art linear probe accuracy at optimal timesteps, with mid-timestep representations providing the strongest semantic signals and dataset-dependent variability in the best timing. The findings support using DDPMs as general-purpose feature extractors while highlighting practical challenges like timestep selection and computational cost for downstream tasks.

Abstract

Generative models capture the true distribution of data, yielding semantically rich representations. Denoising diffusion models (DDMs) exhibit superior generative capabilities, though efficient representation learning for them are lacking. In this work, we employ a multi-level denoising autoencoder framework to expand the representation capacity of DDMs, which introduces sequentially consistent Diffusion Transformers and an additional timestep-dependent encoder to acquire embedding representations on the denoising Markov chain through self-conditional diffusion learning. Intuitively, the encoder, conditioned on the entire diffusion process, compresses high-dimensional data into directional vectors in latent under different noise levels, facilitating the learning of image embeddings across all timesteps. To verify the semantic adequacy of embeddings generated through this approach, extensive experiments are conducted on various datasets, demonstrating that optimally learned embeddings by DDMs surpass state-of-the-art self-supervised representation learning methods in most cases, achieving remarkable discriminative semantic representation quality. Our work justifies that DDMs are not only suitable for generative tasks, but also potentially advantageous for general-purpose deep learning applications.

Figures (13)

• Figure 1: The analogy between a traditional DAE and our proposed DiER, multi-level DAEs derived from DDMs. (Left) A traditional DAE add (multi-scale) noise and predict the original input, the encoder $\varepsilon$ and decoder $\delta$ directly compressing corrupted images into $h^s$ to construct a latent semantic space. (Right) DiER extends DDMs by encoding a vector $v_t^s$ for each noise level, and using a single neural network with parameters $\theta$ to predict noise, thus compressing semantics of different levels into directional vectors $v_t^s$ which provide information that assists ${x_t}$ in indirectly recovering the possible original sample ${\hat{x}_t}$ from its current corrupted state, forming multi-level DAEs.
• Figure 2: The pipeline used to analyze the representational capacity learning of diffusion models, where each block's internal structure remains uniform, serving for embedding through equivalent insertion computations. The encoding $v_t^s$ represents the encoding result of $x_0$ at any randomly sampled timestep $t$.
• Figure 3: Based on the current $t$ in the diffusion sampling process, encode $x_0$.
• Figure 4: The LPA results of DiER across all timesteps are depicted. The x-axis represents the timestep, while the y-axis denotes the LPA. The representations were evaluated every 100 timesteps, with the maximum $t$ being 999. Solid lines indicate the Top-1 Acc., while dashed lines represent the Top-5 Acc.
• Figure 5: The t-SNE visualization of representations extracted from test images of all datasets, where each class is represented by a distinct color.