Understanding Diffusion-based Representation Learning via Low-Dimensional Modeling

Authors: Xiao Li, Zekai Zhang, Xiang Li, Siyi Chen, Zhihui Zhu, Peng Wang, Qing Qu
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: diffusion representation learning, representation learning, diffusion model, denoising auto-encoder

Abstract:
This work addresses the critical question of why and when diffusion models, despite their generative design, are capable of learning high-quality representations in a self-supervised manner. We hypothesize that diffusion models excel in representation learning due to their ability to learn the low-dimensional distributions of image datasets via optimizing a noise-controlled denoising objective. Our empirical results support this hypothesis, indicating that variations in the representation learning performance of diffusion models across noise levels are closely linked to the quality of the corresponding posterior estimation. Grounded on this observation, we offer theoretical insights into the unimodal representation dynamics of diffusion models as noise scales vary, demonstrating how they effectively learn meaningful representations through the denoising process. We also highlight the impact of the inherent parameter-sharing mechanism in diffusion models, which accounts for their advantages over traditional denoising auto-encoders in representation learning.

Decision: reject
Reviewer scores: 5, 1, 8, 3

Reviews and discussion:
<thread>
creation date:
1729549609037

thread_type:
official_review

<thread_object>
creator: reviewer_anhe

<content>
summary:
This work evaluates the potential of representation learning with diffusion models. The authors evaluate the connection between the quality of posterior estimation and the quality of learned representations, extending this analysis to different scenarios.

strengths:
- In Section 4.1, the authors present an interesting and insightful analysis of the effect of diffusion training with multiple levels of noise passed through the same model on the learned representations. It is intriguing to see that diffusion models have relatively steady representations across different denoising timesteps.

weaknesses:
- Several results presented in this work are not novel. For example, the evaluation presented in Figure 1 was already discussed in [1] and [2]. As noted by authors, the fact that representation learning dynamic captures a “fine-to-coarse” shift with the increased amount of noise was already noted in DAE works [3] and [4]
- I fail to see the significance of the contribution “Linking posterior estimation ability of diffusion models to representation learning”. Isn’t this observation a straightforward implication of the fact that samples at the late diffusion steps are, by definition, more noisy, which results in lower quality of the posterior estimation and higher entropy when analyzing predictions from the linear probe?
- The “representation learning” capability evaluation is very limited to the simple linear probing task. This significantly limits the significance of the presented results. Extending the analysis with other SSL tasks would strengthen the submission.
- Editorial: I’m lost in the presentation of this work as in Sec. 2.3 it is written that “since diffusion models tend to memorize the training data instead of learning underlying data distribution when the training dataset is small (Zhang et al., 2023), we focus on the case where sufficient training data is available throughout our analysis in Section 3”, yet there Figure 2 on the second page seems to be a reproduction of work by Zhang et al. The analysis related to Figure 2 is presented on the last page


[1] Xiang, Weilai, et al. "Denoising diffusion autoencoders are unified self-supervised learners." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
[2] Deja, Kamil, Tomasz Trzciński, and Jakub M. Tomczak. "Learning data representations with joint diffusion models." Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Cham: Springer Nature Switzerland, 2023.
[3] Choi, Jooyoung, et al. "Perception prioritized training of diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.
[4] Wang, Binxu, and John J. Vastola. "Diffusion models generate images like painters: an analytical theory of outline first, details later." arXiv preprint arXiv:2303.02490 (2023).

questions:
- As indicated in line 192, in a setup with clean images used being used for classification, the diffusion model is imputed with clear image and timestep t, “where t serves solely as an indicator of the noise level for diffusion model to adopt during feature extraction.” Isn’t this setup out of distribution for a model, as the diffusion model is never trained with a clear image or conditioned with a timestep equal to zero?
- Why limiting the evaluated features to only those encoded in the bottleneck layer of the UNet architecture? A lot of information relevant to classification (such as high-frequency features) might be passed through the residual connections.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730316629919

thread_type:
official_review

<thread_object>
creator: reviewer_bxju

<content>
summary:
I believe the authors try to study the structured emerged from training multi-scale score matching objective. Given different noise scale, for the model to achieve score matching (denosing), it musts learn feature at that scale. For example, for score matching on natural images, if the noise level is low, the model just need to learn low level (texture) features to denoise. But when the noise level is high, the model needs to learn high level (features) concept or also have to guess what's in the image. The author shows different scale of features learnt by multi-scale score matching can be used for semantic task like image classification, as well as for posterior estimation of mixture of Gaussian. The author author shows the performance of doing these task under different noise scale form an inverted u shape curve. Under strict assumption of mixture of Gaussian, the author provide analysis on the posterior estimation power given differer noise level. The author compare the learnt model's performance with the optimal one.

