Exploring Diffusion Time-steps for Unsupervised Representation Learning

Thanks for the in-depth comments and suggestions.

W1 DM for representation learning?

• Regarding computational cost, our DiTi takes slightly longer to train than SimCLR (40 hours vs 28 hours) on our evaluation datasets. However, this may not be the case in large-scale training, because 1) contrastively learning scales poorly w.r.t. data size, 2) our method distils knowledge from a pre-trained DM instead of learning from scratch.
• More importantly, our theory shows that DM has a unique mechanism for representation learning: attributes of different granularities are separated along DM time-steps, where fine-grained attributes are lost earlier than coarse-grained ones. This serves as an effective inductive bias for learning a disentangled representation, enabling robust inference and counterfactual generation, and has huge potential that can be further explored in future, e.g. see Kc2k-W2, Cjby-Q2.
• We also add that in general, learning representation with DM removes the need of data augmentation in contrastive learning, leading to several benefits: 1) As you pointed out, this helps preserve augmentation-related attributes. 2) Practitioners no longer needs to find the most suited augmentations, which would have been a daunting task given the heavy cost of contrastive learning especially in large-scale training. 3) This helps extend representation learning to modalities where data augmentations are not well-defined/work not as well, such as audio or other time-series data.

W2 Comparison with related work. As you rightfully point out, these works tackle representation learning from the perspective of building a more versatile generative model, e.g., to enable image editing or interpolation. We include them in the related work because they essentially still pursue a good representation to enable image editing. As discussed in Section 3.1, a good representation should enable robust inference in downstream tasks, as well as image editing on a targeted attribute (i.e., counterfactual generation). We will revise related work to clarify this.

W3 Similarity with PDAE. We clarify that the similarity with PDAE in implementation details (e.g., network design, hyper-parameters) is deliberate, such that emerged properties of disentangled representation are solely from our theory-inspired algorithm. We follow PDAE for its hyper-parameter choice, and our added hyper-parameter is solely to suit the need for disentangled representation learning, which has a clear-cut choice from our ablation study (Section 5.2).

Thanks for the in-depth comments and suggestions. We will fix the typo and address all points below.

W1 Modular editing consistency. No, we separately interpolate each feature subset corresponding to early/middle/late time-step range, as indicated on top of Figure 5. This is why we expect the edit of eyeglasses to happen in one range instead of multiple ranges. For example, if interpolating the feature subset in the middle range changes the “eyeglasses” attribute, interpolating the subset in the late range should only modify a more coarse-grained attribute without changing “eyeglasses”.

W2 & Q1 Plot of loss against time-steps. We include the plot in Appendix Figure A9. The reason for a small loss at a large time-step $t$ is that DM assigns a weight to each time-step (Eq. 2), and the weight is extremely small at a large $t$. One could also understand this from the $\epsilon$-formulation (equivalent to our reconstruction formulation), where the U-Net predicts the added noise in the input $x_t$. At a large $t$, $x_t$ is close to a pure noise, hence a trivial prediction by directly outputting $x_t$ already achieves a small loss.

W3 Comparison with SimCLR. Sorry for the confusion. We are not suggesting that our proposed DiTi relies on attribute correlation. Instead, we are suggesting possible reasons that SimCLR outperforms our DiTi on the classification accuracy of a few attributes (blue bars in Figure 4). SimCLR’s contrastive objective encourages different samples to have different features (i.e., representation must be an injective mapping). Hence the representation is encouraged to account for all differences between each sample pair to make their features maximally contrastive. This can help representation capture unnoticeable attributes (e.g., earrings). It may be worth exploring in future the merit of contrastive learning in context of diffusion models, e.g., as in [1].

[1] Discrete Contrastive Diffusion for Cross-Modal Music and Image Generation. ICLR 2023

Q2 Extension to conditional models. Good suggestion. We did some follow-up experiments to verify our theory on a text-conditioned DM. Specifically, our theory indicates that only fine-grained attributes are lost at a small time-step $t$. Hence given images from a fine-grained category (e.g., a specific aircraft model), we first obtain their noisy images at a small $t$, then we fine-tune a text-conditioned DM (Stable Diffusion) to reconstruct the original images from noisy ones conditioned on a chosen prompt (e.g., a rare token). If the reconstruction is accurate, the finetuned DM must make up for the attribute loss by associating the prompt with the category’s fine-grained attribute. In this way, to classify if a test image $x_0$ is from this category, we can intervene on the diffusion time-step to only focus on the fine-grained attribute, i.e., if DM can reconstruct it from $x_t$ at a small $t$, then we predict as “yes”, because the fine-grained attribute of $x_0$ matches that of this category. We verify that this intervention strategy on diffusion time-step is extremely effective in removing spurious correlations in the difficult few-shot learning task, e.g., on the challenging FGVC-Aircraft dataset, we significantly outperform previous few-shot learning methods using OpenCLIP (ViT-H/14 trained on LAION-2B) by up to 13%.

Table r3: $N$-shot performance on FGVC-Aircraft dataset.

Thanks for your review. We address all weaknesses and questions below.

W1 Model Training from scratch. We clarify that we only train the feature extractor from scratch, which is standard in representation learning. We use a pre-trained DM and freeze it in representation learning. Our training cost is similar with the recent related work PDAE and much faster compared to previous works, e.g., Diff-AE. Could you kindly clarify what you mean by “the model needs to be re-trained”?

