On the Feature Learning in Diffusion Models
Authors: Andi Han
Wei Huang
Yuan Cao
Difan Zou
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion model
Feature learning
Primary area: generative models
TL;DR: A theoretical framework that shows diffusion models learn balanced features while classification learns discriminative features

Abstract: The predominant success of diffusion models in generative modeling has spurred significant interest in understanding their theoretical foundations. In this work, we propose a feature learning framework aimed at analyzing and comparing the training dynamics of diffusion models with those of traditional classification models. Our theoretical analysis demonstrates that diffusion models, due to the denoising objective, are encouraged to learn more balanced and comprehensive representations of the data. In contrast, neural networks with a similar architecture trained for classification tend to prioritize learning specific patterns in the data, often focusing on easy-to-learn components. To support these theoretical insights, we conduct several experiments on both synthetic and real-world datasets, which empirically validate our findings and highlight the distinct feature learning dynamics in diffusion models compared to classification.

Published as a conference paper at ICLR 2025
ON THE FEATURE LEARNING IN DIFFUSION MODELS
Andi Han"
Wei Huang"
Yuan Cao
Difan Zou
"RIKEN AIP (andi han@r jp, wei. www.g....g.....m vr@riken. ip). Equal contribution. i Depaatmeent of Statistics and Actuarial Science, University of Hong Kong (yuancao @ahunc hk)
! DDeparmment of Computer Science and Institute of Data Science, University of Hong Kong (dzouêcs cs hku.hk)
ABSTRACT
The predominant success of diffusion models in generative modeling has spurred significant interest understanding their theoretical foundations. this work, we propose feature learning framework aimed at analyzing and comparing the train- ing dynamics of diffusion models with those of traditional classification models. Our theoretical analysis demonstrates that diffusion models, due the denoising objective, are encouraged to learn more balanced and comprehensive representa- tions of the data. In contrast, neural networks with a similar architecture trained for classification tend to prioritize learning specific patterns in the data, often focusing on easy-to-learn components. To support these theoretical insights, we conduct several experiments on both synthetic and real-world datasets, which em- pirically validate our findings and highlight the distinct feature learning dynamics in diffusion models compared to classification.
INTRODUCTION
Diffusion models (Ho et al., 2020; Song et al., 2021) have emerged as a powerful class of generative models for content synthesis and have demonstrated state-of-the-art generative performance in variety of domains, such as computer vision (Dhariwal & Nichol, 2021; Peebles & Xie, 2023), acoustic (Kong al., 2021; Chen al., 2021) and biochemical (Hoogeboom al., 2022: Watson et al., 2023). Recently, many works have employed (pre-trained) diffusion models to extract useful representations for tasks other than generative modelling, and demonstrated surprising capabilities in classical tasks such image classification with little-to-no tuning (Mukhopadhyay et al., 2023; Xiang et al., 2023; et al., 2023a; Clark & Jaini, 2024; Yang Wang, 2023; Jaini et al., 2024). Compared to discriminative models trained with supervised learning, diffusion models not only able to achieve comparable recognition performance (Li et al., 2023a), but also demonstrate exceptional transferablity (Li et al., 2023a; Jaini et al., 2024) and improved classification robustness (Chen et al., 2024c).
The significant representation learning power suggests diffusion models are able to extract meaningful features from training data. Indeed, the core of diffusion models to estimate the data distribution through progressively denoising noisy inputs over several iterative steps. This inherently views data distribution as composition of multiple latent features and therefore learning the data distribution corresponds to learning the underlying features. Nevertheless, it remains unclear
how feature learning emerges during the training of diffusion models and whether the feature learning process is different to supervised learning.
Regardless of the ground-breaking success of diffusion models, the theoretical understanding is still in its infancy. Existing analysis on diffusion models has mostly focused on theoretical guarantees in terms of distribution estimation and sampling convergence. Several works have derived statistical estimation errors between distribution generated by diffusion models to ground-trute distribution (Oko et al., 2023; Zhang et al., 2024; Chen et al., 2023a), showing that diffusion models achieve a minimax optimal rate under certain assumptions 2024). Algorithmically. Li et al. (2023c); Han et al. (2024) studied the estimation error of diffusion models trained with gradient descent using kernel methods. Shah et al. (2023); Gatmiry al. (2024);
the true density (Oko al., 2023; Zhang et al.,

---

Published as a conference paper at ICLR 2025
Chen et al. (2024d) introduced algorithms based on diffusion models for learning Gaussian mixture models. In addition, given access to sufficiently accurate score estimation, Lee et al. (2022: 2023); Chen et al. (2023b); Li et al. (2023b) proved the convergence guarantees of sampling in (score-based) diffusion models. Despite showing provable guarantees for diffusion models, existing theories are limited the generative aspects of diffusion models, namely distribution learning and sampling. To the best of our knowledge, no theoretical analysis is performed to elucidate the feature learning process in diffusion models.
Notations. We make use of the following notations throughout the paper. We use || to de- note L2 norm for vectors and Frobenius norm for matrices, unless mentioning otherwise. We use O(-), (2()), ((-), o(-), for the big-O, big-Omega big-Theta, small-o, small-omess notations. We write O-) hide logarithti factors and similar notations hold for (2(.) and ((-). For a binary condition C, we let (C) = 1 if C is true and 1 (CC = otherwise.
1.1 OUR MAIN RESULTS
In this work, we develop a theoretical framework that studies feature learning dynamics of diffusion models and compares with classification. Inspired by the image data structure, we employ multi- patch data distribution x = [My: E] for both classi- fication and diffusion model training. We consider a binary-class data setup with y = #1 as the data label and μ1, f--1 E I2d nal vectors, i.c., 141 1 p.11. representing the sig- nal. On the other hand, is the label-independent noise, which is randomly sampled from a Gaus- sian distribution with standard deviation de.
signal learning
noise learning
w(1)
L
e(1)
two fixed orthogo-
Classification
Diffusion model
<(1
o(1)
@1j
(11
SNR 2
In order to elucidate the difference of feature learning dynamies for the two tasks, adopt a two-layer convolutional neural network with quadratic activation. For diffusion model, we consider a weight-sharing setting for the first and second layer, which is commonly considered for analyzing autoencoders (Nguyen, 2021; Cui & Zdeborová, 2024). For classification, we fix the second layer weights to be +1, following Cao et al. (2022); Kou et al. (2023). In other words, the clas- sifier can be viewed as attaching fixed linear head the intermediate layer of the diffusion model. Given a training dataset of n samples from the multi-patch data distribution, we use gradient descent to minimize the empirical logistic loss for classification and the DDPM loss (Ho al., 2020) with expectation over the diffusion noise.
Figure 1: Illustration of the ratio of signal learn- ing to noise learning when varying n SNR?,, where SNR == [I We show diffu- sion model tends to study more balanced signal and noise while classification has a sharp phase transition and tends to focus on learning either signal or noise.
Under the above settings, we investigate the differences of feature learning dynamics (Allen-Zhu & Li, 2023; Cao et al., 2022: Zou et al., 2023; Huang et al., 2023b; Jiang et al., 2024; Huang et al., 2024a: Lu al., 2024; Meng et al., 2024) between diffusion model and classification. We quantify the feature learning in terms of signal learning and noise learning, measured through the alignment between the network weights to the directions of signal/noise respectively, i.e., [(w, myyi, (w, E We present the following (informal) results for the two learning paradigms.
Theorem 1.1 (Informal). Let SNR == [|pll/(0e dd) be the rignal-to-noie ratio. We can show
For diffusion model, (w, Myyle (w E)1 exhibit linear growth initially and there exists a stationary point along the path of the training dynamics thut satisfies kw,,yYl/le
= O(n SNR²).
For classification, ([w, p.y)i k[,, E) exhibit exponential growth initially and when n SNR2 >
3 for some constant 8 > 1, ((w, p/Ml/ww, E) = w01), and when n SNR? < 1/8.
(w, (w)l//(w. o())
2

---

Published as a conference paper at ICLR 2025
Theorem 1.1 highlights differences in the feature learning process between diffusion models and classification. Especially in the regime where SNR² (1), classification is sensitive to changes in SNR and tends to learn either the signal My or the noise ç. In contrast, diffusion model learns both signal and noise to the same order. Such a claim is visualized in Figure 1.
We believe our framework represents the first attempt to systematically investigate feature learning within diffusion models, potentially uncovering novel insights into the intriguing properties of diffusion models, including but not limited to critical window (Sclocchi et al., 2024; Li & Chen, 2024), shape bias (Jaini al., 2024), classification robustmess (Chen et al., 2024c), feature composition and dependence (Okawa et al., 2024; Yang et al., 2025; Han et al., 2025).
1.2 RELATED WORK
Theoretical analysis of diffusion model. Existing theoretical guarantees for diffusion models focus on distribution estimation and sampling. For distribution estimation, Oko al. (2023) proved that diffusion models achieve nearly minimax optimal estimation error where the true density defined over a bounded Besov space. Zhang et al. (2024) extended the minimax optimality to more general sub- Gaussian densities with sufficient smoothness. When the density is supported oow-dimensional subspace, diffusion models avoid curse of dimensionality with an estimation rate depending only the intrinsic dimension (Oko et al., 2023; Chen et al., 2023a). Furthermore, Shah et al. (2023); Gatmiry et al. (2024); Chen et al. (2024d) introduced algorithms based on diffusion models for learning a mixture of Gaussians. Other works provided guarantees of diffusion model trained by gradient descent (Li et al., 2023c; Han et al., 2024; Wang et al., 2024). For sampling, Lee al. (2022; 2023); Chen al. (2023b); Li al. (2023b) have shown (score-based) diffusion models converge polynomially under sufficiently accurate score estimation. Recent studies also aimed accelerate the convergence via strategies such as consistency training (Song et al., 2023: Li et al., 2024b), advanced design of the reverse transition kernel (Huang et al., 2024b), higher-order approximation (Li et al., 2024a) and parallelization (Chen al., 2024a; Gupta et al., 2024). In addition, Li & Chen (2024) theoretically verified critical windows of feature emergence during the sampling process, provided accurate score estimation.
Theoretical analysis on (denoising) autoencoders. Diffusion models can be viewed as multi-level denoising autoencoders (Xiang et al., 2023). While there extensive research on the theoretical guarantees of autoencoders without denoising, most studies focus on linear autoencoders (Kunin et al., 2019; Oftadch et al., 2020; Steck, 2020; Bao et al., 2020). In contrast, only a limited number of works analyze non-linear autoencoders, primarily in the lazy training regime (Nguyen et al., 2021) the mean-field regime (Nguyen, 2021). Additionally, the training dynamics of non-linear autoencoders have been investigated under population gradient descent (Shevchenko et al., 2023; Kögler al., 2024) and online gradient descent (Refinetti & Goldt, 2022). On the other hand, the training dynamics of denoising autoencoders have been studied in the context of linear networks (Pretorius al., 2018) and in the high- -igheniiensional asymptotic limit (Cui & Zdeborová, 2024).
Diffussion model for representation learning. Pre-trained diffusion models are shown learn powerful representation, which is useful for downstream tasks such as classification (Mukhopadhyay et al., 2023; Xiang et al., 2023; Li al., 2023a; Clark & Jaini, 2024; Yang & Wang. 2023), semantic segmentation (Baranchuk et al., 2022: Zhao et al., 2023; Yang & Wang, 2023). Moreover, many works have found intriguing properties understand shape bias (Jaini et al., 2024) and improved adversarial robustness (Chen et al., 2024c). For more detailed exposition, we refer to the recent survey on this matter (Fuest et al., 2024).
diffusion models used as classifier, including its ability to
2 PROBLEM SETTING
This section introduces the problem settings for both diffusion model and classification, including the data model, neural network functions as well as training objectives and algorithm.
Definition 2.1 (Data distribution). Each data sample consists of two patches, as x = [x(T) (TT)T where each patch is generated as follows:
Sample y (1, uniformly with Ply 11) Ply 1)= 1/2.
3

---

Published as a conference paper at ICLR 2025
Given two orthogonal signal vectors 11, M-11, with μal 1 μ.11 set x()) Hy i.e. x())
N1 y and x(1) 11-1 if y -1. For simplicity, assume || 11p--11 ||
Set x where E~ ~~, of -1 | ||=
1-1 1-11 ?))
This multi-patee data model reflects the structure of image data, where each image consists of multiple patches, and only a subset of the patches are relevant the class label, while the rest contribute as background noise. This data model has been employed in several existing studies (Allen-Zhu & Li, 2023; Cao et al., 2022; Kou et al., 2023; Meng et al., 2024; Zou et al., 2023). difference in our model is the use of two orthogonal signal vectors, in contrast to previous works, that employ single signal vector of the form yj.. Additionally, while our analysis focuses on a two-patch setting for simplicity, it can be readily extended to multi-patch data. We let SNR = |ull/(g//)) denote the signal-to-nois ratio.
Neural network functions. We study two-layer convolutional-tyee neural networks for both diffusion model and classification. For diffusion model, we consider neural network with quadratic activation and shared first-layer and second-laye weights:
f(W, x) = [(( x f2 f2 (W, x
E R
m
where (WW m
wr ( Wr: p == 1,2
where m denotes the network width and represents the neuron index.
For classification, we consider a similar neural network with quadratic activation where second-layer weights are fixed be =1 (instead of w,):
f(W.x) Fi Fi(W111W. (-1(W-1.13
where F(W, = 1 Σ xw.r.xx0 - Σ (W,,,, x (2)
We remark that the use of polynomial activation, such as quadratic, cubic and ReLU with polynomial smoothing is not uncommon in existing theoretical works (Cao al., 2022; Jelassi & Li, 2022; Zou et al., 2023; Huang et al., 2023a; Meng et al., 2023). The aim is to better elucidate the separation between signal and noise learning dynamics.
Training objectives and algorithm. For diffusion model, we employ the objective of denoising diffusion probabilistic model (DDPM) (Ho et al., 2020). We let [x(2), x(2)) E R22d to denote input image. For a given diffusion time step E [0, T), we sample x, = axXp + 3,E, for E, ~NO0 and a pre-determined noise schedule coefficients (a,, 8, T. The aim of diffusion models is to estimate the mean the posterior distribution of the noise E; conditioned on X, This is achieved by training neural network f to predict the noise added each step f. The DDPM loss is given ell up to some re-scaling (Ho et al., 2020). We consider a finite-sample setup
by Exx.e..:ll(xx) given by the training images ) sampled according Definition 2.1 and thus the empirical DDPM loss at time step becomes
1 n Be., (..(..))).)) E2,1l
12
Lr(W,
Σ Ee.. |f(W 1, agXo, + B(E,,) e,,
2m
2n
where we let Xo,i X, and Xt.i 0,x0 + S,e,, Here, we decouple the training of neural network at each diffusion time step with separate weight parameters, a strategy also adopted in (Shah et al., 2023) for simplicity of analysis.
Unlike (Han et al., 2024), where each sample i is associated with a single noise c,, ~ ~O,0 I).
we here consider taking the expectation over the noise distribution, which aligns with the practical setting where multiple noises are sampled for each input data. We use gradient descent to train diffusion model starting from random Gaussian initialization w (O, (f) wh k-1 ww v/V T. (WF), where the superscript l: is the iteration index.
4

---

Published as a conference paper at ICLR 2025
For classification, minimize the empirical logistic loss over the training data {x, V/ V-
(s( 1 Σ (lnf(W,x))) {(2) = log (1 exp(-2)).
The same as diffusion model, we use gradient descent to train the neural network starting from random Gaussian initialization wf.r o
MAIN RESULTS
Our main results are based on the following conditions.
Condition 3.1. Suppose the following holds.
1. Dimension is sufficiently large with m3).
2. The sample size satisfies =
3. The standard deviation of initialization do chosen such that Ō(n'mo 1 11 s 00 S
Ö( min/m (tin/mmm1/6-/8.../0/1 1/3 1/33 mn-1/6d-//12 1/ d d-//4///4 0f n]).
4. The learning rate y satisfies n O0iin(mmmao) - 1-11/2 nmog dd-
5. The signal strength satisfies ||
( and noise variation og satisfies ((max(n97 m//4 nm S o S 1/44)
6. The noise coefficients for diffusion model satisfy at, 31 O(1).
Condition 3.1 requires d to be sufficiently large to ensure learning in an over paramettrrrezedd setting, Furthermore, we require the sample size to be lower bounded by a constant subject to logarith- mic factors. The upper bound on the initialization dg is to ensure random initialization does not significantly affect the signal and noise learning dynamics. The lower bound on og is required to bound the noise inner product at initialization for properly minimizing the training loss of clas- sification. The learning rate is chosen sufficiently small for the convergence analysis for the classification. Lastly for diffusion model, consider the constant order for |||l|| and further re- strict the range of og. Despite these conditions, setting covers broad range of n SNR i.e., Ond 112) S n SNR2 < Óminnn 4m - 2//2 , n -1/ We also consider constant order of ae 3, avoid degeneracy in learning dynamics.
We present the main results for diffusion model (Theorem 3.1) and classification (Theorem 3.2).
Theorem 3.1 (Diffusion model). Under Condition 3.1, suppose m @(1). With probability at least 1 - 5 (for any 8 >> 0), there exists stationary point W; along the training trajectory of diffiusion model, ie., Lw..r
0 that satisfies (1) Wheer
((W" (2)
(w.,: Ei) = ((w" E, and for all = =1, E [m], E [m].
kw //)l//(w", E.)1 E.)+ O(n SNR?).
with (w m1) (()) in SNR" 2(1), and (w E.) ((1) in 1 SNR 2(1).
Theorem 3.1 states that diffusion model training encourages balanced signal and noise learning. i.c.,
the neurons share the same order the directions of signals and noise. Notably, the ratio between signal and noise learning governed by the SNR, with a stationary magnitude
SNR?
Theorem 3.2 (Classification). Let Tμ, = -0nn-mmmlpll (m|lpll1)) and Te = 2mm== -2dd-) and > 0.
Under Condition 3.1, suppose m (n(l))(n/))). There exist two absolute constants C C> such that with probability at least 1 - 6, it satisfies that:
When n 2NR² 2 C, there exists S < Tμ such that Ls ( S 0.1 and
max|(w) u
2, Vj +1,
max (w
j.r {n) o(1).
5

---

Published as a conference paper at ICLR 2025
When SNR* < C, there exists 0 S k Tg such that Ls(W*) 0.1 and
max
(w
Ve.r! Ei) 2 1, Vi E [n],
max ( j,r Myy = o(1)
J.r-2
Theorem 3.2 establishes a sharp phase transition between signal and noise learning under classification training. The transition is precisely determined by n SNR2. That is, when n SNR 2 C for some constant C > 0, the neural network learns signal to achieve small training loss. On the contrary, when n SNR? < C. for some constant C' (0, C), the neural network overfits noise to achieve convergence. With standard techniques, such as in (Cao al., 2022), we can show signal and noise learning corresponds to the regime of benign and harmful overfitting respectively. To the best of our knowledge, this is the first result that shows separation under the constant of n SNR?.
Diffusion model learns balanced features while classification learn dominant features. Compar- ing the learning outcomes of diffusion model and classification, we reveal a critical difference that diffusion models learn more balanced features depending on the SNR conditions, while classification is prone to learning either signal or noise predominately. This can be best understood in the case of n SNR = @(1). By Theorem 3.2, we have either signal learning or noise dominating the learning process in classification, while Theorem 3.1 suggests signal and noise learning are in the same order in diffusion models. The theoretical findings corroborate the empirical observations that the neural network trained for classification is prone to overly rely on learning specific pattern that is casier to learn, a process known as shortcut learning (Geirhos et al., 2020). Mcanwhile, diffusion models tend to learn low-frequency, global patterns (Jaini et al., 2024), which helps improve the classification robustness (Chen et al., 2024b;c).
4 PROOF OVERVIEW
This section outlines the proof roadmap for the main results. For diffusion model, the mean mean-squared loss, the joint training of two layers as well as learning in the direction of initialization, pose significant challenges for the analysis. We adopt a two-stage analysis and characterize the stationary points based on the derived results at the end of the second stage. For classification, the two-stage amalysis is similar as in (Cao et al., 2022; Kou et al., 2023) where the first stage learns signal or noise vector sufficiently fast and the second stage shows convergence in the training loss where the learned scale difference in the first stage is maintained. However for classification analysis, we highlight two critical differences compared to existing works (Cao al., 2022; Kou et al., 2023; Meng et al., 2024), i..., constant SNR² 2 condition and quadratic activation.
4.1 DIFFUSION MODEL
We first simplify the DDPM loss by taking the expectation with respect to the added diffusion noise:
LF(W, d+ 2n
1
E
m
40,
m
| 1l (w,,,, x <b7))
m
m
(p) where for p= 1,2, x = a, + B,e with x = fyy and x E, and e 2 ~ I). We further simplify Jn in Lemma D.2 (in Appendix). We remark that 71 corresponds to regularization term that regulates the magnitude and alignment of neurons, while I2 corresponds to the main learning term. We highlight that apart from the signal and noise directions, the learning term l2 also includes the initialization direction w which further complicates the analysis.
First stage. In the first stage, where all the key quantities, including signal and noise inner products, weight norms and cross-neuron inner products remain close to their respective initialization, we can show the growth of the signal and noise inner products is approximately linear:
(witt (4)) = (w (4) + Ovll l [|a)))
k+1
(w P1 (.) = (wit, { + e (mm- w II (1))
k+
In addition, the change of w along direction w can be properly controlled such that the scale of key quantities remain unaffected and the simplification in (1) is valid throughout the first stage.

---

Published as a conference paper at ICLR 2025
The updates in (1) immediately suggest that once the growth terms of the inner products dominate their initialization, we obtain k M1)l/l(ww; ,,,,, = @( : SNR²). This marks the end of the first stage, as described in the following lemma.
Lemma 4.1. Under Condition 3.1, there exists an iteration Ti = max (u, Te). where Ta Olymoo d (- and Iş @(nymon og -- 3/2 such that for all k S Tı. |w @(o)d), (W wi.)) ((oīd) for all E [m].] = +1, E [n]. Furthermore, we can show for all j.j = =1, E [m], [m].
(wi)) m3) O ((w (g)) (wi), E.) = (((w"), Ex)) and
( /3)l/l + = On SNR*),
Lemma 4.1 verifies that at the end of the first stage, all the neurons are concentrated and the ratio precisely determined by n SNR². This is critically different compared to classification where signal and noise learning exhibits exponential growth as we show later and thus shows clear scale difference at the end of the first stage, even when n SNR" = (1).
Second stage. The second stage aims characterize when the dominant terms in the gradients along the key directions become no longer dominant. To this end, we decompose the gradient into
(Vw., L(W)), = - ([I -1 (l E,
(Vw,. L(W)), = (n [| " + E.
(Vw,, L(W)), wh1)) -((((w)) My + E) - | (ff (w wP,)) + [| (1 1 +E)) + b Bn
where E p,,, 1 By,,, Er,, are the residual terms of the gradients and we let E following lemma shows before Ex,,,yy: E E r,t,uu reach order as the dominant terms, the ratio of signal and noise inner products are preserved.
1 n. Ei= & The
Lemma 4.2. There exists an iteration Tr > Ti with T2 = O(max(1) %0 2
"n noo 2 og)) i E [ni] (1) if SNR² = 2(1), (w Mb)) O(1) E.) = =()); (2) Eii,, = @([["?21PP1 Erg,e =
such that for all
tl,r E [m],i
and if n -1 SNR
-(1,, O(n -| T2 121E,113), ETT E @(((w r.t My E) - [| (w who>+1lw? ||,,, (w2h, + E))
and for all Ti S k S T2. we have
(wf. 14) (((w , (b))). (w E.) = (.)
(ww. W //)l/I(w> <:)1 = O(n SNR*),
Lemma 4.2 characterizes T2 as the point where the dominant terms of the gradients in the first stage become comparable to the residual terms. Meanwhile, we show the scale of signal and noise inner products escape from initialization and reach constant order. Throughout the second stage, the concentration of neurons are preserved and the ratio of signal to noise learning is dictated by n SNR Finally, we identify there exists a stationary point that satisfies the conditions at the end of the second stage (Lemma 4.2).
Theorem 4.1 (Informal). There exists a stationary point W ie.. Vw,, L(W:) = 0 such that the conditions at T2 (in Lemma 4.2) are satisfied, and in particular I(W p/)l/1(w,,, &.) (n SNRR) for all j = +1,r E [m], E [n].
4.2 CLASSIFICATION
For classification, we first let Sy := {i E [ : = y for y +1 and C:t = (y( We can rewrite the gradient descent updates in terms of the signal and noise inner products:
Wir Mgy (w .r My) S/Sy eth (w). Mg)ivileal (1 S/S, ll (all nm
(we
R+1
ju) - Mn), (2)
nm
(wh+1 .Ei) (w) E.) nm
"1 elth (w
nm Σ
e (w)r Ev)jov (3)
7

---

Published as a conference paper at ICLR 2025
Classificatoo
Diffusion model
Clessiiication
Diffusion model
Testure eaming
Festure ure saaming
Feat caminn
gppcilee EH
LP re
mp
mec00e..... PM
masije, 441
shil
Eeeraion
Iteration
Iterotio
Iteratoo
Low SNR (n SNR? 2 0.75)
High SNR (n SNR? 6.75)
Figure 2: Experiments on the synthetic dataset with both low SNR (n SNR? 0.75) and high SNR (n SNR 2 6.75). In the low SNR setting, learning for the classification task and in the high SNR setting, signal learning quickly dominates noise learning. Meanwhile diffusion model converges to stationary point that with signal-to aggal-to-noiss learning ratio respects the order of n SNR?. More experimental results additional SNR values are in Appendix A.2.
noise learning quickly dominates signal
for all j,y +l, E im], [n]. For the signal inner product, its update suggests that for any j = =1, Wjr specializes the learning of μe due to that etth < O, (w j.r Mnyi = (1 nls 2 thh ju)l(w nm
jir ال)س > 1(w), M3)i only when j = y. For the noise inner product, the growth is dominated by the second term where Er), = Ö(d ( |< is negligible. Thus, we show (w (w &i)l grows only for j = ys. In contrast, for j = - Yi. its magnitude cannot increase relative to the scale of initialization. Next, we decompose the analysis into two stages.
First stage. In the first stage before the maximum of signal and noise inner product reaches constant order, the loss derivatives can be lower bounded by an absolute constant, i.e., 1 2 Ce. for all k S Ti. As a result, both signal and noise inner product can grow exponentially and the relative growth rates are determined by n SNR?. A constant order of difference in the growth rates is sufficient to ensure scale separation in the signal and noise learning at the end of the first stage, as shown in the following Lemma.
Lemma 4.3. Under Condition 3.1: (1) When n SNR2 = (((),, there exists T₁ = =Onn (6nn-mmllp -2), T1 such that 1 E (w j.r Myyi 2 2 for all = +1 and maxj,rá (w Ti j.r E.)l = o(1). (2) When
m. 2 = 2(1), there exists Ti = 2m 2d11) 2d-11 such that 1 Em (w Barr Ei) 2 4
n m11, SNN
for all i E [n] and max (( T.r pn)) = o().
Remark 4.1. Different to existing analysis that only shows maximum inner product reaches constant order (Cao al., 2022; Huang et al., 2023a), we also show the average inner product reach constant order at the same time. Such a stronger result is required for the analysis under the constant order of SNR², which reduces the required iteration number in the second stage by an order of m.
Second stage. In the second stage, we show the loss converges while the scale separation established in Lemma 4.3 is maintained. Because n SNR2 can be a constant, we require to carefully bound the loss derivatives in the second stage particularly for establishing the upper bound for (w & when n SNR* 2(1). The naïve bound max, | < max; |eh1 S mLLs(W' used in (Cao et al., 2022) no longer works as it introduces an additional factor of n. To provide tighter bound, we show the ratio of loss derivatives in the case of n SNR² (()), i.e., || with y/ = yv. k 2 Ti, where Ci > O is a constant. This is possible because the network output is dominated by the signal. which is shared across samples with the same label. This allows to bound max; [6,4 = O(ISy.. Eies,, ( I ((Ls)
Ci for all i, [r]
5 EXPERIMENTS
We conduct both synthetic and real-world experiments verify our theoretical claims.
8

---

Published as a conference paper at ICLR 2025
Impur
Etevidiin
100
2015
Classification
5010
2002
(b)
200
2011
Servign (OUF))
Feature ICins)
No (Chane)
Diffusion
Feature
Itanation (Closs
Figure 3: Experiments on NoiSy-MNIST with SNR = 0.1. (First
Figure 4: Experiments on Noisy-MNIST with SNR 0.1. (a) Train loss for classifi- cation. (b) Train loss for dif-
row): Test Noisy-MNIST images; (Second row): Illustration of in- put gradient, i.e., Vx (W, when when y (Third row): denoised image from diffusion model. In this low-SNR case, we see classification tends to predominately learn noise while diffusion learns both signals and noise.
and F= x)
fusion model. (c) Feture learn- ing dynamics.
5.1 SYNTHETIC EXPERIMENT
Setup. We follow Definition 2.1 to generate synthetic dataset for both diffusion model and classification. Specifically, we set data dimension d = 1000 and let μu Ju,0. .0] E Rd and f-1 00, Jr, 0, 0] Rdd We sample the noise patch & NO, Ia). i 'n] (i.e., ag 1). We set sample size and network width to be n 30 and m 20 and initialize the weights to be Gaussian with a standard deviation σe = 0.001. We vary the choice of μ to create two problem settings: (I) low SNR with μ 5, which leads SNR 6.7 We the same two-layer networks introduced in Section 2. For classification, we set a learning rate of 7 0.1 and train for 500 iterations. For diffusion model, we minimize the DDPM loss by averaging over the diffusion noise, following the standard training of diffusion model. In particular, for each sample, we samples by taking an average over the noise. For the noise coefficients, we consider time 0.2 and set ar = exp(-() = 0.82 and 3, /1 exp( - 21) = 0.57. For diffusion model, we set = 0.5 and train for 40000 iterations.
SNR² 0.75 and (2) high SNR with μ = 15, which leads
2000 noise at each iteration and the loss is calculated
Results. Figure 2, we compare signal and noise learning dynamics-visualizee through maximum signal and noise inner product-between classification and diffusion model. In Appendix A.I, also include training loss convergence for both the tasks as well as training and test accuracy for classification. For classification, the training loss converges while diffusion model recovers only stationary point.
In terms of feature learning, noise learning in classification quickly dominates signal learning by exhibiting a significant larger growth in the first stage (up to around 20 iterations). This ensures that noise learning stabilizes second stage, training loss converges and signal and noise learning exhibits logarithmic growth. For diffusion model, in the first stage, where training loss does materially change, both signal and noise learning increase linearly. In the second stage where loss significantly decreases, signal and noise learning start to grow exponentially and the final stage, due to the weight regularization terms, noise and signal reach stationary point that preserves the scale of SNR?.
constant order while signal learning remains relatively small. In the

---

Published as a conference paper at ICLR 2025
5.2 REAL-WORLD EXPERIMENT
Setup. We also conduct experiments on the MNIST dataset (Lecun et al., 1998) to support our theory. In order to better control the signal-to-noise ratio, we create Noisy-MNIST dataset, where we treat each original MNIST image as clean signal patch and concatenate standard Gaussian noise patch with the same size, i.e., 28 28. In addition, we scale the signal patch by constant denoted SNR. Because the noise scale is fixed, higher SNR corresponds to higher SNR. Some sample images with SNR 0.1 are shown in the first row of Figure 3. We select 50 samples each from digit 0 and 1 respectively (i.c., n = 100). We consider the same neural networks as in the synthetic example, where we set m 100 and initialize the weights with dg 0.01. For diffusion model, choose the same (kt, 3r as in the synthetic experiment. In the main paper, we present the results for SNR = 0.1, which corresponds to low SNR setting.
Results. Figure b(a.b) shows that both classification and diffusion model converge in loss. Addi- tionally, Figure 4(c) plots the signal and noise learning dynamics. Because each image is composed of unique signal μå; and noise patch E; for E [m]. we measure the signal and noise learning by computing 1
a >}= max, (w,, (4)) and a Σi= max, (w,, respectively. We notice that due to the low SNK, noise learning in classification dominates signal learning at convergence while diffusion model learns more balanced features. This corroborates our theoretical findings.
To visualize the patterns learned by the neural networks, for classification, we adopt an approach similar to Grad-CAM (Selvaraju et al., 2020) by analyzing the gradient of output with respect to the input. Specifically, for samples of digit O, we plot the gradient of negative function output, Vx (W, x), while for digit 1, we plot VxI++(W x). As shown in the second row of Figure 3, the gradients of six test images indicate that classification primarily learns the noise rather than the signal patch. For diffusion model, we first add diffusion noise to the input images and use the network to predict the added noise. Then we reconstruct the input with the formula &o (x, 3,,6(x,) (x,))/ax, / where é( ((x.) denotes the predicted diffusion noise. The third row of Figure shows that the diffusion model learns both the signal and noise. In Appendix A.3, we present results for high-SNR setting
with SNR 0.5, where we observe the reverse pattern: classification predominately captures signal rather than noise while diffusion model continues to balance the learning of both signal and noise. Additionally, Appendix A.5 presents experiments on all 10 digits of the MNIST dataset, verifying the observed distinctions in feature learning between diffusion models and classification.
CONCLUSION
This work introduces a theoretical framework for analyzing the feature learning dynamics in diffusion models, taking an initial step toward understanding the representation learning in diffusion models. Our findings demonstrate that diffusion models inherently promote a more balanced feature learning. in contrast to classification models, which tend to prioritize certain features over others. This suggests that classification models may be more sensitive to variations in the signal-to-noise ratio compared to diffusion models. Consequently, this may provide explanation for the inherent adversarial robustness of diffusion models in downstream applications, such as classification (Li et al., 2023a; Chen et al.. 2024c;b), as perturbations are less likely significantly affect the learned representations compared to classification models.
Although our study focuses on a two-patch data setup, the proposed framework can be extended to accommodate more complex data settings. For example, our analysis can be extended to multi- feature data distributions, where certain features appear more frequently (Zou al., 2023) or have larger norms than others (Lu et al., 2024). Such extensions could provide deeper insights into the mechanisms of feature learning more realistic scenarios. We hypothesize that, despite the infrequent occurrence or smaller norm of certain features, diffusion models can effectively learn them due to the nature of the denoising objective. This insight has significant implications for downstream tasks, such didto--diitributionn classification, where rare weak features may be the primary distinguishing factors.
We believe our framework holds broader potential beyond the scope of this work and can be adapted to analyze conditional and latent diffusion models, elucidate the mechanisms of various training objectives and optimizers, and examine other generative paradigms, such flow matching.
10

---

Published as a conference paper at ICLR 2025
ACKNOW
We would like to thank the anonymous reviewers and area chairs for their helpful comments. Wei Huang is supported by JSPS KAKENHI Grant Number 24K20848. Yuan Cao is supported part by NSFC 12301657 and Hong Kong ECS award 27308624. Difan Zou is supported in part by NSFC 62306252, Hong Kong ECS award 27309624, Guangdong NSF 2024A1515012444, and the central fund from HKU IDS.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In International Conference on Learning Representations, 2023.
URL https: /tte:://opppee /tpp/pppreeeeeeee revirrevien vew. net./ foet/foo forum? f2 FXGA.
Xuchan Bao, James Lucas, Sushant Sachdeva, and Roger B Grosse. Regularized linear autoencoders recover the principal components, eventually, Advances in Neural Information Processing Systems, 33:6971-6981, 2020.
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.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolu- tional neural networks. Advances in Neural Information Processing Systems, 35:25237-25250, 2022.
Haoxuan Chen, Yinuo Ren, Lexing Ying, and Grant Rotskoff. Accelerating diffusion models with parallel sampling: Inference at sub-linear time complexity. arXiv:2405. 15986, 2024a.
Huanran Chen, Yinpeng Dong, Shitong Shao, Zhongkai Hao, Xiao Yang, Hang Su, and Jun Zhu. Your diffusion model is secretly certifiably robust classifier. arXiv:2402.02316, 2024b.
Huanran Chen, Yinpeng Dong, Zhengyi Wang, Xiao Yang, Chengqi Duan, Hang Su, and Jun Zhu. Robust classification via a single diffusion model. In International Conference on Machine Learning, 2024c.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference Machine Learning, pp. 4672-4712, PMLR, 2023a.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron Weiss, Mohammad Norouzi, and William Chan. Wave- grad: Estimating gradients for waveform generation. In International Conference on Learning Representations, 2021. URL htttps://oppnreviers htt.rviwwwweeeeeee.eeeee.e net forum? dd=NsMLjcF a080.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference Learning Representations, 2023b.
Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning general gaussian mixtures with efficient score matching. arXiv:2404.18800 2024d.
Kevin Clark and Priyank Jaini. Text-to-image diffusion models are zero shot classifiers. Advances in Neural Information Processing Systems, 36, 2024.
Hugo Cui and Lenka Zdeborová. High-dimensional asymptotics of denoising autoencoders. Advances in Neural Information Processing Systems, 36, 2024.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794 2021.
Michael Fuest, Pingchuan Ma, Ming Gui, Johannes Fischer, Vincent Tao Hu, and Bjorn Ommer. Diffusion models and representation learning: A survey. arX7v:2407.00783, 2024.
11

---

Published as a conference paper at ICLR 2025
Khashayar Gatmiry, Jonathan Kelner, and Holden Lee. Learning mixtures of gaussians using diffusion models. arXiv:2404. 18869, 2024.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 1):665-6673 2020.
Shivam Gupta, Linda Cai, and Sitan Chen. Faster diffusion- sampling with randomized midpoints: Sequential and parallel. arXiv:2406. 2024.
Yinbin Han, Meisam Razaviyayn, and Renyuan Xu. Neural network- score estimation in
diffusion models: Optimization and generalization. In International Conference on Learning Representations, 2024. URL hwwp::tttp/..mm htw.//peeeevveeeeeeeeeeeeeeeeeee eew net/ forum d
Yujin Han, Andi Han, Wei Huang, Chaochao Lu, and Difan Zou. Can diffusion models learn hidden inter-feature rules behind images? 0442025502.04255, 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference Machine Learning, pp. 8867-8887. PMLR, 2022.
Wei Huang, Yuan Cao, Haonan Wang, Xin Cao, and Taiji Suzuki. Graph neural networks provably benefit from structural information: A feature learning perspective. 2rxiv:230..1392 2023a.
Wei Huang, Ye Shi, Zhongyi Cai, and Taiji Suzuki. Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations, 2023b.
Wei Huang, Andi Han, Yongqiang Chen, Yuan Cao, Zhiqiang Xu, and Taiji Suzuki. On the comparison between multi-modal and single-modal contrastive learning. arXiv preprint arXiv:241..0283 2024a.
Xunpeng Huang, Difan Zou, Hanze Dong, Yi Zhang, Yi-An Ma, and Tong Zhang. Reverse transition kernel: A flexible framework to accelerate diffusion inference, arXiv preprint arXiv:2405. 16387, 2024b.
Priyank Jaini, Kevin Clark, and Robert Geirhos.
Intriguing properties of generative classifiers.. In International Conference on Learning Representations, 2024. URL net forum idd Ma6=mm000MKRR
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pp. 9965-10040. PMLR, 2022.
Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. arXiv preprint arXiv:2409.199 19345, 2024.
Kevin Kögler, Alexander Shevchenko, Hamed Hassani, and Marco Mondelli. Compression of structured data with autoencoders: Provable benefit of nonlinearities and depth. arXiv:2402.05001. 2024.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021. URL httph:tt//// htttp:://openzevier net wttt/ffrummiicc
x7z3Ymz5J.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer ReLU convolutional neural networks. In International Conference on Machine Learning, pp- 17615-17659. PMLR, 2023.
Daniel Kunin, Jonathan Bloom, Aleksandrina Goeva, and Cotton Seed. Loss landscapes of regularized linear autoencoders. International Conference on Machine Learning, pp. 3560-3569. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2025
Y Lecun, L Bottou, Y Bengio, and Haffner. Gradient-based learning applied to document recogni- tion. Proceedings of the IEEE, 86(1 1):2278-22333 2324, 1998.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score- bose-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 35:22870-22888 2022.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference Algorithmic Learning Theory, pp. 946-985. PMLR, 2023.
Alexander C Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your diffusion model is secretly a zero-shot classifier. In International Conference on Computer Vision, pp. 2206-2217, 2023a.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv:2306.0925/, 2023b.
Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 27942-27954, PMLR, 21-27 Jul 2024a.
Gen Li, Zhihan Huang, and Yuting Wei. Towards a mathematical theory for consistency training in diffusion models. arXiv:2402.00 07802, 2024b.
Marvin Li and Sitan Chen. Critical windows: non-asymptotic theory for feature emergence in diffusion models. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 27474-27498, PMLR, 2024.
Puheng Li, Zhong Li, Huishuai Zhang, and Jiang Bian. On the generalization properties of diffusion models. In Advances in Neural Information Processing Systems, volume 36, pp. 2097-2121 2023c.
Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In International Conference on Learning Representations, 2024. URL https:tt////m /tte:///ppenn /pppneveeeeieeeeeeeeeeeeeeeeeeee net forum? 1d=wYmvN3sQPG.
Xuran Meng, Yuan Cao, and Difan Zou. Per-example gradient regularization improves learning signals from noisy data. arXiv:2303. 17940, 2023.
Xuran Meng, Difan Zou, and Yuan Cao. Benign Overfitting in Two-Layer ReLU Convolutional Neural Networks for XOR Data. In International Conference on Machine Learning. PMLR, 2024.
Soumik Mukhopadhyay, Matthew Gwilliam, Vatsal Agarwal, Namitha Padmanabhan, Archana Swaminathan, Srinidhi Hegde, Tianyi Zhou, and Abhinav Shrivastava. Diffusion models beat gans on image classification. arXiv:2307.0877 2023.
Phan-Minh Nguyen. Analysis of feature learning in weight-tied autoencoders via the mean field lens. arXiv:2102.08873, 2021.
Thanh V Nguyen, Raymond KW Wong, and Chinmay Hegde. Benefits of jointly training autoen- coders: An improved neural tangent kernel analysis. IEEE Transactions on Information Theory, 67(7):4669 4692,
Reza Oftadeh, Jiayi Shen, Zhangyang Wang, and Dylan Shell. Eliminating the invariance on the loss landscape of linear autoencoders. In International Conference Machine Learning, pp. 7405-7413. PMLR, 2020.
Maya Okawa, Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Compositional abilities emerge multiplicatively: Exploring diffusion models on a synthetic task. Advances Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning. pp. 26517-26582. PMLR, 2023.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In International Conference on Computer Vision, pp. 4195-4205, 2023.
Arnu Pretorius, Steve Kroon, and Herman Kamper. Learning dynamics of linear denoising autoen- coders. In International Conference on Machine Learning, pp. 4141 PMIR, 2018.
Maria Refinetti and Sebastian Goldt. The dynamics of representation learning in shallow, non-linear autoencoders. In International Conference on Machine Learning, pp. 18499-18519, PMLR, 2022.
Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv:2402.16991, 2024.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-CAM: visual explanations from deep networks via gradient-based localiza- tion. International Journal of Computer Vision, 128:336-359, 2020.
Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the DDPM objective. Advances in Neural Information Processing Systems, 36: 196499999699,, 2023.
Alcksandr Shevchenko, Kevin Kögler, Hamed Hassani, and Marco Mondelli. Fundamental limits of two-layer autoencoders, and achieving them with gradient methods. In International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 31151-31209. PMLR, 2023.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021. URL https: /wwp///veeeee.e.eeee/eeeee net forum? id=PxTIG12RRHS.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference Machine Learning, pp. 32211-32252. PMLR, 2023.
Harald Steck. Autoencoders that don't I overfit towards the identity. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances Neural Information Processing Systems, volume 33, pp. 19598-19608 Curran Associates, Inc., 2020.
Yuqing Wang. Ye He, and Molei Tao. Evaluating the design space of diffusion-based generative models. arXiv:2406. /2839, 2024.
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein 2023.
structure and function with rfdiffusion. Nature, (1077976)1089-1111
Weilai Xiang, Hongyu Yang, Di Huang, and Yunhong Wang. Denoising diffusion autoencoders are unified self- supervuue learners. In International Conference on Computer Vision, pp. 15802- 15812. 2023.
Xingyi Yang and Xinchao Wang. Diffusion model representation learner. In International Conference on Computer Vision, pp. 18938-18949, 2023.
Yongyi Yang, Core Francisco Park, Ekdeep Singh Lubana, Maya Okawa, Wei Hu, and Hidenori Tanaka. Dynamics of concept learning and compositional generalization. In International Confer- ence on Learning Representations, 2025.
Kaihong Zhang, Heqi Yin, Feng Liang, and Jingbo Liu. Minimax optimality of score-based diffusion models: Beyond the density lower bound assumptions. In International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 601 60-777 PMLR, 2024.
Wenliang Zhao, Yongming Rao, Zuyan Liu, Benlin Liu, Jie Zhou, and Jiwen Lu. Unleashing text-to- image diffusion models for visual perception. In International Conference on Computer Vision, pp. 5729-5739, 2023.
14

---

Published as a conference paper at ICLR 2025
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. The benefits of mixup for feature learning. In International Conference on Machine Learning, pp. 43423 43479. PMLR, 2023.
15

---

Published as a conference paper at ICLR 2025
APPENDIX CONTENTS
Additional experimental results
17
A.I Supplementary results for synthetic experiment
17
A.2 Feature learning comparison under varying SNRs
17
A.3 High SNR setting Noisy-MNIST
17
A.4 Experiments with additional diffusion time step
18
A.5 On the feature learning with 10-class MNIST
20
Preliminary lemmas
21
Classification
C.I Useful lemmas
22
C.2 Scale of inner products
23
C.3 Signal learning
28
C.3.1 First stage
29
C.3.2 Second stage
31
C.4 Noise memorization
35
C.4.1 First stage
36
C.4.2 Second stage
38
D Diffusion model
43
D.1 Useful lemmas
43
D.2 Derivation of loss function and gradient
44
D.3 First stage
46
D.4 Second stage
54
D.5 Stationary point
66
16

---

Published as a conference paper at ICLR 2025
A ADDITIONAL EXPERIMENTAL RESULTS
This section includes additional experiment results.
A. SUPPLEMENTARY RESULTS FOR SYNTHETIC EXPERIMENT
We first include the convergence in loss plots well accuracy for classification under the two SNR conditions considered in the main experiment. The iindistribution) test accuracy is computed with 3000 test samples. We see both classification and diffusion model are able to converge in loss, although diffusion model only finds a stationary point. In the low SNR setting, classification is able to perfectly fit the training samples with 100% classification accuracy. However because it primarily focuses on learning noise, the generalization is poor with test accuracy of around 50%. For the high-SNR case, both training and test sets can be perfectly classified due to the signal learning.
Cassificatio
Duffus on model
Cless fication
Diffuston model
itoratio
360
2000
2007
X/8
20
J8 heration
40
0858
2000
2000
Cassificatio
Classification
Tox: acc
Reration
Iteration
Low SNR
High SNR
Figure 5: Experiments on the synthetic dataset with both low SNR (n SNR² = 0.75) and high SNR SNR- = 6.75).
A.2 FEATURE LEARNING COMPARISON UNDER VARYING SNRS
In this section, compare the feature learning dynamics of classification and diffusion models on additional settings of SNR. Apart from the SNR² = 0.75 and n SNR = -6.75 as shown in the main text, we additionally test on (1) SNR2 = 1.92. (2) n SNR² = 3 (3) n SNR? = 4.32. The feature learning dynamies under the corresponding SNR settings are shown Figure 6.
From the figures, we can see that classification indeed is more sensitive to the SNR scale, where it easily overfit to either signal or noise (except for the case where n SNR = 3 where classification learns signal and noise approximately the same scale). On the other hand, we can verify that at stationarity, diffusion model learns in a more balanced scale for signal and noise.
A.3 HIGH SNR SETTING ON NOIST-MNIST
Here we include experiment results when SNR 0.5, which corresponds to the high SNR setting. The experiment settings are exactly the same as in the main experiment. Figure 8 shows both classification and diffusion model converge in terms of objective. In addition, we see the high SNR encourages classification to learn primarily the signal while ignoring the noise. In contrast, diffusion model still learns both signal and noise to relatively the same order. Figure 7 suggests that classification learns more signal compared to noise while diffusion model still learns more balanced signal and noise. We also plot classification accuracy for both the low and high SNR cases. In the loW case, because classification predominately learns noise, the generalization is poor with test accuracy around 50%. Conversely in the high-SNR case, where the model is able to learn signals, the classification demonstrates effective generalization with nearly 100% test accuracy.
17

---

Published as a conference paper at ICLR 2025
Classificat
Diffusion model
Clessiiication
Diffusion
eaminn
MA00 FM
mystin, 441 mpslins.p mpxiime.p..ll
shil
mmpx
Eerrtion
Iteratiio
Iteration
Iteration
SNR? 2
1.92
SNN
Faature Boming
masiie,.d mpxlim..oll
manlim..-al)
Iteration
Iteration
n. SNR? 4.32
Figure 6: Experiments the synthetic dataset with varying SNRs.
Input
(a)
3:00
140
Classification
200-
100
(b)
2013
2011
ter3con
Diffusion
(Clant)
Feature 33m1
(c)
16 Figure 8: Experiments on Noisy-MNIST with SNR 0.5. (a) Train loss for classifi- cation. (b) Train loss for diffu- sion model. (c) Feature learn- ing dynamics.
Figure 7: Experiments MNNST with SNR 0.5. (First
row): Test Noisy-MNIS" images; (Second row): Illustration of in- put gradient, i.e., Vx (W, when y and V F_, (W x) when y = 0. (Third row): denoised image from diffusion model. In this low-SNR case, we see classification tends predominately learn noise while diffusion learns both signal and noise.
A.4 EXPERIMENTS WITH ADDITIONAL DIFFUSION TIME STEP
Here we also test on additional diffusion time step for learning on noisy-MNIST dataset. In particular,
we consider t = 0.8, which gives 0; = exp( (pp(t) = 0.45 and 3, = /I - exp(- 22t) = 0.89. We include the illustrations of denoised images well loss convergence and feature learning dynamics in Figure 10, 11, 12, 13. We see despite with a larger scale of added diffusion noise, diffusion model still learn both signals and noise unlike for the of classification.
18

---

Published as a conference paper at ICLR 2025
Train acc Teak acc
ation
Itoration
(a) ACC (SNR 0.1)
(b) ACC (SNR 0.5)
Figure 9: Classification accuracy on (a) low-SNR and (b) high-SNR noisy MNIST datasets. This demonstrates that when classification focuses learning noise (as in the low-SNR case), the test accuracy hovers around 50%, thus suggesting failure to generalize. In contrast, when classification focuses on learning signals (as in the high-SNR case), classification generalizes effectively, achieving near-perfect accuracy.
1000
Input
6000
1el ainaggonn (1,0)
Diffusion
Pestur Heaaang
(b)
Makee (14f))
Figure 10: Additional experiments on Noisy-MNIST with SNR
Figure 11: Additional experi- ments on Noisy-MNIST with SNR 0.1 and 0.8. (a) Train loss for diffusion model. (c) Feature learning dynamics.
0.1 and diffusion t = 0.8. (First row): Test Noisy-MNIST im- ages; (Second row): denoised image from diffusion model. We see diffusion still learns both signals and noise even with large diffusion time step.
0.90
Input
(a)
6.97
Iteaatoon (C.I))
Teature (Clams
Diffusion
lamming
Naie (Closs
Teeture (3if))
(b)
Nuse: (24T))
berstio (Claes)
Figure 12: Additional experiments on Noisy-MNIST with SNR 0.5 and diffusion 0.8. (First row): Test Noisy-MNIST im-
Figure 13: Additional experi-
ments on Noisy-MNIST with
ages; (Second row): denoised image from diffusion model. We see diffusion still learns both signals and noise even with large diffusion time step.
SNR = 0.5 and = t- = 0.8. (a) Train loss for diffusion model. (c) Feature learning dynamics.
19

---

Published as a conference paper at ICLR 2025
Impur
Eteridion
1.00
Classification
(b)
3407
2011
Terratoo ID
Feature (Cles)
Molie (Chani
Diffusion
(D4M)
(c)
Boration (Class)
Figure 14: Experiments on 10-class Noisy-MNIST with SNR
Figure 15: Experiments on 10-class Noisy-MNIS, with SNR = 0.1. (a) Train loss for classification. (b) Train loss for diffusion model. (c) Fea-
0.1. (First row): Test Noisy-MNIST images; (Second row): Illus- tration of gradient of output (for the true class) with respect to the input. (Third row): denoised image from diffusion model. In this low-SNR case, we see classification tends to predominately learn noise while diffusion learns both signals and noise.
ture learning dynamics.
A.5 ON THE FEATURE LEARNING WITH 10-CLASS MNIST
In the main paper, we only conduct experiments on Noisy-MNIST restricted to two classes. In this section, we experiment
the 10-class MNIST dataset, which contains more features and is more challenging for both diffusion model and classification.
We adopt the same data processing pipelines as in Section 5.2 except that for each class, we select 10 images. We set the scaled SNR SNR = 0.1, consistent with the main paper. While the diffusion model remains unchanged, the classification model requires modification. Specifically, the second layer's weight matrix has dimensions m x 10, with entries fixed uniformly to values {-1,1} 1). Furthermore, we employ cross-entropy loss for training the classification model.
We plot the visualization of feature learning in Figure 14. We observe that, even with additional features and labels, the similar learning patterns observed, i.e., diffusion model learns both signals and noise in order to reconstruct the input distribution while classification model learns primarily noise for loss minimization. From Figure 15(c), we notice that diffusion model learns features to relatively the same scale while for classification, the growth of feature learning is dominated by noise learning.
20

---

Published as a conference paper at ICLR 2025
B PRELIMINARY LEMMAS
Recall we define S = {i E [m] = 1} and S-1 = {i E [n] yj =-1.
Lemma B.1. Given arbitrary
O. with probability at least - 6. we have
(1 On 1/2) |S-1| 2 ((n-112))
Proof of Lemma B.I. The proof is the same in (Cao et al., 2022: Kou al., 2023) and include
here for completeness. Because |S =i= (1% = 1) and |_- Σi= ( -1) and P(y: = 1) Ply: -1) 1/2 for all E [n). then S S n/2. By Hoeffding's inequality, for arbitrary a > 0, S+1 n/2 2 a) S 2 exp( 2
Setting a vn logllo/4))/ (2 and taking union bound, we have with probability least 1 5,
log((4/6)
S11
Hence the proof is complete.
Lemma B.2. Given arbitrary 0, with probability least 5,
odd(1- - (00l1-ō(d-1/))) < |E:1 5 + ((2-1/2)
(E, 20 2log
for all i, E [n]]
Proof of Lemma B.2. The proof is the same as in (Cao et al., 2022: Kou et al., 2023) and we include here for completeness, By Bernstein's 's inequality, with probability at least 1 6/(2m), we have
|||&/|? of Olo dlog(4n/8)).
which shows the first result.
For the second claim, we can show by Bernstein's 's inequality, with probability at least - 8/(2n²) that for any +
KE. (x>| - 20 <llog(An3/4
Then apply union bound show the results hold for all i, e [n]]
C CLASSIFICATION
We track the inner product dynamics during the training of supervised classification to elucidate the signal learning and noise learning. We first write the gradient descent dynamics as follows.
w
w., Ls (W))
w
aile (w J.r
dl (w J.r Es)iy:E:
j.r
W jr
Ww M1)jm +
n
Σ
e (w) (11))
2)
de w j.r Ediin.
Tum SES-:
num
SESI
Here we restate the Condition 3.1 specific for the case of supervised classification. Condition C.1. Suppose that
1. Dimension d satisfies ②mmax ( mo - n
21

---

Published as a conference paper at ICLR 2025
2. Training sample and network width satisfy 22(log(/66)), ((llog(m/6))
3. The inilialization variation do satisfies O S 0o - {(min 1,
4. The learning rate satisfies 1 s (mmin(m] اهم nmmooo 1/2 - nnmo 2d
We make the particular remarks as follows. The lower bound on m = 2i11 is to ensure the initialization is concentrated and thus provides lower bound the maximum and average inner product. The lower bound on n (2(1) is required such that |Si]; | (n)) and O(2-1/2) remains small. The lower bound on do is required for the noise memorization setting where we need to control the lower bound for the noise inner product at initialization. Thus to ensure the lower bound on is valid, we require further conditions on the dimension apart from = l(n²).
C.I USEFUL LEMMAS
We first provide lemma that bound the inner product at initialization. Lemma C.1 (Cao et al. (2022)). Suppose and that d (llo(mmn/60), m (gloo(1/6)) then with probability at least 1-
Kw J. My)) /2log(8m/8)oolle
1(w j.r E.) s 2ogg(mmn/00jooo
for all j,j' [+1),r r [m]. E [n]. In addition,
max (w J.r my)' 2 00l|m11/2
rE
max re
(w f. 4:) 2 aooe dd/4
for all j,j {+1), i [n].
We decompose the weights into signal components and noise components. Lemma C.2. The weight can be decomposed as
w w + cimm 5
where the noise coefficients Pi.rissatisfy satisfy Pi.ri and
k+1 = P5 nmn
7
( 11)
for all j = #1, E [m]] and i E [n].
Proof of Lemma C.2. The proof follows from (Cao et al., 2022; Kou et al., 2023). First, we recall the gradient descent update as
k+1
#k (w
Σ th (w J.r
p 1-1jju111
tk
(wi.. Ei)iu
Tum
nmm VES-1
nm
iES1
w
(w: Jim Mijim
plke (w j.r 14 111)
==0 S
==0 IES
7
otle ( j.r E.)iv...
By the data model, we have with probability 1, the vectors are linearly independent and thus the decomposition is unique with
&
P
th (w"),
11
YUTT ==0
Then writing out the iterative update for pi... completes the proof.
22

---

Published as a conference paper at ICLR 2025
V1
Lemma C.3. Let x ~ (O Then P(Ja) s c) S erf
(2o
exp(-
2c2
no?
Proof of Lemma C.3. The probability density function for x given by
1
x²
f(x)
exp(
V2Ta
20²
Then we know that By the definition of erf function
x²
P([2] S sc)
exp(
2 )d..
V2TO
VT { expl- dd)d
2
erf(c) =
and variable substitution yields
/20 V2Ta L/ exp( dx..
c
x²
erf(
20
Therefore, we first conclude PP(2)] S c)
2erl(
(-)), Next, by the inequality er) <
VI exp( exp(-4r2/T), we obtain the desired result.
C.2 SCALE OF INNER PRODUCTS
We first derive a global bound for the growth of inner products until convergence, To this end, we let m, d) be the maximum number of iterations considered and " -0 -n,m,
T 7/ poly |=|| 11 ,
let a = 2 log(T"). We also denote 3 == max (Kw J.r: Py)b j.r (.) Then from Lemma
C.1 and from Condition C.1, we can bound
max (oo || aog vd/4) S 3 1/C
for some sufficiently large constant C > Proposition C.1. Under Condition C.I, all 0 < TT', we can bound
(w J.r. (4))b16 v..r, E:)l, 1p S ar
(5)
l(w J.r / --=1)1 S B
(6)
log (4n2/8) 1(w us.r E m.r S S BB+11 12 1
d n1c
for all i [m), E [mi] and j #1.
We will prove the bound by induction and we first derive several intermediate lemmas as follows.
Lemma C.4. Suppose results in Proposition C.I hold at iteration k then we have F, (W) J xx) S 0.5
for all i E [m). j # Vc-
Proof of Lemma C.4. Recall that
1 m
F5 (, xx) = m
Σ
(w (w x: (2)g2 + (w).r> x; (2))22)
r=1
1
(w ( Myy.) 2 (w Ex)?)
m
|lo((422/5) nax
< B +
B 12
d
< 0.5
where the second last inequality is by (6) and (7). The last inequality by Condition C. such that
5 1/C & 0.25 and d > 14 log( 4n. 15).
23

---

Published as a conference paper at ICLR 2025
Lemma C.5. Suppose results in Proposition C.I hold at iteration k, then have
log (n)nAnn)))))))))
kw wirrs Ei) Pf s VV
d na,,
for all j = =1,r E [m]] i E [n].
Proof. By Lemma C.2, recall the decomposition
w " w + ct + c 1 + Σ
slsilll ?
1=1
By the orthogonality, show
(w).r= Ei) = (w3rr Ei) + PJrr. + Enjirol (1:11111111
By Lemma B.2 and suppose 1[(log(n/6)), then (.).,l)) Avlog(4n"/8)d= Thus we have
log Log(4n2/0) na,
kw wG.r Ei) -rrmall s 4/
d
where we the upper bound on 17) S a.
Lemma C.6. For any E [m],j.y = +1, we have sign((w)) jor Hgl) = Sign((wf.rr jarz pn)) for all k O T
Proof of Lemma C.6. We prove the results by induction. First, it is clear at k 0, the results satisfied. Then suppose there exists iteration i such that sign((w), (( (by) sign((w) myy) holds for all k < k 1, we show the sign invariance also holds at k. Recall the gradient descent update as
wh+1
Cthh (wh (w Pi)imi "
Σ nm VES
(W M-1))M
nm
VE
ek (w E.iiE..
71772
Then the update of the inner product is
(w
J. My) = (wi-11 Jrr
Σ oh (w Mu)illlall
(4y)
nm iES,
= -(1
" jy eh [|=1))))))) (w]r² My)
nm
SES
By the condition that 17 1 C-1ml|pa|| for sufficiently large constant C, have I jy Eies, ghh-1 e/r ||| < Thus we can guarantee the sign( (w wir Hy)) sign((w)) an (y))
sign( (wo S. M.)).
Proof of Proposition C.I. We prove the results by induction. For Pf.,, we prove a stronger result
that 1pth & 0.9a S a and l ra S 0.6,3 + 8y log(4n2/8) na. First it is clear at = 0, the results are satisfied based on the definition of / and a 2 B. Now suppose that there exists 7 < T
such that results hold for all 0 < k < T 1. We wish to show the results also hold for k T.
First recall the gradient descent update as
t (wh hwwt
n nm vES-1
Σ cth (w M-1jjm-
NES
24

---

Published as a conference paper at ICLR 2025
c (wS.r E
7177
Then based on the orthogonal data modelling assumption, we have for # j. i.e., y -j,
(w k+1 J.r
(e-5) = (w j.rt (0-5) +
"
t (w" 1 1 |leel?
Tum
iES-;
-(1
(w) j.rt [b
nm
iES
where the second equality by t bign((wjj.. (-3))) and thus
0 for i, k. From Lemma C.6, we have sigu((w)) (wh-- :b-5))
kw j.rt ا-2سس (1 (1- 11
Σ
11)
1
[W T-1 -e (wrrr' T-1 |m sB
nm iES_;
On the other hand, for = j. we have
(w/r M4)
eth (w |
Tu iESy
-(w jum Meg +
Σ
10th (w
nm
iES,
Next, we notice that
10t4
+ exp (F,, (W&
N. x,) F (W -w: x:))
S exp ( Fy, (W U: x;) F-ye (W - x:))
I exp ( Fy. (Wh M x;) + 0.5)
exp (
(wh Bi My.) + (w oarr {i) + 0.5)
(8)
where the last inequality is by Lemma C.4. Let kj.r- be the last time k < T that (w o الموس I 0.5a. Then we have
Mb) + 1 1lllal1 1ett
Jir (4y) + alleall?
(w. j.r Mb) = (w Ryir J.r
Reir (w
1 (w..r:
nm
nm
ky
A2
Without loss of generality, we suppose (wo p,r+ 14g)) 2 O. then by Lemma C.6, (w f My) 2 0 for all Then we can bound
n/l 0.5a S 0.25a
|
nm
where the last inequality is by the condition that 1) S nm |||l= 2/2. Furthermore,
A2 <
(W 0: x,) + 0.5)(w pj) exp(-Fy, (wk
nm
kir T
< 2m 21||pa| Pa T' exp( exp(-a2//4
2n||u|| "a "a exp(
exp(eeepppppp (oop-lllo(TT)
nm
< 0.25a
25

---

Published as a conference paper at ICLR 2025
where the first inequality is by (8) and the second inequality is by upper bound on (w),,14g) S a for all & < T 1. The equality is by the definition of a 2 log( and the last inequality is by the condition r7 S nm/ pa|| 2/8 Thus, we can show
(w jam My) S 0.5a + 0.25a + 0.25a
Next for the noise growth, from Lemma C.2, we have for yi #
p U.r. =
T-1 P-yurr nmm
T-1 (w
Berr 1::12
(9)
When p p=yilna $ 1.3(0.33 + 4 Jog(dn2/0) na), we have
21
log(4n²/ 4n² (2422/)) T-1 v.ral + 0.153 < 0.63 8V
I -nrr S
T-1 Mord
S 10
nim
d
where the second inequality is by triangle incquality and | < and Lemma B.2. The third inequality is by the lower bound on 33 in (4) and the condition that v) S 0.05nmorg -d---------- 1/2a.
Further, because ((w2 w..r E.) & 0.33, when 1.3(3.33 4V log(4n2/4) na) S 1 & 0.63
log (4g( (6) na, we can show from Lemma C.5 that if -B,,r, > O, then
8
3 -
7-1 T-1 -w-iirr Ei) -P--v.n
(w
4 T-1
Then (9) suggests
log (eg(Am2// (8)
p Bara S $11-
F-1
Wiri S P S 0.63 + +8V
no.
3nm
d
--1 If Pamaa
< O, then
4 ύ-1
(w T-1 -wiir
Ex) p -
1 ύ-1
3 birc
Then (9) suggests
El12
log log(4n2/0
Mari 2 (1-
( 1=p
Wirr, 2 p p- 2 -0.63 8/
na
3nm
d
Ing(4n2 Thus this completes the proof that 10th mr.r. S 0.6,3 + 8//
d
na.
Finally, by Lemma C.5 we have for all h: 2 0
log log(4n2/15 (4n2/5)
(og(422/1 28)
lww wurt Ei)l S (ww" wart : &i) + 10th rara + 4
na S 0.93 + 12
na
d
which proves the upper bound for ? -v..r Eill and m
Next, from Lemma C.2, we have for y: j.
k+1 P
6t (w garrs E.
(10)
Let Ky,s the last time < T that 1pt S 0.6a. Then it can be verified that for l: 2 ky
(w r. Ei>l 2 lpgarral - 1( wi.r: Ei) 4/ log 1og(442/// na 2 0.5a
d
where the first inequality is by Lemma C.5 and the last inequality is by Kw tr Ei)
log an2 (8) na < 1 < 0.1a.
26

---

Published as a conference paper at ICLR 2025
We now expand (10) as kire Puitri
Pu...
(w
Σ
14t Kw oe E..1
num
Tum
Ren RAsSSTTT -1
Ag
A4
Then we can bound
2
2n
log /8)
|Aa S
1(
Gi>
(ww (.).,,;i)-- 0.6a + 4V
mm
virt
mm
d
2nd
0.7a
nmm
- 0.15a
where the first inequality is by |,741 < and Lemma B.2 with d = ((log(n/8)) and the second inequality is by Lemma C.5. The last inequality is by the condition S C=-mo -am =dd1 for sufficiently large constant C.
In addition, we bound
2n
|41/ S nm
Σ
exp(-F,, (W x.) + 0.5)
4noRda
T exp(-a2/4)
nm
4no2da
Tum
< 0.15xx
where the first inequality is by (8) and the second inequality is by kww. 4 2 0.6a - 0.1a 0.5a.
The last inequality is by the condition S CC-nnoo dd for sufficiently large constant C.
Combining the bound on |Aal and |Aa], we have
I S 0.6a + 0.15a + 0.15a 0.9a.
Lastly, bound
log( (nnnn
48)
log log(An2/8)
kw wur: Gx)l S 1(w D. G.)l + Ip +41
na < 0.33 - 4V
nax + 0.9g
d
d
< a
This shows the upper bound as kwT w..r I S a.
We require the following lemma that lower bound the loss derivatives in the first stage before the inner products reach constant order.
Lemma C.7. If max,,,.. ((w j.r: My), (w). {i) = O(1), there exists a constant Ct > 0 such that
10 2 for all i [n].
Proof of Lemma C.7. If mmxx, i. ((w fr Mg), (w {.) = O(1), bound for all +1
F. x;) m
(w f ) J.., E.) < S0(1)
Therefore, we can bound [ = (1 + exp( Ey, (W& v. x:) F-y, (W -1 2(1).
We also prove the following upper bound on the gradient norm.
27

---

Published as a conference paper at ICLR 2025
Lemma C.8 (Proof of Lemma C.8). Under Condition C.I, for 0 S l $ T", we can bound
L [I2 O(max [Ilul] of L.s (W*)
Proof of Lemma C.8. The proof adopts similar argument as in (Cao et al., 2022, Lemma C.7) and we include here for completeness. We first bound
2
VVW
x,
Σ m
J.m My + (wf...E:)
3.r
2
2
yerro Myd)llm
kw bi 1:)
m
m7
2
( BUT ||ll
+
2
Σ 1w - E.111
m
m
2
Σ
(K yerr Mad + l( Dirrs Ei)l) max{|l 20 vd)
771
2
m
(1(w* p + lw werr Eiil) max pell, 20 d}
m
5 V Ps (Wh: 8. xi) + (W -y xi) max
s 2
V Fr. (Wk B x.) +
max |||| 20 Vd
where the third inequality is by Lemma B.2 and the fourth inequality is by Jensen's inequality and the last inequality is by Lemma C.4 that F-y. 24 (W& -m. x;) for all E [m]. Then we have
-(( \ x;)11
S (F (W 9: x;) - 0.5) (2 Fee (W.. Ne x;) + 1) max [alee|lllll1 20g v vdt) 2
= -4 (- W: x:) - x0.))
(V Fys (W&. u: xx) + 1)²
I maxx-44f (z - 0.5)( + 1)2) max [|a]] 40
250
(max || ded))
where the last equality is by max: (= (.5) < do because El has an exponentially decaying tail. Then we can bound
LLs
(w.S(w* x))]V/((X x1))
11
Omaxx ||| (W
< 0 (max{]1441 2 ,o( =
[(u x:))
(max 2
where the third inequality is by Cauchy-Schwartz inequality and the last inequality is by -l' S 8 for cross-entropy loss.
C.3
SIGNAL LEARNING
We first analyze the setting, where n SNR2 2 C² for some constant C' > 0, which allows signal learning to dominate noise memorization, thus reaching benign overfitting.
For the purpose of signal learning, we derive an anti-concentration result that provides a lower bound for signal inner product at initialization.
28

---

Published as a conference paper at ICLR 2025
Lemma C.9. Suppose ò and ((log((/6)). Then with probability at least 1 8, we have for all j.y = =1
0o1lue1/2 <
Σ (w (wrr Huyl S solllall
Proof of Lemma C.16. First notice that for any j =1, (w J.r. My) ~ NNO, (bll and thus we have Bll(w).r j.r myY0 = plll By sub-Gaussian tail bound, with probability at least 8/8, for any j.v =1
m2
20 1
log 18/(8/8)
((w
j.r My) ool1 2/ S
m
m
Choosing m (l())(1/)), have
a
Kw j.r Myyi S /
m1 r=1
Then we have //llal1/2 < 1 y +1 completes the proof.
m E 1( {:)l s colloa||- Finally taking the union bound for all
We have established several preliminary lemmas that hold with high probability, including Lemma B.1. Lemma B.2, Lemma C.I, Lemma C.9. We let Eprelim be the event such that all the results in these lemmas hold for given 5. Then by applying union bound, we have P(Eprelim) 2 1 46. The subsequent analysis are conditioned on the event Eprelim.
C.3.1 FIRST STAGE
In the first stage where max,,i.y (ww), My), (w j.r {i)} ((),, show in Lemma C.7 that can lower bound the loss derivatives by constant Ce, i.e., 1 2 Cr for E [m], S Ti.
Theorem C.1. Under Condition C.I, suppose n SNR2 2 C" for some C" 2 Then there T: exists a time Ti = O(1 (1mm-m]||aa) [ull- )2), such that (1) max, ((w j.r: Mty)1 2 2, for all j = +1, (2)
jir p2) 1(w? 2 2, for all j == =1 (3) max,, (( girr E:)l (nnnn/2))
Proof of Theorem C.I. We first upper bound the growth of noise by analyzing inner product dynamics
(w :.)
Ei) (w : Ei)
17
(w W.
(E) Ex)
nm
-w 1 Ei)
"
ef (wh-1
"
Σ
ek- [w
Ep)
k-1 Virt
West
nm
Tm 441
This suggests
l(we Bi.r
Ei)l
S (
k- Bert
<X) +
27
|
k-1
+
[elllw k-1 B1.
Bert
(11)
Next, from Lemma C.7 and Lemma B.2, we have for any 1' + [n] and k < Ti,
(E,,
Ev)
20 /dloo(4n3//)) 2.10.1 log 4n2
1
C,0.99o2d
where we use the lower and upper bound on loss derivatives during the first stage, as well as Lemma B.2. Then taking the maximum of (11) over the neurons and samples, we let B
max,, kw
Barro {.)l and obtain
7 1 2.10 1 log(4n's (og(4n2/0) 2B
Bk < Bk-1
nm
29

---

Published as a conference paper at ICLR 2025
(1
1.01m/15 Eil
Bk-1
nm d.02no² nm
B
where the second inequality is by d = (nn2) sufficiently large and |e/1 S 1. The third inequality is by Lemma B.2.
We then consider the propagation of (w), My). From the gradient update we can show for j =
(w]tt 1+1 say} = (w j.r He))
ge (w".. Mg)llaall
ES
2 (w Jon 14j) +
(w 2.r 14)
nmm
(1 + 0.49 nCeeleaa)
(w M4)
where the first inequality by loss derivative lower bound and the the second inequality is by Lemma B.1 and n = 52(1) sufficiently large. This implies that
Kw ال 2 (1 0.49 Ce Call )(w) j.r -1y)1 2 (1 0.49 0.4999/C CClllell? (
)* (w Hy
m
m
Applying Lemma C.1 and Lemma C.9, we have for all j = +1,
(w J.r 14)س 2 (1 0.49 nC CCellllll 1ll ,ol|wa1//
max
m
1 m Iw H4) 1 0.49
TCCellul1)) oo
12
m
m
Consider
Ti = log(4mo, 11111111111 log (1 + 0.49 0.49 " O1 mm]|pa|]
log 1og(4mo7-1l1n -111=(111)
m
for 7) sufficiently small. Then we can verify that for j +1, we have
m
max [w J. mb)
2,
22 2 2,
m
Now under the SNR condition, we can bound the growth of noise as
B < 1 + 1.01 n
2000 Vdv log(Smn/0)
log(1 + 1.02 po?a
log(40 - 2000 V log(Smn/0)
=exp
log(1 + 0.49 DC |all
(2.1//Cm -SNR ((SSNR?)) log(40, 112001111)1))) 2000 Vev/dd/oo///nn log(Smn/6))
exp
(22.//Cm-ISNR 2 0.01) log(4o "1 11)) 2000 Vdv log(8mm/0))
exp
< 8SNR -1 1 Vlog(8mm//))
:Oōnnnn1 1/2)
where the first inequality is by Lemma C. and the second inequality is by Taylor expansion around 77 0. The third inequality by choosing 7) sufficiently small and the fourth inequality is by the SNR condition that n SNR2 C' > 2.5
30

---

Published as a conference paper at ICLR 2025
C.3.2 SECOND STAGE
First, at the end of first stage, we have
max, (w J.m Mg)) 2 2 for all j = =1.
Er= (w J.n M3) 2 for all = +1.
max,, (ww Wi.r! &i = 1112)
max,, (ww -wi, Eiil S 3 + v
log(4n² /8)
n1..
Next we define
Wrr w + 22lo((/e)siin((ww) (rr (2))) pe + p_j
||
We first show the monotonicity of signal inner product in the second stage.
Lemma C.10. Under the same conditions in Theorem C.I, we have for all j t1, E [m]].
Ti s k < T. kw j.r MyYl kw j.r ال)م 2 2.
Proof of Lemma C.10. From the update of signal inner product, we have for all j = +1, E [m].
Ti < k <T
(wett M4) = (w jur 14))
ptle (w j.r: M
SES,
=(1 ?llla1
(w j.r: M4)
ES
Thus (w pb)س 2 ( (ww-1 p1)) 2 1w m for =1, [m]. S
We then bound the distance between W to W Lemma C.11. Under Condition C.I, can bound WT) W O0/]
Proof of Lemma C.11. Let Pg be the projection matrix to the direction &, i.c. Pe we can represent
GET
Then
[E
w j.r
wh. = -uus (w w9.r) + Puei(W). - w + D Pe. (wh j.r w
12 Pps (w ).
I I-PM Pp
By the scale difference Ti and the fact that gradient descent only updates in the direction of paj. j = =1 and &., we can bound
|WT: - 2
(w w
11
(w
w
3.r
11
11
(w T:
w. Sr E.)
|1012
|//112
| 12
2m
I I-P Pu,
P Pun.
Σ Pg. (w -
w
1=1
maxx
(wT1. Ti fir
12,)2
2(WTh. jim p 1m1))² 2 + 2(w
j.r: 11 1
2(w (.. 11,)22
< 2m
11
||
31

---

Published as a conference paper at ICLR 2025
2(w Ti Jart
61) - 2(w Jor Ei)
)
E j=#1,r/m
(I-PM. k. P PM_ Eps. (w w..r)
1|
50mml 1
where we have use the scale difference at Ti. Therefore,
|W
will S || w + -
- w
I 50(ymll ||م| - 0(/mlog((//))]lll - [0
where we use the definition of W
Lemma C.12. Under Condition C.I, we have for all Ti S k S T
I
w || w
2 (W*) nE
Proof of Lemma C.12. The proof is similar in Cao et al. (2022). We first show lower bound on ye VV f(W,, x.), W*) for any E [n] for all Ti S $ T
W 1
juk(w j.r: My (Marye wirr) + Σ
jus (w j.r (E) wir))
3rr
(w Verrs
Myy (w Bur Myd ww
uerr My (w" -yer Myi
m
+ m
Σ d.r
Su (w jir (E) w
Mi.rt My 12log(4/e)) (w
(w
Derre My (w
Vesro Myy
ww
Mus
(w"
vur und
1
Σ
jui(w Evi w9.T)
m
m 3.r
A7
where the second equality is by definition of W The third equality is by Lemma C.6. We next bound
|eel S Gollu// 2log(Smm))))))) ((0011111))
|Arl < 1 Σ w o Mo (1 mu Mn + 2 log(2//e) ((ooll)4)
|Asl < Ö(ovagVd)
where we use the global bound on the inner product by Ō(1). Next, by Theorem C.I and Lemma C.10,
can show m 5 Kw pyd) 2 for all [n] and lower bound As 2 log(4/e) and thus
9: ( x3), W > llo((//)) 2 log(4/e) 2log(4/c)
(12)
where we bound As + |Arl + As < 2 (4ge) under Condition C.I.
Further, we derive
W 2 W - W0
=2 W W "
2m Σ
2k u. x;) - (V/(W' xx), W - V
32

---

Published as a conference paper at ICLR 2025
21 pk (2th ((W x:) - -20og(2/e) V/
40
(e(ss(W* x:)) - c/4) - VLL
2nLs(W*)) n/
where the first inequality is by (12) and the second inequality is by convexity of cross-entropy function and the last inequality is by Lemma C.8.
Before proving the second stage convergence, we require the following lemma in order to bound the ratio of loss derivatives among different samples.
Lemma C.13 (Kou et al. (2023)). Let g(z) = (z)) = -(1 +1+ exp(2))- Then for any 22 2 21 where e 2 O, we have g(z1)/g(z2) S exp(c).
[W 1 w
Theorem C.2. Under the same settings in Theorem C.I, let T Ti +
g
T + - (1)laalll111 Then we have
there exists Ti S k S T such that Ls(W*) < 0.1.
max (w E.) for all Ti < k S T.
max, (ww j.r 14) for all = 1, 7 < T.
Proof of Theorem C.2. By Lemma C.12, for any T, S k - T, we have
I 2 | W 2 2n (Wh) -7
for all S k. Then summing over the inequality gives
T
WT W 12
Σ Ls (w*)
S
Se
T Ti
nn - Ti + 1) 2
k=
I W
where the last inequality is by the choice T = Ti + On1 -e-- m lolllog(/e//// log(1/e)llaall))) Then we can claim that there exists a k [T,T] such that Ls( (w(k < e. Setting = 0.1 shows the desired convergence.
Ti +
Next, we show the upper bound on max I (. Ei)l for all k E [T, No- tice that by Proposition C.I, we already have maxj.r I(w Ei)l Ei. I 0, where
we let v = 3max/max,, (w Thr T 8.4 log(An"/8) na). Then we only focus on bounding max ( E.)l. From the scale difference Ti. know that Ö(max(n 1/2 o d, ao adl. 1/23) o(1). Next, we can bound
d
T
[T w O( mm log(1/e)
Σ Ls(W
(13)
"
Arm Ti
where we use Lemma C.11 for the last equality.
Then, we first prove max,, Ipf..ral I 20 for all Ti < k < T. First it is easy that Ti, have
log(An2/0)
max 738
Ti Pr..r
max (: w.r Ed) max (ww wi.rs E,) 4v V.2
d na S 8 < 20
Then suppose there T E [Ti, T] such that max,, le....a S 20 for all l: [TiTT 1). Now let *k = m Ipt...a and thus by the update of noise coefficient
log 482
b+t < 0
prk mm
17
O* + 1/3
4
no |1
d
33

---

Published as a conference paper at ICLR 2025
( (o+
+/33 + 4V log(ln2/// na) O(ogg)
max
nm
d
where we use Lemma C.5 in the first inequality. Then taking the summation from T1 to T gives
S
71
max 1 ((ood)
(14)
TATI
where the first inequality is by the induction condition. Next, the aim is bound Exn max; 1.
First, for any i, 5 E [xx] such that ys = yv. we can bound for all T₁ I k S T
I(W xx) V/r f(W, Xy.)
Fu. (W N6? x.) F-v: (Wh 21 x,)) Foy (Wk
No Xx. + F-ye (Wh.
u1r x,,))
m
((w berr 22 + (w w::: {;)
((w uii:: 11y)) 2 (w werr Ex) + 1/C1
m
m
((w Merr: Ex)) (w (w ,, Ex) (5))))+17) 1/C1
S mmxx mi.r E.) 1/C2 12
S max (Ip + max kw wur: Ei)l aV 1og(4n2/0) na)) 12
d
60 < 0
where in the first inequality we notice that F- yi (Wh y. x;)) 2 0, y/ yv and we recall that
F-y. (Wk x,,x) S 32 (B + v log(An? (8) nax) 2 1/C1 for some sufficiently large constant
G1 > The second last inequality is by induction condition and the last inequality by choosing
i S 1/6. Then we can bound the ratio of loss derivatives (based on Lemma C.13) that
E 1/ S exp (W* x:) yr /( xx)) S exp(1)
This suggests 1 - (10) S 1 < 1 + O(0) for all i, E [n]. Ti S k S f 2-11.
Then let
= arg max, | we have
T
Σ max 1
Σ 0
(
Σ 0 Syy
9 SSy.-
Ls(WH)
SyLe
Ti
Ti
iESg..
SES,.
Ö(1 Im log (1/6)//
-2 (15)
where the first inequality by |e'l < and the last equality is from (13) and SSy.- 2 0.49n (based on Lemma B.
This allows to bound (14) as
T
S T1
max E (()2d))
nm
==
< Ti
O(n oo log( |1 -3)
< 0 -SSN log (g/())))
< 20
and the second inequality is by (15) and the last inequality is by setting e = 0.1 and n SNR? > C for sufficiently large constant C'. Thus, have max,,; (w w { S max,; Ie +
4 log(422/8) 2/8) na < 30 o(1). The lower bound on signal inner product is directly from Lemma
1.10.
34

---

Published as a conference paper at ICLR 2025
C.4 NOISE MEMORIZATION
We also analyze the setting where nt SNN -2 2 C' for some constant C" > 0, which allows the noise memorization to dominate signal learning, thus reaching harmful overfitting.
We first require the following anti-concentration result for the noise inner product, which is required ensure the sign invariance of the inner product along training.
Lemma C.14. Suppose 6 >> 0 and do 2 22(log(n2) /6)n
to,)) ma
we have for all j = =1,r E
[m], i E [n], J.:: 6.) V log(4n2 /5) no.
Proof of Lemma C.14. For any j = +l, E [m], E [nn], we have (w Ei) ~ N (0
Then applying Lemma C.3 by setting RHS to 8/ (2mm) and c = Sv Loe(An2// na, we require d
4m2n2
d² ž 42 log (4n /8n2 2000" 05² log( 4m²
52
where we use Lemma B.2 that ||&:11 > 0.99o Finally noticing that 1/ log(4m²n (a/////(mm) (Ammnn 5²)) I Omm2n²) and taking the union bound completes the proof.
An immediate consequence of Lemma C.14 is the following result that allows to derive the sign invariance for (. E.) for all iterations.
Lemma C.15. Under Condition C.I, for any E [n].r E [m]. we have sign( (w yarr Ei) T TT.
(p sign((w) ((w 1i.r E.)) for all 0 < &
Proof of Lemma C.15. First by Lemma C.14 and Lemma C.5, we can bound if (w Biirt Ex) 2 0,
Poenr + 2
(wo (w...r Ei) S (w yeur Ei) S Poerr + (wol De.rt Ei)
and if (wo p..r Ei) S O,
3
Po + 2 ww..r. Ei) S (w w.. Ei) S P + (w!l ......E... Ei)
Next use induction to show the sign invariance. First it is clear when k O, the sign invariance is trivially satisfied. At 1, we have by the iterative update of the coefficients,
Pu..r Pu... +
77
1
(w ..... E.)
_ 1,,00 kw & 1l
nu
and thus sign(p)) r.)) sign((w) (ww (w...: E.)). Further, by Lemma C.5, and without loss of generality that (w.
years Es) 2 O, we have
dog(4n2² na 2 Pyini + (wo 2 wirt Ei) 2 0.
1 (..r:::)
(w)) Mare Ei) 2 Phenn + (wo whr Ei) 4v
d
Similar argument also holds for (wh..s Ex) O. Then we show at k = 1, sign(p) ((.r))
sign((w) wr.r (x)) sign((w!) (eWW...)))))) Ex)). Suppose there exists a time such that for all t < T 1,
the sign invariance holds. Then for l f, suppose sign((?? f-1 9es (.)) sign( 7-1) sign((w ( ( Ex)) +1,
T-1 P = p
7
oT -l W.r
7-1 E.)1E:11
+
nmn
P P
7-1
+
/ -1((p1p (pp o + (wo (..rr.E) -4
7-1
log(An2
d na)) ||
nm
T-1
- T-1 1
c
P P
nmn
2
2
35

---

Published as a conference paper at ICLR 2025
Further,
(w ye.rt Ei) 2 Piier + (w Merro E.) 2v 1ox(4n2200) na 2 P + 1
7(w Mirr Ei) 2 0.
d
2
and thus completes the induction that sign( (w pr Ei)) = sigu(p = sign((w)). 2e.r1 E.)). Similar argument holds when sign((w) (snn
Mi.r (x)) 1.
We also derive the following concentration result for the average noise inner product at initialization.
Lemma C.16. Suppose and m ((log((/6)). Then with probability least 8, we have for all j = +1,1 [n]
1 Σ (w a.r Ei)l S aoogVa
0 o Vd/2 =
Proof of Lemma C.16. notice that for any i [n], (wS, Ex) ~ A/(0, , and thus we have (ll(") j,r E:>l] = ooll</ll 2/ By sub-Gaussian tail bound, with probability at least 1 - 6/(2n),
for any i E [n]
20 <:ll (2g44)/))
kw j.rt Ei) oollE: S
m
Choosing m (logg(n/)) we have
1
GollE:ll 2 S m
(w 2 6.) S aol &
Because from Lemma B.2, we have 0.990 v dd - | E 1.01og by choosing = [(()) sufficiently large. Then we have aodev < Σ (w j.r Et) S oooevd. Finally taking the union bound for all j = +1, i E [m] completes the proof.
We have established several preliminary lemmas that hold with high probability, including Lemma B.I, Lemma B.2, Lemma C.I, Lemma C.14, Lemma C.16. We let Eprelim be the event such that all the results these lemmas hold for a given d. Then by applying union bound, we have P(Eprelim) > 1 56. The subsequent analysis are conditioned on the event Eprelim.
C.4.1 FIRST STAGE
Theorem C.3. Under Condition C.I, suppose n SNR-2 2 2 C¹ for some constant C' Then there exists a time Ti O(1 -i "nmo -2d 1) such that (1) max, (ww Ti Bi.r &.) 2 for all i [n], (2)
m D
(wwT) Wioro {i) for all E [n] and (3) maxj.ry (w]]. Mny) = Ö(
Proof of Theorem C.3. We first bound the growth of signal as follows. From the gradient descent update, we have
S/Syl
1( jor Ha)i = (( k-1 |1)1
(w k-1
0.51
(( |42)
m 7/ 7/llaa11
< 0.51
1( p.r: 14.3 M43ii
(16)
m
where the first inequality by | < 1 and the second inequality is by Lemma B.1 with n = Ō(() sufficiently large.
36

---

Published as a conference paper at ICLR 2025
On the other hand, for the growth of noise, we have from the inner product update, for any i [n]]
(w Mi.r: Ei) = (whrr Nerr
Et)
glk ( (w k-1
E.)(Ev, E.)
Tu
"I
ttth |1::1?
(w 1-
Ei)
e ( (w
R-1 yir :.) Ei)
Mirr
nm
Then this suggests
(ww wi Gi)1 2 (1
" nm
tt [
(w Mirt
Gi)
C 11 kw : KEr,, (17)
nu
We first prove for any [ni]. max, (w G.)l 2 max, (( gi.r. 6t) 2 max, g..: Ei) for all pirr
k S T1- We prove such a result by induction. is clear that at = 0, the result is satisfied. Now suppose there exists an iteration k such that
max (w
21 E.)1 2 max (w Mures Ei) "
ouo vd/4
for all k S k - 1, where the last inequality is by Lemma C.1. Then we can bound based on Lemma C.7 and Lemma B.2, we have for any i' * E [n] and
n2 1 11 KE Ex)1 1( Be.r Ev)
20 / /loo((222/0)
nad 0g 1-''d'/' 1/2
1 11 ||2 max, ((w (
Ce0.990
=8.4C7 na Vlog(4n2/8)
doyog
0.01
(18)
where we use the lower and upper bound on loss derivatives during the first stage, as well as Lemma -1
(oog(42n//)) Then B.2 and Lemma C.I. The last inequality is by do 840nC; d d'' aa have
max (w Dirrs 2i
1
" nm
pk
max (w Ei>1 [IE:1?)
77
1 11
1(w
(WE-1 E-1 2ir
Ev)l 1 4i)
Tum
+
R-1 " 0.99 et -111:111)))) max (w
girr: Ei)
max (w yur+ Ei)
max ((w p..r
(...r:::)) E:)
Let BZ = max, lw wr. {:) and we obtain for any A: < Ti.
no?d
B¹ 2 (1 " -0.99165-1111.111 Bx-1 2
0.98C Bk--1
note
b ".
2 nm
0.98C²
no?d 0.98C/
4: ooog vd/4
nm
where we use (18), which holds for iteration is and Lemma C.1. Consider
Ti = log - log (1 + noon 0.98C;) = (v) - nmo - 2d (og(807 0g 2==20
1---d111 1/2))
nm
for y sufficiently small. Then it can be shown that
BT max
(w 24 wi.r: E 2
37

---

Published as a conference paper at ICLR 2025
In addition, show the average also grows to constant order with a similar argument. In particular, from (17), we have
EX>1
2
(1
"
(w k-1
Ei)
num
kw
Ee) ((Sv.
Wesr
m72
Using similar induction argument, we can show
1 kw Ment
6i) 2
Kw
Sill 3
2 m
Σ
((w
wi.r Eill goog Va/2
m
m
for all & S Ti, where the last inequality follows from Lemma C.16.
Then we can show Ti,
poon
Ti
1.)1
2
0.98C.
voogva/22
m
nm
In the meantime, (16) allows to bound the growth of signal learning as for any +1, max
(w j.r
(1 0.51 11laa13
V 1llg(8mm/))o ||||||
m
log(1 0.51
Mll?
=exp
m noddCe am (0.53C7"nSNRR exp
log (805'0g 1-1)) logg88//8joolll
log(1 0.98
2(1) log (80 0='d (2ll2)))2llll((l1 (2log(8m//)oo))
< 8 8//2log(mm/0))NN
Ö(n 1/2)
where the first inequality is by Lemma C. and the second inequality is by Taylor expansion around
" The third inequality is by choosing 7/ sufficiently small and based the condition that
SNR 2 2 0.55C, The last equality is by the SNR condition.
C.4.2 SECOND STAGE
We choose W* to be
W j.r = W.. + 2log(4/c)
Σ 11y = jissinn((ss))))n j.r Ei)
First we show the invariance of sign of noise inner product after the first stage.
Lemma C.17. Under the same settings as in Theorem C.3, we have max, lww M. Ex) 2 and kw Warrs Ed) 2 for all Ti < k < T* and any i E [n].
1 m
Proof of Lemma C.17. In addition to the two results, we also prove max, Ipt..al 2 1.5 and
1 mm
r+l I 2 3. We prove these results by induction. First, it is clear that at k = Ti.
the bound regarding inner products trivially satisfied by Theorem C.3. Then by Lemma C.5, have
log(4n? (422
max 1PTh Panni 2 max (w Ti Birrt E.) - B 4/ d na 2 2 - 0.5 = 1.5
log 4n 4
Punnal
(w Ti
Harr Eiil - B 41
na > 4 4- =
d
38

---

Published as a conference paper at ICLR 2025
where the last inequalities are by Condition C.1 for sufficiently large constant C.
Now suppose there exists time T1 S 7 I T such that the results hold for all A I T 1. Then at
T, recall the coefficient update as
T-1 P Pyiit
1 ((
T-1
Par.r
+
(19)
nm
If (w N1. Ei) > 0, by Lemma C.15 we have (wT.1 we.r Ei), T-1 > 0. Then
Pyr.r. =
T-1 Puirr
+
1eth
Mer ..illE:11
T-1
nm
Pu.r.
Ā-1
+
peT 11 ye.ri + (wh wirrs E.) - 41 d
log (/n//)/))) /8)
na) ||:
nm
T-1
1 1( T-1
1
(w!) Warr E1:
Ueiris
2
nm
Then taking maximum over r,
nl|E:11
max
max Poraa
12
max | Porste 2 max P 2 1.5
2nm
where the first inequality follows from (wh... /2 s 0.5 S max, -- based on Condition C.I. Similarly, when (w..... Ei) < O, we can obtain the same result. Then, we have
max (w mr E,)l 2 max Ipr.naa B -AV / [og(An2 4m2 (6) (4m2
na 1.5 0.5 =1
Furthermore, we prove the results for the average quantities in similar manner. First, from the (onn((())))) Mirr (.)) and thus taking the average
coefficient update, and by Lemma C.15, (ign(p T-1 sign( (wT T-1 of absolute value on both sides of (19), we get
712
m
7
E
m
1177
m
log log(4n2/5)
T-1 Pyrra
17
P 3 1--44 a
d na) ||E
nm
m71
Po.rral
+
7
Por.ri 111
2nm
r=1
m
p
m
where we use la + b) = |a| + |b] when sign(a) = sigu(b). Then, we have
log 4n² /n²/5)
4 Permall - B - -41
d na 2 3- 3-122
where the inequality is by Condition C.I.
Lemma C.18. Under Condition C.I, we have || W 0 wow.m/oogg////// =
Proof of Lemma C.18. The proof follows similarly in Lemma C.11. Let Pg be the projection matrix to the direction of i.e., PE
SET TEJ
Then we can represent
w j.r j.r
w = -P
(wi jar
wo, + Pu (w
w w.r
+
P (w.
w
jir
39

---

Published as a conference paper at ICLR 2025
" Pe: (w w
(T-P pa.
By the scale difference at Ti and the fact that gradient descent only updates in the direction of pi+ j = =1 and E., we can bound
|WT) w0
Ti Jir
w. J. H1)3
(w
W j.r
11-1)
(w Ti 2.
w E/)
|ea|
||
E
Σ
P - P P - Σ Pg. (w -w].))
[
2(w Ti j.r
13)3 + 2/w 11 j.r -1²²²²² 2² + 2(w j.m 1em1} + 2/ 2(w j.r 14)
< 2m
|
2(w 1.r 6,) + ++((w].,r (1) 5
+ naxx
Σ
(I-PM) m. - P p-
> P&,
(w
Ti
w j.r
1|
i= 1
S <m 2d
where we use the scale difference at T1. Therefore,
| W' < l - w + | -
W
O(v O 1d-1) O0 (og(1
1-1222
where we use the definition of W
Lemma C.19. Under Condition C.I, we have for all Ti S k S T
T
I 1 | W 12 2 2n (W*) n
Proof of Lemma 19. The proof follows from similar arguments as for Lemma C.12. We first obtain lower bound on yo f(W x:), W*) for any E [n]] for all Ti T
ya (WW*, X:), W = 1 Σ
jue(w j.r Mny MMa
Σ
jui jor (i)(E, wj,))
3.r
J.r
Σ
ju (w j.r: (M Wiyh
jus (w j.r: (i) wh.r) Σ
j.r
3,r
(E:, (E,,Er) 2 log (4/e)
jus J.rt E.) vv)
||
P=1
(5i, Ev) 1
Kw W.. (2)22oo((4/)) + m
ΣΣ
wil He E,) log(4/))
m
d.r P
Aa
A1o
jus (wh f.r My Mayy
w
Σ
jus ( f Ei) why
m
A
A12
where the second equality is by definition of w The third equality is by Lemma C.17 and Lemma C.15 on the sign invariance. We next bound based on the scale difference and Lemma B.2,
A11l ((dd-1/2), |An (oullall)) |A1al S ((ooogVa)
40

---

Published as a conference paper at ICLR 2025
where use the global bound on the inner product by ((). Next, by Theorem C.3 and Lemma C.17, we can show 1 Dx= (wr)
Wi.r: Myy) 2 for all E [n], 2 T1 and we can bound
Ag 4lo((4/))
Combining the bound for Ag. A1o, An1 A12, we have
V ( x,), W 2 2lg((4/c
(20)
where we bound |Azol + |Anil + A12l - 2 logg (4/c) under Condition C.I. Further, we derive
W 11 Wh 11
Lss W -W> V (W
2n
4 (2f( x,) - (Vf(W* x,). w - -
20
4th (2 x.) -210g(2/e)) - V
71
4n (e(y.f(W* x.)) c/4) n7 |Ws)1*
(W ne
where the first inequality is by (20) and the second inequality is by convexity of cross-entropy function and the last inequality is by Lemma C.8.
| W 12
Theorem C.4. Under the sume settings in Theorem C.3, let T Ti + T + O "ann'ee-mmnn m 22 ") Then we have
there exists Ti S k S T such that Ls(WW) ( 0.1.
maxjary (w),,: My)) o(1) for all Ti 1 k T.
max, ((w gi.rs &.) 2 for all i E [n), Ti < k < T.
Proof of Theorem C.4. The proof is similar as in Theorem C.2. By Lemma C.19, for any Ti S S. we have
W
| 2 2nLs(W*) nE
8 S k. Then summing over the inequality gives
Σ Ls ( < W W 1 k=
Se
T
(T Ti + 1) 2
where the last inequality is by the choice T Ti -e (o "d11. Then we can claim that there exists a k E [Ti, 7] such that Lss (w*) Setting 0.1 shows the desired convergence.
W
w 11
Ti
n
Next, we show the upper bound on maxj... kw j.r: My)l for all & E [T, Notice that by Proposi- tion C.I, already have maxj.r (ww j,r M4)) O, where we let
log 4n2
0 = 3 max
mmax max (w To j.r My)1. 3, / d na = Ō(max(n- 1/2 aoog V d, ao I|llll nd 1/23)
Subsequently, we use induction to prove maxj,r l(w r. 143)) 20. First notice that
IT Ls(W
W0 12 O(1i nmoo 2d
(21)
"
41

---

Published as a conference paper at ICLR 2025
where the equality is by Lemma C.11 where we choose 0.1.
At k Ti, have maxj,r (w Ti j.r p1)) _ 0 20. Suppose there T E [Ti,T] such that
max,,i Ip g S 20 for all & E [71, T 1|. by the update of inner product
Now we let Th = maxj, kw j. H531 and thus
gk++1
I v
e wr ||]l
mm
iES
I
ph v ||]
iE|n]
20
2n
Ls
where we use |e4| < l in the second inequality. Taking the summation from Tj to T gives
7 < Ti +
20l1
Ls(Wk) m k=
T-1 Σ
4 + +(MSNNR² 20 20
where the second inequality by (21) and the last inequality by SNR large constant C" > 0. The lower bound for noise inner product is directly from Lemma C.17. O
ANNR-2 2 C" for sufficiently
42

---

Published as a conference paper at ICLR 2025
D DIFFUSION MODEL
For the analysis of diffusion model, we restate Condition 3.1 specifically for the case of diffusion model.
Condition D.1. Suppose > and the following conditions hold.
1. The dimension d satisfies d (1(mmxmxx' no=²1)-
2. The training sample satisfies (2(log(m/66) and the network width satisfies (
The initialization dg satisfies Ö(no= 'orr--'d-i/ 5/4) I ou I Ó(min ( 1/0d-1/00/ 1/3 n 1/3
vn m->1oo--/1200 1/3
11/3//2000/////////// n/ d= n1)
The signal strength satisfies ||a||| (1).
-NR-1 =
6. The noise coefficients ag S, satisfy a,, Be = O(1).
We make the following remarks on the conditions. Compared to the conditions required by classifica- tion, diffusion model requires m = @(1) for the analysis of stationary points. The lower bound on sample size n is required for the concentration of |--1 The lower bound on dp is required to ensure the inner products of E; across samples remain small relative to the initialization. The constant order of signal strength | and the bound for n SNR2 are utilized for simplifying the analysis. is also worth mentioning that diffusion does not require small learning rate for convergence.
D.1 USEFUL LEMMAS
any Lemma D.1. Suppose 8 > Then with probability least for any t,
ogd(1 ((d-1/2)) w < od o ((2 (w r. My)1 s //llo/(6mm/00a M1
(w rie: E.) S 2 30og(16mmm/0)onogg
(ww. he: w 1 1og((1mmm20/)00 r+
for all r,r' E [m]] and i [m), and j = 1,2.
Proof of Lemma D.I. The proof is the same as in (Kou al., 2023) and include here for com- pleteness. Because at initialization w ~ (O, I)),, by Bernstein's inequality, with probability at least 1 - 5/ (8)m, we have
2ull2 - ofd = O / dloo(16m//))
Then taking the union bound yields for all E [m]] we have with probability at least 8/4 that
oād(1 1 o(d I wh 11 dn O(d
Further, because (w... 11) ~ No of 14, [1:) for 1, then by Gaussian tail bound and union bound, we have with probability at least 8/4, for all j 1, 2, r [m].
(ww. r. |4) < / logg(1mm00jog
Finally, following similar argument and noticing that ||5.13 - =(o and l|w O(oīd), we have with probability at least - 5/4 that for all i E [n). (( &i) s 2o/og(16mm/Boooaqa and kww... wh. < //ogg(66m2//)22
43

---

Published as a conference paper at ICLR 2025
D.2 DERIVATION OF LOSS FUNCTION AND GRADIENT
We first simplify the objective through taking the expectation over the added diffusion noise. Lemma D.2. The DDPM loss can be simplified under expectation
I(W,)
1
( (W,) (d+
( 12
2n
i=1 jc(2
where
m m Σ
2 (W,)
Wrrel (a xf11** 6o2 (w.,: ())) wr.eel? +33 /w,ll
4. 4ymmar, (Wr,t (1))
10 (W.)
Σ (Wr,t, Wrr (aa((w.., x())² + Bill (at x01 @l]W)r
+ 2.8 (Wr,t, .)) + 40 (Wr,t, Xo,a) (We (W wr.))
corresponding to the learning of r-th neuron and alignment of r-th neuron with other neurons respectively.
Proof of Lemma D.2. Without loss of generality, we consider for single sample Xt, We first write the objective as
f
(p))
(p) E,
1|
E
m
(W
(p)
Wr.t
22
1
m2
((Wer, x (1))
m
m
where we omit the subscript for the expectation for clarity.
First, we can see 1 d. Then
I3
((w x Ga,)
E
((w (2)) )) 2 w., ((
28,
E
(ww,,e wr. 1872
mm
28
/w..
2
(Wr,e, (2))
20
w (w
v m
where the third equality uses Stein's Lemma. Next, we consider lg by writing
m
In Σ E [[( (() (21))) ?
+ Σ Σ E (( (11) x(1)) (w Wrra).
44

---

Published as a conference paper at ICLR 2025
(()1)) respec- tively. For notation simplicity, let or or (Wr,s x Bell ///l and ž, (wr,, E1,3)
Next, compute the two terms E [[((w.,) x;,)) ) and E[((W,,t, ( xx,,,,((,))
first compute [p = 0 and || = 3 W, 11 B| = 3,8 wr,e For the first term.
E E[((w
(p) "" [[ + )+
Eia + 4a , + 6a + -a, + 271
a + [a + B[z!
- a, + ba, + 3b;
= (w,,,,
x (o)) + 6a B8 (Wr,t. x (ply w/,l, 355 |wall"
Next, for We.... X,0 + B,E,,) ( (Wrr,, 0:80, + ),),,) we note that
[[ ( Wr W 2 32 (W..; wrr)),
60 2| = 0
| [2
[[2 + 20(2,,,13 5 ww. w + 2// (wr.. wry))
where the second and third results follow from Isserlis Theorem. Then we can simplify ((W,,,, + B, W ))
=[((a, + 2,) (a, + 2,-') )1
=,,a,, + [2 + a, a, 22 + E[ + B[2
= =o Xo.) (Wr,,, x + 0f o2 Xo.), " + 40 ( (W.,,,,, weye)
+ af (w x7 2 W //t
28 (Wn., wr)a)?
Combining 1,, 12, Is gives
(x( cr, d
|weell? (at( x(())*++ + 67 x Oply 2 + 33/1/ - 3e X
(we.))
+ 2
xx1)) (Wees wyr (wr,, x1,)) + Blll, 12)) (af xx + 3/1lwr
1=lr'r
+ 2,3 (Wr,t Wr. 2 + Xp.)) (Wr Xo. ( Wrr,t)
(w..))
where we respectively denote the two composing loss terms as LP) (corresponding to the learning of r-th neuron) and L (alignment with other neurons).
We next compute the gradient of the DDPM loss in expectation. Lemma D.3. The gradient of expected DDPM loss in Lemma D.2 can be computed
VL(W.) = 2n
(W.) + VL (W.)
2
where
(
(at/Wre, (p) x
12a.33 (w,,,,,,,,, (ply) [| 985 ww,ll* 4/man (Wr,, -
Wr,
45

---

Published as a conference paper at ICLR 2025
2
x |w,,el 112 wwrll (.) - 2/ma, Belc,,llw.... !1? ) xox
VL2. Wr.))
V
Σ
((a)(W,, X..11 2 (, (p)) + wwr, + (w.,,
+ +42 (r.,,
Wors)) Writ
2
(p) + I. Wy Wr, (r (Wr,,, (p) (p) 10))0 + wr.)
m
2 m
Σ
(Wra, We 1) (48) wr, + 80 (wr,, Xo,)
Proof of Lemma D.3. The proof is straightforward and thus omitted for clarity.
D.3 FIRST STAGE
Before deriving the results for the first stage, we derive the following lemma that decomposes the weight norm given concentration of neurons.
Lemma D.4. For any & and r E [m], such that (w)) (((w (kr))) 2 @(oall (w Ei) @((wt, Ev)) 2 ((ongeVd) and (w Mg), (w". Ei) Ó(1), (wf.e: what
(lofd) for any j. = =1, E [n].r [m]. Then we can show
||w",ell = ((w n. 14)) |||||| SNRR(WW,, Ex) +||]||=+++111
and r + r', we have
(w w @((w My) (w |zm| - n SNR? Ei) (w Eillle)1 + (w wh.)))
Proof of Lemma D.4. We decompose the weight w as
wh Pit = w wh + 1 + 2 1 ||| +
(22)
based on the gradient descent updates of w/, starting from small initialization w and the direction of update only involves w and μ+1. &., where 7p = - p.. 0 and of =1 1.
Then given the assumption that (w.,, wh.)) O(o]d), we have of (1) because
(wi ri w = o. wh 112 +
-
(oGddd
where the second equality uses Lemma D.1. This suggests that ob = 0
Then we can see
(wh M4)) = (wo nee 113) + 7
(w", Es) (W Ei) p + (PE ot (w o Ei) ph. + ((nd-1/3)
where the second equality for (wh Ex) by Lemma B.2 and (w Ei) Ō(1), thus p = Ō(1)
Then based on the assumptions that kw (4)) @(Golllo))) and ((w E.>I ( d) and 0f (1)) we can simplify (22) as
w = O(w..) + O(( (1) + ==1)l|=||12) + @((w), Ei) + ((nd-1/2) Σ E | -2
(23)
46

---

Published as a conference paper at ICLR 2025
where use (w p4) ((wwf,: (cy))) 2 of Mly) and Ex) @((w,, Ex))
(fl Es)l in the first equality. For the second equality, we use the assumption that (w., Ei)
((rogeyd)
Then we can show [|wfeel?
= -([lw + (w) +((w,, E.)² ((ndd11/2)] Σ E:1EE,11
+ +((( (g) 7. M3) -- + @(((w E.) (/2) (w
r,e0 Ex)
|| 11
O(a)d) O( 23) ) + @(( {,)² (n2 2 d Ö(n' og dd 3/21)
-=2d
+ + 5:) + Ond 112)) (w 5 2
=Ooodd) + @(( ||||||||| (nog 2 d (w k (1) Onn' 3/2) + 0n2200 'd-1
= 2 || -| nSNRR E1) |||1||||1 + ok)d),
where the second equality uses Lemma D.I. Lemma B.2 and (w (4) = 0((wf,e (yr))
2
(i ug)l- The third equality is by the condition on d = U(n²) and (w Ei)
O((w Ex)) 2 (l Eiil and Lemma D.I. The last equality is by the condition Ø 2
M(max(no 1d-5/4 on )d) = M(no dd-//2) (d5/) given d = =((?))
In addition, we can deduce from (23) that (w r. wh(a)
=O((w) re: w.. + O((w (w) M45) 13)1 + @((w r.: Ei) (i) Σ (w. &M>1E:11
n
+ O((w m3) (w |||| 2) Oliw Ei) Σ
(wo (w Ex) 6.)1
+ Mz) (4) |||| @(n(w/) pic: (w a E.M1l..11
O((w... w w 0 ((w (4) (W ( 2) + Onn(w"... Ex) (w 1.)1
= O((Wh,: w w O r.:: (w) (4) |all 2) + (IISNR"(W Ei) (.) (wh Exllleal -?)
where we use of (1) and the (wi. (4)) = O(ww (. 2 WW Mللil and (wh, Ex)
O(ww Ex)) 2 of kw Ex) for the equalities.
Lemma D.5 (Restatement of Lemma 4.1). Under Condition D.I, there exists an iteration Ti max(T, Teh. where Tμ = @(mmo, d " "m and T/ @ de --1
123/2n11) such that for all 0 < k < Tí, (1) I 2| @aaīd) for all r E j +1, E [n]. (2) (w wh.t (ofd), and (3) the signal and noise learning dynamics can be simplified
(wht) k+1 :/)) = (w.. (4)) + 0
|1
nym
(whtt Ee) (wi, Ex) θ ma
[w w.l.1111..
m
for all j = +1, r [m],i E [n]. Furthermore, we can show (w?)) m4) O((w
Ti 2 )1)))> 0. (W Ex) 0((w") {x))> 0,
|| @(Iw?); (P).
47

---

Published as a conference paper at ICLR 2025
(wil) whoh @(1lw for
(Vw,, L(W?+), w9.))
= M4j E)
Vml /m/lw//ll)) (w wh.)
[wafll (elll?(w" (w My E) where we denote & Σn Et.
(w M3) (w e Ei) D(n-SNR)
for all j.j = #1,r, i, e [m].
Proof of Lemma D.5. We prove the results by induction. To this end, we first compute the scale of the gradients projected to the space of μ1. P-11 and &, for [n] under the initialization scale. For notation clarity, we omit the index k.
As long as //l?,
((o)d) and suppose (w,',t, (b)) O(Wr,s,)5)) (wr.a:s)) O(W,,, (i)), we can identify the dominant terms as follows.
Signal.
First for pjj and for any i E [m]. we compute
1
V 14)))
2n
1
1 Vm
@( 14) + (Wne, 2 ogd + (Wr:Wwr;)
((wr,e (4))))
m
+ - oo wa, 2 + agd (wr,, 1
O(oîd)
m
mm
= 1 (w..)( (y))
mm ((ord)
where the second equality is by (Wr.s, 12 S |w
O(o]d). =(ofd). is clear the dominant term v 3 [|,,, ||/p|1 || The second dominant term comes from ( lI,,ll (Wr,s (4)).
Further, we have due to the orthogonality between signal and noise vectors,
7
1 'O((w,, ( (Wr,ss (5)) + (Wr,..
VL
w.)
14) + od (wr,,
2n
m
/mi Ex) (Wras
In addition, we have
1 n1
VL (1 (wx.).))
0 (( 14)3 (w,,,,,y) ofd (Wras m
2n
m 1 (Wr,,))))
Further,
m - Olwwees Mi) (i) (4) + (o mz) + (Wrs, (Wr.).
2n
m
Then according to the definition of [St1] and μ11. we can simplify the gradient into the dominant terms as
1 @lIlw,,el (Wr,ss Ed) (Waas H4)))
VL(WW).)))
mm - (Wr,r, 4i) (Wress M43) (Wr,,, Ei) ood w1))))))))
48

---

Published as a conference paper at ICLR 2025
Noise.
Similarly, we can also show for the noise learning
(w..).
1
O ( Ei) + (Wr.e. ogd Ei) - (we Ei)
2n 8'
Ei) (1)
mm
1
(lofe Ei) - O(((rees Ei) (1)))
m
where the dominating term is 1 (W (g) E.). In addition,
(wr.),
2n
1
{:) + (wr., Ei)% of (wr.,, 1 [[((Wr,,,)))
Vm
1
(Ol(w,,), ozd + of (wr.). E.) + ((o (O(o (((od) - vmo(o
+
m
1
(( 6.) 5 (w..,. of of (w.,, E:))
mm O((Ww,,,))3
m
+ O(( E.) 2 on (w..,, 3.00
(aao )
where we use Lemma B.2 in the first equality and the second equality is by 0n
Further we can show
1
2, (w..).
(1)
2n
m
O( M4y) (Wr,, (.rreE)) + (5) Ei) + odd(wa., Ei))
m
O E.))
m
Lastly,
m O ((W,,,, E1) (Wr,,, (.),,,)) ad) + of (Wr,), Eil
(w)
2n
2,1
m
m 1 O(( o0 + (Wr.,, 61)0002d"
This suggests we can simplify the gradient along noise direction as
(VL(W.). (.)
vm Olo + (Wr,,, Ei) + (Wr. Es) (Wr,er H3))
+ (fffd Ei + (Wr,t, + (Wr,e og)))
+ +0((wr.:: Ei) of +
of
In summary, as long as |/ d(ffd) and suppose (Wr.), My) O(W (Wrrn O(ww,,t Ed)), we can simplify the gradient as
WV.., L(W.). = -- O((id)) -
1 -wwn (w),
Vm
Vm
+ (W
(wrs, M3) + (weee (i) (Wres my) + (Wres Ei) d( M4j) (24)
49

---

Published as a conference paper at ICLR 2025
V Wr. ((W)), - - ((Wr,s, + (Wr,es E (w),,,
1
Vm
Vm
+ ( (Wr.,, Ei) + (Wr,,, Ei) + ( ozdd)
O(Wr,,, E,) of on (Wr,s, o
(25)
In the initial phase where || = @(a)d), (w My)1 = Ö(oollae])) and (w = ((oogeVe) (by Lemma D.I), we can show (24) reduces to (VL(W.).
( + O + (Wras (i)as 14y) + (Wra, (W
Vm (())d)
where we use the condition that SNR -1 = =((d'/4),ic. ie., og 1/4 o(1) in the first equality, The second equality is by condition on S 00m--1/6d 11/5) and oe (((d1/4. Further, we can show (25) reduces to
(WW.)))
1 Oooo 0 Ei) + E1) + (.),,,E)) ogja) Ewn
-
Vm
0 E.)
1
Oloiood ((Wr, 0000d"n + (Wr,, ( Ed
Vm
1
0loho
vm
where the first equality is by d M(n and d 2 5(n2/3 -2/3). The second equality is by The third equality is by oî - 1/2d d (/
oe <
In summary, we can show as long as l/w? = @(o)d), lww,, الحي = Ō(Golla|]) and l( r.t (i)l = Ō(ovggVa),
(wht) (4)) (w/, r.er (b) + nan (O,dd)
(26)
Vm
(wit) .E.) = (wh rir: Ei)
nar
(27)
nym
In addition, we similarly show that as long as ||w/ 112 = ((o3d), lw Kw Ei) o(1),
(wit), w.t) = =(w whe) + no (lodd + (W Hj E.) (w with wh,) + odd(w + E:))
(28)
aoe d3/3nouff and Ti max Tg)- We prove the results hold for all 0 < < Ti via induction. We partition the proof into two stages, namely when 0 l S min (Tin Te) and when min(T,,, Tς)
Next, let Tμ O (11m(16
cad|||a|| and Tg O "y in log(16mn/6))
Tı.
(1) We first show for all 0 < k I min Tş} that l|w Oooād), (w wh.)) ((oād). l(w pz)مl = ((ooll (4l) and (w &.) = Ö(ana vd) hold and thus (26), (27), (28) are directly satisfied. We prove the claims by induction as follows.
It is clear that at k = O, we have from Lemma D.I that |w w O(oād), (w wha) = /|w (afd) and
kw 2, (w
22 220gg(66m/8)ool/ee (Gn)111))
50

---

Published as a conference paper at ICLR 2025
kw Ei) - 2//og(16mm/00oogg Vd Ö(ovo
Suppose there exists an iteration T S min (Tp: T&} such that || @ (w w..)) ((oād), (w [4y) O(oollal])) and (w E.) Ö( vd) for all 0 S k ST -. Then we have from (26) that
(w (W = (w u3)) + na
nandt
@(ogd)llull= (W,,0113 (4) + ((1)d)
Vm
Vm
(w 7,61 M3) + n
((ard)llu1
vm
(w,, (4)) + (Gollll))
(29)
where we use the Lemma B.1 that |Sy] ((n). In addition, we have from (27) that
T-1 Ex) n/m x (wrr>E)) Ex) (WTT'
3 ( s (w ( Ei)
n/m @
= Vd)
(30)
where we use Lemma D.1 that || - for all [n]] Finally, we deduce from (28) that
(w what (Wii T-1 what + no (lodd + ( My + Ei) (w what + odd( Hy + +Ei)
=w what + no (0800/2(11)) + ocVo) + ofd²)
= = + n ( 3/d +o Vd) + ofdd)
(31)
where we use Cauchy- Schwarz inequality in the first inequality. When T Tpp
10 (0 + ozva) of f ő ((Ium) oeVa)o ogdd² S (Ga)d)
(32)
where the last inequality is by the condition on of (0(dt/4) > and 00 s Odd-1/3), When T Ts.
no (22013(())) + + 000000000011 + +ova)) 20001114
+ no00o= 2012) < (().d)
(33)
where the last inequality is by the condition on d that d = (114a1 and d = i and To I 0( 11133-111/ Hence we have proved the induction on (wh,, whe) and in fact proved a stronger result that (w wh.) O(aīd) for all & < max(T,, Te} as long as || O(oād). (w r.s 13)l. [=)l= o(1)
||/4|l)
Next, we let P&
R be the projection matrix onto the direction of & and we express w = 7 We + (I- Pμ,, - p- - E Pş.) w and due to the
P w + + p_, w + Er Pε. orthogonality of the decomposition, we have
w
(w. r,t: 21)
(w... r.::
2
(w? r.t: Ei)
n
|w w
+
+ (I Pm Pp-, Σ Pe, w
||
11
(w. mer
w.O 12
1
= =Ō(oz) + Ō(no)) +
w
=O(fd)
51

---

Published as a conference paper at ICLR 2025
where we use the induction results that (w
ner uz)1 ((aaala))) and I( c.) Ö vd),
and the ||(I mM - pp_. Ei= Pe. w is dominated by its projection to w
This completes the induction that for all & I min Te). we have ||w @(odd).
(w wax) = (oz]d) (w (4)1 = O pll) and (( E.) Ö(ooog
(2) Next, we examine the iteration min Tc) S k S max Te) = Ti. The magnitude comparison between Tμ and T3 depends on the condition on n- SNNR?. In particular, we can verify that TT/Te= = 1//n-1///nn1555 SSNR 2) = θ(n -SNRR
When Tu S Tτ, i.e., SNR? (111, we use induction to show for all min min(T T&} S k I T2. || @(gdd), (w w.,)) @(ojd), (w Hb)1 kw Ei)l = V(oogeVd). It can be shown that under the condition do I ogdd we have 1(w Ei) o(1), which suggests (w wh.)) ((o]d). Suppose there exists an iteration Tμ < Tg S Te such that the results hold for all Tр S k S T 1. Then we can derive the dominant terms in (24)
[NSN))
(VL(W)).)) vm ((add) (oged
1
+ OO (w.., Ua) + (W),,.)) + (Wr,, ogd Ms)
(caa)) -
1
where the first equality is by = (n) and the second equality is by oă S 01mm1 -td-1/20)) This suggests that we can still leverage (26) to bound
To
T3 ,1, +
nar (a (w., M4)) + man
(w (y) = (w, r.t
Vm
m
(w" 2, (4)) + (0olllel| - Ö )SN))
To r E.)l is the same as (30). Then by the same arguments (31), and (32), (33),
The bound on (w
we can show (w
Te r,t
w .,))) = Oofdd). Thus, we can compute
+ 2|2 = Ö(oğn² SNRR) - Ö(no)) + O(a = O(oād)
where the last equality is by the condition on d that d = "(nllaaloo This verifies the induction on l/w.,, (()2d).
When T3 < <T,, i.e., -1 SNR-2 ē((1), we use induction to show for all min min(T,,. Te} S k < Tı,
|| k F,ell 112 = O(oīd). kw r,1: (1)) ( kw (.) = Ö dn SNR ). Under the condition that op S 0 - dd33/4n) s 06 In), we have ((w E.) = o(1)
which suggests (wh,, wh.) @(ggd). Suppose there exists an iteration Tç < Tμ S Tμ such that the results hold for all Te < k - Tí 1. Thus derive the dominant terms in (25)
(VL(W.). (W)
Lo n 1 O ( Ed) + (Wr., Es) + (w.....)) og Vm
+ O(( (1) 00 ef 2
(Wr. (i)
Looooo + Ei)
(Wr., (
Vm mm (0p
52

---

Published as a conference paper at ICLR 2025
where the first equality due SNR )d and d 11(21//0 -2/3) The second equality
is by of I m/3/)). The third equality is by og = Ō d nm n/31/3).
d-3/4nm-
This suggests that we can still leverage (27) to bound Th (wie:: Ei) (wer E.) + no4
To
( & (wher E1) + na O n/m
nym
Ö dnn -SNR
Tg The bound on (Wiis (4)) is the same as (29). Then following the same argument, we can decompose
w,, -1 2 =
Ō( + + 3) +(ogd) ((ofd))
where the last equality is by the condition that SNR -1 ((d4/4). This verifies the induction on |/ = (ofdd.
Furthermore, at k Ti, we have for all E [m]. j = +1 and i [n]. the growth term dominates the initialization term and thus
(w v.t. (a) = 3m - 1/2 0z (T;) 2 @( 2 ((w My)D
(w/.). r.c Ex) (na 'mm o' 2 0(oogg O(((W,,, E.)l
where the inequality is by the definition of Ti Thus, we verify the concentration of inner products,
ie., (w ne 14y) = @((w") r' (by))) and (w E,) = @((w &.)))))) at the end of first stage well
as the ratio (w p2)/ ///l/(ww)) (w? Ei) = On SNR*) for any r, r' E [m]. Then, we can see directly [[ [w/ = O(Iw") 23
" 2,1l1) ((odd) for all r, [m].
Next, we verify at Ti, we have (w w w = O @[] for all r,r' E [m]] such that Ar. To this end, we first notice that the conditions required by Lemma D.4 are readily satisfied at k = Ti and thus applying Lemma D.4 yields
[| 1|2 = @((w Th 22) ||l| + SNR? (w.). || -2 (!
(w?)) 7.t w @((W (y (y 1y) n SNR" ( Ex) (w?) E.) 11lllll1 -2 (w wh.)))
|||| + SNR? (w?) 6,) ||||| + (wh. wh,,))
e(I[2]] || (wh., w.,,))
= - ogd) + V
= =(([W
where the second equality for (w",, w is due to (w u3) = O((w") o (4y))) and (w Ei) @((w?) Ex))) and the second last equality is by Lemma D.1.
Finally we verify that at Ti,
(V L(WT)) w..))
-(((W)) My + E) (w wh.) + [W (w My + E))
Vm
+0 ((wT) )1)) (w 4x) + Th ) + (w Ti E:) || + [[W" (w wa.))
0 (w (2)) 7:2:1 f.ex M4) 2i || (w Ei)
+
[ Ti (w. T. Mi (w.. (1)) - || ( Ti Eu) (w. E,)
[((((w M + wm
w..) + || (w f. My + E)
Vm
where we use the concentration of neurons along directions Hj at Ti and the scale of (w?) wht (b)) (wi) (.). E.).
n
53

---

Published as a conference paper at ICLR 2025
D.4 SECOND STAGE
For the second stage, we derive an extension of Lemma D.4 given the scale of (w wh,)) can escape initialization. We highlight that unlike (4)) and (w E.) that increase monotonically, the dominant term of (Vw.. L(W?'). whe)) suggests that (wh. whe) can also decrease.
Lemma
D.6. For any & and O(iw re Iw (( (owha)y)))))) (w
[m)) such that (wi Mb) @((wf,)) Ex))
((W
2
@((w", wiiil 3 Eil and (w Mj), (w E.) = Ö(1) (w what
((w wh..)) = Q(minforo 1 11/6mm11///// 1/6 Lov dm 11/6) for any j.j' = +1,i, E
[n]. r,l. E [m]] Then we can show
||w/,a11 @((w M3) 1|lal1 + n ANRR(W"" ( 11) (w w [[ 11)
for + r', we have
(w wh.)) @ (W May) (w (4) pe))llaa1 -nn--SRRMM [t)llll
(w- ner w
+ (wee n,:: what
|w reell
Proof of Lemma D.6. Similar to the proof of Lemma D.4, we can decompose the weight w as
w = &. fi || x 14 ||
First, we show that o5 O(w w w -) as follows. We compute
(w wh.) of ( ,el 11 O((w (w) ( |g)0||||1 (w., Ex) Σ ww Re: E -2)
= = wh 1|2 + Ō + noooo (a
O w 13)
where the second equality is by the assumption that (w ret: M3), (w..,, {.) = O(1) and the last equality by the assumption (wh. wha) = Mminnong in 1/2 m 1/6 oo /1m6 and the condition Then based on the assumption, that oF (4d Ōmm and d 2 Ö 1/6 of whaa)l Wl (l py) and similarly can still bound (w 2y) 2 of kww. الح 0((w T.t+ what
bound (w (w Ee) 2 (w Es)l @((w n.e. w - Ed). This allows to simplify
w rit = O((w,, what || 3) + (ollw.e + p- --))l|
+ + Ex) Σ E-11Ei11-
(34)
Consequently, the assumption that (w wh.) O @(ww wh,))) combined with (34), we can
derive that tº = ( given (w p4)) O ((ww, (1))) and (w .. {i) = O((w,, {.) Thus, we can compute
[/w (1 + (wh. 14) || nSNR? (w E,) 61/2110011 + @(Ilw
In addition, we can derive for r of *
(wit whi
= -O((o)) (w/,, whr + (wh (w 14)) 13)0 uSNRR?(W: (w) & (0)1|
which completes the proof.
54

---

Published as a conference paper at ICLR 2025
Lemma D.7. Let T 2 Ti and suppose for all Ti k < Tí. it satisfies that
for all j +1, E [n].r E [m]) (whtt rt (4)) (wh+1 ..)) Ö(1), (wh. w w.))) A(min n
1/ m- 1/6 116 2ov dm ddm-//6) and
(wrx" (ej) = -(w rst: M3) + θ
7
(35)
vm
(wit) Eii = (w Et) + θ
"
||w 1111_:11
(36)
mym m
(witt k+l whtt) = (w wi.e)
+ 11 0
((wW,, Mg + + - /mml // (w what llw (l My + E)) (37)
Vm
Then we have for all Tj S k $ T
(1) (w r.t m4) = @( :h , Mp'
(2) (whee: Ei) ((w ri:: Ev))
(3) (wh. Fits w re) @((w). r(, wo.a))
(4) (wh
Fits 14) 2 @((w no wh.)) (t 1z))),
(w r,b. E:) 2 @(w r,0 whet // 1 (w 6x) ).
(5) |w 112 +[ 1?)
(6) (w w @ [[ (1) for *r
(7) (ww, v. (3) //)l/((w (.)l = O(n SNR?)
for all j = r E [m], [n].
Proof of Lemma D.7. The proof is by induction. First, when k TT. claims (1-8) are satisfied by Lemma D.5 with (wi)) wh.) = O(aād), (w?) Ilwh a. ||- O(1). Now suppose there exists
T < T such that for all Ti S k S T.t (1-6) are satisfied. We aim to show for it is also satisfied for k + I. By the assumption that for any r E [m]]
(wA+) 14) (w r. M4) + θ
w 1 (11111))))))
(wht) Ei) = (w" (w,,, E.) + e
1
ny m
(wht) wha) = (W w..))
@ ((((((())
((ww r. (: Ei) / (w wh.) + [/ (w (W ui + Eci
+ vm
we can show
(wht h+1 pej) = (w/.t. 143) + 0
11 / / (11e11111)111) ( + vm mm
(((btt) 11::))
where the second equality is by induction condition, thus verifying the induction for claim (1).
Similarly, we can use the same argument for verifying claim (2). For the claim (3)
(witt) wh.)) = -wwf why) + T _0( (((wh, Mi + E) - /m wht w.:,l?) (w whee + (.f::ll?(w?, (W:.: 14 E))
Vm
O((w wh
55

---

Published as a conference paper at ICLR 2025
+ 77 0
(( (t E Vmllww/.ll (w wprre + ll (all (wll re Me + E)
vm
@(ww
h+1 r's (P.)))
where the second equality is due to induction claim that (w n,r: w whe) = 0 ((W>,,, wp(,)) (wh wie By <) = @((w , M4j + E) and we can show that @((w), r.:: v)) = (w-, v) holds for any v, and any r' E [m]] with constant probability due to m = (1). Next, we verify claim (7)
.))
1++1 pi Ex)
143) (w 7. (4))
2 (we E.) + θ
O(n SNRR)
12
where the last equality follows from the induction condition and || @(SNR²) by Lemma
B.2 and w I| 2 = @[] ,.ll") by induction condition. Thus the induction for (7) is verified.
Next in order to verify (4), we only need to show the growth of (w (4)). (w E:) is larger than the growth of (w wh. ? ( this end, we consider upper bounding the update of (wh, whee)
([w...+ (i) respectively. To (4y) and (w",., w [/w 1o 3
kw e M4 + E) y.e: w - (lI 1 (11aa1) vd)
(w
[|l (l M4 + E) s @(Il (11 + og well.
Then we consider two cases depending on the magnitude of |ا| and og vd:
When || 2 oe ie.,
Then
(1 + m)vd))l | 1: M5Y1 = Ö( k F.z
(2d-1/3) S @ (/ll |||||]]]) (38)
|w ( + ddll ((w {:) = Ō(llw (l/))))=== - [1 |E:11)) 13) (39)
where we use the condition on d = ( and ||&/|| ((o for (39).
When || S Vdv i.e, we have d 22)1) Then
w.w.,, ([llllll + oevd)ll "K
= = oe (4))) = "(l = 1(llw,,,.'dd--/)3 s O(Iw" (1)l(()))))
(40)
1121111 + of aevdd) w - Ei)l
= = l (e,. E.)) = Ō( 1l2l Vo
= -[_1| 11,,11111:1 (2))
(41)
where the second last equality of (40) is by the condition that SNR 1 = Ö(m) which implies that og = 1/2)) The second last inequality of (41) is by d = {(n2)).
This suggests that
(w (w wo w < ([[ (1luu1))))))
[Iw. (W Mj 1 < θ | || |||:: 112)
which verifies the claim (4) by combining with the update (35), (36), (37).
Next, in order to verify (5,6), we leverage Lemma D.6. First, it is easy verify that at k + 1, the conditions for Lemma D.6 are satisfied by the induction claims (1-4) at k + 1. Then we have
[ 2 = O ((w 1+1 . - + n SNR² (whtt [ 2 + (wh, w w. 212 (((W/t) 11; |||||| SNR (wit) |||l |||||| (wh wh 32 [w -2)
56

---

Published as a conference paper at ICLR 2025
= [[ (!
Finally, to verify (4) for l + 1, we have from Lemma D.6 that (wize, w./t))
(w.. w
n- = -((w (4) (w k+1 (4511a1) 2 + SNR* (w 1ot1 (.) (W k+1 lllll - + (witt we1 w wo,))
((w m4) | n SSNR , |1|1 + whit w (wo Tits whet)
I| /|
= =((IwA 112 (w w w [ 21 + (wir)) w 1 (wh r. wP
[w ww
=
where we use the induction claims (1-2) for + and Lemma D.1. The last equality is by (wit1, who
w W[ & @(] F.t 11? which completes all the induction.
From Lemma D.5 and Lemma D.7, know that for Ti S ls s TT we can decompose the gradient into two parts, the dominant term and the residual term:
I(W.) ),n)) =
+ (I (rcll//1lea)) + Pirtas
(42)
mm
(V,, ((W n/m
+
( [1.,,11111 + F F
(43)
(V w... I(W?), wh.))
0 ((iwtt (/ 5) /mllwf (w", what + ll (W My + E) + E E
mm
(44)
where let Er, Ers,,, E E Ek w. denote the residual terms. Therefore, before Ex,,,,, Ek Dr
grow to reach Dr. = O vm ww (l)),, ER = O m(mm w ( 1% 1. = n.c. wha) + ||w (wh Hy + E)), it can be veri- fied that (35), (36), (37) are satisfied respectively. If further, (wh,,, My)) (wf. Ei) = Ō(1) 22min a -1 1n -10. vooommm1/)) are satisfied, then we readily have kw a My)l/WWW &d)l O(n SNR?) by Lemma D.7.
v. O (an My + E) /m r,e "ll,,ll") (w,,
(wo, ri:: what
The next lemma characterizes the end of second stage where the residual term reaches the same order as the dominant term.
Lemma D.8 (Restatement of Lemma 4.2). Consider the gradient decomposition defined in (42), (43) and (44). There exists T2 Ti with T2 ((max(1) -Imm mm000000 2d 1 . m m oz)) such that for all j = =1,r E [m], E [m).
(I) Ifn -NR2 = 2(1),
(w" = (m (m10,, ( = (((w Hz)). (w w w.W.ll < 0( SNR2 (N (w.T..
211).
(w mj) = O(n SNR2 m -1/6),
(w (w E.) = 10000mmmmp///)
(w?3. what) I/ -1 < @(V SNR? (.)
(2) T,,,
o 0 11? F.i
O @ mm) /
E.a
O((w. w. M4 (w: wh.)) + /|l (wf.e: Mo + E.)).
mm
57

---

Published as a conference paper at ICLR 2025
In addition. for any Ti S k & T2 and all j = =1, [m], i [n].
(3) (w (a) = (( Hp)) and (w E.) = (( .. Er)),
(4) [I @[I , .113) and (wf,,, w = e [[[ (1
(5) (w M53/ ( E.) = O(n SNR²).
Proof of Lemma D.8. Here we let Tz be the first time such that 1, = 0 ( | 1 ام or Er.... = 0 1
mym w || or E Eyst.
J ((((" r.t. M4 + E) mi // (| w +
[ (Wr, My + E)). In order to prove the results, we use induction k to simultancously prove the following conditions (k), (k), & (k)) (T2), for T1 < k Tzy
of (k)) (w (4)) (w (.) = Ö(1, (w what) Q0minfooo 1/3 mm1mm-116, aoodmm / for all =1, [m]. i [n].
1/6,
(k)) (Vw... ((W)), m2)) = @
11 (Vw., I(W?), Ei) = 0 -
1 |w. (115:1) and (Vwr. I (WW), wit) = "mm (
O(w p4j + (W whit)
nym
(wf, Mo + E)) for all j = =1,re E [m].iE [n]]
(kk)): Claims (3-5), i.e., (w 14)) = O((w es (yy)), (wh, E.) ( Ev)), ||
e[IIw), I13),, (w., e wh.)) (/ll and (w 14)) (w (. E.) = O(n SNRR).
(T2): Claim (1), i.e., n SNR² (1)), (W.t, (2) 10000mm///))) (wit Es) O(n
SNR 2 - -1/6) and n SNR -2 2(1), (w..) 14)) = -Onn- SOR--NRR-... -m -166,, (Wi.) Ex)
Omm-1/))
& (T2): Claim (2), i.e., p,,,,, = 0 vm
ET = 0
( mm 7 || 11²), and
T
(((( To vm
r.t. M4 + - - m
w To
To (w
P,2 w
+
[w To
(wwo
(ww
T.tx H + E)
r,t.u
The initial conditions of (Ti) (T,). 6 (Ti) are satisfied by Lemma D.5 at the end of the first stage. In order to show (kk), (Tx), 6 (T2), we show the following claims respectively.
Claim D.1. sf (k), (k) = (k), for any Ti S k S T2.
Claim D.2. 6 (Ti), ..... (Tp) (T2),, (T2).
Claim D.3. (T2), 8 (Ty) = o (Ti), of (T)) H (T,),..,
Proof of Claim D.1. Claim D.I directly follows from Lemma D.7.
Proof of Claim D.2. First, when ((k) is satisfied, can simplify |/w from Lemma D.6
[/w,,l? = @((wf, |1)) + SNR² (. (1)) (wh, ? 2?)) *((n2NNR* nSNR |A|| -2 (wwh.. Ei) + (wh w (ll w.lllll? @((x² + ? E.) + vh.)
where we temporarily denote x == n SNR2 and w/th = (wit wh. Wl -l for notation clarity.
Then, for the update of (w Mt)), we can compute
( 145)
2n 1=1
@(w k r.es 14,)0 + (w. (4,) W/ + (w.;, (4y) wl Vmn Hy)?
m
58

---

Published as a conference paper at ICLR 2025
1 @((w r,81 |1 + (w mg)) l| lllll wr, 11al1 - /mall "r.",, (12,011))))))
O (w (1) + (x + x ( ||l|| + (x + x= xx=+
E,) ||ell vmxx (w Ei)
((x" + (w 5a) (x* + ( <;) |||1 -2 vml(2 2 + x)(w Ei)
+ m
O(D 2 ( 41) u1l (wi
(.)1 12- /mmott,
m
O vml(22 x) E.)² (x + E1) (x" + x²) ) Ex) ||
m
+ ( ( 11 + (wh.)) x( 1 -
/mmss (1)))
Similarly, we obtain
[VL (w.)). my)
_ @(w E1)* (w W Uz) + (w Ei) (3) ms)ll (w (e) wwll
Vm(w Ei) p1)))
(x(w 41) + (x 3 (W 11 -2 (x + x²) Ei) |lpal
1 0 4:) x( (w Ei)
mmx(w n. 2) +
m
L e - v Ei)² x(w 6.) (x* + (w [ -2 - (x + (w (:)
m
O(xus ( E,)a + (2 Ei)
m
1
V
(1) 2,2
(Wr.),,
2n
1=1
@((w : 23)8 + (w wh 14)) [ // 11? + (w h. (4) w/, (w p1)ll 1l
m
+ (w :: mو) || - 1111)))
m
((x2 + x () 6,) + (x + x (w Ex) ||ell -2
m
1
O (ww 4. plll + (utte) xx(w 1i)11 /mneth 1111122)
+ m
1
2,1 ( (Wr.)). 143)
2n
1=1
((w (b) (w:, Ei)" (wh. www + (w (4y (w E1) 11 112)
m
m 0 (xiw 6,) + (x2 + +x²) (w 61) 1| (xx x² ( (1) [1
m
+ m xxx (w 6.)3 + x( (Wff Ei)
Combining the above results, we have
1 (V wr, L(W!) (z) Vm 0 (1111111111))) + E
(1
0(xx _-__ + +X) Ex) w wh 11al(2)
59

---

Published as a conference paper at ICLR 2025
+ 0 ((x + +x Xw Ei) + (x" + X 11al1
+ 0 (x 61) (x2 + (w r. E1) ||lll 2 (x + (W Ei) Ilall -1)
+ O (W 61) ||all/ + (1e ) X Gi)llll + (w Si)
Similarly, derive for the update of (w Ei) as follows:
2n
1=1
1
(w =((ww) ( M3)) (w 7:0 Ei) + (e (W Ei) + +ll))++www,,,, [lw. Ei)
m
mm r.t (w 1. Ei)
0(x m
(x* + 2 (w |1 -2 (x* (.) vmx(w Ei)
Ei) (w E.)
O Vmx(w"..cm E1)² (w E1) + (x* + x²) ( E1) + (x2 + +x²) (x+a
+ m
O ( <x) (w E.)
(2) (w: E.)
O (w : 5i) (w 51)% whlll [/wh (w Ei) Vm(ww. {:)²)
(w 61) || + [| (1
( 1.)16/112 - Vmllw 1. 1111)111))
+ nm
e (wwr 6i) (x2 + (w |l1l (x )) ) ||lll| - vm(ww, {3)²)
m
1 0 (xx2 x) 4 (x* + ( {,) ||||||||||| vmlx² x) Et)P)
xm
O (e 63)3 + (wi.) (w E.) W ( Ex) $1000 + x (wttt) (w &illl
+ m
@ Vmix + () 2 + (x2 + x)(w 7.
Ei) (0) 1111 -2 + (x* + x²) 1111
(x + w) r. Ei)s (x2 + () Ex) 111111111
+ 0 p (w r. E1) 3 + (w (2 {.) +
(w ||0|2 + (xsk, (w (1)11
x V 11
where the second equality follows from Ev= (E)) = (1 ((nddd//))lle:1 ((1ls/11) by Lemma B.2 and condition on d. Further,
x
( 2,3'
w..)). Es)
2n
θ (w f. (1,) (wh w (i) + [| (w Ex) + (w ?lw? f/wl (weee E.)
m
m
0 7. 62) + (x* + x2 ( [|]| (x2 + x ( Ex) 11101|11111-
m
0 0 Ex) (with ( (w Ei)
m
60

---

Published as a conference paper at ICLR 2025
(2) 2,2
(w..):: Et)
1=11
m
θ (w r.t+ 5:) + ll (( 2,81 E.) + (w E1) l1 //
m
3 (w r,t, (1) 1|w 1 || + //w" (w 6t) 1
mm
11-1 -(w E1) + (x" + ( |||||| 6x2 + x( Ex) 11lllll1
m
(x* + ( r.: Ei) ||ea| -2) - @((x (w w 61)
xm
1 1 (w 6,)5 + ( (w Ex) + x w (w ||al + x (vbte)) Ex) 1u11
m
"Vmut.llel??
@((x² + (w 11)l11=111 [|]laall+++ (x* + +x²) (w [1)_= (x 1) Et)
m
+ (xx + x) 4:) 21
m - 1 θ (v (w {i) + (57,)) (w Ei) + +x (w, + x (( Ex) ||1l
X Vmmpt =1
Combining the above results, have
1 O - (121/11))))))1 + E
(V w... I(W.). Ei)
n/m
Vm o (((x1)((w Ei) + + whilleel13)
+ 0 x (x 5:) (x' + x2) /w n. {x) ||1|| (x" + (w 11)
+ 0 0(x2 + x)(w & ||ll|| + (x* xxx2) (w,, "1all1+
+ 0 (x + 1) (w 4 + (x3 + (w Ex) |laa| [1xa)1-))) 22))
((x
+ θ bh (wh 6,) + (1btt)) ( (w {i) + wht (w E.) + x tus (W ries [.)
+ (W/ r. <.) (111))
(ww
In summary, we finally arrive at
1 (Vw.,, L(W)) W.)). M3) vm 0 1ll (1/llall)) + Ei
Vm Lo((x (W f,, Ei) wh. (1ul))
+0((x² + 6.) + (x + x (w :: {;)& [401 2
+ (x( (x3 + x3 (w (,) ||l||||||| 22 + (x" + (w Ei) 1||pl|lll 4)
+ 0 ( Ex) |||1|1 2 + (witt)) xx (w Ei) ||ll + xush (w
(45)
0 --
1
V w., ((W.). Ex)
E.
m/m""
0 (x 1)(w 41) + x - (1
+ 0 xxx((w" (:) (x* + (w {5) =|||| -2 + (x* + x E,) =1121
+ 0 0((22 + xw |1111 - (x* + (w 1/)5/11a|]1+
61

---

Published as a conference paper at ICLR 2025
+ θ
(x ((x
1)(w E.) + 6x" + x) Et) 110111111
+ θ b
Vrxex" (w rif. 6i) + (1btee) ( Si) + (w E.) + x - W r.ts (.)%%
+ x (wtt.e) P(w r::1 S.) 1
(46)
We also examine (V wr. L(W?). wheh follows. We first upper bound for r' *
(w r, wP .)) = O((w", w / (w whe @((w (y) ",
+ @(w r, (1) Σ Pl 6.)
((a) + + + O ' 1/2)
=(0o) + O'noro 1-1122
where we use (34) in the first equality and Lemma D.I, Lemma D.7 in the second equality. Next we simplify the gradient as
(V... w.., L( L(W)), wh.))
θ (w r,t. M4 + Ei) (w whyt + (ww (w. Ma + Eii)
Vm
θ
((W n.07 p2y)* + (Wh.es (x)") (w what ((wf, 1ay) (.),)))))))) Illl)llw [/feell (w/ whe))
m
θ
| [w (w F. what + ((w".e M4) (w. M4) + (w ( (i)))
θ
((w (4y) (wh pe)) + (wh, Ei) (wh, wr E. Ilwllllllll
+ hhe 14))* + (w (w..e.s)) 63) 4 + [[w (1 14;) + (wh.. Ei)
m
+ ||w/tall 11 Ō(oo + nooooa -a 1/2)
θ ((w 7,6s 24 (w M4)) + (w" {.) 2 (w... {.) ll
m
θ
((w r,61 (4) (w.) r, Mb) + (w (i)(w) fe {i) / f I1*)
m
m
θ
(w r.t1 H2)) + (w Ei)?) [l w. (w. whee))
m
m
+
w wlll (w rees whi))
m
O (w. ye pej + Ee) (w w + ll (we Hej Ei)
mm
+ 0 (wi. /4/)* + (wht (i)* (ww what
+ ((Ww 21 2 (wh.r E1) w 112 + llw (1() (w e whot
+
((ww, (4 (we..: m4)) (w.. (.) (w" Exd l/ IP) θ ((w/: (1g) (wo R.ex 113) + (wf, Ex) (w (.)) (.)) ww
(47)
where use that (w w 2 @oofd) 2 Ö(oo + nooo 1d21 due to the condition that
Go 2 Omaxx(no 1d d)-)).
In order to identify the dominant terms of (45) and (46), we separate the analysis for three cases depending on the scale n SNR For each case, we also consider two sub-cases depending on the scale of (we
( Ex). Recall that we define wit (wk w 2 wf.ll
62

---

Published as a conference paper at ICLR 2025
When X x= SNR" (1)),
If (I 61)² @(vit,))) we can identify the dominant terms as
1
1
| ||all
O (w E,)
Eniips =(( E1)
v Vm
vmm
| 111,11
((w r,t Ei)?), E = @((w Ei)
nym m
Hence, we can see at k T2, when (w Ei) O(m-1/0) and (w
(xm (m 1/G) (by the condition of x), we have Enins O V. | 22 ||| and E ( (//m//l /| 1
It remains to show that E r, = V (((W j4j + - vml mm/lw
wh.t) +
| / ( f....P My + E)). We compute that w = 6:)² + + = 6 1/33)
which implies
((ww? r,6 Mi + vm/ll (w 23 r, wie + || (WP.
(w 14 + E)
O( 1/0 Wr wh.) m 1/3 (wo W, M4 + +
vm
= (mm 2/3 (w r,t: w(t) + m 5/6 (wo p. Mi + E))
-ETT -
based on the derivation in (47). In this case, we can show with I (w",, E:)² = mmm31/3)
feasible given the lower bound on (w r.:: wha) = Di -166 in ol (k)). This verifies
(T2),, & (T2) in the case where (w. r, Ex)² 2 (W.,)) - If (w =2 = ofv))))) we show (Vw...L L(W)), wha) # 0 and thus cannot converge.
More specifically, we can identify the dominant terms as
11l-ulll=--- (whte) |/
|111e1122
E Franu = (( (w Ei)
vm
Vm 1
1
E - O((V),)) )2(())( E.)).
Lo(tt))
nymm"
Vm
Thus, we see Fie F.r.t.ps
O 1 || and F. O( mmmm | || only when (w",,: Ei) = 0mm -1) (w (y)) This implies that wy (mm-1/1)))))))))
and thus l/w/flll
O(v).) 2 0m We show in this case, the gradient along direction wh, is negative. Examining (V.... L(WE), whi)., we see
((ww (Ww My + E /mil| (w wh.) + l/ (l ( w.:: Hy + +E))
(0(m- 11/2 ( - /m () whet + O (w Whe 14j + E
Vm
Omm -( ) (w. wh.) - (v+t,)² (w wh.)) + m 1/2 w wh (w.. M4 + E))
E& E O((m -3 (VA,.) -4 +m +m (vbt,)) ( + (utt, )2) (ww w
+ -0(m 3/2 (1/k:) -2 + m 1/2 ( By + t))
= = 3 (w . wh.t + 3mmm1/2 1/2 b/2 (w. r. ,b + E))
where we use that w/dk 2 O(m O(m-1/3). It is clear that that due to the condition wht 2 (m/33,) it satisfies (Vw.., L( LWW)), wf.)) < O, which then concludes that w/tk would decrease and cannot reach stationary point.
When X n 2 Ā(1).
63

---

Published as a conference paper at ICLR 2025
If x2 (w 7, = (WW).)) can simplify (45) and (46) to
1 Vm /w/ = 1
( E
(
E,i,, +xx" (w E,)
Vm
Bym | 115:11 = Vm Vm (( E)) E =0(x* (w Ei)
1
Hence, we can see at l = T2, when (w Th Ei) = 0(x - 'mm 1/6) and thus (w H4y) =
(mm=//0), we have Ey,,,n O mm wrilllla|| and Eniss O mm ?| [E:1)))
Next we show ET2 rd
7m θ ((( My + E) vmillw [3 (w wi.)
[w ( r fyy + E)). We first compute that [[ = @(x2(w7)) Ei) + wit))
Om 1/5), which implies
1
((w77) Th r, pej + Vm//l 7il (
3)
T3 r, w + [|w (w My + +E)
O(m 1/6 (w y w wh + m mm13////--
1/3 (w.. ( Mj + E))
Vm
= (m 2/3 (w (w..:
r.t: w..) + m 5/6 (w - M4j + E)
=ET
based on the derivation in (47). In this case, we can show e/th S xx2(wx E:) 1m is feasible given the lower bound on (w what = R(ooo dm 1/00 in of (k). This verifies
(T2)), 6 (T2) in the case where x² (w p,es 6.) 2 O(v
If ( E2)
o(v)t1), we can follow similar argument to show that (Vw,, L I(W)), w wha) * Specifically, we can simplify (45) and (46) to
1
1
wolllllllll =
vm (wtte) Prim = -(( xx( r,c: Ei)
mm
1
"
0(x-1xt)
P @((v:,, 2 E:))
nym
Vm
Thus, we see ERan, = 0 Vm / 1111) and E = 0 1 |
p.m r,t 12 only when x(w 7.t. Et) = ((()).)) ''mm 1/22,, which implies v.A. 2 O 1/3). However, by similar argument as the case n SNR² = @(1), we can show (V.... w.., L(W)), wha) < O, which then concludes that v would decrease and cannot reach stationary point.
When x =n=-RR - == =Ô(1)
If 2 O(xx=1p (1r,))) we can simplify (45) and (46) into
1
l/w | = 1
( E,)?),
Ehaws @(((w)) 3)
Vm
mm
m wwt..1111llll
mm @((w
r.:: {; ERae: = O((w Ei);)
Hence, we can see at k = T2, when (w (w Ei) Om (1/6) and (w To mوم = Ox 1/6).
we have ETT = O( Amm || ?) and ETa O // ||w 2 11) In this case, we can show 4/, < x(w E,)² is feasible given the lower bound on (wh: wha) 20ro0=" 122mmmmmm11)) in (k) when x = Ā(1).
1
Next we can check that at Th. max ( (w Th Hg), (w (.) = (wT3. E.) = 90mm1/00 Further,
WWlll = (((Wittn))) + (W?3, (4y) (w Ti E.) + (WZ), wh.) w. = O 01/3-1133
64

---

Published as a conference paper at ICLR 2025
where the second equality is by (w..). 23 wh.
// 12 23= , @ 1/3). This leads to
1
((w p.t3 Mij + E
w.,) w ?. (w My + F)
Vm
= (m 2/3 (w r.t w + xm
5/6 (w r,t: رíj + E))
T r,d,w0
6 6((m 11 + x)m 2/3
To r.ts
r.1)
||
+ xm
5/6 (w r,c: Méj + E)
where ET² o, /mn
1
((((((1)) ((w
r. M4
m | || (w Ta w +
Mj + +5)
holds due to m = =1 This verifies 2 (Tz), & (T2).
If (w 1 61)² < x0xx (bxx),,,) :/,e), we can simplify (45) and (46) into
| |wtllllal|? =- 1
(vvr,),
Ext, @(( 2 (w 6.)x + fffeww Er) x)
v vm
vm
|| 1|/.11
(x E = 0(x - ( E1) + x 2 EX)
mym
vm
If (w
6,)² 2 O Ovvk.))) the equalities become
1
(v..e):
Exttan = =V,,, E.) <,))²x)
m
vm 1
1
|w ||
E -1i (w -.. E))
nym
vm
and E Ei,t, 0( Vin
when (wf, Ei) 9mm-116x-11 and E P.
O mym w 11? when (w E.) 0(m-1/0) Due to the scale that x (2(1).
two equalities cannot hold at the same time.
If (w 4, (vvh)))) the equalities become
| |ll =
Vm (vtt.)
Exttay = ((V+t,)) (ww Ex)x)
m
nmm || (w2 1//1/ vm O wWtt),
F (v (w E()
Thus Ehan = 0 4mm (/w (1)) when (wh,, E.) = (v+t)) 1) and D
0 ( || 12211, when (w/ Ei) = (((U)))))) which clearly cannot be satis- fied at the same time given x-= ((1)).
To conclude, verify that when (wf,,, E.) < O(x Ex,,,n O( (m / [ffelllaa() || cannot be satisfied.
or E,,.s. = 1 ay
In summary, we obtain the following results:
When n SNR² = @()), we can show E
0 ( Vm
|w (( and Ek Eh
O m
w (.lll when (w Ei) (mm-1/6)
When n SNR? ((()), we can show Ek Erst, O
1
w and Et Ehe.e.
vm
O( 1 w w 112 when (w Ex) O(x m 1/6).
hym
When -1 SNR-= 2 (()), we can show EY,,,
O
(ulull) and E.
vm
0 1 nym
[1 I| when (w Ex) (m/-6.).
Combining the definition of T2, we complete the proof for Claim D.2.
65

---

Published as a conference paper at ICLR 2025
Proof of Claim D.3. By the definition of T2 and Lemma D.5, we know that for all Tj < k S T2, the gradients can be written as
(V w... I(W = 0
Vm ll rell (1
1
LwW.,.W))...... = ( nymm Iw (11:1111)/)13)
and thus (35), (36) are satisfied, which verifies (k). In addition, this suggests, for all Ti I 7,to k S T2, the increase in (w Ei) (w (,,t,/))) (45) is monotonic. Combining with (T2),, we have
(wf, My), (w",, E.) = =0 (00)= = ((1 for all Ti s le $ T2, thus verifying of (k))
Hence, the proof completes the induction (k), (k), ((T)),((I)) Ti < k T2.
k and verify the claims
Finally, we derive upper bound on T2. Because for all Ti S k S T2, we can decompose l! from Lemma D.4 as
||w/,ell @( (b) || -2 + x(w n. E.) | 2 + | 11] ((ggd)
Therefore, we can upper bound the update in (35), (36) for Tj - k - T2 by a liner growth as
(witt) (85) = (w r. 14)) + 0
| 1 2 (w M.) + θ m 1/2 1/200/)
vm
(w/t), .Ei) = (wh, yir: Ei) +
| [w (11:)1 2 (w Ei- + θ 17X 11x mm -//2
nn/m m
, Therefore, we can upper bound as T3 v (m 2d ,,- m 13n00==200000 20
D.5 STATIONARY POINT
This section analyzes the stationary point with the conditions at the end of the second stage.
Theorem D.1. Under Condition D.I, suppose (1) (w M₂) ((
Ö(1. (2)
(wi,e: {i) = O((w,,e (v)) ((:Eee))) = Ö(1). (3) l|w = O ((l and (4) (wi, w = (l] (1,) (5) (w wh.)) O((W,, wx,))) = [[minn(ooo - 2n2//2 1/6 Doodm -1/o]) hold
for all j = =1, E [m], E [m]] Then there exists a stationary point W;, iee Vw.,I(W L(W;) = 0
that satisfies
Iw:,,, (w,,, -I -SNR))
with (w;, Ei) O(n SNR m 1/6), (Wi,e) wh.a)
(..l)//w wh S O((w", Hy)) if n
SNR2 2 R(1), and (wi, Ex) (mm)), (wi,,s wh. w @(Vn SNR? SNR? E()
nn -1. SNR = (1).
Proof of Theorem D. The analysis mostly follows from Lemma D.8. Due to the concentration of neurons, we can derive
WW.., N...L((i)..))
1
vm @ ((w + llw,,, 11l + (w,,, Ex) (W (4)))
+ ((Wits My))" + (w ( "| + [[wfell (W 14y) + (WWres (ay)ersy))) 24,) "1w 1 || 21?
+ +Ilw (WS 11 (W Ex) (wi, M5) + (WG,)) (e)"ll well (I (Wires
V w.., Sw..L(W77..k.
- (W., (4)))))) (wi,es Ex) + (Wi, E.) + - [l
-
Vm @
+ +0((; (Wi.e) 5:) + (wi,es 41) /w, + (WF u4)) (wi,es es Ei) + (WF u4y) [Il E.)
66

---

Published as a conference paper at ICLR 2025
11 + /| "ll" Ei) (wi,e Ex) Ilw 11 + 1 - (w
W... L(W)), w9,))
0 (w te My + Ex) (wi,s) whet + || (wo Ex)
vm
@ (Wi M4)) + (w Ei)) (w what + 0 ((w,, 1e)) + (W Ei) |/ (w whe)
m
+ θ
[ (w wit + (( H4) (w H4) + (wi, ( (w Es) [/,3)
+ +
((W (WW My) 1, 14y) + (Wi,e: Ei) (w 6:) (well"
And we can verify that W; is stationary point if and only if for all j = =1,r E [m]. i E [n]. (Vw., L(W7), = (V (V.,, L(Wi), Ei) = (V (Vw.. L(WI), w = O. This leads to the following equation system:
Vm (wW p2)) + (wiss M3) WI !l? + (w Ei) (w M4) + [| (o (w.,,)
+ (Wr 61) wT, (w (w 15) + (Wa, (3) I| /i. 11l,, [[w (W (4) 11)1)
0 ((wier Ei) M3) + (wh pey) + W[ (11111))
(48)
(w Vml
6,)3 + (wi,s 6,) we, + (w (4) (w Ei) + all
+ (wi,, Mg) [W, ( (W Ei) + (W,, Ei) Il 1| [IW,,e (w ( 6))1):
o ((w (4))( (wi, Ei) + (wi, Ex)²
(49)
v 2 ((w) pay)* + (wh, Ex) (Wia) wi.)) + ((w 14) + (wi, Ei) Iwr PR
wi.))
+ wWi,ell" ( what + ((w), e (4)) (w (wh.: (4) + (wi, 6,) (w/ (:))
+ (Ww pay) (w ao mg) + (wf 4.) (w w. E.) [|
= (W My + Ez) wh.) [[ ( (whe My + Ei)
(50)
(w)) In order to solve the system, we let Ti. = (1 for any i E [m],] = +1. We let 7 = O(T..))
We first consider solving (48) and (49) and then analyze (50).
Furthermore, because the claims (I-4) are assumed, we can leverage Lemma D.4 to decompose
|| 2| = ((w",, /2)) |||00 + SNR (w", E.) (1) + (w",,: wh.) 2l -1 = O(77² +n SNR*) (wi,,, E.) [|ml] + (Wi,es wh.) // //w.elll
where the third equality is by the scale of (wi,e> (15) and (w", Et).
Next, we separately consider three SNR conditions, namely (1) n- SNR" n 2(1): and (3) - SNR-- 52(1)
1. When n 2SNR2 = O(1) we first can bound (w"..: wh.) (/ || 1l
w.. 11 & O((wi, Mz)) and derive
u))) ((w ( ||
Next, we can simplify (48) and (49) depending on the scale of
When r r= Ō(()), we have |/, = O((w), (451 M,3l and the equations reduce to ( (1) = O(7 (W) (1)
10(Vmr {1) ((72 9)
It is clear for = 5((1, the equations cannot be jointly satisfied.
67

---

Published as a conference paper at ICLR 2025
When -- == ((11,, have l/ 1l? (((w,,, (1) ||= and the equations reduce to
O @ (w (,) = (ww,,, E1)
(@(ym(wr... (,) @((w E.)
which cannot be satisfied simultaneously f = (2(1). When ((1), //w,,ell @( 14 (u -2 = ((w ( E.) (1) and thus
can simplify the equations to
[0(mm(w :: E1 O((wi,c E)) @ 5 O((wi,e: {1)
which has solution with (w {. @ (wi,,) uz), thus verifying the scale and (n SNRR). Then we can verify (50) holds under the scale of (w Ei), (W M4) =
(m 1/0) and /|w", = Om 1/3). With the same argument as in Lemma D.8, we can show (V... w... L(W)), wha 0.
2. When SNNR (11): we first can bound (wi,,, w., w -o S (((w (z)) and derive || 1| = max (m(x()(()))))))))))))))))()))) 0(SSNR ) E1> 11111
We only consider the scale when r = On SNR?), where we can simplify ||w;,,l?? ((7")(W) 2 and thus the system becomes
{
[ 5 = 0(7 {i) (W/ 61) = O(T(W {1)²)
order to satisfy both equations, we require (wi, E:) 0 1/6) and (w M4) (mm which verifies the scale. We can then verify (50) holds under such condition. With the same argument in Lemma D.8, we can show (Vw., I(W:), what =0.
3. When n SNR @(vn SNR? (wi,,r {i) and derive
2(1):
we first can bound (wie
S
|| max max(O(77) ((NSNN"))WW 4:) 1,,
We only consider the scale when r O(n SNR where we can simplify || = O(n SNR3) & 1,1l= and thus the system becomes
0 @(VmT(ww E,) (nSNN (w E.)
@0(m((;; 5 O((w,, E))
which can be satisfied when (w".., E.) 0/m0 and (w My) 10mm00(mm))) and thus verify the scale. In this case, can also verify that (50) holds under the condition that O(1). With the same argument as in Lemma D.S, we can show WV...L(W) L(W.), w = 0.
This concludes the proof that suppose the scales and concentration are the same as the end of second stage, then there exists stationary point where (w",, (t)) (w",, E.) O(n SNR?).
68


Reviews:
[{"review_id": "t3bLsezOXf", "reviewer": "Reviewer_nXhk", "summary": "This work presents a novel theoretical framework for analyzing the feature learning dynamics in diffusion models, marking the first such contribution to the existing literature. Through rigorous analysis, the authors demonstrate that diffusion models inherently promote the learning of more balanced features, in contrast to traditional classification methods, which tend to prioritize certain features over others. Additionally, the framework can be adapted to accommodate more complex data settings. This theoretical insight has significant implications for downstream tasks, such as out-of-distribution classification, where only these rare or weak features may be present.", "questions": "1. My main concern is how applicable the theoretical findings are to a broader range of classifiers and diffusion models. To be honest, the shallow neural network function defined in this manuscript appears unsuitable for real-world data.\n2. Is the requirement for orthogonality in data setup common in practice?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The theoretical findings are solid and  innovative.\n2. The contribution appears valuable for comprehending the internal mechanisms of diffusion models and neural networks.", "weaknesses": "1. The presentation requires enhancement. Some notations are introduced without prior definition upon their first occurrence, e.g., what does $m$ refer to in line 188?\n2. The classifier considered in this manuscript lacks sufficient generality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6wWdnaySeK", "reviewer": "Reviewer_VpM3", "summary": "This work study the feature learning mechanism behind the performance of diffusion models in fitting complex distributions and their out-of-distribution transferability. The authors attribute these capabilities to the denoising objective, which results in balanced training and comprehensive representations, in contrast to a classification objective that focuses on specific patterns. A theoretical and empirical analysis are conducted which show that diffusion models (DMs) tend to learn both the signal and noise in a balanced manner, in contrast to supervised classification, which focuses either on signal or noise. These claims are supported using a synthetic dataset and an edited MINTS dataset.", "questions": "1) It is interesting to analyze the learning regimes of diffusion models in comparison to classification objectives. However, it is somehow expected that a generative learning objective (as in diffusion models) results in more balanced learning, unlike the classification objective, which is designed to focus on learning features related to specific labels. A natural question that arises is how DMs behave in comparison to other generative approaches, such as GANs or VAEs. Could the authors discuss this point or provide experimental results?\n\n2) In the diffusion case: Since the neural network input is a noisy version of the data sample, the SNR should be altered as additional noise is introduced. Consequently, the SNR of the neural network input differs between the diffusion model and the classification case. Could you please elaborate on this point?\n\n3) DMs can also be viewed as maximizing the Evidence Lower Bound (ELBO) with a form of data augmentation [1]. Since the neural network input is a noisy version of the data sample, unlike in the classification case, it may be unclear which factors are really influential. Could you elaborate on this point? If data augmentation were applied in the classification case, what behavior would be expected?\n\n4) Redundant content between Sections 1.1 and 2 should be revised.\n\n5) The notation is sometimes difficult to follow in several parts of the paper and needs to be revised:\n\n   * In 1.1, some elements are used without being introduced, such as \\(d\\). I assume \\(d\\) and \\(n\\) refer to vector dimensions, but it would be better not to rely on guesswork. Other variables like \\(r\\), \\(\\gamma\\), and \\(k\\) are first used without any introduction.\n\n   * In Section 2, \\(W\\) and \\(x\\) sometimes have one or two subscripts, which is confusing.\n\n   * Figures 2 and 4 should be improved (Y Axis labels missing ).\n\n6) In Section 5.1, Figure A.1 should be moved to the main text for better cohesion with the section.\n\n[1] Kingma, D., & Gao, R. (2024). Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This work addresses a highly interesting problem: elucidating the underlying mechanisms in diffusion models that lead to their performance and transferability. \nThe authors have done a good job in building a framework where the mechanism and learning dynamics can be analyzed both empirically and theoretically. \nThe problem is well-motivated, and the paper is generally pleasant to read.", "weaknesses": "1) The paper’s objective is to analyze feature learning in diffusion models, which could explain their success. A detailed comparison with supervised classification is provided, however a comparison to other generative approaches such as GANs or VAEs is missing. Therefore, it is unclear if the balanced learning regime is due to DMs' specific denoising based learning, or due to a more abstract or fundamental difference between generative and discriminative approaches.\n\n2) The paper could benefit of more exhaustive experimental setting such a case with larger feature space rather than one label.\n\n3) The organization and the notation of the paper could be improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z7RDbdtUwH", "reviewer": "Reviewer_cG9c", "summary": "The paper introduces a new theoretical framework specifically for analyzing the feature learning dynamics in diffusion models compared to classification models. This is a significant contribution to the existing body of literature, which has primarily focused on the generative aspects of diffusion models.", "questions": "- The theoretical analysis relies on specific assumptions about high-dimensional spaces and signal-to-noise ratios. How would these theoretical insights hold up when applied to real-world datasets like ImageNet or COCO? Is there any plan to test these results on more practical benchmarks?\n- The analysis is based on a simple two-layer convolutional network. If the architecture changes—say, using deeper networks or different types like ResNets or Transformers—would the feature learning dynamics of diffusion models change? Are there any experimental or theoretical results to explore this?\n- The paper suggests that classification models tend to overfit noise when the SNR is low. However, in certain tasks, could learning noise or background features also be beneficial (e.g., in robustness or generalization)? Is there any discussion or analysis of cases where noise learning might be advantageous?\n- Diffusion models are typically slower to train due to their iterative denoising process. The paper highlights the balanced feature learning in diffusion models, but are there any methods to accelerate training while maintaining this advantage? Would techniques like accelerated sampling or reduced time steps affect the feature learning dynamics?\n- Besides classification, how do diffusion models fare in other tasks like object detection, semantic segmentation, or even reinforcement learning? Are there any theoretical or empirical results that support the use of diffusion models in these more complex tasks?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors not only provide a theoretical analysis of the feature learning process in diffusion models but also validate their findings through experiments. This strengthens the credibility and practical relevance of the proposed framework.\n- A key contribution of the work is the comparison between diffusion models and classification models in terms of feature learning dynamics. The results show that diffusion models tend to learn more balanced representations of both signal and noise, which contrasts with the classification models' tendency to focus on one or the other, depending on the signal-to-noise ratio (SNR).\n- The paper provides clear mathematical results (e.g., Theorem 1.1 and Theorem 3.1) that demonstrate key differences between diffusion and classification models, particularly in terms of feature learning speed and final outcomes. These results are well-supported by rigorous proofs and visualization (e.g., Figure 1).", "weaknesses": "- While the theoretical insights are valuable, the paper does not sufficiently discuss how these findings might impact practical applications of diffusion models. The focus is more on theoretical analysis, and less on how these insights could be used to improve real-world tasks like image generation, classification, or other downstream applications.\n- The theoretical results depend on several strict assumptions (e.g., Condition 3.1), particularly regarding the dimensionality and network initialization. In real-world scenarios, these assumptions may not hold, limiting the practical applicability of the theoretical findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dM1hfPWEFh", "reviewer": "Reviewer_T68k", "summary": "This paper provides a theoretical analysis of feature learning in diffusion models and contrasts that with feature learning in classification models. This work finds that diffusion models tend to prioritize learn balanced representations of the data whereas classification tends to prioritize learning easy to learn features of the data. \n\nAs an applied researcher with no experience in the theory behind feature learning or theoretical analysis of neural networks, I honestly can’t explain how I was assigned this paper. As such it is difficult for me to judge the novelty, impact, and soundness of this work.", "questions": "* Would this analysis apply to the training of a classifier on the same levels of noise as a diffusion model? I.e. my question is whether these results apply specifically to diffusion models, or if they also apply to any model trained on all levels of (Gaussian?) noise. \n* Does this analysis also apply to flow-based models used in flow-matching? I would guess that this is likely and it may increase the impact of this work to study more general frameworks.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The motivation and essential results were clear even to a non-expert in this subfield\n* To the best of my knowledge no one has characterized the differences in feature learning in these settings.", "weaknesses": "* I don’t find it particularly surprising that when training on Gaussian noised samples that the features learned would be broader than those trained on clean samples where small details that will be destroyed by Gaussian noise will still appear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s759liSk3p", "reviewer": "Reviewer_2j1t", "summary": "This paper focuses on the theoretical analysis of the feature learning framework of diffusion models. While previous studies have primarily analyzed the generative capabilities of diffusion models, this work examines how diffusion models differ from traditional supervised learning in terms of feature learning. Specifically, they show that, unlike classification models that tend to focus dominantly on either signal or noise, diffusion models learn features that are balanced with respect to signal and noise. They experimentally validate their theoretical claims using synthetic datasets and the MNIST dataset.", "questions": "* Could the authors provide their thoughts on the issues raised in Weaknesses?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* They provide a theoretical analysis of the features learned by diffusion models, a relatively unexplored area. While diffusion models are often expected to have good representational capabilities due to their powerful generative capabilities, they provide a preliminary theoretical grounding for this expectation.\n\n* The comparison between diffusion models and classification models, showing how they differ in feature learning, is interesting and meaningful. This work could contribute to the long-standing research on the discriminative vs. generative perspective in machine learning.", "weaknesses": "* Comparison with other generative models, such as VAE, could be more informative. Diffusion models have unique features, such as the diffusion timesteps, that distinguish them from other generative models; an analysis of how these differences affect feature learning would add significant value.\n* It would be beneficial to discuss how the theoretical results can be applied to real-world applications, even if only briefly, to provide practical relevance.\n* The meaning of formulations could be further explained. For example, in Definition 2.1, it would be helpful to explain the rationale or significance of setting the variance of $\\xi$ as defined.  Additionally, in Condition 3.1 (5), it would be useful to clarify whether the assumption of constant-order $\\alpha_t$ and $\\beta_t$ in diffusion models differs from conventional diffusion model formulations.\n* A detailed explanation of the theoretical differences between diffusion and classification models would be valuable. For example, diffusion models tie two layers, while classification models fix the second layer to ±1; it needs to be explained whether this difference affects the theoretical results.\n* It would be helpful to know whether the theoretical analysis could also be applied to the score matching loss [1] and the reconstruction loss [2] in diffusion models.\n\n[1] Song et al., Score-Based Generative Modeling through Stochastic Differential Equations, ICLR 2021   \n[2] Karras et al., Elucidating the Design Space of Diffusion-Based Generative Models, NeurIPS 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vn2REI8Uon", "reviewer": "Reviewer_Cb9H", "summary": "This paper investigates feature learning for diffusion models through a theoritical approach. Under assumptions of two-class label setup and using certrain two-layer neural networks, the authors reveal an interesting difference between diffusion models and classification models.  Diffusion models learn both signal and noise in a \"balanced way\", where the \"learning\" grows proportionally to $\\text{sample size} \\times SNR^2$. In contrast, classification models tend to focus on certain features and exhist a \"jump\" learning behavior when $\\text{sample size} \\times SNR^2$ exceeds some threshold.\n\nThis paper provides theoritical supports for their findings as well as experiments on synthetical data and MNIST data.", "questions": "1. In Figure 6 of Appendix A.3, the seeting is described as high SNR, but the caption states \"In this low-SNR case, we see classification tends to predominately learn noise while diffusion learns both signal and noise.\" Seems the caption is not consistent with the setting in the figure. Is my understanding correct?\n\n2. Comparing the plot of difuusion model in left and right panel of figue 2, the $n \\, SNR^2$ change significantly from 0.75 to 6.75. However, the ratio between  $max(w_r, \\mu)$ and $max(wr, \\xi)$ does not appear to change as dramatically much. Could you explain how this aligns with your theory results?  What is the motivation behind selecting 0.75 and 6.75 as low and high SNR settings? I am a bit curious if you have experimental results for other SNR values?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper provides a noval perspective to understand the behavior of diffusion model: the feature learning. Defusion model learns feature over noise in a balanced manner. This addresses a gap in existing theoritical study that focuses mainly on distribution learning and convergance. This provides insights on why diffusion model could learn weak features.\n\n2. The paper provides a rigourous theoritical framework and validate its results with numerical experiments. Experiments results clearly illustrate the feature learning patterns. \n\n3. The theoretical results are well-organized and explained.", "weaknesses": "1. Using SNR to demonstrate the feature learning pattern may oversimplify other factors. Real-world data can be heterogeneous, where signal structure and noise can vary significantly, making it unlikely to have a conistent SNR.\n\n2. The assumptions of orthogonal signals and a single feature seem quite strong. Is there any justification provided for these choices?\n\n3. The synthetic data experiment settings are somewhat limited. Including more SNR settings could provide a better illustration of how feature learning behavior changes gradually with different SNR values.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)