strengths:
1. It relates multi-scale score matching to posterior estimation to representation learning.
2. It make (realistic?) assumption on the data so the connection in 1 can be studied and analyzed theoretically.
3. The theoretically analyze under simple data assumption is thorough.

weaknesses:
1. I felt like the paper is missing lots of detail. I have to make many assumption outside the paper.
2. I wish I see a cleaner method. Despite these's extensive experiment and theoretical analysis under the strict "mixture of low-rank Gaussians" setting. I don't really see the key message. What I get is that different level of features learnt at different scale can be used for posterior estimation of mixture of Gaussian, as well as doing classification. How are these two related specifically? One statement is that they both has the inverted U shape curve. But is that it? I hope the author can make this more clear. 
3. Figure 4b tells us the learnt multi-scale score matching function (with neural network?) behave different from the optimal one on a simple dataset. I think this is really interesting result. But I also think the author should provide some insight on why is this the case, because it might reveal the nature of bias in neural network. For example, in [1], they find cnn based denoiser cannot learn optimal solution for simple toy example like global planer wave because it has spatial locality as bias. 

[1]  Z, Kadkhodaie, et al, 2023, Generalization in diffusion models arises from geometry-adaptive harmonic representation

questions:
1. For image denosing case, I'm assuming the work is based on the idea of latent diffusion. So x_0 is the clean data embedded into the latent space using a VAE? Or it's the raw image? 
2. How is U_k calculated for cifar10? Is it calculated by doing PCA on all examples (in vae latent space) with class label k?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730598636052

thread_type:
official_review

<thread_object>
creator: reviewer_aaut

<content>
summary:
This paper makes an attempt to explain *representation quality* in a diffusion model using a simple model of Gaussian data over low dimensional subspaces. They observe a uni-modal trend in *representation quality* as a function of noise level, which indicates that the best representation of the clean input image can be obtained at a certain noise level. It seems that the representation quality is empirically measured by classification accuracy over test data using the internal representation at the bottleneck of a UNet. 

To explain this observation, they 1) assume image data lies on a union of manifolds, where each manifold corresponds to a class. 2) They approximate manifolds with linear sub-spaces. 3) They assume each subspace of a class contains features that are relevant for high quality representation ($U_k$) and anything that is not relevant is noise and lies on orthogonal complement ($U^T_k$). Moreover, the assumption is that data is Gaussian distributed over the image subspaces. 4) They define Class Signal to Noise Ratio to measure the goodness of representation using denoised images (i.e. posterior mean estimate) and the class sub-spaces. 5) They show both of their measures for quality of representation (i.e. classification test accuracy from bottleneck and CSNR) have a uni-modal trend as a function of noise level. 6) Due to the theoretical analysis, for the low rank Gaussian model, the highest representation quality is a function of the ratio of added Gaussian noise level and the *noise* level of portion of image that lies in ($U_k^T$).

strengths:
1. Important and interesting topic: The question of representation learning through diffusion models is interesting and worthwhile for more investigation. 
2. This paper could potentially be interesting with major improvement in writing and clarity. Overall, the paper needs more refinement to be ready for publishing.

weaknesses:
1. The paper is very poorly written. It's hard to follow the logic and it's not clear what are the contributions. It's also not clear what parts are borrowed from other works and what parts are novel. For example, a good portion of the assumptions, modeling, and toy experiment setup are borrowed from [1] which is not obvious from the text. 
2. The key concept are not explained clearly. For example, what do you mean by "representation quality"? It's been referred to throughout the text and figures without any definition. How do you connect your two measures of quality to each other? Overall, many of the key concept in the text are never defined! How can one assess the results without knowing what the experiments are trying to measure?! 
3. The linearity assumption (approximating each class nonlinear manifold with a subspace) is obviously too simplistic for real data. As a result the theoretical result, expressed in theorem 1 does not extend beyond the simple low rank Gaussian model. 
4. Even within the union of low-rank Gaussian model, a major drawback is the arbitrary split of $U_k $ and $U_k^T$. How does one decide what it relevant to representation and what is *noise* or irrelevant perturbations? To answer this question and correctly decide where the image subspace is, one needs to solve the representation learning problem first. To decide what is relevant information to representation, one needs to define the task the representation will be used for. Features that are noise with respect to one task are relevant information w.r.t another task. For example: for a coarse level classification task more information is irrelevant hence $U_k$ is lower rank. That means the mode in your uni-modal plots will be at higher noise levels. For a more granular classification task, more information about details are needed, so the mode will be at lower noise levels. Thus, the whole notion of "representation quality"as a function of noise level  is not well-defined. 
5. The authors are using two separate notions of representation: the ave-pooling responses in the bottleneck layer, and the strength of projections of posterior mean onto the class subspace (CSNR). Merely showing that both of these are uni-modal as a function of noise level doesn't prove anything! Specially given that the maximums happen at different noise levels! 

