Understanding Representation Dynamics of Diffusion Models via Low-Dimensional Modeling
Authors: Xiao Li
Zekai Zhang
Xiang Li
Siyi Chen
Zhihui Zhu
Peng Wang
Qing Qu
Venue: NeurIPS 2025 poster ()
Keywords: diffusion model
representation learning
Primary area: deep_learning

Abstract: Diffusion models, though originally designed for generative tasks, have demonstrated impressive self-supervised representation learning capabilities. A particularly intriguing phenomenon in these models is the emergence of unimodal representation dynamics, where the quality of learned features peaks at an intermediate noise level. In this work, we conduct a comprehensive theoretical and empirical investigation of this phenomenon. Leveraging the inherent low-dimensionality structure of image data, we theoretically demonstrate that the unimodal dynamic emerges when the diffusion model successfully captures the underlying data distribution. The unimodality arises from an interplay between denoising strength and class confidence across noise scales. Empirically, we further show that, in classification tasks, the presence of unimodal dynamics reliably reflects the diffusion model’s generalization: it emerges when the model generate novel images and gradually transitions to a monotonically decreasing curve as the model begins to memorize the training data.

Understanding Representation Dynamies of Diffusion Models via Low-Dimensional Modeling
Xiao Li" University of Michigan xlxiaolumich ch. edu
Zekai Zhang"
Xiang Li
University of Michigan zzekai @umi ch. edu
University of Michigan forkobedumi ch, edu
Siyi Chen University of Michigan siyichEumich edu
Zhihui Zhu Ohio State University zhu 34408osu. edu
Peng Wang
University of Michigan pengwaQumic ch edu
Qing Qu
University of Michigan qingquêumi ch. edu
Abstract
Diffusion models, though originally designed for generative tasks, have demon- strated impressive self- supersueer representation learning capabilities. A particu- larly intriguing phenomenon in these models is the emergence of unimodal repre- sentation dynamics, where the quality of learned features peaks at an intermediate noise level. In this work, we conduct comprehensive theoretical and empirical investigation of this phenomenon. Leveraging the inherent oinddimeniionality structure of image data, theoretically demonstrate that the unimodal dynamic emerges when the diffusion model successfully captures the underlying data dis- tribution. The unimodality arises from interplay between denoising strength and class confidence across noise scales. Empirically, we further show that, classification tasks, the presence of unimodal dynamics reliably reflects the diffu- sion model's s generalization: emerges when the model generate novel images and gradually transitions to monotonically decreasing curve as the model begins memorize the training data.
Introduction
Diffusion models, a class score-based generative models, have achieved great empirical success in various tasks such image and video generation, speech and audio synthesis, and solving inverse problems [5]. In general, these models, consisting of forward and backward processes, learn data distributions by simulating the non-equilibrium thermodynamic diffusion process 2 [6 17). Specifically, the forward process progressively adds Gaussian noise to training samples until they fully corrupted, while the backward process involves learning score-based model generate samples from the noisy inputs [17 [8].
Beyond their strong generative capabilities, recent studies have revealed that diffusion models also possess remarkable representation learning capabilities [19244]. In particular, the internal feature extractors of trained diffusion models have been shown to serve powerful self-supervised
learners, achieving strong performance on downstream tasks such as classification [20 21, semantic
'The first two authors contributed equally to the work.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

CIFARIO
CIFAR 100
FFHQ
CelebA
94.0
90.0
0.52
0.53
CIFAR Test 80c.
CIFARIO0 Test acc.
FFHQ Test mloU.
Test mioU.
85.8
77.0
0.512
0.525
64.0
0.505
0.52
69.4
51.0
0.49%
0.515
61.2
38.0
53.0
0.0
0.01
0.02
0.06
0.14
0.3
01.58
1.09
1.92
0.01
0.03
0.1
0.28
0.48
Time step (ot))
Time step (or))
(a) Classification
(b) Segmentation
Figure/ 1: Unimodal representation dynamics in diffusion- representation learning tasks. This unimodal representation pattern has been previously observed in diffusion- representation learning tasks; [99 20 23]. To verify this, train diffusion models various datasets and evaluate downstream performance using noisy images 2G at different timesteps In both classification and segmentation tasks, the performance consistently follows unimodal trend, peaking at intermediate noise levels. In (b), "mloU" denotes mean Intersection over Union, a standard metric used segmentation tasks.
and image aligmment 233 In many cases, these representations match or even surpass those from specialized self-supervised methods. implying that diffusion models have the potential to unify generative and discriminative learning vision.
segmentation 19
Alongside the empirical success, a prevalent phenomenon, which we refer to as the unimodal representation dynamics, has been widely observed in diffusion models for representation learning 19 20 23]: the quality of learned representations, as measured by downstream task performance, follows unimodal trend across noise levels. Specifically, the most effective features consistently appear at an intermediate noise level, while performance degrades as inputs become either fully noisy entirely clean (see Figure TT. Despite being widely observed, the underlying cause of this unimodal representation dynamic remains poorly understood.
Our contributions. In this work, we investigate the emergence of unimodal representation dynamics in diffusion models and its cause. Our analysis characterizes how representation quality varies across noise levels and offers new perspectives on how diffusion models, despite their generative design, can excel at representation learning. Specifically, we conduct a comprehensive theoretical and empirical study of unimodal representation dynamics in diffusion models. Motivated by the well-established observation that natural image data typically reside on a low-dimensional manifold 27 we examine the representation dynamics through how diffusion models learn from noisy mixtures of low-rank Gaussian (MoLRG) distributions. We show that the emergence of unimodal representations is intrinsically linked to the model's ability capture the underlying low-dimensional structure of the data. Our key contributions are as follows:
Mathematical framework for studying representation learning in diffusion models. analyze the representation dynamies, provide mathematical study of how diffusion models learn from MoLRG distributions, where data lie on a union low-dimensional subspaces. We adopt a simplified yet analytically tractable network architecture that mimics key structural properties of Noise U-Net. Under this setting, we quantify the quality of learned representations via the Signal- -oiise Ratio (SNR) within the target subspace, enabling to characterize how representation quality evolves across timesteps in the diffusion process.
Theoretical explanation for the emergence of unimodal dynamics. Leveraging the structures of the MoLRG model, we prove that the unimodal pattern representation quality naturally arises when the model effectively captures the wow-dimensiona data distribution. We show that this unimodal pattern is driven by an interplay between denoising strength and class separability that varies with noise levels. There exists an intermediate diffusion timestep where class-irrelevant components are maximally suppressed and class-relevant features are best preserved, resulting in optimal representation quality.

---

Empirical connection between unimodal dynamics and generalization. Empirically, we demon- strate that the presence of unimodal representation dynamics serves reliable indicator of model generalization in classification tasks. Specifically, the unimodal pattern consistently emerges when the model generalizes well, but progressively vanishes as the model shifts toward memorizing the training data.
Relationship to prior results. Recent empirical advances in leveraging diffusion models for down- stream representation learning have gained significant attention. However, theoretical understanding of how diffusion models learn representations across different noise levels remains largely unex- work and defer
plored. Here, we focus the results most relevant survey to Appendix A A recent study by [24] takes initial steps in this direction by analyzing the optimization dynamics of two-layer CNN trained with diffusion loss binary class data. Their focus is on contrasting the learning behavior under denoising versus classification objectives, without examining how representation quality evolves across timesteps. In contrast, our work characterizes and compares representations learned at different timesteps, provides deeper understanding of diffusion-basee representation learning and also extends to multi-class settings. recent study by representation learning, focusing
comprehensive
28 also investigates the influence of timesteps in diffusion- on attribute classification and counterfactual generation. contrast, our work provides theoretical explanation and practical metric for the emergence of unimodal representation dynamics and shows its relationship with data and model complexity and training iterations.
Problem setup
Basics of difflusion models. Diffusion models are class of probabilistic generative models that aim to reverse progressive noising process by mapping an underlying data distribution Pdata to Gaussian distribution. The forward diffusion process begins with clean data 2p Paaa and adds Gaussian noise over time E [0, 1]. which can be described by the following stochastic differential equation (SDE):
dæ; f() dt + g(t) dw,,
where f(() is the drift coefficient, g(() is the diffusion coefficient, and {w, is standard Wiener process. Then, one can verify that the noisy data satisfy 2;, 8,20 + StOTE with e ~ (O, I). where 8, and o2 scaling factors determined by (tt) and g()). For ease of exposition, let pa denote the probability density function (pdf) of the noisy data D; for each / E [0, 1]. particular, Po = Pdata To simplify the analysis, we assume throughout the paper that s; 1.
The reverse process transforms noise back into clean data by leveraging the score function V log pe (xx)) and is governed by the reverse SDE 229
dæ; = (f(()) (t) log dt g())d
where {wi is a Wiener process independent of {w:}. This enables diffusion models to generate new samples from the underlying data distribution Pdata by initializing from pure Gaussian noise and iteratively denoising via the score function.
Training loss of diffusion models. Note that the score function v (a.) depends the unknown data distribution Pdata. According to Tweedie's 's formula (30, i.c.,
p [x,] 23 + ofV log (x(,))
we can alternatively estimate log by training a network (0 t), which is referred to as denoising autoencoder (DAF), approximate the posterior mean E[arglay 20 22 31. To learn the network parameter 0, we minimize the following empirical loss:
NN
min L())
Σ 2m1
de BE.N(O...)
pa
a tore,)) a
(0)
i.i.d. Pdtta for all i 1...., N denote the training samples and A represents the weight associated with each noise level.
where mo
3

---

3 Good distribution learning implies unimodal representation dynamics
Building on the setup introduced in Section 2 this section provides a mathematical framework for analyzing the representation quality of diffusion models, and theoretically we characterize the unimodal representation dynamics exhibited by diffusion models across noise levels. We validate our theoretical findings both synthetic and real-world datasets.
3.1 Assumptions of data distribution
In this work, we assume that the underlying data distri- bution is cow-dimensional, which follows noisy version of the mixture of low-rank Gaussians (MoLRG) 32 defined as follows.
Noka MoLBG = 0.3)
Neky MoL3GG 16 0.3)
Assumption 1 (K-Class Noisy MoLRG Distribution). For any sample 2n drawn from the noisy MoLRG distribution with K subspaces, we have
Figure 2: An illustration of MoLRG with different noise levels. We visualize sam- ples drawn from noisy MoLRG with noise levels 5 = 0.1, 0.3 and K = 3.
To = =Uja+ toUee with probability Tu 2 O,
(3)
where (i) U E xnxdi denotes an orthonormal basis of the k-th subspace for each k 1...., K. Ug Onx is orthonormal basis for the DE- dimensional subspace spanned by {Ut 1 + k); (ii) "M N(O, Id.). Lo,)) e ~ N ID.) are independent standard Gaussian vectors representing the latent signal and noise components of an. respectively; (iii) 5 > 0 controls the data noise level, and (iv) Thk 2 0 is the mixing proportion satisfying Σ Tk 1.
One can verify that any sample æn drawn from the above distribution satisfies:
K
ZD TW (O. + +8ūūū;]
As shown in Figure 2 data from MoLRG resides union of low-dimensional subspaces, each following a Gaussian distribution with low-rank covariance matrix representing its basis perturbed by some noise. For simplicity, bases (i.e.. U T U. = 0 for k * l). and uniform mixing weights (лı ... = а 1/K) Additionally, the study of noisy MoLRG distributions is further motivated by the following facts:
assume equal subspace dimensions (dN
dx d), orthogonal
MoLRG captures the intrinsic low-dimensionullly of image data. Although real-world image datasets are high-dimensional in terms of pixel count and data volume, extensive empirical studies (25/27) demonstrated that their intrinsic dimensionality is considerably lower. Additionally, recent work [35] 36) has leveraged the intrinsic oow-dimensional structure of real-world data to analyze the convergence guarantees of diffusion model sampling. The MoLRG distribution, which models data a low-dimensional space with rank dę < n, effectively captures this property.
Latent diffusion models encourage the latent space toward Gaussian distribution. State-of-the-ae large-scale diffusion models [37 38 typically employ autoencoders [39] to project images into a latent space, where KL. penalty encourages the learned latent distribution to approximate standard Gaussians [ Furthermore, recent studies
22 40 show that diffusion models can be trained to leverage the intrinsic subspace structure of real-world data.
Modeling the complexity real- weal-world image datasets. The noise term ~Ü,e, captures perturbations outside the k-th subspace via the noise space Ut. analogous to insignificant attributes of real-world images, such as the background of an image. While additional noise may be less significant for representation learning, plays a crucial role enhancing the fidelity of generated samples.
'The assumptions of subspace orthogonality, equal subspace rank and uniform mixing weights are adopted for analytical simplicity. Empirical results in Appendix demonstrate that the unimodal dynamics persist when these assumptions are relaxed.

---

In other words, our data model assumes that the overall dataset lies in a union of now-dimensional subspaces, where each class is associated with a dominant class-specific subspace, and there exists a shared subspace common across all classes that captures class-independent fine-grained details. This seared-specffie decomposition has been empirically supported in the context of subspace clustering and representation learning 41 42
3.2 Network parametrization
To analyze how feature representations evolve across noise levels in diffusion models, we study the following nonlinear network parameterization inspired from the MoLRG data assumption... Specifically, we parameterize the DAE ze (2 t) and its corresponding latent feature representation hg (z,, as follows:
oo(0, 1) = U2 f), ho(a,, D(,,, U)U St, D = diag (3f ,.............. Ia (4)
where A = }U} denotes a set of learnable parameters with U = [U,,... Ux) E O Let 8a
1 and & 84077 where of the noise scaling in Correspondingly, we parameterize 1+o?
31 & + (Gr - (a)) () in with
exp (xt, t))
Se (1(,, t) = 20;
& |0 2ill2+
wr( t)
00 xel??.
(5)
Σ exp t))²
20
This network architecture can be interpreted as a shallow U-Net
43 composed with blockwise mixture-of-experts [44] mechanism, or equivalently, a restricted form of self-attention, with the following components:
Low-dimensional projection. The input æ; is projected into a latent space via learned orthonormal basis U æt, which is partitioned into K groups of dimension d. Each block U/ can be viewed as an individual expert operating on distinct subspace of the input.
Expert weighting. Each projected latent group is then reweighted by a coefficient 8/ which depends on the input æ; and timestep / via softmax function wr (xx,,.) These coefficients form the block-diagonal matrix ( l). which scales each group independently. This structure can be interpreted as a mixture-of-experts model (44) where the contribution of each expert is modulated by the input. It may also be viewed as restricted self-attention mechanism [45), where attention is applied the group level rather than individually, yielding the feature representations of the input as in
Symmetric reconstruction. The modulated feature representation is projected back to the input space via the same expert blocks U, forming symmetric encoder-decoder architecture.
Moreover, this parameterization in 4 induces a time and data-dependent feature representation he h2, t), which enables systematic analysis of representation quality across noise scales.
3.3 A metric for measuring representation quality
To understand diffusion-based representation learning under the MoLRG data model, we define the following signal-to-noise ratio (SNR) to measure the representation quality follows.
Definition 1. Suppose the data 2n follows the noisy MoLRG introduced Assumption | Without loss of generality, let k [K] denote the true class of ang For any trained DAE ths parameterized as
in we define:
Ex, U ho (z) 1) k)
SNR(
Ex. @1 (x,, U.he ho (z,, t) 1
where do ( () U ho () denotes the decoded reconstruction, and Už is the basis matrix for the true class data subspace.
This formulation measures how well the learned feature emphasizes the signal components aligned with the true class subspace versus those aligned with irrelevant or confounding directions. Intuitively, the numerator captures the energy of the feature projected onto the correct class subspace, while the
5

---

denominator measures the residual energy after removing this component from the reconstructed signal. A higher SNR value at given noise level ż indicates that the representation hgg encodes more discriminative structure with respect to the true class, implying better alignment with the downstream classification objective.
3.4 Main theoretical results
Based upon the aforementioned setup, we show the following results explain the unimodal representation dynamics.
Proposition 1. Suppose the data Bo is drawn from noisy MoLRG data distribution with K-class and noise level ó introduced in Assumption [ Then the optimal {U} minimizing the loss (2 is the ground truth basts defined in (3). and the optimal DAE 2g (x t) admits the analytical form:
K
1 (2 1) = Σ wi (at, (GUFUT + -GÜūūūT) Tt.
(1t,
where wi (Iz, f) are the coefficients in when (U) = (U
Link to the fine-to-coarse generation shift. Since the -rerlteeed component captures low- dimensional class-relevant attributes and the -relttee component captures small-scale, class- irrelevant attributes, the optimal DAE exhibits fine-to toeerr-- transition [8] in its output, where the class-irrelevant attributes are progressively removed as the noise level d, increases. Specifically, Se T+o?
quantifies the reduction rate of U,: term while & = कभッ quantifies the reduction rate of
8²
U/ term, or grows, & decays much more rapidly than Çt, indicating that the output retains more class-related coarse information while discarding fine-grained, irrelevant details.
This phenomenon correspond to an important observation formalized in the next theorem: there exists a balance timestep during the diffusion process. at which Mlass-rrreevaan components are maximally suppressed while class-relevant component is preserved, yielding peak classification accuracy from the feature. Substituting the optimal DAE formulation into (6), we can approximate the SNR the following theorem and analyze the unimodal dynamics via the approximation:
Theorem (Informal) Suppose that the data Fo follows the noisy MoLRG introduced in Assumption Then the SNR of the optimal DAE in can be approximated as follows:
o² hi (1+ [ 1 SNR(# (K (K-1)
2
Ct
(1+ ( k
(8)
Here, Ct is a monotonically decaying constant that has minimal impact to the overall unimodal shape. The function h(w, 5) (1 82)w + 52 is monotonically increasing in w, where h(w ,5) and h h(w, 8) denote positive and negative class confidence rates with
w (az, 5) = E& Eu,, ( t) k:],
W (a,, 5) = E& Eu, [w] (I,, t) le + Z]].
which are the softmax coefficients assigned to the correct class k and the other classes 1 o k.
We defer the formal statement of Theorem | and its proof to Appendix E.2 In the following, we discuss the implications of our result.
The unimodal curve of SNR across noise levels. Intuitively, our theorem shows that a unimodal curve is mainly induced by the interplay between the "denoising rate" /82 and the positive class confidence rate h(0t, 5) as the noise level increases. As observed in Figure the "denoising rate" (o? (02) increases monotonically with o,, while the class confidence rate h ( monotonically declines. Initially, when o, is small, the class confidence rate remains relatively stable due its flat slope, and an increasing "denoising rate" improves the SNR. However, as indicated by (7). when o, becomes too large, h(wr, 5) quickly decays and approaches h(ū, 8), leading
drop in the SNR. It is worth noting that though we are mainly characterizing the unimodal dynamics of an approximation of SNR, in practice it closely mimic the trend of the actual SNR as in Figure 3b

---

46.0
40.0
SNR(Gg.
SNR(Eg.
32.0
SNR(
32.0
Approx. (Thm. I)
24.0
Feature Ace.
Fosture
SNR
24.0
8.0
8.0
0.05
0.1
0.28
0.48
0.7
0.97
0.03
0.24
0.48
0.7
0.97
Time step (dt,)
Time step (or)
(a) SNR aligns with feature probing accuracy
(b) Tightness of the derived approximation
Figure 3: Feature probing accuracy and associated SNR dynamics in MoLRG data. In (anel(a) we plot the probing accuracy and SNR with the feature obtained from a learned DAE žg. both of which exhibit consistent unimodal pattern. The DAE is trained on a 3-class MoLRG dataset with data dimension n = 50, subspace dimension = 5, and noise scale = 0.2. Additionally, in pbnel(b) include the optimal SNR calculated from the optimal DAE #t and the derived approximation in Theorem as reference.
3,80
2.75
Feature Tost Acc.
0.00
Feature Acc.
$N
Feamir Test Acc.
3.50
2.25
0.04
2.0
SNR[#g:
40
175
30
0/0
8:14
0.59
1.92
532
Time step (Oi)
Time step (0x)
(a) CIFARIO
(b) TinyImageNet
Figure 5: Dynamics of feature accuracy and associated SNR on CIFAR1O and TinyImageNet.
Feature accuracy is plotted alongside SNR(2g. t). Feature accuracy is evaluated on the test while the empirical SNR is computed from the training set. Both exhibit an aligning unimodal pattern. We released EDM models 49) trained on the CIFARIO (50) and ImageNet (51) datasets, evaluating them on CIFAR 10 and TinylmageNet [52], respectively. To compute SNR we apply PCA on the CIFARIO/Tinylaage Appendix D
features to extract the basis Ujs. Further details can be found in
Alignment of SNR with representation learn- ing performance. As shown in Figures 3 and[5] our theory derived from the noisy MoLRG distribution effectively captures real-world phe- nomena. Specifically, we conduct experiments on both synthetic (i.e., noisy MoLRG) and real- world datasets (i.e., CIFAR and ImageNet) to measure SNR(&g, t) as well as the feature prob- ing accuracy. For feature probing, we use fea- tures extracted at different timesteps inputs for linear probing. The results consistently show that SNR( follows unimodal pattern across all cases, mirroring the trend observed in feature probing accuracy as the noise scale increases. This alignment provides formal justifica- tion for previous empirical findings [19] (0) (33), which have reported a unimodal trajectory in the representation dynamics of diffusion models with increasing noise levels. Detailed experimental setups are provided in Appendix D.
0.9
M(u .6)
5.0
0.8
9.0
0.05
0.24
0.43
0.65
0.9
1,21
1.6
2.12
stop (ot)
Figure 4: Illustration of the interplay between the denoising rate and the class confidence rate. The settings follow Figure 3
Unimodal representation dynamics predicts model generalization
In the previous sections, we theoretically showed that when diffusion model successfully captures the low-dimen www-dimensional distribution of the data, the unimodal representation dynamics emerge. In
7

---

Generaization
10
80
UNa-54
80
Fealure Test acc.
70
600
0.
0.6
30
04-
10
30
UNer-32
30
UN
20
0.0 0.01
0.02
0,56 0.14 0.3 0.59
1.09
192
3.26
24
2"
0.0 0.51 0.02 0.06 0.14 055 0.59
1:59
1.92
3.26
Time steps (17)
Training dotanet size
Time stops (o))
(a) UNet-32
(b) Generalization score
(c) UNet-64
Figure 6: Representation dynamics across model and data sizes. We train DDPM-based UNet-32 and Unet-64 diffusion models on the CIFARIO dataset using different training dataset sizes, ranging from 2 to 215 The unimodal representation dynamics across noise levels consistently emerges in the generalization regime (sufficient data size) and gradually disappears in smaller data settings.
this section, we investigate the opposite direction: can the presence of the unimodal representation dynamics serve reliable prediction of good generalization of diffusion models?
Answering this question sheds light on the distribution learning capabilities of diffusion models and closely related recent studies on their generalizability, which reveal that diffusion models operate in two distinct size, model capacity, and training duration captures the underlying data distribution and generates diverse, novel samples. In contrast, in the memorization regime, it overfits to the training data and loses the ability to generate novel samples. Further discussion on the generalization of diffusion models are provided in Appendix A
and cetmnorizationddeeeeeee
on factors such as dataset [32 53) In the generalization regime, the model
In this section, using classification tasks as case study, empirically demonstrate that the presence of unimodal representation dynamics reliably indicates generalization, while its gradual shift to a monotonically decreasing trend indicates memorization. Specifically, study the effects of data size and model capacity in Section 4.1 and the effects of learning dynamics in Section 4.2
4.1 Effects of dataset size and model capacity on representation dynamics
Recent studies [31 have shown that diffusion models exhibit phase transition from memorization to generalization the number of training samples increases. Specifically, when the network size significantly exceeds the number of training samples, diffusion models tend to memorize rather than capture the underlying wiwdimensional data distribution 332 53). For fixed model capacity,
generalization typically emerges when the number training data is larger than a certain threshold; Figure (b)) Here, demonstrate that the representation dynamics undergo similar transition with data scaling. Specifically, we train UNet-32 and UNet-64 models using varying training dataset sizes to examine how their representation dynamics change
regimes. The generalization score metric introduced in [53 57 is used as a reference for quantifying model generalization.
As shown in Figure 6 we observe that reducing the training dataset size leads decrease in the generalization score. Correspondingly, the unimodal representation dynamic becomes less obvious, eventually transitioning into monotonically decreasing curve (i.e., monotonic representation dy- namics). These observations highlight a strong connection between representation and distribution learning in diffusion models specifically, the emergence of unimodal representation dynamics aligns with the ability of the model to capture the underlying data distribution for achieving good generalization.
4.2 Effects of learning dynamics on representation dynamics
Second, we investigate how learning dynamics influence representation dynamics and also general- ization performance, particularly in the limited data regime (e.g., training size N 212 as shown in Figure (b)). In this regime, recent studies (54-56) have shown that early stopping improve

---

Peak Feature Acc.
Peak Feature Acc.
40
Peak Acc.
Peak Acc.
30
20
FID
FID
15
30-
70
20
10
2.5m
5.0m
7.5m
10m
15m
2.5m
5.0m
7.5m
10m
15m
Training Iteration (imgs)
Training Iteration (imgs)
(a) CIFARIO
(b) CIFAR100
Figure 7: Negative correlation between peak classification accuracy and FID. We train UNet-128 diffusion models on N = 2 training samples from CIFARIO and CIFARI00, As training progresses, the peak representation accuracy across noise levels shows consistent negative correlation with FID.
generalization performance. Specifically, these works observe an early learning phenomenon, where generalization improves during the initial phase of training but deteriorates the model begins to memorize. As illustrated in Figure 7 this effect is reflected in the evolution of the FID score (58), which initially decreases (indicating better generative quality) and then rises as memorization starts. Notably, we find that this trend negatively correlates with the linear probing accuracy of learned representations. This observation implies that representation quality could potentially serve as an carly-stopping criterion prevent memorization in diffusion models trained on limited data without relying on external models.
Moreover, we show that this early learning behavior can be captured by the transition between unimodal and monotonic curves of representation dynamics during training. Experimentally, we demonstrate this by training EDM-based diffusion models on subsets of CIFAR[50 using N = 2 training samples. Specifically, from Figure we observe that the evolution of representation dynamics during training can be divided into two distinct phases:
Early phase of unimodal representation dynamics. In the early stage of training, the model learns the underlying low-dimensional data distribution and is able to generalize. As predicted by our theoretical analysis in Section 3 representation follows a unimodal dynamic across noise scales. This unimodal dynamic is clearly observed before training iteration Iter < 7.5M in Figure (a). The generalization behavior is further supported by the new outputs of the model as shown Figure 8 (b), which more closely resemble those from reference generalized model than the nearest neighbors in the training set. Moreover, the peak representation quality improves steadily during this phase as the model better captures the data distribution.
Late phase of monotonic representation dynamics. However, as training progresses toward conver- gence, the model begins to memorize the training samples, resulting in reduced ability to capture the underlying data distribution. This transition is obvious the outputs of the model, which increasingly replicate training examples (see Figure g (b) Iter 15M, 100M). During this phase, the unimodal representation dynamics give way to a monotonic representation dynamics, where representation quality consistently degrades the noise level increases. Furthermore, as shown in Figure 7 the learned features become less informative, and the peak probing accuracy begins decline.
Conclusion
In this work, we developed a mathematical framework for analyzing the representation dynamics of diffusion models. By introducing the concept of SNR under mixture of low-rank Gaussians, we showed that the widely observed unimodal representation dynamic across noise scales emerges naturally when diffusion models capture the underlying data distribution. This behavior arises from trade-off between denoising strength and class confidence across noise levels. Beyond theoretical insights, our empirical results demonstrate that the emergence of unimodal representation dynamics closely linked to the model's distribution learning and generalization ability. Specifically, this unimodal pattern consistently appears when the model generalizes and gradually fades as the model starts to memorize. Our findings take step toward bridging the gap between generative modeling and representation learning in diffusion models.

---

Representation Dynamics at Different Training Itemaiins
Generalized EDM Reference
Training lteaatin 7.5m memonizatioo
Gen NN
Feature Test acc.
30
Training lieration 15m [Slight mncrmmizatiion
Gen NN
24mm: Unimeral
Training Iteation 100m [Sevare memorizationl
7.5m: Unimscaa
Gen
Uninocal X
Kiler Unimocal
NN
0.0
0.02
0.14
0.59
192
steps
(0) Representation dynamiis
(b) Generated images neighhors at different training iterations
at different training iterations
Figure 8: Representation learning and generative performance across training iterations. We train a UNet-128 diffusion model on N = 2 training samples in CIFARIO, monitoring both representation learning and generative performance training progresses. clear phase transition is observed: early in training, the representations exhibit a unimodal pattern, and generated samples resemble those of generalizing EDM model, with no signs of memorization. As training continues, the unimodal pattern gradually transitions to a monotonically decreasing trend, aligning with the model's ss shift toward memorizing the training data. "NN" denotes nearest neighbor in the training dataset.
While our analysis captures key aspects of representation dynamics, it relies on simplified data assumptions and model formulations that facilitate tractable derivations. Extending this framework to more realistic data distributions and complex model architectures remains an important direction for future work. Moreover, establishing rigorous theoretical connection between the representation geometry of diffusion models and their phase transition between generalization and memorization remains an important and promising direction for future work.
Acknowledgment
We acknowledge funding support from NSF CAREER CCF-2143904. NSF CCF-2212066, NSF CCF- 2212326, NSF IIS 2312842, NSF IIS 2402950, NSF IIS 2312840, NSF IIS 2402952, ONR N00014-22-1-25299 ONR 000014251239, and MICDE Catalyst Grant. We also thank all the anonymous reviewers for their valuable suggestions and fruitful discussions.
References
[1] Ismail Alkhouri, Shijun Liang. Rongrong Wang, Qing Qu, and Saiprasad Ravishankar. Diffusion- based adversarial purification for robust deep mri reconstruction. In ICASSP 2024- 2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 12841-12844 IEEE, 2024.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances Neural Information Processing Systems, 33:6840-6851. 2020.
[3] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEF/VV Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[4] Huijie Zhang, Yifu Lu, Ismail Alkhouri, Saiprasad Ravishankar, Dogyoon Song, and Qing Qu. Improving training efficiency of diffusion models via multi-stage framework and tailored multi-decoder architectures. In Conference on Computer Vision and Pattern Recognition 2024,
2024.
10

---

[5] Omer Bar-Tal, Hila Chefer, Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Guanghui Liu, Amit Raj. et al. Lumiere: A space-time diffusion model for video generation. In SIGGRAPH Asia 2024 Conference Papers, pages 1-11, 2024.
[6] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David Fleet, et al. Imagen video: High definition video generation with diffusion models, arXiv preprint arXiv:2210.02303, 2022.
[7] Jungil Kong, Jachyeon Kim, and Jaekyoung Bae. Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis. Advances in neural information processing systems, 33:1 7022-17033, 2020.
[8] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DIFFWAVE: A versatile diffusion model for audio synthesis. In International Conference Learning Representations, 2021.
[9] Daniel Roich, Ron Mokady, Amit H Bermano, and Daniel Cohen-Or Pivotal tuning for latent-based editing of real images. ACM Trunsactions on Graphics (TOG), 42()):1-33 2022.
[10] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aber- man. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 22500-2233 22510, 2023.
[11] Siyi Chen, Huijie Zhang, Minzhe Guo, Yifu Lu, Peng Wang, and Qing Qu. Exploring low- dimensional subspace in diffusion models for controllable image editing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[12] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683-25555 256962 2022.
[13] Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. In The Twelfth International Conference on Learning Representations, 2024.
[14] Xiang Li, Soo Min Kwon, Ismail R Alkhouri, Saiprasad Ravishanka, and Qing Qu De- coupled data consistency with diffusion purification for image restoration. arXiv preprint urXiv:2403.00054 2024.
[15] Ruihan Yang and Stephan Mandt. Lossy image compression with conditional diffusion models. Advances in Neural Information Processing Systems, 36:64971-64995, 2023.
[16] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256 22256.2265.222 PMLR, 2015.
[17] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
[18] Aapo Hyvärinen and Peter Dayan. Estimation of non non-normllized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[19] Dmitry Baranchuk, Andrey Voynov, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. Label-efficient semantic segmentation with diffusion models. In International Conference on Learning Representations, 2022.
[20] Weilai Xiang, Hongyu Yang. Di Huang, and Yunhong Wang. Denoising diffusion autoencoders are unified self-supervised learners. In Proceedings the IEEE/CVF International Conference Computer Vision, pages 15802-15812, 2023.
[21] Soumik Mukhopadhyay, Mathew Gwilliam, Vatsal Agarwal, Namitha Padmanabhan, Archana Swaminathan, Srinidhi Hegde, Tianyi Zhou, and Abhinav Shrivastava. Diffusion models beat gans on image classification. arXiv preprint arXiv:230..08002 2023.
11

---

[22] Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing denoising diffusion models for supersuue learning. arXiv preprint arXiv:2401.14404, 2024.
[23] Luming Tang, Menglin Jia, Qianqian Wang. Cheng Perng Phoo, and Bharath Hariharan. Emergent correspondence from image diffusion. Advances Neural Information Processing Systems, 36: 13611-36919 2023.
[24] Andi Han, Wei Huang, Yuan Cao, and Difan Zou. On the feature learning in diffusion models. arXiv preprint arXiw:2412.01021, 2024.
[25] Sixue Gong, Vishnu Naresh Boddeti, and Anil K Jain. On the intrinsic dimensionality of image representations. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3987-3996, 2019.
[26] Phil Pope,
Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Representations, 2021.
[27] Jan Stanczuk, Georgios Batzolis, Teo Deveney, and Carola-Bibiane Schönlieb. Your diffusion model secretly knows the dimension of the data manifold. arXiv preprint arXiv:2212.12611, 2022.
[28] Zhongqi Yue, Jiankun Wang. Qianru Sun, Lei Ji, Eric I-Chao Chang, and Hanwang Zhang, Exploring diffusion time-steps for unsupervised representation learning. In The Twelfth International Conference on Learning Representationss 2024.
[29] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-322 1982.
[30] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(496):1002-1614 2011.
[31] Zahra Kadkhodaie, Florentin Guth, Eero P Simoncelli, and Stéphane Mallat. Generalization in diffusion models arises from geomeury-daptive harmonic representations. In The Twelfth International Conference on Learning Representations, 2024.
[32] Peng Wang, Huijie Zhang, Zekai Zhang, Siyi Chen, Yi Ma, and Qing Qu. Diffusion models learn low-dimensional distributions via subspace clustering. arXiv preprint arXiv:2409.02426, 2024.
[33] Ehsan Elhamifar and René Vidal. Sparse subspace clustering: Algorithm, theory, and applica- tions. IEEE transactions on pattern analysis and machine intelligence, (5(11):27652273 2013.
[34] Peng Wang, Huikang Liu, Anthony Man-Cho So, and Laura Balzano. Convergence and recovery guarantees of the k-subspaces method for subspace clustering. In International Conference on Machine Learning, pages 22884-22918. PMLR, 2022.
[35] Zhihan Huang, Yuting Wei, and Yuxin Chen. Denoising diffusion probabilistic models optimally adaptive to unknown low dimensionality. arXiv preprint arXiv:2410.18784 2024.
[36] Jiadong Liang, Zhihan Huang, and Yuxin Chen. Low-dimensional adaptation of diffusion models: Convergence in total variation. arXiv preprint arXiv:2501.12982 2025.
[37] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
[38] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: Improving latent diffusion models for high-resolution image synthesis. The Twelfth International Conference on Learning Representations, 2024.
[39] Diederik P Kingma. Auto-encoding variational bayes. arXiv preprint arXiv: 1312.63111111 2013.
12

---

[40] Bowen Jing, Gabriele Corso, Renato Berlinghieri, and Tommi Jaakkola. Subspace diffusion generative models. In European Conference on Computer Vision, pages 274-289. Springer, 2022.
[41]
Konstantinos Bousmalis, George Trigeorgis, Nathan Silberman, Dilip Krishnan, and Dumitru Erhan. Domain separation networks. Advances in neural information processing systems, 29, 2016.
[42] Tao Zhou, Changqing Zhang, Xi Peng, Harish Bhaskar, and Jie Yang. Dual shared-specific multiview subspace clustering. IEEE transactions on cybernetics, 2019.
[43] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assiste assisteed intervention-MICC 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer, 2015.
[44] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- mixture- of-experts layer. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[45] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[46] Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pultern Recognition, pages I 1142--11481, 2022.
[47] Binxu Wang and John J Vastola. Diffusion models generate images like painters: an analytical theory of outline first, details later. arXiv preprini arXiv:2303.02490 2023.
[48] Mason Kamb and Surya Ganguli. An analytic theory of creativity in convolutional diffusion models. arXiv preprint 241222241,.,029,, 2024.
[49] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Proc. NeurIPS, 2022.
[50] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[51] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition. pages 248-255. Ieee, 2009.
[52] Serena Yeung Li Fei-Fei, Justin Johnson. Cs231n: Convolutional neural networks for visual recognition, 2015, http://cs231n stanf orddfrrdeoo edu
[53] Huijie Zhang, Jinfan Zhou, Yifu Lu, Minzhe Guo, Peng Wang. Liyue Shen, and Qing Qu. The emergence of reproducibility and consistency in diffusion models. In Forty-first International Conference on Machine Learning, 2023.
[54] Puheng Li, Zhong Li, Huishuai Zhang, and Jiang Bian. On the generalization properties of diffusion models. Advances in Neural Information Processing Systems, 36:2097-2127, 2023.
[55] Xiang Li, Yixiang Dai. and Qing Qu Understanding generalizability of diffusion models requires rethinking the hidden gaussian structure. Advances in Neural Information Processing Systems, 17:57499-57538. 2024.
[56] Ricardo Baptista, Agnimitra Dasgupta, Nikola B Kovachki, Assad Oberai, and Andrew M Stuart. Memorization and regularization generative diffusion models. arXiv preprint arXiv:2501.1578. 2025.
13

---

[57] Ed Pizzi, Sreya Dutta Roy, Sugosh Nagavara Ravindra, Priya Goyal, and Matthijs Douze. A self-supervised deseriptor for image copy detection. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14532-14542, 2022.
[58] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochre- iter. Gans trained by two time-scale update rule converge to local nash equilibrium- supplementary material. Advances in Neural Information Processing Systems, 2017.
[59] Pascal Vincent, H. Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and composing robust features with denoising autoencoders. In International Conference on Machine Learning, 2008.
[60] Pascal Vincent, H. Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion. J. Mach. Learn. Res., 11:3371-3408, 2010.
[61] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequenee-to-sequene pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 2020.
[62] B. Chandra and Rajesh Kumar Sharma. Adaptive noise schedule for denoising autoencoder. International Conference on Neural Information Processing, 2014.
[63] Krzysztof Geras and Charles Sutton. Scheduled denoising autoencoders. In International Conference on Learning Representations (ICLR) 2015, 2015.
[64] Qianjun Zhang and Lei Zhang. Convolutional adaptive denoising autoencoders for hierarchical feature extraction. Frontiers of Computer Science, 12:1 140 1148, 2018.
[65] Arnu Pretorius, Steve Kroon, and Herman Kamper. Learning dynamics of linear denoising autoencoders. In International Conference Machine Learning, pages 4141-4150. PMLR, 2018.
[66] Harald Steck. Autoencoders that don't 't overfit towards the identity. In Neural Information Processing Systems, 2020.
[67] Daniel Kunin, Jonathan Bloom, Aleksandrina Goeva, and Cotton Seed. Loss landscapes of regularized linear autoencoders. In International conference on machine learning, pages 3560-3569. PMLR, 2019.
[68] Michael Fuest, Pingchuan Ma, Ming Gui, Johannes S Fischer, Vincent Tao Hu, and Bjorn Om- mer. Diffusion models and representation learning: A survey, arXiv preprint arXiv:2407.00783, 2024.
[69] Kamil Deja, Tomasz Trzciński, and Jakub M Tomezak. Learning data representations with joint diffusion models. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 543-559. Springer, 2023.
[70] Yujun Shi, Chuhui Xue, Jun Hao Liew, Jiachun Pan, Hanshu Yan, Wenqing Zhang, Vincent YF Tan, and Song Bai. Dragdiffusion: Hamessing diffusion models for interactive poiet-based image editing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8839-8849, 2024.
[71] Chandramouli Sastry, Sri Harsha Dumpala, and Sageev Oore. Diffaug: A -inuuse-and-denoise augmentation for training robust classifiers. arXiv preprint arXiv:2306. 09192, 2023.
[72] Xingyi Yang and Xinchao Wang. Difffusion model as representation learner. In Proceedings of the IEEF/CVE International Conference on Computer Vision, pages 18938-18944) 2023.
[73] Daiqing Li, Huan Ling, Amlan Kar, David Acuna, Seung Wook Kim, Karsten Kreis, Antonio Torralba, and Sanja Fidler. Dreanteacher: Pretraining image backbones with deep generative models. In Proceedings of the IEFE/CVE International Conference on Computer Vision, pages 16698-16708. 2023.
14

---

[74] Nick Stracke, Stefan Andreas Baumann, Kolja Bauer, Frank Fundel, and Bjöm Ommer. Cleandift: Diffusion features without noise. arXiv preprint arXiv:2412.03439, 2024.
[75] Grace Luo, Lisa Dunlap, Dong Huk Park, Aleksander Holynski, and Trevor Darrell, Diffusion hyperfeatures: Searching through time and space for semantic correspondence. Advances in Neural Information Processing Systems, 36, 2024.
[76] Korbinian Abstreiter, Sarthak Mittal, Stefan Bauer, Bernhard Schölkopf, and Arash Mchrjou. Diffusion-based representation learning. arXiv preprint arXiv:2105. 2021.
[77] Yingheng Wang. Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusion models. In International Conference on Machine Learning, pages 36336-36354, PMLR, 2023.
[78] Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James L McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. Soda: Bottleneck diffusion models for representation learning. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23115-23127, 2024.
[79] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakom. Diffusion autoencoders: Toward meaningful and decodable representation. In Proceedings of the IEEE/CVF conference on compuler vision and pattern recognition, pages 10619-10629, 2022.
[80] Yujin Han, Andi Han, Wei Huang. Chaochao Lu, and Difan Zou. Can diffusion models learn hidden inter-feature rules behind images?, 2025.
[81] Matthew Niedoba, Berend Zwartsenberg, Kevin Patrick Murphy, and Frank Wood. Towards mechanistic explanation of diffusion model generalization. In Forty-second International Conference on Machine Learning, 2025.
[82] Artem Lukoianov, Chenyang Yuan, Justin Solomon, and Vincent Sitzmann. Locality image diffusion models emerges from data statisties. Advances in Neural Information Processing Systems, 2025.
[83] Gabriel Raya and Luca Ambrogioni. Spontaneous symmetry breaking in generative diffusion models. Advances in Neural Information Processing Systems, 2023.
[84] Giulio Biroli, Tony Bonnaire, Valentin De Bortoli, and Mare Mézard. Dynamical regimes of diffusion models. Nature Communications, 2024.
[85] Antonio Sclocchi. Alessandro Favero, and Matthieu Wyart. phase transition in diffusion models reveals the hierarchical nature of data. Proceedings of the National Academy of Sciences, 2025.
[86] Luca Ambrogioni. In search of dispersed memories: Generative diffusion models are associa- tive memory networks. Entropy, 2024.
[87] Carlo Lucibello and Mare Mézard. Exponential capacity of dense associative memories. Physical Review Letters, 2024.
[88] Bao Pham, Gabriel Raya, Matteo Negri, Mohammed Zaki, Luca Ambrogioni, and Dmitry Krotov. Memorization to generalization: Emergence of diffusion models from associative memory. arXiv preprint arXiv:2505.21777 2025.
[89] Anand Jerry George, Rodrigo Veiga, and Nicolas Macris. Analysis of diffusion models for manifold data. arXiv preprint 2502:22502.03339 2025.
[90] Beatrice Achilli, Enrico Ventura, Gianluigi Silvestri, Bao Pham, Gabriel Raya, Dmitry Krotov, Carlo Lucibello, and Luca Ambrogioni. Losing dimensions: Geometric memorization in generative diffusion. arXiv preprint arXiv:2410.087277 2024.
[91] Omkar M. Parkhi, Andrea Vedaldi, Andrew Zisserman, and C. V. Jawahar. Cats and dogs. In IEEE Conference on Computer Vision and Pattern Recognition, 2012.
15

---

[92] Ya Le and Xuan S. Yang. Tiny imagenet visual recognition challenge. 2015.
[93] Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over large number classes. 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing, pages 722-729, 2008.
[94] Anh-Dzung Doan, Bach Long Nguyen, Surabhi Gupta, Ian Reid, Markus Wagner, and Tat-Jun Chin. Assessing domain gap for continual domain adaptation in object detection. Computer Vision and Image Understanding, 238: 103885, 2024.
[95] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[96] M. Cimpoi, S. Maji, I. Kokkinos, Mohamed, and Vedaldi. Describing textures in the wild. In Proceedings of the IEEE Conf. on Computer Vision and Puttern Recognition (CVPR), 2014.
[97] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[98] Tero Karras, Samuli Laine, and Timo Aila. A Style-Based Generator Architecture for Generative Adversarial Networks IEEE Transactions on Pullern Analysis & Machine Intelligence, 43(12))4217-4228, 2021.
[99] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. International Conference on Learning Repre- sentations, 2018.
[100] tanelp. tiny-diffusion. http::://ggthuu...cmmmmmaee////p
2022.
[101] Diederik P Kingma. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
[102] Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016.
[103] Pascal Vincent. connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-166 2011.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Later sections follow the structure of our main contributions as outlined in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We note that we have carefully discussed the analysis tools, assumptions and the scope of experiments in our paper, certain assumptions and design choices (e.g., synthetic data models, classification-fouus focused benchmarks) may not capture the full complexity of all applications. We highlight several promising directions for future work, including broader task coverage and theoretical extensions, in the Appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
17

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We discussed the assumption used in our analysis and provide the full set of proof in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, we discuss the experimental details in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
18

---

reproducible, they 5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: Our work uses publically available datasets (such as CIFAR) and codebases (such as EDM).
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/nips. wttpcc/pcccccccccccooooooop httis//gguideee/eeeuubbiiiiii seriissssoooooiii
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We primarily use publicly available pre-trained models for evaluation. For models we train ourselves, the training setup is detailed in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The main empirical results in our paper focus on revealing trends and dynamics, rather than reporting relative accuracy difference.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
19

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, we discuss the computing resources and other details in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: :tttlps://neuuiiii hww.rpsu...../ wwps.ccc/put hccppullic/tthhccccidee........e
Answer: [Yes]
Justification: Yes, the research conducted in the paper conforms the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of Machine Learning. There's societal consequences of our work we feel must be highlighted here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
20

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We carefully cite the data/code/model assets used in the paper and discuss the license (if applicable) in Appendix D
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
21

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets is released in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No such experimentssresearch is involved in the paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No such experiments/research is involved in the paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
22

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
23

---

The Appendix is organized as follows: in Appendix A we discuss related works; in Appendix B we present auxiliary findings that complement the main discussion; in Appendix we provide additional complementary experiments; in Appendix [ we present the detailed experimental setups for the empirical results in the paper. Lastly, in Appendix we provide proof details for Section
A Related works
Denoising autoencoders (DAEs) are trained to reconstruct corrupted images to extract semantically meaningful information, which can be applicd to various vision [59] 60] and language downstream tasks [6]]. Related to analysis of the weight-sharing mechanism, several studies have shown that training with a noise scheduler can enhance downstream performance 64).64] On the theoretical side, prior works have studied the learning dynamics 655 66 and optimization landscape [67] through the simplified linear DAE models.
Denoising auto-encoders.
Diffusion-based representation learning.
ceftusion-basete representation learning 68 has demon-
strated significant success in various downstream tasks, including image classification 20 21 (69] segmentation [19], correspondence [23], and image editing [70]. Using diffusion models for data augmentation has also been shown to improve robustness against covariate shift 77. To further enhance the utility of diffusion features, knowledge distillation (72-75) methods have been proposed, aiming to bypass the computationally expensive grid search for the optimal in feature extraction and improving downstream performance. Beyond directly using intermediate features from pre- trained diffusion models, research efforts has also explored novel loss functions [76] and network modifications 78 [99 develop more unified generative and representation learning capabílities within diffusion models. The work (80) investigates whether diffusion models are capable of learning latent inter-featuu dependencies underlying image data. Unlike the aforementioned efforts, our work focuses more on understanding the representation dynamic of diffusion models and its relationship with model generalization.
Generalization of Diffusion Models. In this paper, we examine the relationship between a uni- modal curve in representation-prole
accuracy and the generalization of diffusion models, a topic of There has also been notable work on the mechanisms that enable diffusion models to learn the underlying score function from discrete empirical samples [8] 82 thereby gen-
active interest 31
48 53
erating novel in-distribution samples (i.e., generalizing). We propose a representational perspective on the generalization diffusion models, together with corresponding measures (unimodality in representation dynamics) for evaluating generalization.
The unimodality of representation-prooinn accuracy across time steps is also related to symmetry breaking [3] and carly-stopping generalization during sampling 84
85 when the model learns empirical score functions. This line of work further connects empirical DMs and dense associative memory (AM) networks through their objectives and sampling behavior 86 88). where generalization interpreted as novel attraction sinks in AM, which is also related to Gaussian and spherical patterns in the data. However, in this paper our focus is on learning a parameterized model with data from dinensnnsiooaa distribution, an assumption also used in [89 90) (which links generalization to adapting to fow-dimensional data manifolds and learning local geometric components), rather than on sampling.
B Auxiliary results
B.1 Extended results from Section 4.2
In Section 4.2
we showed that when training on limited data, the representation dynamics undergo clear phase transition: from a unimodal pattern to monotonically decreasing trend, where the diffusion model also exhibits a transition from generalization to increasingly memorize the training data. Here, provide additional empirical evidence supporting this insight.
We train UNet-based diffusion models using the DDPM+ architecture with the EDM configuration [49 on the Oxford-IIIT Pet [91] and TinylmageNet 992 datasets, using training subsets of 3680 and 2048 images, respectively. Throughout training, we monitor both the evolution of representation
dynamics and generative outputs. As shown in Figure and Figure 10 consistent with our findings
24

---

Training Tteration 5.0m No memonization
Representation Dynamics
Gen NN
at Different Training lterations
27.5
25.0
Training Iteration 10.0m [No memorization]
Gen
Feature Test sec.
NN
20.0
Training Iteration 25.0m [Slight mecmorivation
Gen
15.0
NN
Unimodal
Training Iteration - 60.0m [Severe meniorization
10o:: Uninodal
Gen
25mM: Uninodal
60m: Unimodal
NN
0.0
0.02
0.14
0.59
1.92
Time steps (o,)
(a) Representation dynamies
(b) Generated images and nearest neighbors al different training ilerations
al different training iterations
Figure 9: Representation learning and generative performance across training iterations for Oxford-IIIT Pet dataset [91]. We train UNet-128 diffusion model on 3680 training samples in Oxford-III' Pet dataset, monitoring both representation learning and generative performance training progresses. clear phase transition is observed: early in training, the representations exhibit a unimodal pattern, and generated samples resemble those of a generalizing EDM model, with signs of memorization. As training continues, the unimodal pattern gradually transitions to monotonically decreasing trend, aligning with the model's 's shift toward memorizing the training data. "NN" denotes nearest neighbor the training dataset.
in Section 4.2 we observe that in the early stages of training, the model exhibits clear unimodal pattern and generalizes well. As training continues, this unimodal structure gradually shifts into monotonically decaying trend, and the models start increasingly replicate training examples.
B.2 Disentangling the role of input noise in representation dynamics
One might argue that the declining portion of the unimodal curve is simply due the increasing noise level or, which makes the input Φ; progressively noisier, leading to a natural drop in classification accuracy. However, show that this noise-induced degradation alone does not account for the observed representation dynamics.
Our theoretical analysis in Section 3 attributes the unimodal pattern fundamental trade-off between denoising rate and class confidence rate across noise levels. To validate this explanation and disentangle the effect of additive Gaussian noise, conduct experiments where feature extraction performed directly on clean inputs Ey rather than noisy inputs a;, We show that the unimodal behavior remains clearly observable even in the absence of injected noise, indicating that the dynamics are not solely a consequence of input corruption. In fact, one can verify that SNR (mr)), () also exhibits unimodal trend- this can be shown through analysis analogous that of Theorem [ SNR( (xx)), ().
To support this claim, train EDM (with the VP configuration) 49) CIFAR datasets and classical DDPM 2 on CIFARIO and Flowers- 102 [93], and evaluate the representation quality using both noisy inputs 2, and clean inputs &g. shown in Figure II For Figure ITa)), a). select the layer that yields the best accuracy, while for Figure bTb),, we directly use the bottleneck layer to demonstrate that the observed unimodal behavior is not sensitive to layer choice. Across all settings, the unimodal representation dynamics remain clearly visible, reinforeing that additive Gaussian noise not the sole factor responsible for this phenomenon.
also visualize posterior estimation dh across noise scales using both noisy and clean inputs. Since diffusion models are trained approximate the posterior mean different noise levels.
25

---

Training lteration 5.0m NNo memorization
Representation Dynamics at Different Training Iterations
Ginn
NN
Training Iteration 12.5m NNo memorization)
Gin
Feature Test acc.
NN
Training Iicration 50.0m [Slight memorization
Gen
NN
5.0m: Unimodal
Training Iteration
95.0m [Severe memorizzttion
12.5m: Unimodal
Gen NN
50.0m: Unimodal X
95 Omc Unimodal X
3,2
0.14
Time steps (ar,)
(a) Representation dynamics
(b) Generated images and nearest neighbors al different training iterations
at different training ilerations
Figure 10: Representation learning and generative performance across training iterations for TinylmageNet [92]. We train UNet- 128 diffusion model on N 2 training samples in TinylmageNet, monitoring both representation learning and generative performance training progresses. A clear phase transition is observed: early in training, the representations exhibit unimodal pattern, and generated samples resemble those of generalizing EDM model, with signs of memorization. As training continues, the unimodal pattern gradually transitions to a monotonically decreasing trend, aligning with the model's shift toward memorizing the training data. "NN" denotes nearest neighbor in the training dataset.
........
CIFARIO (Ize)
CIFARR00 (Ite)
CIFARIO (Z;)
Flowers (m:,)
CIFARIO (ata)
CIFAR100 (ato)
CIFARIO (zo)
Flowers102 (x²g)
40.0
CIFARIO Test acc.
86.2
78.6
Lest acc.
CIFARIO Test acc.
72.4
34,0
77.4
65.2
28.0
68.2
22.0
59.8
38.4
46.6
DDPM model
16.0
$1.0
25.0
10.0
0.0
0.01
0.02
0.14 0.3
0.58
1.09
1.92
0.0
0.01
0.02
0.06
0.14
0.3
0.58
1.09
1.92
Time step (ax)
Time step (a))
(a) EDM
(b) Classical DDPM
Figure 11: Unimodal representation dynamics persist when using clean inputs æt. We train EDM and DDPM models on CIFAR and Flowers- 102 datasets and evaluate feature quality using both noisy inputs æ; and clean inputs æg. Across settings, the unimodal trend consistently observed.
their representation features emerge intermediate products of this denoising process. As such, improvements or degradations in representation quality should be mirrored the posterior estimates.
We visualize the posterior estimation results for clean inputs (ie((z, t)) and noisy inputs (0o(2,, t)) across varying noise scales de Figure [2 both cases, the posterior outputs undergo clear fine-to-coars, transition as de increases. This supports our theoretical claim that as noise grows, class-irrelevant attributes are gradually removed. The peak representation quality occurs at an intermediate point where class-essential structure is preserved while irrelevant details are suppressed. When o; becomes too large, class confidence rate drops significantly, resulting poor representations. The additive noise o;E merely accelerates this degradation but is not the root cause.
26

---

Og 0,023
R 0.96
d- 0.14
7 0,256
7 0.385
01. 1,038
1923
3.207
5315
Cr 8.401
Network Inguls
Posterior Extimation
(a) te ar, f): Posterior estimation using noise image inputs.
dh 3,257
6g 5,15
6 6.401
Network Inputs
Posterior Estinasion
(b) te (wo, l): Posterior estimation using clean image as inputs.
Figure 12: Fine-to-coarse shift in posterior estimation. We use a pre-trained DDPM diffusion model on CIFARIO visualize posterior estimation for clean inputs and noisy inputs across varying noise scales og We can observe seemingly fine-to-coarse shifts in both figures.
Text (Diffasion)
SWD (Diffasion)
Test acc. (Diffisiio)
SWD (Diffssion)
Test scc. (DAF)
SWD (DAF)
aec. (DAR)
SWD (DAF)
Peature Test ucc.
Feature Test aCc
60
60
SWD
40
SWD
40
20
20:
Time stop (a,)
Time step (ox)
(a) CIFARIO
(b) CIFARI00
Figure 13: Diffusion models exhibit higher and smoother feature accuracy and similarity compared to individual DAEs. We train DDPM-based diffusion models and individual DAEs on the CIFAR datasets and evaluate their representation learning performance. Feature accuracy, and feature differences from the optimal features (indicated by *) are plotted against increasing noise levels. The results reveal inverse correlation between feature accuracy and feature differences, with diffusion models achieving both higher/smoother accuracy and smaller/smoother feature differences compared DAES.
B.3
Weight sharing in diffusion models facilitates representation learning
While our theoretical analysis captures the emergence of unimodal representation dynamics under an idealized network parameterization, an important future direction extend this framework to deeper and more complex architectures. Real-world diffusion models often involve highly complex feature transformations, and understanding how these interact with noise scales to influence representation quality remains open and valuable avenue for exploration.
In this section, we present some interesting preliminary results found that may potentially explain why diffusion models outperform classical denoising autoencoders (DAEs) in representation learning: although both share the same denoising objective (2), diffusion models demonstrate superior feature learning capabilities largely due to their inherent weight-sharing mechanism. Specifically, by minimizing the loss across noise levels, diffusion models enable parameter sharing and interaction among denoising subcomponents, effectively creating an implicit "ensemble" effect. This interaction enhances feature consistency and robustness across noise scales, leading to significantly improved representation quality compared to DAEs [22], as illustrated in Figure 13
27

---

Method
MinilmageNer" Test Acc. %
Label Noïse
Clean
20%
40%
60%
80%
MAE
73.7
70.3
67.4
62,8
51.5
EDM
67.2
62.9
59.2
53.2
40.1
EDM (Ensemble)
72.0
67.8
64.7
60.0
48.2
DIT
77.6
72.4
68.4
62.0
47.3
DIT (Ensemble)
78.4
75.1
71.9
66.7
56.3
Table 1: Comparison test performance across different methods under varying label noise
levels. All compared models are publicly available and pre-trained on ImageNet-IK [51], evaluated using MinilmageNet classes. Bold font highlights the best result in each scenario.
To test this, we trained 10 DAEs, each specialized for single noise level, alongside DpPMMbaseed diffusion model on CIFARIO and CIFAR100. We compared feature quality using linear probing accuracy and feature similarity relative to the optimal features at oe 0.06 (where accuracy peaks) via sliced Wasserstein distance (SWD) 994
The results in Figure 13 confirm the advantage of diffusion models over DAFs. Diffusion models consistently outperform DAES, particularly in low-noise regimes where DAEs collapse into trivial identity mappings. In contrast, diffusion models leverage weight-sharing to preserve high-quality features, ensuring smoother transitions and higher accuracy as noise increases. This advantage is further supported by the SWD curve, which reveals an inverse correlation between feature accuracy and feature differences. Notably, diffusion model features remain significantly closer to their optimal state across all noise levels, demonstrating superior representational capacity.
Our finding also aligns with prior results that sequentially training DAEs across multiple noise levels improves representation quality 64) Our ablation study further confirms that multi-scale training essential for improving DAE performance classification tasks in low-noise settings (details in Appendix
C Table 3).
Beyond the implicit feature ensembling effect, we further introduce straightforward method that explicitly ensembles features from multiple noise levels to enhance downstream task performance. Our experiments demonstrate that this approach significantly improves robustness against label noise in classification tasks, both in pre- training and transfer learning settings. For detailed methods and results, refer interested readers Appendix B.4
B.4 Feature ensembling across timesteps improves representation robustness
Our theoretical insights in Section 3 imply that features extracted at different timesteps capture varying levels granularity. Given the high linear separability of intermediate features, we propose simple ensembling approach across multiple timesteps to construct more holistic representation of the input. Specifically, addition to the optimal timestep. we extract feature representations at four additional timesteps- two from the coarse (larger σ,) and two from the fine-gran fine-grained (smaller o,) end of the spectrum. We then train linear probing classifiers for each set and, during inference, apply soft-voting ensemble by averaging the predicted logits before making final decision (experimet details in Appendix D)
We evaluate this ensemble method against results obtained from the best individual timestep, as well as self-supervised method MAE [95, on both the pre-training dataset and a transfer learning setup. The results, reported in Table I and Table 2 demonstrate that ensembling significantly enhances performance for both EDM 49 and DiT 37,, consistently outperforming their vanilla diffusion model counterparts and often surpassing MAE. More importantly, ensembling substantially improves the robustness of diffusion models for classification under label noise.
C Additional Experiments
Validation of - ippprox approximation in Appendix E.2 In Theorem we approximate the optimal
posterior estimation function á using d= appmor by taking the expectation inside the softmax with respect to Et. To validate this approximation, we compare the SNR calculated from ag and from appprox using the definition in Proposition I and (0) in Appendix E.2 respectively. We use fixed
28

---

Method
Transfer Test Acc. %
CIFAR100
DTD
Flowers102
Label Noise
Clean
20%
40%
60%
80%
Clean
20%
40%
60%
80%
Clean
20%
40%
60%
80%
MAE
63.0
58.8
54.7
50.1
38.4
61.4
54.3
49.9
40.5
24.1
68.9
55.2
40.3
27.6
9.6
EDM
62.7
58.5
53.8
48.0
35.6
54.0
49.1
45.1
36.4
21.2
62.8
48.2
37.2
24.1
9.7
EDM (Ensemble)
67,5
64.2
60.4
55.4
43.9
55.7
49.5
45.2
37.1
22.0
67.8
53.9
41.5
25.0
10.4
64.2
58.7
53.5
46.4
32.6
65.2
59.7
53.0
43.8
27.0
78.9
65.2
52.4
34,7
13.3
DIT (Ensemble)
66.4
61.8
57.6
51.3
39.2
65.3
60.6
56.1
46.3
30.6
79.7
67.0
54.6
36.6
14.7
Table 2: Comparison of transfer learning performance across different methods under varying
label noise levels. All compared models are publicly available and pre- pre-trained on ImageNet- IK [5I], evaluated on different downstream datasets. Bold font highlights the best result in each scenario.
K-4
K-6
K-8
K-19
104
10
SNIROL SNRO
SNIIEL SNROK
SNROKG D SNIKK
314 sto
sta
aic cls Ttes d=
Timeanep (0)
Time rape (0)
Time rape (0) ==:
d
df
SNIVNN
SNIING
SNRING 0
SNIROL SNROK
SNRRR
SNR
SNRo 19
sips
Tiis (0) 6w
Tinserigsio
0.0
fu 0.2
6..0.5
340
SNRANG SNR. 5)
SNIINN SNR ?)
SNRING SNRI
SNBAA 4)
S
Figure 14: Comparison between SNR calculated using the optimal model di and the SNR cal- culated with our approximation in Theorem I We generate MoLRG data and calculate SNR using both the corresponding optimal posterior function ag and our approximation a Hipprox from Theorem Default parameters are set as n = 100, d 5, K = 10, and 8 = 0.3. In each row, we vary one parameter while keeping the others fixed, comparing the actual and approximated SNR.
dataset size of 2400 and set the default parameters to n 100, d d= 5, K = 10, and d 0.3 to generate MoLRG data. We then vary one parameter at a time while keeping the others constant, and present the computed SNR in Figure 14 As shown, the approximated SNR score consistently aligns with the actual score.
Visualization of the MoLRG posterior estimation and SNR across noise scales. In Figure 3 we show that both the classification accuracy and SNR exhibit unimodal trend for the MoLRG data. To further illustrate this behavior, we provide a visualization of the posterior estimation and SNR at different noise scales in Figure [5] In the plot, each class is represented by colored straight line, while deviations from these lines correspond to the 8-related noise term. Initially, increasing the noise scale effectively cancels out the 6-related data noise, resulting in cleaner posterior estimation and improved probing accuracy. However, as the noise continues to increase, the class confidence rate drops, leading an overlap between classes, which ultimately degrades the feature quality and probing performance.
Mitigating the performance gap between DAE and diffusion models. Throughout the empirical results presented in this paper, we consistently observe performance gap between individual DAES and diffusion models, especially in low-n regions. Here, we use a DAE trained on the CIFAR 10 dataset with a single noise level 0.002, using the NCSN+ architecture [49. In
29

---

SNR 12:
SNR: 15:23
SNK: 29.36
SNK 12:76
SNR: 0.96
Diflissin
Input
SNR
Optimal
0.02
0.09
0.28
0.58
1.71
Figure 15: Visualization of posterior estimation, higher SNR correspondings to higher clas- sification accuracy. The same MoLRG data is fed into the models; each row represents different denoising model, and each column corresponds to different time step with noise scale (T₂).
Table 3: Improve DAE representation performance at low noise region. vanilla DAE trained on the CIFAR1O dataset with single noise level of a 0.002 serves the baseline. We evaluate the performance improvement of dropout regularization, EDMM-bsed preconditioning, and multi-level noise training (a {0.002, 0.012, 0.102)). Each technique is applied independently to assess its contribution to performance enhancement.
Modifications
Test acc.
Vanilla DAE
32.3
Dropout (0.5)
35.3
+ +rropout (0.9)
36.4
+rropout (0.95)
38.1
+DDM preconditioning
49.2
Muulti-lvee noise training
58.6
the default setting, the DAE achieves a test accuracy of 32.3. We then explore three methods to improve the test performance: (a) adding dropout, as noise regularization and dropout have been effective in preventing autoencoders from learning identity functions [66]; (b) adopting EDM- based preconditioning during training, including input/output scaling, loss weighting, etc.; and (c) multi-level noise training, in which the DAE trained simultaneously on three noise levels [0.002, 0.012, 0. 0.2. Each modification is applied independently, and the results are reported in Table 3 As shown, dropout helps improve performance, but even with a dropout rate of 0.95, the improvement minor. EDM-based preconditioning achieves moderate improvement, while multi- level noise training yields the most promising results, demonstrating the benefit of incorporating the diffusion process DAE training.
Empirical robustness to assumption relaxation The data assumptions in Assumption D were made simplify the analysis and derive closed-form results. Empirically, we find that the unimodal SNR trend remains robust when these assumptions are relaxed. To demonstrate this, we train our parameterized DAE under binary MoLRG data while systematically violating each assumption. Unless otherwise stated, each experiment uses the DAE network as introduced in Section 3.2 and a dataset contains 30,000 samples generated with 50, d 5, and 8 0.2.
control the principal angle (0) between class subspaces and observe that overlap tends to reduce SNR across timesteps while the peak remains stable. We conjecture that overlap can be viewed as introducing additional intrinsic noise beyond the ò-related term, thereby affecting the SNR value.
Overlapping class subspaces. As shown in Table
Varying subspace ranks. As shown in Table we set the class subspace dimensions to dg = 10 and di 2. Intuitively, the higher-rank class retains
signal and is less sensitive to noise, yielding higher and later-peaking SNR, while the low-rank class decays earlier.
Non-uniform mixing weights. As shown in Table 6 we set mo = 0.8 and क = 0.2 and observe consistently higher SNR for the majority class. We conjecture that this may stem from both the
30

---

score function of the distribution and DAE training being biased toward denoising more frequent samples.
Training Setting
SNR @
0.008
0.023
0.060
0.140
0.296
0.585
1.088
1.923
0 90% (Non-overlapping)
24.88
25.15
26.95
35,84
58.27
20.84 22.88
4.05
1.57
i 30 (Overlapping)
16.70
16.92
18.39
25.97
38.54
14.56
12.55
Table 4: Effect of overlapping class subspaces on SNR across noise levels.
Training Setting
SNR @ a₁
0.002
0.008
0.023
0.060
0.140
0.296
0.585
1.088
1.923
Class 0 (dg- = 10)
68.90
69.00
69.83
74.87
102.24
195.44
215.31
34.17
8.44
Class 1 (d) 2)
7.75
7.75
7.80
8.08
8.79
6.50
1.80
0.46
0.27
Overall SNR
24.12
24.13
24.28
25.21
28.19
23.34
8.18
2.84
1.45
Table 5: Effect of subspace rank variation SNR across noise levels.
Training Setting
SNR @
0.002
0.008
0.023
0.060
0.140
0.296
0.585
1.088
1.923
Class 0 (Tg = 0.8)
52.35
52.43
52.98
56.74
75.61
134.76
72.38
13.56
4.50
Class 1 (Ii 0.2)
12.15
12.16
12.31
13.10
17.49
23.91
6.05
1.17
0.52
Overall SNR
38.69
38.74
39.17
41.85
55.77
90.54
35.28
7.40
2.81
Table 6: Effect of non-uniform class mixing SNR across noise levels.
Well-separated clusters. In our main analysis, we considered zero-mean Gaussian clusters. In this section, we show that when the clusters are well separated, the unimodal trend disappears: both classification accuracy and SNR decrease monotonically, peaking at = 0. In this regime, the dataset remains linearly separable throughout the denoising process, and both metries are primarily determined by the class means. Adding noise (increasing f) simply blurs the class boundaries, thereby reducing separability. To verify this, we conduct an experiment on well-separated data and report results in Table 7 which clearly exhibit this monotonic behavior. These observations highlight the importance of overlapping clusters in giving rise to the unimodal representation dynamics.
Unimodal dynamics beyond pixel space. The unimodal dynamics are not restricted to pixel-space denoisers. In the table below, we analyze feature accuracy and SNR dynamics using the DiT-XI./2 37 model on mini ImageNet, and observe in Table & that the unimodal trend persists in this latent diffusion model setting as well.
Projection-based classification analysis. In the main paper, our linear probing experiments always logistic regression classifiers with inner-product-based logits. Here, re-ran the experiments in Figure 5) using projection- based classifier: for each class, we compute its principal subspace [V....., Vxl from the training features, and classify each test sample by identifying the class whose subspace captures the most projection energy, computed as arg max; ( V
The results shown in Table 9 While the mojectioo-bbsse classifier yields lower overall accuracy, reveals more pronounced unimodal trend that aligns closely with the SNR dynamics.
D Experimental Details
In this section, provide technical details for all the experiments in the main body of the paper.
31

---

0.030
0.053
0.098
0.189
0.282
0.379
0,480
0.588
0.704
0.830
0.989
1.492
Acc
100.0
100.0
100.0
99.8
97.7
92.4
85.7
78.5
71.8
66.0
61.6
$2.6
SNR
0.540
0.539
0.535
0.524
0,513
0.502
0.484
0.465
0.436
0.406
0.370
0.297
Table 7: Classification accuracy and SNR on well-separated clusters across diffusion noise levels.
0.010
0.026
0.044
0.076
0.108
0.140
0.205
Acc SNR
73.53
74.64
75.17
75.54
75.83
75.93
76.10
0.0369
0.0372
0.0375
0.0380
0.0382
0.0384
0.0387
0,271
0.339
0.719
1.234
2.031
3,424
6.135
Acc SNR
75.92
74.85
66.47
49.31
29.28
16.41
7.81
0.0384
0.0387
0.0350
0.0321
0.0282
0.0255
0.0211
Table 8: Feature accuracy and SNR dynamics for DiT-XL//2 on mini ImageNet.
Assets license information. We primarily utilize the codebase of EDM [49]. which is released under the Creative Commons Atrbuttion-NnnCommrriall license. We also use code from the GitHub repository accompanying [9 which is licensed under the MIT License.
4.0 (CC BY-NC-SA 4.0)
For the datasets, CIFAR datasets [0). ImageNet [51], Oxford 102 Flowers 93,, and DTD [96 are publicly available for academic use. The CelebA dataset [97] is released for non-commercial research purposes only under a custom license. Oxford-IIIT Pet 9 is available under the CC B BY-SA 4.0 license. The FFHQ dataset [98] is distributed by NVIDIA under the CC BY-NC-SA 4.0 license.
Computational resources. Most experiments are conducted on single NVIDIA A40 GPU, except for training subsets of images (e.g.. Figure 8 which is performed using two A40 GPUs.
Experimental details for Figure
Experimental details for Figure (a). We train diffusion models based on the unified framework proposed by 49). Specifically, we use the DDPM+ network, and use VP configuration for Figure Ma). a). (49 has shown equivalence between VP configuration and the traditional DDPM setting. thus we call the models as DDPM* models. We train two models on CIFARIO and CIFARIO0 respectively. After training, we evaluate the learned representations via linear probing. For each noise level (t), we corrupt the clean input æo with Gaussian noise to obtain æ; Var(arg + n), with ~ N (0,oāTT and vou 1/ 1. We then extract features from the decoder's 166x16 block1' layer, which consistently yields the highest classification accuracy. A logistic regression classifier is trained on these features using the training split and evaluated on the test split. We perform the linear probe for each of the following noise levels: [0.002, 0.008, 0.023, 0.060, 0. 14. 0.296, 0.585, 1.088, 1.923].
Experimental details for Figure ()). We exactly follow the protocol in 19 using the same datasets which are subsets of CelebA 97 [99] and FFHQ (98), the same training procedure, and the same segmentation networks (MLPs). The only difference that we use newer latent diffusion model (3] pretrained on CelebAHQ from Hugging Face and the noise are added to the latent space. For feature extraction we concatenate the feature from the first layer of each resolution in the UNet's decoder (after upsampling them to the same resolution as the input). We perform segmentation for each of the following noise levels:[0.010, 0.015, 0.030, 0.053, 0.098, 0.189, 0.282, 0.379, 0.480, 0.766, 1.123].
For the MoLRG experiments, we train the our parameterized model 4 following the setup provided in an open-source repository [100] The model trained on d 5, n = 50, K = and 0.2 MoLRG dataset containing 12000 samples. Training is conducted for 200 epochs using DDPM scheduling with T 500, employing the Adam optimizer
Experimental details for Figure and Figure 15
32

---

Experimental details for Table T For EDM, we use the official pre-trained checkpoints on ImageNet 64 64 from 49 and for DíT, we use the released DIT-XL/2 model pre-trale pre-trained on Classifier Type
Dataset
Ø, = 0.008
0.023
0.060
0.140
0.296
0.585
1.088
1.923
CIFARIO
Logistic Regression Projection
93.72
94.94
95.20
94.11
91.36
85.25
72.95
56,21
CIFARIO
86.95
90.36
90.93
90.78
88.48
82.82
70.53
53.84
Dataset
Classifier Type
0,
0.008
0.023
0.060
0.140
0.296
0.585
1.088
1.923
TinylmageNet TinyImageNet
Logistic Regression
17
34.70
43.09
53.58
50.78
42.13
27.56
15.10
Projection
12.89
14.38
24.10
34.78
34.88
30.02
20.99
12.64
Table 9: Logistic linear probing vs. projection- -aaee classification accuracy on CIFARIO (Top) and Tiny ImageNe (Bottom) across diffusion noise levels.
with learning rate of 5 x 10-4, For SNR computation, we follow the definition in Section 3.3 since we have access the ground-truth basis for the MoLRG data, i.e., Ui, U2, and U3 E R 2o For probing we simply train a linear probe on the feature.
For both panels in Figure E] train our probe the same training set used for diffusion and test on five different MoLRG datasets with 9000 samples generated with five different random seeds, reporting
the average accuracy and SNR time steps [5, 10, 20, 40, 60, 80, 100, 1 120, 140, 160, 180, 240, 260, 280]. In Figure we visualize the posterior estimations at time steps [5. 20, 60, 120, 260] by projecting them onto the union of the first columns of Ui U2. and Ug (a 3D space), then further projecting onto the 2D plane by a random 3 x 2 matrix with orthonormal columns. The subtitles of each visualization show the corresponding SNR calculated explained above.
Experimental details for Figure 5 We use pre-trained EDM models 49] for CIFARIO and ImageNet, extracting feature representations from the best-performing layer each timestep. For the ImageNet model, features are extracted using images from classes in TinylmageNet [92]. Feature accuracy is evaluated via linear probing. To compute the SNR metric, we first normalize the extracted features by dividing each by its norm and subtracting the global mean. At each timestep, we perform class-wise SVD on the normalized features and compute SNR defined in Section 3.3 Specifically, we use the top 5 right singular vectors of each class to form U
Experimental details for Figure 6 We use the DDPM++ network and VP configuration to train diffusion models [49] on the CIFARTO dataset, using two network configurations: UNet-32 and UNet- 64, by varying the embedding dimension of the UNet. Training dataset sizes range exponentially from 2 to 215 For each dataset size, both UNet-32 and UNet-64 are trained on the same subset of the training data. All models are trained with duration of 50M images following the EDM training setup. After training, we calculate the generalization score as described in [53] using 10K generated images and the full training subset to compute the score.
Experimental details for Figure and Figure g
We use the DDPM++ architecture with the HDM
configuration to train UNet- 128 diffusion model 40 CIFAR10 and CIFAR 100, using 4096 image training subsets. We track the evolution of representation dynamics throughout training. For Figure 7 FID [58) is computed using 50K generated samples compared against the full training dataset. Classification accuracy is obtained by extracting features from the training subset and evaluating linear classifier on the full test set. Nearest neighbors identified by computing the smallest & distance between each generated image and the training subset.
Experimental details for Figure 13
We train individual DAEs using the DDPM+ network and VP configuration outlined in at the following noise scales:
[0.002, 0.008, 0.023, 0.06, 0.14, 0.296, 0.585, 1.088, 1.923, 3.257]
Each model is trained for 500 epochs using the Adam optimizer [0o1 with fixed learning rate of 1 x 10-4, The sliced Wasserstein distance is computed according to the implementation described in 94
33

---

ImageNet 256 x 256 from [7 As baseline, we include the Hugging Face pre-trained MAE encoder (ViT-B/16) [95].
For diffusion models, features are extracted from the layer and timestep that achieve the highest probing accuracy, following (20) After feature extraction, we adopt the probing protocol from (2), passing the extracted features through probe consisting of a BatchNorm Id layer followed by a linear layer. To ensure fair comparisons, all input images are cropped or resized to 224 224, matching the resolution used for MAE training.
For ensembling, extract features from two additional timesteps either side of the optimal timestep. Independent probes are trained on these timesteps, yielding five probes total. At test time, we apply soft- ensemble by averaging the output logits from all five probes for the final prediction. Specifically, let W. E R xd be the linear classifier trained on features from timestep 1, and let h, Rdd denote the feature representation of sample at timestep t. Considering neighboring timesteps t 2, 1 1, + 1, and t + 2, our ensemble prediction is computed as: 2 W.hi)
y = arg max
We evaluate each method under varying levels of label noise, ranging from 0% to 80%. by randomly mislabeling the specified percentage of training labels before applying linear probing. Performance is assessed on both the pre-training dataset and downstream transfer learning tasks. For pre-training evaluation, we use the images and classes from MinilmageNet [102] to reduce computational cost.
For transfer learning, evaluate CIFAR100 [50] DTD 96 and Flowers [93].
E Proofs
We first provide some auxiliary results for proving Theorem
E.1 Groud truth posterior mean
We begin by deriving the ground truth posterior mean [ro [e] under the MoLRG distribution. When Zo follows the noisy MoLRG assumption, the optimal solution 27 (æ,, t) to the training objective in Eq. (2) is exactly the posterior mean E (xo xz).
Proposition 2. Suppose the data æo is drawn from noisy MoLRG data distribution with K K-class and noise level introduced in Assumption 7 Then the optimal {U} minimizing the loss (2) is the ground truth basis defined in (3 and the optimal DAE z. (2,, admits the analytical form:
K
(b t) =
Σ
wi (x (GU;URT G,Ü;ū;T) xe.
(9)
where C,
5a 1 and E1 .4077 and Too7
exp (g) (axt, t))
Se |Ui xz 20
& 2
- ūū w.l?? 20?
wi (x);t)
gi (x
EK 1= exp (gt (a,, t))
Proof. By 103 we see that the ground-truth score/posterior estimator calculated from the pdf is the global minimizer of the denoising score matching loss in (2), so here we first calculate the ground- truth score of the noisy MoLRG distribution and the corresponding posterior. We follow the same proof
ax and Ü = [U: sū:] where d
steps as in 32 Lemma with a change of variable. Let
ex
and denote the dimensions of signal space and noise space as in definition, we first compute the conditional pdf
Pt (x Y =k) k)
=f
/ pe
0, I2+D) des
= -/nic = - - (cx: 0, Ia+D) dck
34

---

- N T WW (exi 0, La+D) dcc
1
s, s:Uscal?) exp = dec
(2T) (d4
exp
2%
(2T) 2/))12(220) (d4
exp
1
1
TeT
a
2
ÜRCH + t U ÜNc* - -hick ex) dek
87 +
vi
-d/2
$200
2
-D/2
s788
(2T)2//228
11
exp
2%
In
ot + 7t
URUK
$78² + 78
(2T)
2
d/2
exp
2
2
+
1
2
8t
+ 78
U
T
2
da
(27)D/2
1
8702 + 7f
11
DD/2
exp
8
27
1t
2
Cx
s482 + 7
8,6
U x
dej
1
$/00
(2T)a/2
(5)
+ 78
22)d/2 (5,002 + 1/2 exp
2%
In
8 + 78
U
8282 + 7
(2T) det 1/2 US:URUR STÜūūØ In)
exp (-12" (BUUU + ,, +
= = SURU + + + Ia)
where we repeatedly apply the pdf of U ;;;UUÜ;;; In) = (s) (5 + and (S!U, U;;
Gaussian and the second last equality uses
I I, 2/(5 + U; - -5 + U //8 because of the Wood-
bury matrix inversion lemma. Hence, with PP (V k) = my for each & E [K] we have
K
K
(x) = Σ pr(a = k)P(Y = k)
T SUUU + ~ + I1)
Now we can compute the score function
Σx K
0, SUUU + ,,,ūūū,, +
7 xf ,, a 28
VP((x)
V Po
(x)
Eam UU;U + ::U;Ü I1
1= TaN(2; U. :: I.)
25 UFÜ x))
x))+ प,ा
2
x
K= E Ti + ,T I
According to Tweedie's 's formula, we have
It + 22 (xx))
E 2 x
St 9t
K K
+ +
87 7
0( + +
Σ K
82 + 7
Mix;0,3U.UK +
35

---

Ek=1 Th exp (0e1lU% x 12) exp dh U 2.113) U æ
k
8/
7t
2
E The exp (0: | z 121 exp 4/ (1:1 x; 11)
k= ThK exp [ x 1 exp dh UK x211?) U
8182
8,62 + 7
E Th exp (6: U T 2+133) exp w | T æz11²)
with o; /2/ (s} + ))) and 1/
+ (2) The final equality uses the pdf of multi-variant Gaussian and the matrix inversion lemma discussed earlier. Since т is consistent for all k and 8, we have
K
82
EE [00 = Σ we
ULU
82 x000
exp
201(1+of)
UU x:11 2 + +01
6²
U T a:e|1
where wxs()) ==
k= exp
2of[1+of]
|U T x 1|2 + 20
|
xx 1?)
Finally we show the equivalence between ground-truth posterior and our parameterized DAE when its weights are just the ground truth basis. We begin by rewriting the optimal posterior function by leveraging the fact that Ui-
[Ui
Ut
Ut+
...Uk]
K
E Jaon [ =
Σ
(2 1) (S.URURT + +GÜūŪNT) x
1=1
K
Σ
wft tx),
GUU;U,T + ΣUjUūT
371
K
K
( t) (GUFFUT + Σ (wf (ze, E) N UUJU,TD )ar
K
xz + & Euj( UU;U; xz
171
K
U,GUU;;""" on + &e (1 - wi (xr. 1)) U.U; x))
(Soo E - (x,, UT x
K
Σ 1=1
Ke + (St - (2), (,;; xt.
Now if we let U 1U, U2 Ukl E Onx Kd and (' diag Ia 3% Ia) to
be a block- diagonal matrix. Each 87 defined as 87 = E. + (Gr wf (I,, t), we can then write: U D* (*,, I)U (*,,
E [
|x.]
Thus, the optimal solution for our network parametrization as defined in (4) is exactly E [x:). And by such equivalence, the optimality of the DAE is induced.
E.2 Proof of Theorem 1
We first state the formal version of Theorem
36

---

To simplify the calculation of SNR as introduced in Section 3.3 on feature representations, which involves the expectation over the softmax term w we approximate di as follows:
K
e apprext (z) t) = Σ Si (SURUNT+&ūüÜ)" x,
+
where 20 == exp (Ex, (2 t)l)
8²
(10)
Se =
& =
Ek=1 exp (Bu, ( &
(2,, t))
1 1+
82 + o
2
St | x0l2 + 20? l x x01?
and (x,
20?
In other words, we use rig in T0) to approximate w (#t, 6) in Proposition I by taking expectation inside the softmax with respect to Tt. This allows to treat ugg constant when calculating SNR.
making the analysis more tractable while maintaining E|] hi (z,, +] = Ell aopnor (Zo, t)]l"]
(
for all / E [K]. We verify the tightness of this approximation Appendix C (Figure 14). With this approximation, we state the theorem as follows:
Theorem 2. Let data By be any arbitrary data point drawn from the MoLRG distribution defined in Assumption D and let k denote the true class ao belongs to. Then SNR introduced in Section 3.3
depends on the noise level ot in the following form:
of 1+ g2 h(wg.
2
1 oz SNR( appprox: () (K - 1)(02 + of)
(11)
1+ g h(w;
82. Since 8 is fixed, h(w, 8) is a monotonically increasing function with respect to w. Note that here 5 represents the magnitude of the fixed intrinsic noise in the data where or denotes the level of additive Gaussian noise introduced during the diffusion training process.
where (
Proof. Following the definition of SNR as defined in Section k Mult(K, K, Tx) with ग
3.3
Lemma 2 and the fact that
n 1/K, we can write
ERz, ||U t: Hapmou (xt, 1)]]
UUU E appees (at, 1)]13] ER
SNR ipppeer 1)
Rag, Expr 10;U, f approx (3ee 1
Σler Eee. 107U 5: Kappra (1,, [
wig 1+o7
+
(1 (1+07)d
82407
(K 1)
82 2))0;))
(82 + of)d
1ta7
8t
1
v - o + (K 1)0² + (K - 1)62 o²
(K- 1)(82 + +
1778 82 प 0
8 u. (K 2)02 + ouig
(K 2) 2
1+0?
82 + + (K (K 1)82(07)
(K 1)) + +of)
82 + of + 8 + (K- - 2)
1+
+ 23 (11 - 52) + 00 + (K- 1)(0t)
(K 1) + +
1 1 62 ((1 - - + 8² (1o + oz + (K- 2)
2
1 1
1 # ((1 0²) + 88)
[K 1) + tzf
1 / - ((1-62) + 82)
2
1 1 o
+ و
(K 1) of
+ o h 5)
where 1(,, 6) == (1 - -82)0 + 52, and we set C. =
1+02
(8=
37

---

Lemma 1. Let 91, .2.............................. 9K be sequence of real-valued inputs. and fix an index k E [K]. Assume that 9i 93 for all 1,j & k; that all entries except gse share the same value. Let the softmax weights be defined as
exp(gj)
Wi
j== exp(gj)
Then we have:
exp(gx - 9t)
Wk K 1 + exp(s: 9t)
and wn = K - exp(xx 9i)
for all + k.
Proof. The softmax weight for class k is:
exp(ga)
Wk 2ー=1 exp(gi)
We simplify the denominator by factoring out exp(gi):
exp(o,) exp(gk)
Σ
exp(()) exp(gx) (K exp(())
Therefore.
exp(gk)
exp - 9t)
Wh exp(gk) (K - 1)) exp(gi)
1 (K 1)exp(a 9k) (K 1) + exp((x 91)
For any + ki we similarly have:
exp(gi)
Wi = exp(sx) + (K - 1) exp(()) (K 1) + exp(g: gi)
This simple fact can also be proven by applying the log-sum-exp trick.
Lemma Given the setup of a K K-class MoL&G data distribution defined in 3 consider the following approximate posterior mean function:
K
Sx x (I, t) =
Σ wi km
+ ( + x,
cxp (En, 1ot (x,, t)l)
1
8²
where wi
&e
Σx=1 exp (Eu, l9t (It, t)]) &. U x1?
1 o+
82
oo
Gt | T x +
and gh (x, = 20
20?
That is, we consider a simplified form of the expected posterior mean from Proposition where the expectation is taken inside the softmax argument (i.e., over gV (zt, t)) to obtain tractable approximate weights ut
Under this approximation, for any sample 2o from class k, i.e., 2o = U,a, + -U we have:
mig
(K (K-1) -1)0P(0)) ) (d od+0,3 o7d)
Em, | x= pppmee (x),, t)][9]
+
(12)
of
82 of
wij
82 + (K - 2)(Ot) (62d o,7)
Em, U.UUT * Epprror (wt, t)][3] =
+
(13)
1 of
82 of
38

---

(K 1)62un
Eig x apppeo (Ft, t)]1
+
(d (d+
1 + o²
+ oz
BllUge Apoe (xr..)) 2)
(14)
(0g + (K - 2)
(K - 1)
(8 (8'd+
1 1
82 of
B UF & &igoe (aee 13]
and
exp (D₂)
ot K 1 + exp (D,)
(15)
1
oi
+ (D,)
for all class index 1 & k where D;
(1-8 (1-5")d (1+o0=
Proof. Throughout the proof, we use the following notation for slices of vectors.
e:(a: b] Slices of vector e, from ath entry to lih entry.
We begin with the softmax terms. Since each class has its unique disjoint subspace, it suffices to 6²
consider gi and 9 (2o, 1) for any + k. Let 2070+075 and C2 207(6'+075)'
have:
Elgh(x t)] Elael U x +c x1
+ Elail| (Uga, e Hoo (( + 80 e, + o
Elaalla: +
+ o E|calee + +o &
= a. + + a (5 1)d
1)d).
where the last equality follows from a; itd N Ia), e, iid NOO, ((K- and e, itd NO Without loss of generality, assume the j = k: + 1, we have:
Elgi (æe, t)] = Bja, |0, Tu, x; c 1 |
(U E[a,,lU, 0Ü +,,6 + Elal (Uaa + +ūūe, +,,6,)11 O Rd
=Eja,lde d + ., < [ [0 0 RD-
[ DI] +o, e: [
+
a: (o'd ct (d 5K(K 2)d + OK 1)d)
Hence have
(agh [)] Eggg (x,, 1)]
a, (d + o,d 82d - oidd + a - 1)d
of (
- 1)d d (K 6K - 1)d)
(1 - 6d)d (62 - 1)d
61-
dd2 -)
-
21 + o?? + 202 of?) 2(1 (02
This, together with Lemma i yield 15
Now we prove 12)
(T+
82
U &pppna (x =
+
82 oo
39

---

8²
+
EOTUR
(T+o. UiU;
8t
1#
8
UjU 2t)
+ +
U
01
1 1+007 0z
8 R+
(K 1) 1)
+
U (Uja, + +0 +..
+ o
82 + of
(K 1) 1)
+
(URa, + oU E.)
1 + of
8² + oî
Since U3 E Onxd:
U E
appros (a:, ()1,))
(1
wk 11 0z
+
(K - 1) 10
(d+ odd)
82 + oo²
and similarly for 13):
1
8²
Ui x= Taprror (rt, 1) un UFU
U
Un
1o
82 + o
8
UUU
+ F
x1
1+0,
Σ By UFUTT
1
8²
ú,ū,T)
1+
+
82 o
jFL1
0
82
UFUU 81)
+ uij
1 1 - m.) + Σ
E0j'
( 82
UJU at
82 + of
sth,
wf
8 ( + (K - -2)mt)
UNU;T(UAa + oÜge, t,,,))
1 of
82 + +
wf
(1g + (K -2)(tt)
(SU) eill : d] U,U;U;T ai)
1+ of
82 + of
where the third equality follows since 0Z =wi for all j # k, I. Further, we have:
( (K 2))
2 (6 oid).
UT 2 Epppex (I t)]
+
1 o?
82 + o;
Lastly, we prove (4). Given that the subspaces of all classes and the complement space are both orthonormal and mutually orthogonal, we can write:
(appaa(x1,100 B|| appnos (Ze, 1 = E i Sppprax (xt, 1)] + EE E
s apprec (0,, 1) + ElIIU LUSTTATT a (xx,, 1)
where we define the noise space as U, = = Uz E (nx(n-Kd) Kd), representing the directions orthogonal to all class subspaces. Since U, is orthogonal to each Uf. the third term vanishes. Combining the remaining terms, we obtain:
(K - 1)-0rp
E + opprox (It; t)
+
(d oid)
1 + of
82 + o
20
8 8( + (K- - 2)(0))
+ (K -1))
odd+o?d)
1+o7
8 + oo
40


Reviews:
[{"review_id": "KXoasK0VOc", "reviewer": "Reviewer_cQRG", "summary": "The paper shows that if images lie on a noisy mixture of low-rank Gaussians and the diffusion denoiser is idealized as a block-diagonal “mixture-of-experts,” then one can derive a closed-form signal-to-noise ratio (SNR) for learned features that rises at low noise, peaks at an intermediate noise level, and then falls, explaining the empirically observed “unimodal” representation quality in diffusion models and validates this theory on both synthetic and real vision benchmarks, also demonstrating that the peak SNR reliably indicates generalization versus memorization and can serve as a principled early-stopping criterion.", "questions": "- How well does the Mixture of Low-Rank Gaussians (MoLRG) actually capture the structure of real high-resolution image data (e.g. ImageNet)? What happens if class subspaces overlap or have different ranks?\n- To what extent does the block-diagonal “mixture‐of-experts” abstraction preserve the key representational behaviors of full U-Net architectures with attention and skip‐connections? Could cross‐block interactions materially change the unimodality result?\n- The derivation of the approximate SNR curve hinges on swapping the expectation and softmax (Eq. 10). Under what conditions might this approximation break down, and how sensitive is the peak location to that error?\n- The SNR metric is evaluated for classification accuracy. Would an analogous unimodal behavior hold for other downstream tasks, e.g. segmentation, detection, or self‐supervised objectives?\n- The paper suggests that losing unimodality signals onset of memorization. How robust is this indicator across dataset sizes, architectures, and training regimes?\n- Do similar unimodal representation dynamics appear in latent diffusion models (e.g. Stable Diffusion) where the denoiser operates in a lower‐dimensional embedding space?", "limitations": "I'm happy to raise my score if my concerns are addressed properly.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The assumption that images lie on a noisy mixture of low-rank Gaussians (MoLRG) is supported by a growing body of empirical work on intrinsic dimensionality in vision data. It lets the authors derive closed-form expressions for the optimal denoiser and an explicit SNR metric (Eq. 6) that quantifies feature quality.\n- By modeling the denoising autoencoder as a block-diagonal “mixture-of-experts” layer (Eq. 4–5), they capture key behavior of U-Net–style architectures while keeping the math clean. This lets them prove (Prop. 1) that the optimal basis aligns with the true subspaces and derive an analytical form for the denoiser.\n-  The informal Theorem 1 shows the SNR follows a unimodal curve because as noise $\\sigma_t$ increases, a “denoising rate” term grows while a “class confidence” term decays, yielding a single peak in their ratio (Eq. 8).\n\nWeaknesses:\n- Real data rarely satisfy exact orthonormal subspaces or uniform mixing weights. The analysis assumes all subspaces have the same dimension and are mutually orthogonal—conditions that hold only approximately in practice.\n- Actual diffusion models use deep, highly nonlinear U-Nets with skip connections, attention, and non-diagonal feature mixing. Reducing them to a single block-diagonal layer may miss interactions that affect real-world representation dynamics.\n- The analysis treats a decaying constant $C_t$ as negligible (“minimal impact to unimodality”), but in practice its behavior across $t$ could shift the peak. Furthermore, the metric is derived for classification accuracy—its relevance to other downstream tasks (segmentation, alignment) may require additional work.\n- Some missing related works should be added in the discussion [1, 2, 3]. There might be others but I suggest the authors examine the literature carefully.\n\n[1] Infodiffusion: Representation learning using information maximizing diffusion models\n\n[2] Lossy image compression with conditional diffusion models\n\n[3] Soda: Bottleneck diffusion models for representation learning", "was_revised": true, "final_justification": "The authors properly addressed most of my concerns. I'll raise the score to 4 as promised.", "rebuttal": "We thank the reviewer for the thoughtful and constructive feedback. We're glad the MoLRG assumption and theoretical results were appreciated. Below, we address the reviewer’s concerns one by one.\n\n> Real data rarely satisfy exact orthonormal subspaces or uniform mixing weights. The analysis assumes all subspaces have the same dimension and are mutually orthogonal—conditions that hold only approximately in practice. What happens if class subspaces overlap or have different ranks?\n\n**Answer**: We appreciate this important point. As the first theoretical study on diffusion-based unimodal representations, we adopt a simplified setting to isolate the root cause of unimodality. That said, we agree it is valuable to test the robustness of our findings under relaxed assumptions. Empirically, we observe that the unimodal SNR trend persists even when class subspaces overlap, differ in rank, or follow non-uniform mixing:\n\n1. Overlapping Class Subspaces: We control the principal angle ($\\theta$) between class subspaces and observe that overlap tend to reduce SNR across timesteps while the peak remains stable. We conjecture that overlap can be seen as introducing additional intrinsic noise beyond the $\\delta$-related term and thus affect the SNR value.\n\n| Training Setting |SNR @  $\\sigma_t = 0.008$|$0.02$|$0.06$|$0.14$ | $0.30$ | $0.59$ | $1.09$ | $1.92$ |\n|:------:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|\n| $\\theta = 90^{\\circ}$ (Non-overlapping)| 24.9|25.2|27.0|35.8|58.3|20.8|4.1|1.6|\n| $\\theta = 30^{\\circ}$ (Overlapping) |16.7|16.9|18.4|26.0|38.5|22.9|14.6|12.6|\n\n2. Varying Subspace Ranks: We set class subspace dimensions to $d_0=10, d_1=2$. Intuitively, the higher-rank class retains more signal and is less sensitive to noise, yielding higher and later-peaking SNR. The low-rank class decays earlier.\n\n| Training Setting | SNR @  $\\sigma_t = 0.008$ | $0.02$ | $0.06$ | $0.14$ | $0.30$ | $0.59$ | $1.09$ | $1.92$ |\n|:------:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|\n| Class 0 ($d_0 = 10$) | 69.0 | 69.8 | 74.9 | 102.2 | 195.4 | 215.3 | 34.2 | 8.4 |\n| Class 1 ($d_1 = 2$) | 7.8 |  7.8 |  8.1 |  8.8  |  6.5  |  1.8 |  0.5 |  0.3|\n| Overall SNR | 24.1| 24.3 | 25.2 | 28.2 | 23.3|  8.2 | 2.8| 1.5 |\n\n3. Non-Uniform Mixing Weights: We set $\\pi_0=0.8, \\pi_1=0.2$ and observe consistently higher SNR for the majority class. We conjecture that this may stem from both the score function of the distribution and DAE training being biased toward denoising more frequent samples.\n\n| Training Setting | SNR @  $\\sigma_t = 0.008$ | $0.02$ | $0.06$ | $0.14$ | $0.30$ | $0.59$ | $1.09$ | $1.92$|\n|:------:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|\n| Class 0 ($\\pi_0 = 0.8$) |52.4 | 53.0 | 56.7|75.6|134.8|72.4|13.6| 4.5|\n|  Class 1 ($\\pi_1 = 0.2$) |12.2 |12.3|13.1|17.5|23.9 |6.1|1.2|0.5|\n| Overall SNR| 38.7| 39.2|41.9|55.8 |90.5 |35.3|7.4| 2.8|\n\nIn summary, these new results demonstrate that the unimodal SNR behavior generalizes to more complex and realistic data distributions beyond our theoretical setup. We will include this expanded study in the revision.\n\n> How well does the Mixture of Low-Rank Gaussians (MoLRG) actually capture the structure of real high-resolution image data (e.g. ImageNet)?\n\n**Answer**: We thank the reviewer for this question. As noted in Section 3.1, high-resolution datasets like ImageNet exhibit low intrinsic dimensionality [1], which our MoLRG is designed to model.\n\nMoreover, prior work [2] used MoLRG to simulate real data and produced qualitatively reasonable samples resembling those from pre-trained diffusion models. While no analytical model fully captures high-resolution image complexity, MoLRG offers a reasonable abstraction of the key low-dimensional structure.\n\n[1] Pope et al; The Intrinsic Dimension of Images and Its Impact on Learning; ICLR 2021.\n\n[2] Wang et al; Diffusion Models Learn Low-Dimensional Distributions via Subspace Clustering; 2024.\n\n>Actual diffusion models use deep, highly nonlinear U-Nets with skip connections, attention, and non-diagonal feature mixing. Reducing them to a single block-diagonal layer may miss interactions that affect real-world representation dynamics. \n\n**Answer**:  We agree with the reviewer that U-Net architectures are more complex, but this also makes theoretical analysis being intractable. Our simplified model aims to isolate the core mechanisms behind unimodal representation dynamics under the MoLRG assumption, which we find sufficient for explaining key unimodal representation behaviors.\n\nAdditionally, many prior works (see for example [3–5]) have similarly adopted simplified architectures when modeling Gaussian-like data, enabling theoretical insights into otherwise complex systems. A more complete theoretical treatment of U-Nets is an important future direction, and we view our analysis as a first step in that direction.\n\n[3] Shah et al; Learning mixtures of Gaussians using the DDPM objective; NeurIPS 2024.\n\n[4] Ji et al; The Power of Contrast for Feature Learning: A Theoretical Analysis; JMLR.\n\n[5] Han et al; On the Feature Learning in Diffusion Models; ICLR 2025.\n\n> To what extent does the block-diagonal “mixture‐of-experts” abstraction preserve the key representational behaviors of full U-Net architectures with attention and skip‐connections? Could cross‐block interactions materially change the unimodality result?\n\n**Answer**: Despite its simplified block-diagonal structure, our DAE reproduces the key unimodal SNR trend observed in full diffusion models when trained on CIFAR10. While complex U-Net components may influence feature evolution or shift the peak, the core dynamic is captured by our simplified model.\n\n|Network|Dataset|$\\sigma_t = 0.008$|$0.02$|$0.06$|$0.14$|$0.30$|$0.59$|$1.09$|$1.92$|$3.26$|\n| :-----: | :---: | :---: | :---: | :---: | :----: | :----: | :----: | :---: | :---: | :---: |\n| Simplifed DAE|CIFAR10|1.39|1.41|1.49|1.88|2.54|2.78|2.63|2.20|1.33|0.36|\n\n>The analysis treats a decaying constant $C_t$ as negligible (“minimal impact to unimodality”), but in practice its behavior across t could shift the peak.\n\n**Answer**: In light of the reviewer's suggestion, we ablate the effect of $C_t$ in Theorem 1 under the MoLRG setting and find that while it alters the SNR magnitude, it has minimal impact on the peak location.\n\n| Data Setting|SNR Type| $\\sigma_t = 0.008$|$0.02$|$0.06$|$0.14$|$0.30$|$0.59$|$1.09$|$1.92$|\n| :-----: | :---: | :---: | :---: | :---: | :----: | :----: | :----: | :---: | :---:|\n| $n=50,d=10,K=3,\\delta=0.1$|w/ $C_t$|50.3| 52.6 |67.7|144.5|449.9|1310.9|30.6|1.1|0.6|\n|  $n=50,d=10,K=3,\\delta=0.1$|w/o $C_t$|0.5|0.6|0.9|4.2|40.5|344.3|16.7|0.9|0.5|\n| $n=50,d=2,K=10,\\delta=0.2$|w/ $C_t$|2.78|2.81|3.02|4.05|8.00|2.35|0.30|0.15|0.12|\n|  $n=50,d=2,K=10,\\delta=0.2$|w/o $C_t$|0.11|0.12|0.13|0.24|0.94|0.67|0.17|0.12|0.11|\n\n>Some missing related works should be added in the discussion [1, 2, 3]. There might be others but I suggest the authors examine the literature carefully.\n\n**Answer**: We thank the reviewer for these relevant references. Due to space constraints, we postponed the related work section (including [1] and [3]) to Appendix A. Following the reviewer’s suggestions, we incorporated [2] and other recent works on diffusion-based representations. These changes will be updated in the manuscript.\n\n>The derivation of the approximate SNR curve hinges on swapping the expectation and softmax (Eq. 10). Under what conditions might this approximation break down, and how sensitive is the peak location to that error?\n\n**Answer**: We appreciate the reviewer’s careful attention to the derivation. To assess the robustness of the approximation, we varied $K$, $d$, and $\\delta$ and found the peak may shift slightly, but the unimodal SNR trend remains stable (Figure 14, Appendix C).\n\n>The SNR metric is evaluated for classification accuracy. Would an analogous unimodal behavior hold for other downstream tasks, e.g. segmentation, detection, or self‐supervised objectives? the metric is derived for classification accuracy—its relevance to other downstream tasks (segmentation, alignment) may require additional work.\n\n**Answer**: We agree that our current theory is framed in a classification setting via subspace-aligned SNR, and extending it to other tasks may require task-specific SNR formulations which is a promising direction. We believe that the core SNR principle: trading off signal preservation and noise suppression should be applicable broadly. Empirically, unimodal trends also appear in segmentation (Figure 1b) and correspondence [6,7], suggesting the phenomenon generalizes beyond classification.\n\n[6] Baranchuk et al; Label-Efficient Semantic Segmentation with Diffusion Models; 2022.\n\n[7] Tang et al; Emergent Correspondence from Image Diffusion; 2023.\n\n>The paper suggests that losing unimodality signals onset of memorization. How robust is this indicator across dataset sizes, architectures, and training regimes?\n\n**Answer**: The indicator is quite robust. In addition to the results in the main paper, we include experiments on subsets of Oxford-IIIT Pet (3680 images) and TinyImageNet (2048 images) in Figures 9 and 10 of the Appendix. In both, loss of unimodal dynamics signals the onset of memorization.\n\n> Do similar unimodal representation dynamics appear in latent diffusion models (e.g. Stable Diffusion) where the denoiser operates in a lower‐dimensional embedding space?\n\n**Answer**: Yes, unimodal dynamics persist beyond pixel-space denoisers. Using DiT-XL/2 on miniImageNet, we observe similar trends in feature accuracy and SNR, confirming the effect holds in latent diffusion models.\n\n| $\\sigma_t$ |0.03|0.04|0.08|0.11| 0.14| 0.21|0.27|0.34|0.72|1.23|2.03|\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| Acc|74.6|75.2|75.5|75.8|75.9|76.1|75.9|74.9|66.5|49.3|29.3| \n| SNR |0.0372|0.0375|0.0380|0.0382|0.0384|0.0387|0.0384|0.0387|0.0350|0.0321|0.0282|"}, {"review_id": "qsuaCNGMyj", "reviewer": "Reviewer_95ka", "summary": "This paper develops a mathematical framework to analyze the expressiveness of features learned by denoising diffusion probabilistic models (DDPMs), with a focus on how this expressiveness varies across timesteps. The analysis is grounded in an assumed data distribution, a noisy low-rank Gaussian mixture, and a parametrization of the denoising network resembling eigen-decomposition. The theoretical findings align with empirical observations: features extracted at intermediate timesteps tend to be more expressive, as measured by downstream task performance. Additional experiments further demonstrate that this behavior correlates with the generalizability of the diffusion model.", "questions": "- Could you clarify and justify the data distribution assumptions used in the theory?\n- How is the proposed network parameterization connected to U-Net architectures used in practice, and how does this impact the validity of the theoretical conclusions?", "limitations": "I do not see any major ethical or societal concerns raised by this work. However, a more explicit discussion of the limitations and the realism of the theoretical assumptions would be helpful for readers and could strengthen the paper.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n- To the best of my knowledge, this is the first work that provides a theoretical analysis of the relationship between the feature expressiveness of diffusion models and the denoising timestep.\n- Under the stated assumptions, the derivation offers a sound theoretical explanation for the empirical phenomenon under investigation.\n- The experiments reveal interesting and potentially useful connections between feature expressiveness and practical aspects of diffusion model training.\n- The experimental evaluation is extensive, with validation across multiple model variants, datasets, and evaluation metrics.\n\n**Weaknesses**:\n- Several critical assumptions underlying the theoretical framework are insufficiently justified. The authors should consider discussing these assumptions further or citing prior works that support them.\n  - While the data distribution assumption is discussed to some extent, others remain unaddressed. For instance, is it reasonable to assume that data from different classes lie in orthogonal subspaces? Similarly, is the modeling of class-independent fine-grained details as belonging to subspaces of other classes theoretically or empirically supported? \n  - The parameterization of DAE is underexplained. It is claimed to resemble a shallow U-Net, but this connection is not clearly established. Furthermore, most diffusion models in practice use deeper U-Nets. The feature extractor is modeled as a linear projection followed by a self-attention-like reweighting, which differs from standard U-Net architectures. Additionally, the feature is defined in the theory as a linear projection from the input data itself, whereas in the experiments (as described in the supplementary material), features are extracted from intermediate layers of the network. These intermediate features are not obviously related to a direct linear projection from the data space, raising questions about how well the theoretical definition aligns with the practical implementation.\n- Proposition 1 assumes that the latent feature dimensionality equals the intrinsic dimensionality of the data. This assumption is difficult to justify, especially given that deep networks are often overparameterized and operate in much higher-dimensional latent spaces. The authors should either justify this assumption more explicitly or discuss how relaxing it might affect the conclusions.\n- There is a notational inconsistency regarding the DAE. It is denoted as $x_\\theta$ in Section 3.2, but as $\\hat{x}_\\theta$ in Definition 1, without explanation. Clarifying this would improve readability.", "was_revised": true, "final_justification": "Please see my response to the authors' rebuttal.", "rebuttal": "We thank the reviewer for the thoughtful and detailed review. We are glad to see that the reviewer recognizes the novelty and contributions of our theoretical analysis, as well as the extensive experimental validation. We also appreciate the reviewer's recognition of the connection we draw between feature learning and generalization in diffusion model training. Below, we address the reviewer’s comments point by point.\n\n>While the data distribution assumption is discussed to some extent, others remain unaddressed. For instance, is it reasonable to assume that data from different classes lie in orthogonal subspaces? Similarly, is the modeling of class-independent fine-grained details as belonging to subspaces of other classes theoretically or empirically supported? \n\n**Answer**:  We thank the reviewer for raising this important point. Below we address the reviewer’s concerns: \n\n**The assumption that different classes lie in othogonal subspaces.** This assumption is primarily made to simplify the analysis and enable closed-form derivations. However, this is not essential for the validity of our main findings. In practice, subspaces corresponding to different classes may not be perfectly orthogonal, but as long as the principal angles between them are sufficiently large, our theoretical results remain valid with appropriate technical refinements. \n\nTo support our claim beyond this assumption, we conduct an ablation study where we control the principal angle $\\theta$ between class subspaces. As shown below, the unimodal SNR trend remains present even with substantial subspace overlap, suggesting that it is not sensitive to the orthogonality assumption.\n\n| Training Setting                           | SNR @  $\\sigma_t = 0.002$ | $0.008$ | $0.023$ | $0.060$ | $0.140$ | $0.296$ | $0.585$ | $1.088$ | $1.923$ | $3.257$ |\n|:-----------------------------------------:|:------------------:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|\n| $\\theta = 90^{\\circ}$ (Non-overlapping)   |       24.85        | 24.88 | 25.15 | 26.95 | 35.84 | 58.27 | 20.84 |  4.05 |  1.57 |  1.13 |\n| $\\theta = 30^{\\circ}$ (Overlapping)       |       16.67        | 16.70 | 16.92 | 18.39 | 25.97 | 38.54 | 22.88 | 14.56 | 12.55 | 12.13 |\n\n**The modeling of class-independent fine-grained details as belonging to subspaces of other classes.**  Our data model assumes that the dataset lies in a union of subspaces, where each class is associated with a dominant, class-specific subspace, and there exists a shared subspace common across all classes that captures class-independent fine-grained details. This type of shared–specific decomposition has been empirically supported in the literature on subspace clustering and representation learning (e.g., [1,2]). \n\nWe will include these discussions, along with supporting ablation results, in the revised manuscript to clarify and validate this modeling assumption.\n\n[1] Bousmalis et al., Domain Separation Networks, NeurIPS 2016\n\n[2] Zhou et al., Dual Shared-Specific Multiview Subspace Clustering, IEEE Transactions on Cybernetics, 2019\n\n>The parameterization of DAE is underexplained. It is claimed to resemble a shallow U-Net, but this connection is not clearly established. Furthermore, most diffusion models in practice use deeper U-Nets. The feature extractor is modeled as a linear projection followed by a self-attention-like reweighting, which differs from standard U-Net architectures. Additionally, the feature is defined in the theory as a linear projection from the input data itself, whereas in the experiments (as described in the supplementary material), features are extracted from intermediate layers of the network. These intermediate features are not obviously related to a direct linear projection from the data space, raising questions about how well the theoretical definition aligns with the practical implementation. How is the proposed network parameterization connected to U-Net architectures used in practice, and how does this impact the validity of the theoretical conclusions? \n\n**Answer**:  Our primary goal with the simplified model was to isolate the core mechanism behind unimodal representation dynamics. While the parameterization does not fully match a modern U-Net, it is sufficient to reproduce this key behavior on both synthetic data (Figure 3) and real datasets.\nTo validate this, we train our simplified DAE on the full CIFAR10 dataset and report the SNR curve below. The unimodal trend persists, suggesting that this behavior is not an artifact of complex functionalities within a modern U-Net, but rather rooted in the intrinsic structure of the data, and hence can be captured with a simple model like ours.\n| Network                          | Dataset |  $\\sigma_t = 0.002$ | $0.008$ | $0.023$ | $0.060$ | $0.140$ | $0.296$ | $0.585$ | $1.088$ | $1.923$ | $3.257$ |\n| :--------------------: | :-------------------: | :---: | :---: | :---: | :---: | :----: | :----: | :----: | :---: | :---: | :---: |\n| Simplifed DAE  | CIFAR10 |         1.39 |   1.39 |   1.41 |   1.49 |  1.88 | 2.54 | 2.78 | 2.63 |    2.20 |    1.33 |   0.36 |\n\nThat said, we conjecture that the more complex architecture of full U-Nets may influence feature evolution and shift the peak location. We agree with the reviewer that a more comprehensive theoretical treatment of U-Nets remains an important direction for future work.\n\n>Proposition 1 assumes that the latent feature dimensionality equals the intrinsic dimensionality of the data. This assumption is difficult to justify, especially given that deep networks are often overparameterized and operate in much higher-dimensional latent spaces. The authors should either justify this assumption more explicitly or discuss how relaxing it might affect the conclusions.\n\n**Answer**: We thank the reviewer for this valuable comment. The equal-dimensionality setting ($d' = d$) for feature dimension ($d'$) and data intrinsic dimension ($d$) is used only to keep the theoretical results clean and to isolate the essential factors behind the unimodal representation dynamics. The result extends naturally to an overcomplete feature space. Specifically, for any $d' > d$, we may take $U_{l,\\text{new}}^{\\star}=\\begin{bmatrix} U_{l}^{\\star} & 0_{d'-d} \\end{bmatrix} \\in \\mathbb{R}^{n\\times d'}$ and any permutation or orthogonal rotation applied within each class block of this form would make the loss, the soft-max weights and the SNR remain unchanged, and preserve the optimality, although the solution is no longer unique. The conclusions of Proposition 1 remain intact.\n\nWe appreciate the reviewer’s attention to this detail and will clarify this extension in the revised manuscript.\n\n> There is a notational inconsistency regarding the DAE. It is denoted as $x_{\\theta}$  in Section 3.2, but as $\\hat{x_{\\theta}}$ in Definition 1, without explanation. Clarifying this would improve readability.\n\n**Answer**: We appreciate the reviewer for raising this concern. We want to clarify that we intend to use $x_{\\theta}$ to denote an arbitrary DAE for any parameter set $\\theta$ where we use $\\hat{x_{\\theta}}$ to denote a trained DAE. Because Definition 1 evaluates SNR with the trained network, we use $\\hat{x_{\\theta}}$ here.\nWe will follow the reviewer's suggestion to add a sentence in Section 3.2 that explicitly states this convention and ensure the same notation is used consistently throughout."}, {"review_id": "v4WBYAAEsk", "reviewer": "Reviewer_9v41", "summary": "This paper analyzes the representations of diffusion models from two perspectives.\n\nFirst, it theoretically explains why *unimodal representation dynamics*—the phenomenon where the quality of learned features peaks at intermediate diffusion timesteps—arises. Specifically, the paper shows the following:\n\n* It assumes the data is a low-dimensional mixture of Gaussian distributions (MoLRG), and the model is a very simple linear U-Net combined with a mixture-of-experts-like architecture.\n* Under this setting, it derives the model’s optimal parameters (Proposition 1).\n* It defines the SNR (Definition 1) as a metric to measure the quality of the model’s representations for classification and provides an analytic expression for SNR (Theorem 1).\n* The paper shows that this analytical SNR matches both empirical SNR and feature accuracy in synthetic experiments (Figure 3) and also explains trends on real datasets like CIFAR-10 and TinyImageNet (Figure 5).\n\nSecond, the paper empirically investigates the relationship between diffusion model representations and generalization. For example:\n\n* It shows that unimodal dynamics are linked to generalization as measured by feature accuracy: when unimodal dynamics are absent, feature accuracy is very low.\n* It demonstrates that feature accuracy correlates with the memorization behavior of generative models: the trend of high FID – low FID – high FID (underfitting – generalization – memorization) observed during training iterations also appears when measuring feature accuracy.\n\nOverall, this is a strong paper that provides an analytic understanding of representation learning in diffusion models and connects it empirically to memorization, one of the interesting phenomena in generative modeling.", "questions": "* I assume feature accuracy was measured using a linear classifier with logits defined as the inner product between the class embedding vector and the feature vector. Have you considered trying an L2 similarity-based classifier? Since the SNR is defined via the L2 norm, I wonder if this might reveal even more consistent trends, although I understand that training might be more challenging.", "limitations": "* The paper does a good job describing limitations, such as the fact that the analysis tools, assumptions, and scope of experiments (e.g., synthetic data models and classification-focused benchmarks) may not capture the full complexity of all applications.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n\n* Despite being theoretically heavy, the paper is very easy to understand. The writing is clear, and the visualizations are excellent.\n* To my knowledge, this is the first work to provide an analytic treatment of diffusion model representations. This is a very impressive contribution. Beyond simply describing trends, the paper demonstrates that the derived model can accurately predict the peak in unimodal dynamics of feature accuracy (Figures 3 and 5), which is particularly impressive.\n\n**Weaknesses**\n\n* This is not a major issue, but the connection between the first part (theory) and the second part (generalization and memorization) is not very tight. Fundamentally, this may be hard to address, as the theoretical analysis assumes an almost-linear model. Linear models, by their nature, have limited expressivity and regularization effects, which makes them less appropriate for analyzing the generalization–memorization behavior observed in neural networks. It would be nice if the first and second parts were linked a bit more naturally.", "was_revised": true, "final_justification": "I think this paper successfully builds a theoretical framework to explain the representation learning of the diffusion model. There are some downsides, such as (1) logical flow of the paper (2) over simplified setting, but I believe overall this paper contributes significantly to the theory of understanding the diffusion model.", "rebuttal": "We thank the reviewer for the detailed and generous review. We are pleased to hear that the reviewer finds our theoretical analysis of diffusion model representations clear, novel, and impactful, and appreciates our empirical visualizations. Below, we address the reviewer’s comments in detail.\n\n> This is not a major issue, but the connection between the first part (theory) and the second part (generalization and memorization) is not very tight. Fundamentally, this may be hard to address, as the theoretical analysis assumes an almost-linear model. Linear models, by their nature, have limited expressivity and regularization effects, which makes them less appropriate for analyzing the generalization–memorization behavior observed in neural networks. It would be nice if the first and second parts were linked a bit more naturally.\n\n**Answer**: Thank you for the thoughtful feedback. We agree that the current theoretical analysis does not directly cover the generalization–memorization transition described in the latter part of the paper. \n\nOur main theoretical contribution is to show that unimodal representation dynamics emerge when the diffusion model effectively captures the low-dimensional data distribution. The second part of our study aims to reinforce this view by demonstrating that, under limited data, this distribution learning fails, leading to a breakdown of the unimodal behavior. \n\nAs the reviewer noted, a complete theoretical understanding of this transition will likely require analyzing a more nonlinear model. In particular, such a model would need to capture both the fitting of the true data distribution when data is abundant, and the overfitting to the empirical distribution when data is limited. Interestingly, our experiments suggest that diffusion models first align with the true distribution before adapting to the empirical one (as discussed in Section 4.2), hence a learning dynamic based analysis would also be interesting. \n\nWe are actively investigating this phenomenon as a future work and would greatly welcome any further insights or suggestions from the reviewer.\n\n> I assume feature accuracy was measured using a linear classifier with logits defined as the inner product between the class embedding vector and the feature vector. Have you considered trying an L2 similarity-based classifier? Since the SNR is defined via the L2 norm, I wonder if this might reveal even more consistent trends, although I understand that training might be more challenging.\n\n**Answer**: Thanks for the insightful suggestion. Yes, our linear probing experiments use logistic regression classifiers with inner-product-based logits. \n\nInspired by the reviewer's comment, we re-do the experiments in Figure 5 using a projection-based classifier: for each class, we compute its principal subspace $[V_1, ..., V_K]$ from the training features, and classify each test sample by identifying the class whose subspace captures the most projection energy, computed as $\\text{argmax}_k ||h(x_i) V_k||^2$. The results are shown below. While the projection-based classifier yields lower overall accuracy, it indeed reveals a more pronounced unimodal trend that aligns closely with the SNR dynamics. We will discuss these results in the updated manuscript.\n\n| Dataset       | Classifier Type                  | $\\sigma_t = 0.002$ | $0.008$ | $0.023$ | $0.060$ | $0.140$ | $0.296$ | $0.585$ | $1.088$ | $1.923$ | $3.257$ |\n|:-------------:|:--------------------------------:|:------------------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n| CIFAR10       | Logistic Regression (Figure 5)   | 91.32              | 93.72   | 94.94   | 95.20   | 94.11   | 91.36   | 85.25   | 72.95   | 56.21   | 41.98   |\n| CIFAR10       | Projection                       | 78.93              | 86.95   | 90.36   | 90.93   | 90.78   | 88.48   | 82.82   | 70.53   | 53.84   | 37.84   |\n\n| Dataset       | Classifier Type                  | $\\sigma_t = 0.002$ | $0.008$ | $0.023$ | $0.060$ | $0.140$ | $0.296$ | $0.585$ | $1.088$ | $1.923$ | $3.257$ |\n|:-------------:|:--------------------------------:|:------------------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n| TinyImageNet  | Logistic Regression (Figure 5)   | 31.13              | 32.17   | 34.70   | 43.09   | 53.58   | 50.78   | 42.13   | 27.56   | 15.10   | 8.39    |\n| TinyImageNet  | Projection                       | 12.02              | 12.89   | 14.38   | 24.10   | 34.78   | 34.88   | 30.02   | 20.99   | 12.64   | 7.89    |"}, {"review_id": "OCuDgQ2Bvu", "reviewer": "Reviewer_DRKN", "summary": "The paper offers an interesting theoretical and experimental analysis on the dynamics of latent representations in generative diffusion models. The theoretical analysis solely considers the case of mixtures of gaussians with zero mean and different covariance structure. The analysis is performed by considering a simplified network architecture that can be optimized exactly in the mixture of Gaussians.\n \nThe quantity of interest is a form of feature SNR that quantifies the ratio between the norm of the projection of the score on the true subspace with the norm of the projection on its orthogonal complement (noise space). The authors argue that the presence of a single mode in this feature SNR curve is an index of generalization. This particular claim is unfortunately only supported by experimental data and it is not proven theoretically.\n\nIt is interesting to understand why the SNR curve has a unimodal shape in the mixture of gaussians cases considered by the authors. The intuition is that the score function s_t(x) at time t is a (approximately) local average over the data points in a ball of radius sigma(t). For large noise levels, the ball is so large to include all components of the mixture, leading to a low feature SNR. On the other hand, for very small noise levels the average is very local and the particles are very close to the target distribution. In this case, for a datapoint close to the shared mean the covariance structure is averaged over the covariance structures of all components, which leads to a somewhat paradoxical reduction of feature SNR. \nHowever, the feature SNR would have a second peak at t = 0 if the different components are allowed to have different means since, in this more general case, the score at the ‘typical points’ will have the same covariance structure of the component-specific score when sigma(t) is small enough to properly separate the mean vectors (assuming that the distance between the mean vectors is larger than the sds of the individual components).\nWhile the authors do not make this point, this can provide a connection with memorization since in the memorization case each datapoint becomes the center of its own Gaussian component [1,2,3,4].\n\n[Edit: I increased my score based on the response to my and other reviews.]", "questions": "1) My intuition is that you should see a second peak of the SNR at t = 0 if the centers of the gaussians are well separated. Am I correct? Can you elaborate on this point?\n\n2) Is your network parameterization optimized by the exact mixture of gaussians score? This score can be obtained in closed form and it can be directly analyzed. If not, can you elaborate on the differences and similarities between the constrained optimum and the exact score?\n\n3) Can you connect the peak in the SNR with the speciation/symmetry breaking time identified in [1] and [5]? It would also be useful to discuss the connection with the transitions observed in [6]. \n\nMinor observations:\nI do not know what an informal theorem is. It would be better to either express Theorem 1 in a more rigourous way or to not present it in theorem format. \n\nIn line 133, you claim that the latent space of diffusion models is approximately Gaussian. This is just not true, if it were the case, we would not need fancy diffusion models and we could just sample by fitting the covariance matrix. In reality, the autoencoder has a very mild 'gaussianization' effect even thought it uses a KL regularization term.", "limitations": "The discussion of limitations and future directions is rather limited. I would have loved to see a more substantial discussion of the effect of changing their distributional assumptions away from the mixture of gaussians case.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Generally speaking, I do think that this paper contains good ideas and interesting theoretical and experimental analysis.  The link between the feature SNR and generalization/memorization is insightful and seems to be well supported by the experiments on real images. The theoretical part is somewhat limited but it contains interesting results. \n\nA weakness is the lack of theoretical results connecting the SNR with the memorization phenomenon. A second related weakness is the very limited coverage of very relavant existing literature on generalization, memorization, manifolds and class speciation [1-8]. It is important for the authors to reference and discuss these papers and I believe that the methods used in [1,2,4] could be very useful to connect this analysis to the theory of memorization, which is the main concern of the authors in this paper.\n\nReferences:\n\nMemorization\n[1] Biroli, Giulio, et al. \"Dynamical regimes of diffusion models.\" Nature Communications 15.1 (2024): 9957.\n[2] Ambrogioni, Luca. \"In search of dispersed memories: Generative diffusion models are associative memory networks.\" Entropy 26.5 (2024): 381.\n[3] Pham, Bao, et al. \"Memorization to generalization: Emergence of diffusion models from associative memory.\" arXiv preprint arXiv:2505.21777 (2025).\n[4] Lucibello, Carlo, and Marc Mézard. \"Exponential capacity of dense associative memories.\" Physical Review Letters 132.7 (2024): 077301.\n\nClass speciation:\n[1] Biroli, Giulio, et al. \"Dynamical regimes of diffusion models.\" Nature Communications 15.1 (2024): 9957.\n[5] Raya, Gabriel, and Luca Ambrogioni. \"Spontaneous symmetry breaking in generative diffusion models.\" Advances in Neural Information Processing Systems 36 (2023): 66377-66389.\n[6] Sclocchi, Antonio, Alessandro Favero, and Matthieu Wyart. \"A phase transition in diffusion models reveals the hierarchical nature of data.\" Proceedings of the National Academy of Sciences 122.1 (2025): e2408799121.\n\nGeneralization on manifolds:\n[7] George, Anand Jerry, Rodrigo Veiga, and Nicolas Macris. \"Analysis of Diffusion Models for Manifold Data.\" arXiv preprint arXiv:2502.04339 (2025).\n[8] Achilli, Beatrice, et al. \"Memorization and Generalization in Generative Diffusion under the Manifold Hypothesis.\" arXiv preprint arXiv:2502.09578 (2025).\n[3] Achilli, Beatrice, et al. \"Losing dimensions: Geometric memorization in generative diffusion.\" arXiv preprint arXiv:2410.08727 (2024).", "was_revised": true, "final_justification": "The author addressed some important points concerning clarity, connection with pruor literature and assumptions. I think that this work provides a novel view on the theory of generative diffusion that provides insights and complements existing approaches.", "rebuttal": "We thank the reviewer for the thoughtful and detailed feedback. We are glad to hear that the reviewer found our theoretical and experimental analyses interesting, and appreciated the insights linking representation learning to generalization and memorization. Below, we respond to the reviewer’s comments in detail.\n\n> A weakness is the lack of theoretical results connecting the SNR with the memorization phenomenon. A second related weakness is the very limited coverage of very relavant existing literature on generalization, memorization, manifolds and class speciation [1-8]. \n\n**Answer**: We are grateful to the reviewer for pointing us to this body of work, and will include a thorough discussion of all the papers above in the updated version of our manuscript. We have carefully reviewed the references provided and summarize their main contributions as follows:\n\n    [1] Studies the dynamics of diffusion model (DM) sampling using an empirical score. The authors define speciation (when the denoised sample falls into one of the data classes) and collapse (when it converges to a discrete training sample).\n    [2] Draws a connection between empirical DMs and dense associative memory (AM) networks through their objectives and sampling behavior.\n    [3] Identifies spurious samples that appear as a DM transitions from memorization to generalization, from an AM perspective.\n    [4] Analyzes the capacity and retrieval behavior of AMs for Gaussian and spherical patterns.\n    [5] Investigates the symmetry breaking in DM sampling, as in [1], and proposes a sampling strategy that starts from a smaller noise scale.\n    [6] Uses the RHM model to show that different levels of features are recovered at different timesteps during DM sampling.\n    [7] Extends the analysis in [1] to data lying on a low-dimensional manifold defined by a linear subspace and an activation function.\n    [8] Analyzes the discrepancy between the empirical and true score during sampling, illustrating how memorization leads to the loss of local geometric components.\n    \nWe note that [7] and [8] in particular adopt data assumptions closely aligned with ours. We also appreciate the reviewer’s suggestion to relate our SNR analysis to the speciation and symmetry breaking behavior observed in [1] and [5], as well as to the phase transitions analyzed in [6]. We will expand on these connections in the revised manuscript.\n\n> Can you connect the peak in the SNR with the speciation/symmetry breaking time identified in [1] and [5]? It would also be useful to discuss the connection with the transitions observed in [6].\n\n**Answer**: \nIn [1,5], the authors characterize the emergence of speciation by tracking the probability that a denoised sample aligns with a class cluster, mostly using gaussian mixtures with well-separated means. Our approach differs in that we model speciation through alignment to subspaces corresponding to different classes, under a low-dimensional data structure. The SNR we define reflects the signal strength along these subspaces, which can be interpreted as a subspace-based analog of speciation. Furthermore, the class confidence metric introduced in our Section 3.4 may be viewed as a probabilistic surrogate analogous to the metrics proposed in [1].\n\nRegarding [6], their analysis reveals that features emerge hierarchically across timesteps, starting from coarse to fine-grained structures. This matches our own observations in Section 3.4, where we note a fine-to-coarse shift on the model output caused by the representations learned by DMs during denoising. However, our focus is on internal representations rather than generated outputs.\n\n>It is important for the authors to reference and discuss these papers and I believe that the methods used in [1,2,4] could be very useful to connect this analysis to the theory of memorization, which is the main concern of the authors in this paper.\n\n**Answer**: We thank the reviewer for this constructive suggestion. We agree that the energy-based perspective in [1, 2, 4] is highly relevant, as these papers offer a powerful description of memorization by analyzing the dynamics of the empirical score, leading to phenomena like 'class speciation.' Our work provides a representation learning perspective on a similar problem, where we extensively show that generalizing DMs effectively learns the low-dimensional data structure and how this learning breaks down when data memorization occurs.\n\nReferences:\n\n\n[1] Biroli, Giulio, et al. \"Dynamical regimes of diffusion models.\" Nature Communications 15.1 (2024): 9957. \n\n[2] Ambrogioni, Luca. \"In search of dispersed memories: Generative diffusion models are associative memory networks.\" Entropy 26.5 (2024): 381.\n\n[3] Pham, Bao, et al. \"Memorization to generalization: Emergence of diffusion models from associative memory.\" arXiv preprint arXiv:2505.21777 (2025).\n\n[4] Lucibello, Carlo, and Marc Mézard. \"Exponential capacity of dense associative memories.\" Physical Review Letters 132.7 (2024): 077301.\n\n[5] Raya, Gabriel, and Luca Ambrogioni. \"Spontaneous symmetry breaking in generative diffusion models.\" Advances in Neural Information Processing Systems 36 (2023): 66377-66389. \n\n[6] Sclocchi, Antonio, Alessandro Favero, and Matthieu Wyart. \"A phase transition in diffusion models reveals the hierarchical nature of data.\" Proceedings of the National Academy of Sciences 122.1 (2025): e2408799121.\n\n[7] George, Anand Jerry, Rodrigo Veiga, and Nicolas Macris. \"Analysis of Diffusion Models for Manifold Data.\" arXiv preprint arXiv:2502.04339 (2025). \n\n[8] Achilli, Beatrice, et al. \"Memorization and Generalization in Generative Diffusion under the Manifold Hypothesis.\" arXiv preprint arXiv:2502.09578 (2025).\n\n> My intuition is that you should see a second peak of the SNR at t = 0 if the centers of the gaussians are well separated. Am I correct? Can you elaborate on this point?\n\n**Answer**:  We thank the reviewer for raising this insightful point. When the Gaussian class centers are well separated, both classification accuracy and SNR tend to decrease monotonically with increasing noise, the unimodal dynamic disappears.\n\nIn this setting, the data remains linearly separable at low noise levels. As $\\sigma_t$ increases, the added noise begins to blur the class boundaries, reducing both accuracy and SNR. Thus, instead of a unimodal shape, we observe a monotonic decline.\n\nTo validate this, we conduct an experiment with well-separated class means. The results confirm this behavior:\n\n| $\\sigma_t$ | 0.030 | 0.053 | 0.098 | 0.189 | 0.282 | 0.379 | 0.480 | 0.588 | 0.704 | 0.830 | 0.989 | 1.492 | 1.717 | 1.978 |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n| **Acc** | 100 | 100 | 100 | 99.8 | 97.7 | 92.4 | 85.7 | 78.5 | 71.8 | 66.0 | 61.6 | 52.6 | 50.1 | 48.0 |\n| **SNR** | 0.540 | 0.539 | 0.535 | 0.524 | 0.513 | 0.502 | 0.484 | 0.465 | 0.436 | 0.406 | 0.370 | 0.297 | 0.266 | 0.240 |\n\n> Is your network parameterization optimized by the exact mixture of gaussians score? This score can be obtained in closed form and it can be directly analyzed. If not, can you elaborate on the differences and similarities between the constrained optimum and the exact score?\n\n**Answer**: Thanks for the good catch. Our parameterization is indeed inspired by the closed-form posterior mean of the MoLRG model, since a denoising auto-encoder should estimate that mean. Nonetheless, we note that our network is not fitted by directly plugging in the closed-form MoLRG score; instead, it is trained with the standard denoising objective (Eq. 2). \nIn essence, our network acts as a constrained function approximator for the true score. The architectural constraints enforce the correct functional form, while the standard denoising training finds the parameters of that function that best fit the data. The success of this approach is validated by Figure 3, where the SNR of our trained network closely tracks the theoretical optimum. We will make this connection more explicit in the updated manuscript.\n\n> Minor observations: I do not know what an informal theorem is. It would be better to either express Theorem 1 in a more rigorous way or to not present it in theorem format.\n\n**Answer**: Thanks for pointing this out. We will follow the reviewer's suggestion to rename “Theorem 1 (Informal)” in the main text to “Key result (informal)” and add a pointer to the formal version of the Theorem in Appendix E. We hope this can prevent any confusion and maintain rigor.\n\n> In line 133, you claim that the latent space of diffusion models is approximately Gaussian. This is just not true, if it were the case, we would not need fancy diffusion models and we could just sample by fitting the covariance matrix. In reality, the autoencoder has a very mild 'gaussianization' effect even thought it uses a KL regularization term.\n\n**Answer**: Thanks for pointing out this problem. We will soften the wording from “The latent space of latent diffusion models is approximately Gaussian” to “Latent diffusion models encourage the latent space toward a Gaussian distribution”."}]

Author rebuttal:

Meta-review: This paper develops a low-dimensional analytical account of unimodal representation dynamics in diffusion models and tests its implications empirically. Under a mixture-of-low-rank-Gaussians data model and a simplified denoiser, the authors derive an explicit metric for representation quality that rises with denoising strength at small noise and falls with waning class confidence at large noise, yielding a single peak. They show that this metric predicts the empirical peak in feature quality across timesteps on synthetic data and standard vision benchmarks, and they further observe that the presence of a unimodal curve tracks generalization, while its erosion accompanies memorization. Reviewers found the problem important and viewed the analysis as a first theoretical treatment of diffusion-model representations; they also noted practical value in using the peak as an empirical indicator (e.g., for early stopping).

Strengths include (i) a clear theoretical mechanism explaining unimodality that aligns with measurements, (ii) a unifying SNR formalism that connects denoising and class confidence, and (iii) broad empirical support across datasets, probing schemes, and architectures. The main weaknesses are the reliance on stylized assumptions (orthogonal class subspaces, matched ranks, simplified denoiser), gaps in formally tying SNR to memorization, and some missing or deferred related work and notation/clarity issues. During rebuttal the authors ran additional ablations (overlapping subspaces, unequal ranks, non-uniform mixing), added alternative probing that strengthened the SNR link, and showed the trend in latent diffusion models; they clarified approximations and notation, softened claims, and committed to incorporating suggested literature. All reviewers vote for acceptance. I recommend acceptance.
Decision: Accept (poster)