On the Feature Learning in Diffusion Models

We are grateful for the reviewer's positive comments on our work. We also appreciate the reviewer for raising constructive comments and questions. Below are our responses to each comment and question.

---

1. (W1) On the heterogeneous data structure where signal structure and noise can vary significantly.

We believe our analysis can be extended to multiple features by following a similar proof strategy. For instance, consider three feature patches $\mu_1, \mu_2, ..., \mu_k$ along with a noise patch where the SNRs are different (i.e., $||\mu_1|| \neq ||\mu_2|| \dots\neq ||\mu_k||$), we can still use the same analysis to write the gradient descent updates (in early stages) for diffusion models as follows: $\langle w_{r,t}^{k+1} , \mu_j \rangle = \langle w_{r,t}^k , \mu_j\rangle + \Theta( \frac{\eta}{\sqrt{m}}||w_{r,t} ||^2 ||\mu_j||^2 ), j = 1,2,\dots,k$ $\langle w_{r,t}^{k+1} , \xi_i \rangle = \langle w_{r,t}^k , \xi_i\rangle + \Theta( \frac{\eta}{n \sqrt{m}} ||w_{r,t} ||^2 ||\xi_i||^2 )$ Then following a similar analysis, we can show even with varying SNRs, the feature learning outcomes of diffusion models would still be determined by their relative SNRs. Thus diffusion models still learn all the features. In contrast, existing studies [5,6] found that under the presence of multiple features with varying magnitude, classification models only lean partial features that are more significant.

Therefore, we believe the conclusions drawn in this paper that diffusion models learn more balanced features than classification models would not change with more diverse features. We have conducted additional experiments on 10-class Noisy-MNIST dataset (in Appendix I of the revised manuscript) to numerically confirm this claim. See more details in Reply 2 below.

---

2. (W2) On the assumption of orthogonal signals and a single feature.

We have followed prior analysis on classification models [1,2,3,4] to introduce the orthogonality assumption, in order to facilitate theoretical analysis. The orthogonality assumption is used to simplify the analysis and enables clear insights into the training dynamics of diffusion models and classification.

Furthermore, our analysis can be extended to address multiple features, as discussed in Reply 1 above.

Numerically, in Appendix I of the revised manuscript, we conducted additional experiments on Noisy-MNIST dataset with 10 classes (compared to 2 classes in the main paper). In this case, the features are more diverse and are not orthogonal. From the results, we see despite with more features, diffusion models learn all the features to relatively the same level whereas classification models predominantly learn subset of features. This verifies the applicability of the theoretical results in this paper to more realistic settings.

---

3. (W3) Experiments on different SNR values.

We have now conducted additional experiments over a broader range of SNR values, as detailed in Appendix G of the revised manuscript. These include $n SNR^2 = 1.92$, $n SNR^2 = 3$, and $n SNR^2 = 4.32$, complementing the main experiments on $n SNR^2 = 0.75$ and $n SNR^2 = 6.75$.

From the results, we observe that classification models are more sensitive to changes in $n SNR^2$, often overfitting to one feature while ignoring others. In contrast, diffusion models consistently learn more balanced features across the range of SNR values. These findings are consistent with our theoretical analysis and further support our conclusions.

---

4. (Q1) Caption not consistent with setting in Figure 6.

Thank you for pointing this out. We have now corrected the caption for Figure 6 in the revised version.

---

5. (Q2) On the ratio of signal and noise learning in relation to $nSNR^2$.

This is consistent with our theory because we only claim the ratio of signal to noise learning is on the same order of $n SNR^2$ (rather than exactly matching it). The absolute magnitude of changes in signal and noise inner products of diffusion models when switching from $nSNR^2 = 0.75$ to $6.75$ is not meaningful without benchmarking against the changes for classification models. In comparison, we see diffusion models are more aligned with the order of $n SNR^2$ compared to classification models.

---

Reference

[1] Cao et al. Benign overfitting in two-layer convolutional neural networks. In NeurIPS 2022.

[2] Chen et al. Why does sharpness-aware minimization generalize better than SGD?. In NeurIPS 2023.