[1] Wang, Peng, et al. "Diffusion models learn low-dimensional distributions via subspace clustering." arXiv preprint arXiv:2409.02426 (2024).

questions:
1. Why does it make sense to replace $x_t$ with $x_0$ in $x_\theta(x_t, t)$?! What is $x_\theta(x_0, t)$ computing? This ad-hoc replacement creates an inconsistency between the noise level on the input image and the noise variance given to the network. Now it is not clear what the output of the function $x_\theta$ is after creating this inconsistency, and is not clear how to interpret empirical results that involve this (which seems to be all the empirical results). 
2. What does figure 1a convey? How did you measure *posterior accuracy* for CIFAR images? What is called posterior accuracy throughout the paper is equivalent to denoising performance. Indeed you can measure denoising performance, but how did you measure its accuracy for **real** images?! This is an unsolved problem, and it is not explained in the paper how the authors solved it. 
3. Figure b shows a result that has been around for a long time. That is denoising at higher noise levels results in the loss of details. It's not clear how this is part of their results, given that this has been known for more than 100 years and also has been rediscovered in deep learning era again and again. There should be at least some citations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730756569174

thread_type:
official_review

<thread_object>
creator: reviewer_pyve

<content>
summary:
The paper presents a low-rank decomposition theory of diffusion representation learning, which supports the empirical result that using a medium denoise level t achieve best classification performance.

strengths:
1. Theorem 1 provides direct connection from the noise level sigma_t to the empirical results. 

2. The paper is well-written and easy to read.

weaknesses:
1. The results are not rigorous.  theorem 1 is obtained by many approximations: 1)  x_theta(x_t, t) defined in proposition1 is replaced with x_theta(x_0, t) with some informal arguments. The gap between these two terms cannot be easily ignored. 2) eq 9 replaces eq 6 for simplicity. It should be put in theorem instead. 

2. The empirical results. Since authors do not provide new empirical method, I am wondering why no larger experiments are conducted, e.g., ImageNet with some pretrained diffusion models. Evidendece on Cifar are not strong enough to support all arguments in this paper.

3. The K-space. Do you use K=number_of_classes? Is it too strong assumption? Human classify objects into different classes, but it would be overly strong to assume that K classes subspaces are independent?

questions:
as above
</content>
</thread_object>

</thread>
<thread>
creation date:
1732252392221

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
Dear AC and all reviewers:

We thank all the reviewers for the detailed review, the appreciation of our work, and constructive feedback. Specifically, Reviewer BxJU acknowledged the thoroughness of our theoretical results under the MoLRG data assumption, Reviewer PYve noted the alignment of our theoretical findings with empirical observations, and Reviewer ANHe found our empirical comparison between DAEs and diffusion models interesting.

In summary, our work presents the first theoretical analysis of the unsupervised representation learning capabilities of diffusion models. We characterize the unimodal representation dynamics across timesteps, grounded in the strong relationship between posterior estimation and representation learning. 

Below, we address a common concern raised by the reviewer regarding the use of clean images as inputs for diffusion-based representation learning. While this approach may seem unconventional, it is well-justified by the following considerations:

* **Comparable performance with using noisy image**. Our experimental results in Figures 1 and 8 demonstrate that using clean image $x_0$ as input for inference achieves linear probing accuracy on par with or better than using noisy image $x_t$ as inputs. Additionally, in the updated manuscript, we include a visualization of the denoising results when using $x_0$ or $x_t$ as inputs. These results highlight that, in the posterior estimation setting, using clean images as inputs also outperform using noisy inputs particularly in high-noise regimes. For more details, we refer the reviewers to Appendix A.5 (Figure 11).

