Exploring Diffusion Time-steps for Unsupervised Representation Learning

Abstract: Representation learning is all about discovering the hidden modular attributes that generate the data faithfully. We explore the potential of Denoising Diffusion Probabilistic Model (DM) in unsupervised learning of the modular attributes. We build a theoretical framework that connects the diffusion time-steps and the hidden attributes, which serves as an effective inductive bias for unsupervised learning. Specifically, the forward diffusion process incrementally adds Gaussian noise to samples at each time-step, which essentially collapses different samples into similar ones by losing attributes, e.g., fine-grained attributes such as texture are lost with less noise added (i.e., early time-steps), while coarse-grained ones such as shape are lost by adding more noise (i.e., late time-steps). To disentangle the modular attributes, at each time-step t, we learn a t-specific feature to compensate for the newly lost attribute, and the set of all 1,...,t-specific features, corresponding to the cumulative set of lost attributes, are trained to make up for the reconstruction error of a pre-trained DM at time-step t. On CelebA, FFHQ, and Bedroom datasets, the learned feature significantly improves attribute classification and enables faithful counterfactual generation, e.g., interpolating only one specified attribute between two images, validating the disentanglement quality. Codes are in https://github.com/yue-zhongqi/diti.

• The paper introduces DiTi, a method that leverages diffusion time-steps in denoising diffusion models to capture and disentangle progressively lost data attributes.
• The approach partitions the latent space into time-step specific feature subsets that compensate for reconstruction errors incurred during each diffusion phase.
• Empirical results indicate that modular feature vectors enhance targeted manipulation of attributes, outperforming methods that rely on a single feature representation.

The paper "Exploring Diffusion Time-steps for Unsupervised Representation Learning" (2401.11430) introduces a method called DiTi that utilizes diffusion time-steps as an inductive bias for unsupervised representation learning. The central concept revolves around the connection between diffusion time-steps in Denoising Diffusion Probabilistic Models (DMs) and the hidden modular attributes inherent in the data.

The application of diffusion time-steps for unsupervised representation learning and the model's approach to disentangling attributes are as follows:

1. Time-Steps and Attribute Loss: The forward diffusion process in DMs progressively adds Gaussian noise to the data across multiple time-steps ($T$). A core observation is that this process systematically eliminates different data attributes based on their granularity. Fine-grained attributes, such as texture, are lost earlier in the diffusion process (smaller $t$), while coarse-grained attributes, like shape, are lost later (larger $t$). Section 4.1 formalizes this in the paper's theory section and is supported by the theorem that relates attribute loss to the time-step. Relevant background can be found in "Denoising diffusion implicit models" [(Song et al., 2020) [cs.LG]] and "Denoising Diffusion Probabilistic Models" [(Ho et al., 2020) [stat.ML]].
2. Learning Time-Step-Specific Features: DiTi learns features that compensate for the attributes lost at each time-step. An encoder maps the original data ($x_0$) to a feature vector ($z$), which is partitioned into $T$ disjoint subsets ($z_1, ..., z_T$). At each time-step $t$, the model uses the first $t$ feature subsets ($z_1, ..., z_t$) to compensate for the reconstruction error of a pre-trained DM. The model learns a $t$-specific feature to compensate for the newly lost attribute, using the cumulative set of features up to time $t$ to recover from the reconstruction error made by the pre-trained DM at that time-step.
3. Disentanglement through Reconstruction Error Compensation: The training objective encourages each feature subset $z_t$ to capture the attribute(s) lost specifically at time-step $t$. By making each $z_t$ responsible for a specific set of attributes (those lost at time $t$), the model learns to disentangle the modular attributes. If the feature enables perfect reconstruction, it must capture the cumulatively lost attributes. The expanding set of features captures the expanding cumulatively lost attributes, which differs from methods like PDAE that use a single feature vector to compensate for the reconstruction error at all time-steps, potentially leading to entangled representations.
4. Modular Feature Vector Space: The DiTi approach results in a modular feature vector space where different dimensions (or subsets of dimensions) of the feature vector correspond to different underlying generative attributes. This modularity allows for targeted manipulation of specific attributes during counterfactual generation. The authors empirically validated that this design enables the disentanglement of attributes.