[3] Chen et al. Understanding and improving feature learning for out-of-distribution generalization. In NeurIPS 2023.

[4] Lu et al. Benign Oscillation of Stochastic Gradient Descent with Large Learning Rate. In ICLR 2024.

[5] Zou et al. The benefits of mixup for feature learning. ICML 2023.

[6] Shen et al. Data augmentation as feature manipulation. ICML 2022.

We sincerely appreciate the reviewer for acknowledging our contributions and below we provide detailed responses to all the comments and questions.

---

1. (W1) Comparison to other generative models.

We would like to emphasize that the goal of this work is to understand the training dynamics of diffusion models, rather than all generative models. While we agree that comparing diffusion models to other generative models is interesting, it is beyond the scope of this work.

Nevertheless, we believe the theoretical frameworks developed in this work can be extended to other generative models, such as VAEs. We have added Appendix F in the revised version discussing the feature learning of VAEs, both theoretically and numerically.

We find that under the current setup, the objective of VAE simplified to include a reconstruction term and some regularization terms. The reconstruction compels the model to learn both the signals and noise present in the data, thus promoting a balanced feature learning. To test this claim, we conducted an experiment, examining the feature learning dynamics of VAEs. We notice that similar to diffusion model, VAEs learn more balanced features compared to classification. However, a thorough comparison between diffusion models and VAEs would require dedicated efforts and additional analysis, which are beyond the scope of this paper.

---

2. (W1) On the effect of diffusion timesteps on feature learning.

In this work, we consider a fixed range of timesteps rather than all timesteps in order to focusing on the training dynamics of diffusion model at specific timesteps and their effects on feature learning. Nevertheless, extending the analysis to cover broader scales of timesteps is indeed our next step. This extension would allow us to further investigate how features are reconstructed throughout the entire reverse diffusion process.

---

3. (W2) Discussion on how the theoretical results apply to real-world applications.

We believe the theoretical results in this paper have real-world substance. For instance, the more balanced feature learning of diffusion models may be able to explain the inherent adversarial robustness of diffusion models used for downstream classification (as shown in [3,4]). Because the learning outcomes of diffusion models are less sensitive to changes in SNR compared to models trained for classification, the feature learning of diffusion models tend to be more resistant to adversarial perturbation.

We have now added such discussions in the conclusion section (Section 6 of the revised version).

---

4. (W3) Explanation on the meaning of formulations.

(1) The noise $\xi$ is generated from the Gaussian distribution such that it is orthogonal to the signal vectors $\mu_1, \mu_{-1}$ for simplicity of analysis. Such an orthogonal signal-noise setup is common in existing works, such as [5,6].

(2) We believe that the constant order of $\alpha_t$ and $\beta_t$ reflects the standard practice in diffusion models. For instance, in DDPM [7], the noise schedules are pre-defined, and the number of timesteps is typically fixed. This ensures that $\alpha_t$ and $\beta_t$ remain of constant order, independent of the data or model size. Nevertheless, we agree that extending the analysis to cover broader range of $t$ is an important future direction as we have discussed in Reply 2.

We have added the explanations in the revised version accordingly.

---

5. (W4) On the effect of different networks for classification and diffusion models on feature learning.

Due to the differences in the objectives and output sizes between diffusion models and classification, it is difficult to employ exactly the same model setup. The current model setup, i.e., weight-tying for diffusion model and weight-fixing for classification model ensures the first-layer and the trainable parameters are consistent across two models.

We believe such a difference in networks does not affect the theoretical results. To validate such a claim, we have conducted an additional experiment in Appendix H of the revised manuscript, comparing feature learning dynamics between diffusion model and classification on a three-layer networks. The results indicate that while adding more layers, the similar learning patterns still maintain, i.e., diffusion model learns more balanced features while classification focuses on learning partial features.

We are grateful for the reviewer’s acknowledgment of our work’s contributions. We also appreciate all the comments and questions raised by the reviewer. Below are our detailed responses to each comment.

---

1. (W1 + Q1) On the feature learning when training on noised inputs.

