Paper: Exploring Diffusion Time-steps for Unsupervised Representation Learning (ICLR-2024)

Reviewer comment (Novelty):
First, I echo with my previous NeurIPS reviewers folks on the concern: whether it is the direct direction to actually integrate DMs into the current representation learning paradigm? Based on our previous discussions, the proposed method has a higher computational/time cost compared to its contrastive counterparts. And if the argument here is because the contrastive paradigms lose attribute information, then it seems to me the proposed DiTi may not be the optimal method to introduce the attribute inductive bias.

Author response:
**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.

Reviewer comment (Novelty):
Second, in the related work section, the authors include a comparison with the current DM for representation learning methods. I find this new part slightly biased, because the approaches in 2) and 3) are not really proposed under the context of representation learning. For instance, (Kwon et al., 2022) proposed the bottleneck feature of the U-Net under the context of image editing; similar to (Preechakul et al., 2022), in which separate autoencoders are learned to achieve editing purposes. I think while these works do touch the latent features/attributes in DMs, they tackle different scenarios other than representation learning. In other words, the first concern is not resolved by comparing these works.

Author response:
**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.

Reviewer comment (Novelty):
Going towards the technical level, one concern (less urgent in my sense) is that the proposed DiTi method shares quite similarities with PDAE, with the main difference in the disentanglement, with the cost of more hyper-parameters.

Author response:
**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).

Reviewer comment (Experiments):
This work trains models from scratch. Due to the introduced term for disentangled feature learning, the model needs to be re-trained.

Author response:
**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”?

Reviewer comment (Experiments):
The authors only evaluated the proposed method with datasets that are domain-specific (e.g., CelebA, FFHQ). The method has not been evaluated on large models trained on more diverse and general image datasets (e.g., LAION). Therefore, it does not indicate whether the method can scale to more complicated settings.

Author response:
**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.

Reviewer comment (Experiments):
Do the subsets have to be uniform? Does it benefit the method if we have more features for coarser or finer features?

Author response:
**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.

Reviewer comment (Presentation):
The qualitative results shown in figure 5 do not show consistent improvements wrt to modular editing. In the second row, for example, the eyeglasses appear in the middle range and disappear again in the final ranges. Are the time steps plotted in a cumulative fashion within the early/middle/late “t” range?

Author response:
**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”.

Reviewer comment (Experiments):
On the use of timesteps [100-300] and it corresponding to the maximum value of the loss: One would expect the loss to be greater at the later stages when the image becomes complete noise. Eg: it should be difficult to construct image from t = 1000 to t = 700. Can a plot be included for different samples to validate this?

Author response:
**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.

Reviewer comment (Evaluation):
Comparison to SIMCLR: SimCLR does obtain a higher accuracy on attributes affecting local appearances (e.g., “Hat”). We postulate that its contrastive training has some effects in regularizing f as an injective mapping. This is not clear. What is being implied in these findings? Does this mean that the proposed approach depends on attribute correlations?

Author response:
**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

Reviewer comment (Presentation):
From point 2 above, it will be great to include the plot to validate the time steps and loss.

Author response:
**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.

Reviewer comment (Experiments):
The results presented in the manuscript hold for unconditional models. Do the findings also extend to conditional models?

Author response:
**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.  
| Method                      | 1        | 2        | 4        | 8        | 16       | Avg      |  
| --------------------------- | -------- | -------- | -------- | -------- | -------- | -------- |  
| Zero-Shot OpenCLIP          | 42.3     | 42.3     | 42.3     | 42.3     | 42.3     | 42.3     |  
| Zero-Shot Stable Diffusion  | 24.3     | 24.3     | 24.3     | 24.3     | 24.3     | 24.3     |  
| Tip-adapter (on OpenCLIP)           | 48.4     | 53.9     | 57.0     | 62.0     | 67.4     | 57.7     |  
| **DiTi ( on Stable Diffusion)** | **48.5** | **55.8** | **64.2** | **74.2** | **79.9** | **64.5** |

Reviewer comment (Writing):
The equations are missing parenthesis on expectation.

Author response:
Thank you for the review and comments. We have added the clarification for W3 and parenthesis on expectation in revision.  
For Figure 5, we believe the editing is consistent, as from left (w/ sunglasses) to right (w/o sunglasses), the editing should be adding sunglasses (Figure top), and from right to left, the editing should be removing sunglasses (Figure bottom).