W2 Training on LAION. We follow the standard evaluation datasets in this domain following PDAE and Diff-AE. Due to limited GPU access at the point of submission (4$\times$A100), it is too expensive to train on LAION. However, our method is grounded on a provable theory, which still holds on more diverse datasets. In fact, we are currently doing a large-scale pre-training based on this method.

Q1 Non-uniform Subsets. Yes, our method benefits by having more dimensions on the feature subset corresponding to time-step range 100-300. Please refer to the implementation details in Section 4.2, the ablation results in Section 5.2 and additional details in Appendix Section E.

Thanks for the rebuttal. I still recommend acceptance of this work, as reflected in my rating.

As a clarification for W1 and W2, different from works that take information from pretrained diffusion models that only require a very small amount of compute workload to realize or improve their downstream (e.g., [1]), this work presents another term in the loss for better representation learning that contributes to improved downstream, which makes it need to be re-trained (i.e., not using the information from the pre-trained models from large datasets such as LAION). The reviewer accepts the theoretical foundations for this work, but the need for retraining still leads to barriers to users who do not have the compute to train the representation on a large dataset but have access to public weights of diffusion models. The reviewer is glad that the authors are training a model on large datasets with the method presented in their work and encourages the public release of the trained model.

[1] Diffusion Hyperfeatures: Searching Through Time and Space for Semantic Correspondence. NeurIPS 2023. https://arxiv.org/abs/2305.14334

W3 ImageNet-style recognition. Unfortunately, as our representation is trained on face or bedroom dataset, it cannot be directly used for ImageNet evaluation. As we have limited GPU access at the point of submission, it would also be too expensive to train a representation on ImageNet. However, our theory and method can be extended to a larger and more diverse dataset, and we are currently working on large-scale pre-training with this method.

W4 Background information. As rightfully pointed out, our DiTi will disentangle background information in the feature subset corresponding to a late time-step $t$. However, we highlight that this is in fact desirable, because:

1. The background information may be useful in downstream tasks, e.g., in semantic segmentation on road images, the model needs to classify the type of each background pixel into road, pavement, etc.
2. If background is irrelevant for a downstream task, one can perform human intervention to discard the feature subset corresponding to late time-steps, such as the above few-shot learning examples on fine-grained categories.

Q1 Discussion for [1]. The paper shares a similar high-level idea, i.e., learning a feature as an input condition for DM to enable perfect reconstruction. Key differences: 1) [1] relies on empirical evidences and hypothesis for the concept of attribute loss, while we provide a rigorous theoretical analysis. 2) [1] learns an infinite-dimension feature, while we learn a compact feature that is easy for downstream leverage.

Thank you for the detailed review and insightful suggestions.

W1 Counterfactual generation results. We are sorry that our presentations might have confused you. We believe that our counterfactual generation results are sufficient to support that our DiTi learns a more disentangled representation compared to baselines. Specifically, we have two main types of experiments: 1) By interpolating a feature subset between a pair of images (Figure 5), and 2) by manipulating a feature subset guided by the weight of an attribute classifier.

• For 1), only our DiTi allows meaningful interpolation between image pairs as shown in Figure 5 and A11 in Appendix. Hence, we include the baseline results in Appendix (Figure A12-A13), where the generated counterfactuals are either identical to the original image, or contain substantial artifacts, validating the baseline features are still entangled.
• For 2), a disentangled representation should enable manipulating a single attribute without affecting other ones. As shown in Figure 2(a), 6 and A14-A16 in Appendix, baseline methods typically entangle multiple attributes (e.g., manipulating age and gender together), and are still prone to generating visual artifacts. In contrast, our DiTi mainly manipulate only the target attribute. Note that the results can be quite subtle when the attribute is fine-grained, e.g., “wearing lipstick”. Following your suggestion, we will adjust the claim about bedroom experiments to reflect the aforementioned point.
• We have other supporting experiments to verify disentanglement quality, e.g., visualization of classifier weight (explained in the last paragraph of Section 3.1), strong inference results as noted by other reviewer (Section 5.2), and the results of interpolating the whole feature instead of a feature subset (explained in the last paragraph of Section 5.3). Overall, our results are from multiple evaluation methods, on different datasets, with many examples, and we also compare with diverse baselines (e.g., including conventional disentanglement methods based on VAE and GAN). Hence we believe that they are sufficient to validate our disentanglement quality.

W2 Mitigating spurious correlation by disentangled representation. Great suggestion. We evaluate our method and baselines on a few-shot learning task for the fine-grained attributes in CelebA. This is particularly challenging, as the classifier needs to isolate fine-grained class attributes from visually prominent yet spurious ones. As our method disentangles fine-grained attributes on the feature subset corresponding to a small $t$ (e.g., early $t$ in Figure 5), we can perform human intervention to only use that subset to train a classifier, which discards the spurious attributes. As shown in Table r1 and r2, using the full feature of DiTi for few-shot learning does not bring much improvements, as the classifier can still exploit the spurious correlations on the disentangled coarse-grained attributes. However, after intervention, we observe significant improvements over baselines.

Table r1: $N$-shot accuracy on "smiling" attribute averaged over 10 runs.

Table r2: $N$-shot accuracy on "wearing lipstick" attribute averaged over 10 runs.

In fact, we have performed similar follow-up experiments using a text-conditioned DM (Stable Diffusion), where our method beats few-shot learning methods using OpenCLIP by up to 13%. See Cjby-Q2 for details.