We would like to clarify that the difference between diffusion model and classification models in feature learning lies in their different training objectives, rather than the additive Gaussian noise to the data. In fact, when performing the classification on the Gaussian noisy data, the feature learning results will not be changed. In particular, this can be justified both from both theory and experiments:

• Theoretically, in Appendix J of the revised manuscript, we show that if we replace $\mu$ and $\xi$ with $\mu + \epsilon$ and $\xi + \epsilon$, and take expectation of the gradient with respect to the noise added, we find that $E_{\epsilon_{i} } [\nabla_{w_{j,r}} \tilde L_S(W^k)] \approx \nabla_{w_{j,r}} L_S(W^k) + \frac{1}{nm} (\sum_{i=1}^n \ell_i'^k) w_{j,r}^k$ where $\tilde L_S(W) = \frac{1}{n}\sum_{i=1}^n \ell(y_i f(W, x_i + \epsilon_i))$ and $L_S(W) = \frac{1}{n}\sum_{i=1}^n \ell(y_i f(W, x_i))$ denote the noise-added objective and noise-free objective respectively. The above expression of gradient implies that the Gaussian noise perturbation is equivalent to (in expectation) adding an $L_2$-type regularization. This regularization will perform the same penalization for both learning signals and noises, thus will not alter our current results for classification models.

• Numerically, we have conducted an additional experiment in Appendix J of the revised manuscript, examining the feature learning dynamics of classification models with the same Gaussian noise injected (as for diffusion model). The results suggest that even with the presence of Gaussian noise in the inputs, classification still overly rely on partial features, rather than the balanced feature learning performed by diffusion model training. This validates our claim that training classifiers on Gaussian-noised inputs does not promote the learning of more balanced features.

---

2. (Q2) Applicability of analysis to flow-based models.

Thank you for your question. We agree that extending our analysis to flow-based models is an interesting direction for future work, particularly given the importance of comparing diffusion models with flow-based models. The key differences between these approaches lie in (1) diffusion paths and (2) their prediction targets, which consequently result in different training objectives. This naturally aligns with our next step, which aims to investigate the role of training objectives in feature learning dynamics of diffusion models.

We sincerely appreciate the reviewer’s positive feedback on our work, as well as the comments and questions. We would like to take this opportunity to further clarify the concerns raised and highlight our contributions. Below are our detailed responses to all the comments.

---

1. (W1 + Q1) How these findings impact practical applications of diffusion models. How the insights hold up when applied to real-world datasets.

We believe our findings provide valuable insights into the properties of diffusion models. For instance, the results in the paper potentially offer an explanation for the inherent adversarial robustness of (pre-trained) diffusion model when used as classifiers (as in [1,2]). Since the learning outcomes of diffusion models are less sensitive to changes in SNRs compared to classification models, diffusion models may be more resistant to adversarial perturbation. We have now added such discussions in the conclusion section (Section 6 of the revised manuscript).

While experimenting on large-scale datasets, such as ImageNet, is difficult within the limited rebuttal period, we conducted an additional experiment with 10-class Noisy-MNIST dataset, which contains more diverse features (than the 2-class Noisy-MNIST used in the main paper). We have included the experiment settings and results in Appendix I of the revised manuscript. The results indicate similar learning dynamics to those observed in the 2-class setup: diffusion models learn more balanced features, whereas classification models predominantly rely on a subset of features. This suggests the insights in the paper hold for more real-world scenarios.

In theory, such a more complicated dataset can be modelled as the data distribution with multiple signals and we plan to extend our theoretical analysis to this more challenging setting. Intuitively, our main technical analysis can still be adapted. For instance, let $\mu_1,\mu_2,\dots,\mu_k$ be the signals that have different strengths, we may still obtain their early-phase learning dynamics as follows: $\langle w_{r,t}^{k+1} , \mu_j \rangle = \langle w_{r,t}^k , \mu_j\rangle + \Theta( \frac{\eta}{\sqrt{m}}||w_{r,t} ||^2 ||\mu_j||^2 ), j = 1,2,\dots,k$ $\langle w_{r,t}^{k+1} , \xi_i \rangle = \langle w_{r,t}^k , \xi_i\rangle + \Theta( \frac{\eta}{n \sqrt{m}} ||w_{r,t} ||^2 ||\xi_i||^2 )$ When the first-order terms dominate the learning dynamics, we may still prove balanced feature learning in this setting. Such a more general result can further trigger a number of follow-up works on understanding more detailed training designs of diffusion models.

Finally, we believe that this first attempt at analyzing the training dynamics of diffusion models establishes a theoretical foundation for advancing our understanding of diffusion model training and paves the way for exploring more complex problem setups. For instance, we expect the theoretical framework can inspire research on the roles of the loss function, model architecture, and optimizers in diffusion models. This can improve our understanding of the strengths and weaknesses of current practices, such as the Adam optimizer, U-Net architecture, and flow-matching versus score-matching losses.

---

2. (W2) On the assumptions being strict, particularly the dimensionality and initialization.

Our theoretical setups follow a series of prior theoretical works on classification models [3,5,6,9,10], where similar assumptions were widely adopted to simplify the analysis and highlight key phenomena.

Regarding dimensionality and initialization, our work seeks to provide a theoretical comparison of the feature learning dynamics of diffusion models and classification in the over-parameterization regime. The high-dimensional setting we adopted ensures sufficient over-parameterization, as the number of trainable parameters is $d \times m$. Additionally, we require small Gaussian initialization to enable successful feature learning and convergence analysis, as demonstrated in [1]. Using a large initialization would place the model in the lazy/NTK regime [4], which does not allow for meaningful comparisons of feature learning dynamics.

We are grateful for the positive feedback and recognition of the strengths of our work. We also thank the reviewer for providing feedback and comments. Below, we provide detailed responses to all the comments and questions.

---

1. (W1 + Q1) Comparison to other generative models.

We would like to emphasize that the focus of this paper is to understand the training and feature learning dynamics of diffusion models, rather than all generative models. While characterizing the feature learning capabilities of other generative models is an interesting direction, it is beyond the scope of this paper.

Although the focus is on diffusion models, we believe our analysis in this paper can be generalized to other generative models, such as VAEs. We have now included a discussion in Appendix F of the revised manuscript, where we explore the feature learning dynamics of VAEs. Under the current model setup (where we take expectation over the latent sample), we show the objective of VAE includes a reconstruction term with a specific regularization: $L = \frac{1}{n} \sum_{i=1}^n \sum_{p=1}^2 || x_i^{(p)} - W^\top (W x_i^{(p)})^2 ||^2 + L_{\rm reg}.$ The reconstruction loss (the first term of the above equation) forces the network to learn both the signals and noise. Specifically, if we compute the gradient of the reconstruction loss, we can see in the early stage, the updates can be written as $\langle w_{r}^{k+1} , \mu_j \rangle = \langle w_r^k , \mu_j\rangle + \Theta( \eta \langle w_r^k, \mu_j \rangle^3 )$ $\langle w_{r}^{k+1} , \xi_i \rangle = \langle w_r^k , \xi_i\rangle + \Theta( \frac{\eta}{n} \langle w_r^k, \xi_i \rangle^3 )$ Then we can follow a similar analysis as for diffusion models to analyze the feature learning dynamics of VAEs and show that VAEs also learn balanced features.

To validate the claim numerically, we conducted additional experiments training a VAE on the same synthetic dataset used in the main paper. The results, presented in Appendix F, demonstrate that under both low and high SNR settings, the VAE behaves similarly to diffusion models, learning more balanced features compared to classification models.

---

2. (W2) More exhaustive experimental setting, such as with larger feature space rather than one label.

We have now included an additional experiment on 10-class MNIST dataset in Appendix I of the revised version. This dataset contains more diverse features compared to the 2-class MNIST dataset (in the main paper). Even with a larger feature space and more labels, we observe similar learning patterns as for the two-class setup, i.e., classification models tend to learn some partial features over the others, while diffusion models learn all the features.

---

3. (W3 + Q4 + Q5 + Q6) "Organization and notation of the paper could be improved."

Thank you for the suggestions on improving the clarity of the paper. We have incorporated your suggestions in the revised version.

(Q4): We have now removed some redundancy in Section 1.1 and Section 2. Nevertheless, we believe Definition 2.1 in Section 2 is necessary to formalize the data distribution, as briefly introduced in Section 1.1.

(Q5): (a) In Section 1.1, the notation $d$ is already introduced in Line 69, which corresponds to the feature dimension and the notation $n$ is introduced in Line 85, which refers to sample size. Variables $r$ are now introduced in Line 190. (b) We have now removed the subscript $t$ when introducing neural network functions for diffusion model for clarity purpose. (c ) We have now added y-axis label to all figures thanks to your suggestion.

(Q6): Thank you for the suggestion, we have now moved the Figure in Appendix A.1 to Section 5.1.

We hope we have addressed you concerns regarding organization and notation.

---

4. (Q2) SNR of input differs between the diffusion model and classification case.

In our work, SNR is defined as the ratio between the signal and noise in the clean data, which is used as a criterion to quantify the difficulty of a learning problem. It is not correct to claim that SNRs of inputs differ between diffusion and classification models. This is because the noise addition process is part of the diffusion model, known as the forward diffusion process. As a result, the input to diffusion models is the clean data, which is consistent with the classification models. Thus, we ensure the SNR of inputs is consistent across both settings.

We sincerely thank the reviewer for all the positive comments and general appreciation of our work. Below, we address the questions and concerns raised in detail.

---

1. (W1) Presentation requires enhancement. Meaning of $m$ in Line 188.

We have proofread the paper and have added introduction to the notations, including $m$ in Line 188, which refers to the network width, $k,r,t$ in Line 220, $\omega(\cdot)$ in Line 60.

---

2. (W2 + Q1) Classifier considered lacks sufficient generality. Applicability of theoretical findings to broader range of models.

Given there is no prior work analyzing the feature learning dynamics of diffusion models, we follow existing theoretical works on classification models [1,2,3,4,5] to establish the model setups. Although we focus on simplified models, they successfully capture the primary factors necessary to understand and compare the training dynamics of both classification and diffusion models. Furthermore, we believe the model setup is suitable for theoretical studies, given the non-linear, two layer structure and the non-convex nature of neural network optimization.

To further support the applicability of our theoretical findings to broader model classes, we have conducted additional experiments on three-layer classifier and diffusion model with ReLU activation. Details of the setup and the results are included in Appendix H of the revised version. Notably we observe a similar pattern as for the two-layer network setup. That is, classification tends to bias learning towards one of the features while diffusion model learns more balanced features. These findings verify the generality of our findings.

As the first study on diffusion model training dynamics, we believe our work lays a foundational framework that can inspire future research into more general problem setups and more complex model architectures.

---

3. (Q2) "Requirement on orthogonality in data setup common in practice?"

The orthogonality assumption is introduced to simplify the analysis, enabling clear insights into the training dynamics of diffusion models and classification. While perfect orthogonality is often not common in practice, this assumption has been commonly adopted in prior theoretical works [1,3,4] to facilitate the study of complex dynamics in high-dimensional spaces.

We believe extending the analysis to account for non-orthogonal features could follow a similar approach (as done in this paper), like the previous work [6] that studies the non-orthogonal XOR data models for classification problems. However, such an extension would require a more rigorous and detailed analysis, which we leave for future work.

---

References

[1] Cao et al. Benign overfitting in two-layer convolutional neural networks. In NeurIPS 2022.

[2] Chen et al. Why does sharpness-aware minimization generalize better than SGD?. In NeurIPS 2023.

[3] Chen et al. Understanding and improving feature learning for out-of-distribution generalization. In NeurIPS 2023.

[4] Lu et al. Benign Oscillation of Stochastic Gradient Descent with Large Learning Rate. In ICLR 2024.

[5] Kou et al. Benign overfitting in two-layer ReLU convolutional neural networks. In ICML 2023.

[6] Meng et al. Benign Overfitting in Two-Layer ReLU Convolutional Neural Networks for XOR Data. In ICML 2024.