* **Aligning with standard supervised/self-supervised training and inference**. Here, we utilize clean images solely during inference to extract representations, while adhering to the standard denoising procedure during model training. This pipeline aligns with conventional self-supervised learning approaches in contrastive learning. Training on noisy images can be interpreted as a form of data augmentation, akin to techniques like cropping, color jittering, or masking, commonly employed in self-supervised learning to enhance model performance. During inference, clean images are typically used to evaluate representation quality. Such an approach bridges the practices in diffusion models with established conventions in both supervised and self-supervised representation learning.

We will make those points clear in the revision of our work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732550957632

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Looking forward to feedbacks

comment:
Dear Reviewers,

Thank you once again for your thoughtful reviews and the time you have dedicated to evaluating our work.

As the **author-reviewer discussion period is reaching to an end soon**, we would greatly value your feedback on whether our responses have successfully addressed your concerns. Your input is crucial to the continued improvement of our work.

If you have any remaining questions or require further clarification, we are more than happy to provide additional information. If our responses have adequately resolved your concerns, we would sincerely appreciate it if you could consider updating your score.

Thank you for your time and consideration.

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732672812032

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Regarding Malicious Review from Reviewer aaUT

comment:
Dear AC, PC, and All Reviewers:

We would like to respectfully raise concerns regarding the **malicious** review provided by **Reviewer aaUT**. We strongly suspect that **Reviewer aaUT's review of our paper was biased and unprofessional, potentially driven by personal retaliation rather than an objective evaluation**, based on the following evidence:

1. **False Ethical Accusation.** Initially, Reviewer aaUT falsely accused us of plagiarism and dual submission without providing convincing evidence. Following our clarification, Reviewer aaUT removed the ethical flag but continued to make unfounded claims, alleging that "authors *may* have borrowed theorems from [1]" without concrete evidence. Moreover, the only similarity with [1] lies in the model assumption. If this is plagiarism by Reviewer aaUT's logic, then any paper using the same model assumptions would be as well, which is clearly false.

2. **Superficial Reading of Our Paper:** Reviewer aaUT withdrew the ethical flag by stating: "*I removed the flag as the authors clarified that the paper is cited and acknowledged **in the appendix**. However, note that reviewers are not required to read the appendix. Given the level of influence of [1], it should be explicitly mentioned in the main text.*" Reviewer aaUT's claim is obviously false, as **[1] is cited over four times in the main body** of our paper (spread in Sections 1, 2, 3). This reflects a superficial reading of our paper. 

3. **Unjustified Score Reduction:** After his removal of the ethical flag, Reviewer aaUT maliciously lowered the score **without any justification**, claiming: "*I do not find the counter-arguments compelling, as they do not go beyond reiterating statements in the submission.*" This dismissal disregards **our detailed and professional rebuttal**, which included point-to-point responses and extra experiments provided in Appendix A.5 (Figure 10-12). Such behavior and superficial statements demonstrate a lack of basic respect for our efforts, and we suspect **Reviewer aaUT is retaliating against us due to the prior false ethical accusation.**

4. **Baseless and Exaggerated Criticism of Our Contribution.** In Reviewer aaUT’s comments, quote “*Figure 1b shows a result that has been around for a long time. That is denoising at higher noise levels results in the loss of details. It's not clear how this is part of their results, given that **this has been known for more than 100 years** and also has been rediscovered in deep learning era again and again.*”  In our paper, we have never claimed this as our original finding and have properly cited the relevant work. Instead, we present it as the motivation for our theoretical study.

We would appreciate your intervention in such malicious behaviors.

Thank you for your time and attention.

Best,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper received mixed reviews, with the majority of reviewers (PYve, aaUT, ANHe) recommending rejection, and one reviewer (BxJU) suggesting acceptance but with low confidence. The AC has thoroughly examined the submitted materials, including the reviewers’ concerns and the authors’ responses, and has carefully considered all factors in the decision-making process.

While the paper offers theoretical insights into the representation learning capabilities of diffusion models under the assumption of a low-dimensional mixture of Gaussians, its empirical contributions remain unclear. In particular, addressing the raised question (by BxJU) regarding the interpretation of higher CSNR values at a particular noise level and their implications for representation quality could significantly enhance the paper’s empirical impact.

Given the mixed reviews (amongst all ICLR submissions) and the absence of a strong advocate among the reviewers, the work is not yet deemed ready for publication at this status. However, given the potential, the authors are encouraged to revise the paper based on the feedback and resubmit for consideration in the next cycle.