U-Nets as Belief Propagation: Efficient Classification, Denoising, and Diffusion in Generative Hierarchical Models
Authors: Song Mei
Venue: ICLR 2025 Poster (2025)
Keywords: diffusion model
denoising
belief propagation
generative hierarchical models
Primary area: learning theory
TL;DR: This paper presents a novel theoretical interpretation of U-Nets, showing how their architecture can efficiently approximate the belief propagation denoising algorithm in generative hierarchical models.

Abstract: U-Nets are among the most widely used architectures in computer vision, renowned for their exceptional performance in applications such as image segmentation, denoising, and diffusion modeling. However, a theoretical explanation of the U-Net architecture design has not yet been fully established. 

This paper introduces a novel interpretation of the U-Net architecture by studying certain generative hierarchical models, which are tree-structured graphical models extensively utilized in both language and image domains. With their encoder-decoder structure, long skip connections, and pooling and up-sampling layers, we demonstrate how U-Nets can naturally implement the belief propagation denoising algorithm in such generative hierarchical models, thereby efficiently approximating the denoising functions. This leads to an efficient sample complexity bound for learning the denoising function using U-Nets within these models. Additionally, we discuss the broader implications of these findings for diffusion models in generative hierarchical models. We also demonstrate that the conventional architecture of convolutional neural networks (ConvNets) is ideally suited for classification tasks within these models. This offers a unified view of the roles of ConvNets and U-Nets, highlighting the versatility of generative hierarchical models in modeling complex data distributions.

Published as a conference paper at ICLR 2025
U-NETS AS BELIEF PROPAGATION: EFFICIENT CLAS- SIFICATION, DENOISING, AND DIFFUSION IN GENER- ATIVE HIERARCHICAL MODELS
Song Mei
University of California, Berkeley songmei @berkeley. edu
ABSTRACT
U-Nets are among the most widely used architectures in computer vision, renowned for their exceptional performance in applications such as image seg- mentation, denoising, and diffusion modeling. However, a theoretical explanation of the U-Net architecture design has yet been fully established.
This paper introduces novel interpretation of the U-Net architecture by studying certain generative hierarchical models, which are tree-structured graphical mod- els extensively utilized both language and image domains. With their encoder- decoder structure, long skip connections, and pooling and up-sampling layers, we demonstrate how U-Nets can naturally implement the belief propagation de- noising algorithm in such generative hierarchical models, thereby efficiently ap- proximating the denoising functions. This leads to an efficient sample complex- ity bound for learning the denoising function using U-Nets within these models. Additionally, we discuss the broader implications of these findings for diffusion models generative hierarchical models. We also demonstrate that the conven- tional architecture of convolutional neural networks (ConvNets) is ideally suited for classification tasks within these models. This offers unified view of the roles of ConvNets and U-Nets, highlighting the versatility of generative hierarchical models modeling complex data distributions.
INTRODUCTION
U-Nets are one of the most prominent network architectures in computer vision, primarily employed for tasks such as image segmentation, denoising (Ronneberger et al.. 2015; Zhou et al., 2018: Sid- dique et al., 2021; Oktay et al., 2018), and diffusion modeling 2020; Song Ermon, 2019; Song et al., 2020). These networks are structured as encoder-decoder convolutional neural networks equipped with long skip connections, and their input and output typ- ically maintain the same dimensions. While U-Nets have demonstrated exceptional performance across a variety of applications. the theoretical foundations of their key components-includde encoder-decoder
et al., 2015; Ho et al.,
the encoder- decoder structure, long skip connections, and the pooling and up-sampling layers remain inadequately understood. Notably, long skip connections have significant impact performance as shown in empirical studies (Drozdzal et al., 2016; Wang et al., 2022). Existing explanations, often anecdotal, suggest their efficacy stems from improved information propagation and reduction of the vanishing gradient issue, but thorough theoretical exploration still lacking.
In this paper, introduce novel interpretation of the U-Net architecture, viewing it through the lens of neural network approximation. We posit that:
U-Nets naturally approximate the belief propagation denoising algorithm in certain generative hierarchical models.
The generative hierarchical model discussed herein is tree-structured graphical model, which has been widely employed in language and image generative modeling (Chomsky, 1959; Lee, 1996; Allen-Zhu & Li, 2023; Li et al., 2000; Willsky, 2002; Jin & Geman, 2006). We detail the precise definition of such generative hierarchical models in Section 2. A series of recent work (Mossel,

---

Published as a conference paper at ICLR 2025
2016; Sclocchi et al., 2024; Tomasini & Wyart, 2024; Petrini al., 2023; Kadkhodaie et al., 2023a) have pioneered the use of generative hierarchical models in studying classification tasks and dif- fusion models. Kadkhodaic et al. (2023a) has empirically shown that U-Nets can effectively learn the denoising function for these models. Furthermore, as noted by Sclocchi et al. (2024), the belief propagation denoising algorithm, which computes the denoising function in these models exactly, includes a downward process and an upward process, with the latter reusing the intermediate results from the downward process. In Section 4, we demonstrate how the belief propagation algorithm naturally induces the encoder-decoder structure, the long skip connections, and the pooling and up- sampling operations of the U-Nets. This gives rise to an efficient sample complexity bound for learning the denoising function in generative hierarchical models using U-Nets.
In addition to our main findings, in Section 3, we demonstrate that the standard architecture of con- volutional neural networks (ConvNets) is well-suited for classification tasks within the same gener- ative hierarchical model. We provide efficient sample complexity results to support this assertion. This offers unified perspective on the role of both ConvNets and U-Nets in image classification and denoising tasks, and also highlights the versatility of generative hierarchical models in modeling data distributions across language and image domains.
2 THE GENERATIVE HIERARCHICAL MODEL
To define the generative hierarchical model, we start by introducing some key notations. Consider a tree T (V, with a height of L, where we conventionally designate the root of the tree r as layer 0. For each node E V, we denote pa(v) as the parent of u, C(v) as the children of v, and (wv) as the siblings of v. We denote v(0 the set of nodes at layer f. We assume that for any v E ((-11), the number of children is precisely (n(() for é E [L]. The leaf nodes E V(L) have no children. Additionally, for each E D(O), we assume an ordering function (a bijection) e C(u) + [m ensuring that any child v' E C(v) possesses a unique rank c(v') E [m]9]. We denote the number of nodes at layer as d(e) and the number of nodes at layer L as d = d(4), We
p= m!)) By these definitions and assumptions, we have d( IIis.se (+)), and ( S (22) S ..... 5...5 = d.
further denote m = (m(()) and ||m|] = Σ/_
For each layer l ..... L. every tree node v e (( is associated with a variable 2! [S] for some S E N22. (For simplicity, we use the same variable space [S] across all layers, although our framework can accommodate variations across different layers f.) We denote (n()) = (x!")
vever E [S]d') as the variables at layer L. The variables at the leaves, æ (L() E [S] are considered the observed covariates, exemplified by the pixel representation of an image. Conversely, the root node variable .τ E [S] is treated as the associated label. Variables for the intermediate layers (x( h remain unobserved.
(0)
The generative hierarchical model. (GHM)' which is a joint distribution μa, over variables
We consider a specific type of generative hierarchical model
(y (0) e [S],
x(1)
x (f.-1) [s]di-- x(4) x [S]"),
associated with a set of functions (1b(0) [S] x m0) + BBooeetL defined as
He (v, (1) ...., ((. (L--1) x) x (1(0) (y. ()). (Ilvevas ( x =()))) (.Ievve
(GHM)
(Invvva (lll) (x (1-1) ).(w))
The formula specifies that any two nodes U1, E v(1) within the same level uses the same function 20, thereby embedding specific invariance properties into μa, Consequently, this GHM ensures that (#a)eze, d (Ie)ezez for any U1 02 E v(O) The notation ž U1 denotes that is either
We define "generative hierarchical models" as general probabilistic models with hierarchical structure. The specific model discussed in this paper is an instance of such generative hierarchical models. These mod- els are also known by various other names, including "hierarchical generative models", "latent hierarchical models", "Bayesian hicrarchical models", "hierarchical Markov random fields", among others.
2

---

Published as a conference paper at ICLR 2025
Lager
MN
y
Layer
Figure 1: Left: The generative hierarchical model with 3 layers and ml)) 3, m2 3, and m 2 children in each layer. Right: A 3-layer convolutional neural network.
identical descendant of un We will explore later how this invariance property interacts with the convolutional structure of the neural networks to be introduced.
Throughout the paper, impose specific assumption the de functions in GHM, namely that each ab function can be factorized into product of functions that depend solely on the ordering of the child. This assumption essential for enabling convolutional neural networks to approximate the associated belief propagation algorithm.
Assumption (Factorization of v/). For each layer E [L] and node v E D(e-1), we have
(xb (-((-)) - Ierrcct)
((-1)
(x(f-1) (0))
We also state a technical assumption concerning the boundedness of these & functions of GHM. Assumption 2 (Boundedness of a/). For any layer E E [L] and child node rank i E [m["], the transition probabilities are bounded as follows:
1/K S min vy(() (x,x²) max (1(1 (I,x²) S K. a.z.
is helpful to think about the joint distribution μ4, as tree-structured Markov process, which admits the factorization
p+ (y (1)
(1-11) ,x)
(m(()))( (IIeeves Un un (= (m(")) (e ... (Ilvvve-- M (x(
(3)
where abuse the notation denote n. (2) (= | ((-1)) the conditional probability of TC(')
(2)
given 2 (f(--)) and n. (y) as the marginal probability of y. Indeed, any graphical model specified as in Eq. (1) be cast into the form of (3). Furthermore, It can be checked that Eq. (3) coincides
C(v) | (2(f-1)) for & 2 1, and vi(1) x (1)) = 14 (x))) ly). this scenario, the factorization assumption (1) implies that (x())vec(()) given
with (GHM) upon taking (x (x(f-1) C(m) = n. (=
Te (6-1) are conditionally independent. We include a schematic plot of the generative hicrarchical model with 3 layers as in Figure 1(eeft).
GHMs as natural models for languages and images. In the field of linguistics, GHMs are very similar to context-free grammars (CFGs) (Chomsky, 1959; Lee, 1996; Allen-Zhu & Li, 2023). The generative process a context-free grammar involves creating valid strings or sentences based on given set of production rules (the & functions) that dictate how symbols can be extended to form new strings. Starting with an initial symbol (the label y), the generation proceeds iteratively by applying production rules until all symbols belong the terminal set, thus forming complete sentence (the covariate 2:).
In computer vision, GHMs often utilized to model natural images (Li et al., 2000; Willsky, 2002; Jin & Geman, 2006), where they are sometimes referred to as multi-resolution Markov models. The hypothetical image generation process begins with high-level concept of the image (represented
3

---

Published as a conference paper at ICLR 2025
by the label y), which is then iteratively refined using production rule (the &b function). This refinement continues through successive resolution levels until the detail reaches the pixel level, resulting in the final image (the covariate x).
We remark that a series of studies (Sclocchi et al., 2024; Tomasini & Wyart, 2024; Petrini et al., 2023; Kadkhodaie al., 2023a) have provided both theoretical and empirical evidence supporting the efficacy of GHMs as powerful tools for modeling the properties of combinatorial data.
3 THE WARM-UP PROBLEM: CLASSIFICATION IN GHMS
In this section, we consider the warm-up problem of classification within the GHM. In the subse- quent section, we will investigate the denoising task, where the results and intuitions will be similar and parallel to those presented in this section.
In the classification task, consider the scenario where we observe set of iid samples { y1) hein) drawn from μs+ under the GHM. Our objective is to learn probabilistic classifier (ayl (z) from this dataset. With suitable loss function, the optimal classifier is the Bayes classifier μ., (ylæz) which
represents the true conditional probability of given æ. We aim to examine the sample complexity of learning this classifier through empirical risk minimization over the class of convolutional networks.
The ConvNet architecture. We here introduce the convolutional neural network (ConvNet) chitecture used for classification, represented as (N E A((S)) for input x E [S] Initially, set b!
Zy E [S] for each node v E D(L) The operational flow of the network unfolds as
follows:
9 ((0)
(E) NN (0) (h!")
E ERS,
( E [4,, E v(o)
Σ Evec(e) ! E RRS MNN (+|x)- softmax(h" (0)) A((S)).
h, (8-1) = normalize
L [2), We p(t-!)
(ConvNet)
The functions NN(())
S + R2 (which adjust their input dimensionality 5/60 based layer ((v)
depth, s(() (5+1) le + L) 2 1fe L}) and normalize: RR - R are defined as follows:
o (h) = w (2) ReLU(W! ReLU(W!!! [)),,
(4)
normalize(h). ()., = ht eis) hag.
(5)
The dimensions of the weight matrices within the ConvNet specified follows:
W = (ww(") e RSSD) RR (eE(mo)) (Wi1) E DxD1 (RDDDhemmmn
hejmooj (w(f) 3,4 E RDx heimte
[E(T]
(6)
Furthermore, we denote INN as the ConvNet classifier ANN parameterized by W. A schematic illustration of ConvNet with 3 layers is provided in Figure 1(right).
Remark 1 (An explanation of the "ConyNet" architecture). We remark that the neural network layers described in (ConNNet) are different from the "convolution operations" typically seen in practice. The convolution operations used in practice involve computing the inner products be- tween convolutional filters and image patches, whereas in (ConNNet) and (4), point-wise product is employed instead. Despite this, these layers are still referred to as convolutional layers because the mapping from {h Joevie to {q(h)) Jevio, as per the first line of (ConvNet), preserves the translation-invan invariance property. Specifically, we use the same same function NN(!) across different imputs h!? and (f) as long as ((11) i(v2)
Additionally, the "normalization operator" defined in (5) differs from commonly used ones. We adopt this specific form for technical reasons, to effectively control the approximation error.
Despite these differences from standard convolutional networks, (ConyNet) represents an iterative composition of convolutional layers, pooling layers, and normalization layers, aligning closely with the architecture of convolutional networks used in practice. Figure 1(right) shows the sequence of these operations in detail.
4

---

Published as a conference paper at ICLR 2025
The ERM estimator. In the classification task, we employ empirical risk minimization ConvNets as outlined in the following equation:
(R(UNN) = loss(ys, PNN (NN(- (-1zz)))
arg
min
(7)
WeW W.Wam.L
where, for simplicity of analysis, we opt for the square loss rather than the more commonly used cross-entropy loss:
loss(y, y, NN (-(=)) E (14 (oN(())a)))
(8)
The parameter space for the ConvNets is defined as:
Wd,m,L,SDD.. = W as defined in (6) |W| max max max
|| llop - BJ.
JE EE (E[m(G]
We anticipate that the empirical risk minimizer, INN could learn the Bayes classifier μ4, as the global minimizer of the population risk over all conditional distributions yields the Bayes classifier:
W
(+1-) arg min (R(p) = Elloss(y u [(-[x))1]
W In theoretical analysis, we measure the discrepancy between UNN and μ. using the squared Euclidean distance:
(2(,,,)) = Ennn ES (((s)z)-(((sl))) I
(10)
Sample complexity bound. The subsequent theorem establishes the bound of the D3 disdistance between the ConvNet estimator nN and the true Bayes classifier f4..
Theorem (Leaming classify using ConvNets). Let Assumption and 2 hold. Let Wa,m. W.,mLL.S.D L.D.D.B
be the set defined as in Eq. (9), where D 2 k2 3 and B = Poly(d, S, K, 3 D). Let W be the empirical risk minimizer as Eq. (7). Then with probability at least 1 1}, we have
S 32L
LD(D + 28 + 1) log(d d m DS- 3 + 1og(1/1))
(UNNN p.) & C.O
D2
n
(11)
The proof Theorem detailed Section
Remark To ensure the -i-distance is less than e2, Theorem requires to take
D D= O /e), n = 223 |m mmll
(12)
3- (T))). The dependency on any of the param- eters (S, d, 3h, K. ||m]]], e) could potentially be refined by imposing additional assumptions on the d functions through more detailed analysis of approximation and generalization. This question of improving rates remains open for future work.
where θ hides a logarithmic factor log(d]|m]]
Consider a simplified scenario where S 2, K is constant, and (!) leading to d = m : In this setup. the sample complexity gives
for each E E [L],
n = Ó (2+21o%,, 3 3/e0 Ó(L'm d /c))
(13)
exhibiting a polynomial dependence on and 1/c. Such polynomial scaling aligns with existing literature (Poggio et al., 2017; Malach & Shalev-Shwartz, 2020; Schmidt-llieber 2020; Allen-Zhi & Li, 2022; Petrini et al., 2023), which indicates that learning hierarchical models using multi- layer networks avoids the curse of dimensionality. Theorem I serves warm-up result in the classification context. In Section we aim to extend similar methodologies to address denoising problems, employing analogous proof strategies.
5

---

Published as a conference paper at ICLR 2025
Proof strategy. The proof strategy begins by decomposing the squared distance between the learned model and the true Bayes classifier into approximation and generalization error terms. The generalization error is bounded using standard chaining argument, leading to a rate ((vdd D(vdp/m) dp/n), where dp denotes the number of ConvNet parameters. The focus then shifts to controlling the ap- proximation error. This is done by first introducing the belief propagation and message passing algorithm for computing the Bayes classifier, and subsequently showing that ConvNets can approxi- mate this algorithm effectively. A detailed outline of the proof strategy is provided in Appendix A.I, with the complete proof presented in Appendix E.
DENOISING AND DIFFUSION IN GHMS
In this section, we consider the denoising task within the GHM. Consider the joint distribution of noisy and clean covariates (z, æ), generated from the following: ~ μ., represents the clean covariates, and = x + g where g ~ (O, La) denotes the independent isotropic Gaussian noise. For simplicity in notation and with a slight abuse of notations, we continue to refer to the joint distribution (z, x) as μ4.
We consider a scenario where set of iid samples {(z,, m:) Wid M4 is drawn from the distri- bution. Our objective is to learn a denoiser m(z) from this dataset. With suitable loss function, the optimal denoiser is the Bayes denoiser (z) (x,,))) [z]|z], which calculates the poste- rior expectation of æ given z. We aim to examine the sample complexity of learning this denoiser through empirical risk minimization over the class of U-Nets. The approaches and results of this
section closely align with those discussed in Section 3 on the classification task.
The U-Net architecture. We here introduce the U-Net architecture used for denoising, repre- sented as mzN(=)) E I& for input E Rd. Initially, we set ((.)) = (-(x z,.)² /2)xe(1)) E R$ for v E (LL) for each node v E y(D) The operational flow of the network unfolds as follows:
( NN
Li(e) (normalize (h(.)) RR,
E [4], VE V(O
(e-1) Emec(v) L,,' E E
Le [4], v v((-1)
1.0
u t. = =b((-1)
(0)
T.pa(r) ERS, (with b = h (0))
L [4], ve v(o,
(UNet)
b (8) NN((1)
7, (u(l)) (normalize (8) 7. (..) + h1 E R
(
LE [4], ve v())
(2) EeE[s S- softmax(b] (b
V (())
The functions (N
1.c' NN R
+ R heetz and normalize R + R are defined as follows:
NN!!)(h). (! 1,0. Re ReLU d,0. [h: 1)), S [S], (4,1), (T) [m(9].
2,,
(14)
normalize(h). = h eax hyr
(15)
The dimensions of the weight matrices within the U-Net are specified as follows:
W = { (w/(0)
1,0.1 E RRDx) (R hefm/oj
(w-l)
2,0,4 E E DDD DDDD [e([m()) (w(() 3,0.c E DX(+++)) he(mlo) (ee(mt))
oE(L)
(16)
Furthermore, we denote mNN as the U-Net denoiser mNN parameterized by W. A schematic illustration of U-Net with 3 layers is provided Figure 2.
Remark 3 (An explanation of the "U-Net" architecture). As noted in our discussion of the con- volutional network as in Remark 1, the "convolutional layers" and the "normalization operator" described in (UNet) are different from practical implementations. However. we continue to use these terms because they retain core characteristics of their practical counterparts.
An important feature of (UNet) its encoder-decoder architecture and the inclusion of long skip connections, which closely mirror practical implementations. Specifically, the encoder sequence (L-1)
in (UNet) progresses (L) q._ h
9 h consisting of series

---

Published as a conference paper at ICLR 2025
NM4
MM
MNN
Figure 2: A U-Net with L = 3. "AP" stands for average-pooling. "US" stands for up sampling.. "SKIP" stands for long skip connections.
of convolutional, average pooling, and normalization layers. This part of the architecture is the same (ConvNet) for the classification task. The decoder sequence ascends as by²) u bi ..... + bb bb.-1 + u. mNN. consisting of series of convolutional, up-sampling, and normalization layers. Moreover, the computation of b.)) utilizes ui from the upward sequence, and (h() .9. (0) from the downward process, which is enabled by the long skip connections. This encoder-decom decoder architecture and the long skip connections are effectively visualized in Figure
The ERM estimator. In the denoising task, employ empirical risk minimization over U-Nets as outlined the following equation:
W arg
min W.Wa..L.s.DD
R(m
W
Σ ==1
x m NN (2 (2,)112
(17)
nd
The parameter space for the U-Nets is defined as:
W.S.m..S.DDFF W defined in (16): |W]
max max max oE (4.1) 2E(3] rE[r. (((m())
| llop
B].
(18)
We anticipate that the empirical risk minimizer, m m could learn the Bayes denoiser m., as the global minimizer of the population risk over all functions yields the Bayes denoiser:
m. (+) arg min (R(m) (x,,)) (d |d= m(=)]2) m
In our theoretical analysis, we measure the discrepancy between NNN and m, using the squared Euclidean distance:
W
D2 m.) Ez-w. [d-l|/m (2)
(19)
Sample complexity bound. The subsequent theorem establishes the bound of the Disditancce between the U-Net estimator mAN and the true Bayes denoiser m...
W
Theorem (Learning to denoise using U-Nets). Let Assumption 1 and hold. Let W.m.L.SS.D.F be the set defined in Eq. (18), where B = Poly (d, S, K, 184, D). Let W be the empirical risk minimizer as in Eq. (17). Then with probability at least 1 - 17. we have
S S*Kt& 182L.
LD(D + 28 + 1) m
log
m DS 184) log (og())))
W W m.) S C.
+
D2
The proof of Theorem 2 is detailed in Section H.

---

Published as a conference paper at ICLR 2025
Remark 4. To ensure the D2 disttnnee is less than 2,, Theorem 2 requires take
D = O K2 K2d1188//) /t), n = ( 1 ||m]| /e");
(20)
where & hides logarithmic factor log d S (ne)), Similar to Theorem 1, the depen- dency on any of the parameters (S, d. 184, K, fl.lm//aaa could potentially be refined by imposing additional assumptions on the và functions or through more detailed analysis of approximation and generalization. This question of improving rates remains open for future work.
Consider a simplified scenario where S 2, K is constant, and ml() = =m 2 3 for each E E [L]. leading to d = me, In this setup, the sample complexity gives
n = ( m. d2+2lo8m 18 (8)e) S @ d² (e),
(21)
exhibiting a polynomial dependence on d and l/c. Although the degree of the polynomial is substan- tial and potentially improvable, this gives the first polynomial sample complexity result for learning the Bayes denoiser in a hierarchical model using the U-Nets.
Connection to the task of diffusion generative modeling. The denoising task examined in this section is closely related to the diffusion model approach (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song & Ermon, 2019; Song al., 2020) generative modeling. Diffusion models involve learning a generative model from a dataset of n independent and identically distributed samples {a,];=1 drawn from an unknown data distribution μ., P([S]t). The goal is generate new samples â ~ j that match the distribution of steps that are closely related the denoising task described earlier. Here, we illustrate how they work using a variant of diffusion models, the stochastic localization process (Eldan, 2013; El Alaoui et al., 2022; Montanari & Wu, 2023; Celentano, 2022; Montanari, 2023):
Most diffusion model formulations involve a series
Step 1. Fit approximate denoising functions my Rd + Rd for t E T]. This is done by minimizing the empirical risk over a class of neural networks F (i.e., the denoising task discussed in this section):
m = arg min NN.EF nd E1a; (NN Vt. g.)l1.. gi Sed (N
(ERM)
Step 2. Simulate a discretized version of stochastic differential equation (SDE) starting from zero, whose drift term gives the approximate denoising function:
dz; mz))))) dt + d,,,, t E [0, 20 0,
(SDE)
and generate an approximate sample â 2r /T Rd at the final time T.
Standard analysis shows that by replacing the fitted denoising functions m2,(z) with the true denois- ing functions mz))) in Eq. (SDE) and allowing T + oc, we can effectively recover the original data distribution fu.. Consequently, the quality of samples generated from diffusion models hinges on two critical factors: (1) How well the fitted denoising functions m, (ERM) approximate the true denoising functions m,; (2) How accurately the SDE discretization scheme approximates the continuous process as in (SDE).
Recent work has made substantial progress in addressing these two theoretical questions: controlling the SDE discretization error, assuming reliable denoising function estimator is available (Chen et al., 2022a; 2023a; Lee et al.. 2023; Li et al., 2023; Benton et al., 2023); and controlling the denoising function approximation error through neural networks (Oko al., 2023; Chen et al., 2023b; Mei & Wu, 2023). However, these works have not explained the benefit of employing U-Net in image diffusion modeling, which is the primary focus of the current work. Indeed, by integrating the sample complexity bounds for learning denoising functions, as established in Theorem 2, with standard SDE discretization error bounds, such as the result established in Benton et al. (2023), is straightforward to derive end-to-end error bound for the sampling process of diffusion models in GHMs, similar to the strategy of Oko et al. (2023); Chen et al. (2023b); Mei & Wu (2023)?.
We note that stochastic localization formulation is equivalent to the DDPM diffusion model, differ- ing only in parametrization (Montanari, 2023). In the DDPM model, U-Nets serve to approximate the score function. The score function is a linear combination of the denoising function with an identity map, per Tweedie's 's formula. Consequently, Theorem 2 can be readily adapted to establish sample complexity bound for learning this score function.
8

---

Published as a conference paper at ICLR 2025
Proof strategy. The proof strategy for the denoising task parallels that of the classification task. The squared distance between the learned model and the true Bayes denoiser is decomposed into approximation and generalization error terms. The generalization error is bounded via standard parameter counting argument. The approximation error is controlled by first introducing the belief propagation and message passing algorithm for computing the Bayes denoiser and then showing that U-Nets can effectively approximate this algorithm. A detailed outline of the proof strategy is provided in Appendix A.2, with the complete proof presented in Appendix H.
FURTHER RELATED WORK
Generative hierarchical models. Hierarchical modeling of data distributions has been proposed in series of works (Mossel, 2016; Poggio et al., 2017; Malach & Shalev-Shward 2020; Schmidt- Hieber, 2020; Allen-Zhu & Li, 2022; Petrini et al., 2023; Sclocchi et al., 2024; Tomasini & Wyarr, 2024; Cagnetta & Wyart, 2024; Garnier-Brun et al., 2024; Kadkhodaie et al., 2023a;b). While the hi- erarchical models in Poggio et al. (2017); Malach & Shalev-Shwartz ww0-SSaart22200 Schmidt-Hieher (2020); Allen-Zhu & Li (2022) remain deterministic. Mossel (2016); Petrini et al. (2023); Sclocchi et al. (2024); Tomasini & Wyart (2024); Cagnetta & Wyart (2024); Garnier-Bron et al. (2024) studied the generative version of hierarchical models. The diffusion model for multi-scale image distribution representations has been empirically examined in Kadkhodaie et al. (2023a;b), which demonstrated that U-Nets are effective in modeling denoising algorithms. The theoretical and empirical evidence presented in Sclocchi et al. (2024); Tomasini & Wyart (2024); Petrini et al. (2023) underscores the effectiveness of generative hierarchical models in capturing the combinatorial properties of image datasets. Given their significant relevance to this work, we delve deeper into these studies.
Contributions of Petrini et al. (2023); Sclocchi et al. (2024); Tomasini & Wyart (2024). The series of works on hierarchical generative models (Petrini et al.. 2023; Sclocchi et al., 2024; Tomasini & Wyart, 2024) inspired the current study. Sclocchi et al. (2024) first pointed out that the belief prop- agation denoising algorithm of hierarchical models consists of downward and upward processes. Through mean-field analysis on random generative hierarchical model, they identified a phase transition phenomenon, aligning with empirical observations in diffusion models, thereby providing strong evidence of the efficacy of these models in handling combinatorial data properties. Petrini et al. (2023), on the other hand, first introduced these models in a classification context. Petrini et al. (2023); Tomasini & Wyart (2024) demonstrated that learning hierarchical models using multi- networks circumvents the curse of dimensionality. Specifically, they theoretically and empirically characterized the sample complexity, showing that it remains polynomial in dimension when learn- ing convolutional networks under random generative rules. On the other hand, in the absence of correlations, they showed that the sample complexity again exponential in the dimension, even for hierarchical generative models. This learning incapability is not captured by our analysis which does not consider optimization.
ConvNets and U-Nets and their implicit bias. Convolutional networks (LeCun et al., 1989; 1998; Krizhevsky et al., 2012; Szegedy et al., 2015; He et al., 2016) have become the state-of-the- art architecture for image classification and have been the backbone for many computer vision tasks. U-Nets (Ronneberger et al., 2015; Zhou et al., 2018; Siddique et al., 2021: Oktay et al., 2018) have been particularly well-suited for image segmentation and denoising tasks (Ronneberger al., 2015), and have served as the backbone architecture for diffusion models (Sohl-Dickstein et al., 2015; Ho al., 2020; Song & Ermon, 2019; Song al., 2020). A series of theoretical works has explained the inductive bias of CNNs (Bruna& & Mallat, 2013; Gunasckar et al., 2018; Bietti & Mairal, 2019b;a; Scetbon & Harchaoui, 2020; Li et al., 2020: Bietti, 2021; Mei et al., 2021; Misiakiewicz. & Mei, 2022: Cagnetta et al., 2023; Favero et al., 2021; Bietti et al., 2021: Xiao, 2022; Petrini et al., 2023; Tomasini & Wyart, 2024; Wang & Wu, 2024). However, they mostly focused on the classification and regression setting and were not concerned with the role of U-Nets in denoising tasks. The implicit bias of U-Nets has been theoretically investigated in Williams et al. (2024); Falck et al. (2022), where they found that the U-Nets conjugate the ResNets. In contrast, we demonstrate that U-Nets can effectively approximate the belief propagation denoising algorithms of GHMs. We note that Cui et al. (2023) analyzed the learning dynamics for a simple U-Net in diffusion models.

---

Published as a conference paper at ICLR 2025
Neural networks approximation of algorithms. A recent line of work has investigated the pressiveness of neural networks through an algorithm approximation viewpoint (Wei et al., 2022; Bai et al.. 2024; Giannou et al., 2023; Liu et al., 2022a; Marwah et al., 2021; 2023; Lin et al.. 2023; Mei & Wu, 2023). In particular, Wei et al. (2022); Bai et al. (2024); Giannou et al. (2023); Liu al. (2022a); Lin et al. (2023) demonstrate that transformers can efficiently approximate several algorithm classes, such as gradient descent, reinforcement learning algorithmss, and even Turing ma- chines. In the context of diffusion models, Mei & Wu (2023) shows that ResNets can efficiently approximate the score function of -igh-dimensional graphical models by approximating the varia- tional inference algorithm. Our work is closely related to Mei & Wu (2023), except that we study neural network approximation in a different statistical model and network architecture.
From a practical viewpoint, line of work has focused on neural network denoising by unrolling iterative denoising algorithms into deep networks (Gregor & LeCun, 2010; Zheng et al., 2015; Zhang & Ghanem, 2018; Papyan et al., 2017: Ma et al., 2021; Chen et al., 2018; Borgerding et al., 2017; Monga et al., 2021; Yu et al., 2023a;b). While this literature has primarily focused on devising better denoising algorithms, our work leverages this perspective to develop neural network approximation theory and explain existing network architectures.
Related theory of diffusion models. In recent years, diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song & Ermon, 2019; Song et al., 2020) have emerged as a leading approach for generative modeling. Neural network- -ased score function approximation has been recently studied from the function approximation viewpoint in Oko et al. (2023); Chen et al. (2023b); Yuan et al. (2023); Shah et al. (2023); Biroli & Mézard (2023), and from the algorithm approximation viewpoint in Mei & Wu (2023). Theoretical studies of other aspects of diffusion models include Liu et al. (2022b); Li et al. (2023); Lee et al. (2023); Chen et al. (2022b; 2023d; 2022a; 2023c;a); Benton et al. (2023); El Alaoui et al. (2022); Montanari & Wu (2023); Celentano (2022); Ghio et al. (2023); Biroli & Mézard (2023); Biroli et al. (2024); Cui et al. (2023); Fu et al. (2024); Wu et al. (2024). For comprehensive introduction the theory of diffusion models, see the recent review (Chen et al., 2024).
CONCLUSIONS AND DISCUSSIONS
In this paper, we introduced a novel interpretation of the U-Net architecture through the lens of generative hierarchical models. We demonstrated that their belief propagation denoising algorithm naturally induces the encoder-decoder structure, the long skip connections, and the pooling and up- sampling operations of the U-Nets. We also provided an efficient sample complexity bound for learning the denoising function with U-Nets. Furthermore, we discussed the broader implications of these findings for diffusion models. We also showed that ConvNets are well-suited for classification tasks within these models. Our study offers a unified perspective on the roles of ConvNets and U-Nets, highlighting the versatility of generative hierarchical models in capturing complex data distributions across language and image domains.
The results presented in this paper offer considerable scope for enhancement. We initially assumed that the covariates x lie in the discrete space [S]d, and extending these results to continuous spaces would be an intriguing direction for future research. Additionally, the dependencies of the sample complexity bound on d and 1/c may be amenable to improvement through more careful analysis. Moreover, the convolution operations employed in this paper are different from those commonly employed in practical settings. It would be worthwhile to explore graphical models where the be- lief propagation algorithm aligns more naturally with ConvNets and U-Nets that utilize standard convolution operations.
On the practical side, our theoretical findings generated a hypothesis of the functionality of each layer of the U-Nets. Verifying these hypotheses in pre-trained U-Nets, such as those used in stable diffusion models, using interpretability methods, could yield valuable insights. Furthermore, ex- tending these results to include conditional denoising functions represents exciting direction for future research. Finally, we hope that the insights provided in this paper could guide the design of innovative network architectures.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported by NSF DMS-2210827, CCF-2315725, CAREER DMS-2339904, ONR N00014-24-S-B001. Sloan Research Fellow Award, an Amazon Research Award, Google Re- search Scholar Award, and an Okawa Foundation Research Grant. The author would like to thank Hui Xu, Yu Bai, and Yuchen Wu for their valuable discussions.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. How can deep learning performs deep (hierarchical) learning. 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint ar Arvi:2330 13673, 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. arXiv preprint arXiv:2308. 03686, 2023.
Alberto Bietti. Approximation and learning with deep convolutional models: kernel perspective. arXiv preprint arXiv:2102.10032 2021.
Alberto Bietti and Julien Mairal. Group invariance, stability to deformations, and complexity of deep convolutional representations. Journal of Machine Learning Research, 20(25): 2019a.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kemelss. Advances in Neural Information Processing Systems, 32, 2019b.
Alberto Bietti, Luca Venturi, and Joan Bruna. On the sample complexity of learning under geometric stability. Advances in neural information processing systems,
2021.
Giulio Biroli and Mare Mézard. Generative diffusion in very large dimensions. arXiv preprint 03X1:22006033188 2023.
Giulio Biroli, Tony Bonnaire, Valentin de Bortoli, and Mare Mézard, Dynamical regimes of diffu- sion models. arXiv preprint arXiv:2402. 18491, 2024.
Mark Borgerding, Philip Schniter, and Sundeep Rangan. Amp-inspired deep networks for sparse linear inverse problems. IEEE Transactions on Signal Processing, 65(16))4293-4300 2017.
Joan Bruna and Stéphane Mallat, Invariant scattering convolution networks. IEEE transactions on pattem analysis and machine intelligence, 35(8): 1886, 2013.
Francesco Cagnetta and Matthieu Wyart. Towards a theory of how the structure of language is acquired by deep neural networks. arXiv preprint arXiv:2406.0000 00048, 2024.
Francesco Cagnetta, Alessandro Favero, and Matthieu Wyart. What can be learnt with wide con- volutional neural networks? In International Conference on Machine Learning, pp. 3347-3379.
PMLR, 2023.
Michael Celentano. Sudakov-femique post-amp, and proof of the local convexity of the tap free energy. arXiv preprint arXiv:2208.09555) 2022.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, pp. 9735-4763. PMLR, 2023a.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estima- tion and distribution recovery of diffusion models on low-dimensional data. arXiv preprint 1.Xiv:3302.071944 2023b.
11

---

Published as a conference paper at ICLR 2025
Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. An overview of diffusion models: Ap- plications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771, 2024.
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. arXiv preprint arXiv:2209. 11215, 2022a.
Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast, arXiv preprint arXiv:2305. 2023c.
Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-dderadaation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. In International Conference Machine Learning. pp. 4462 4484. PMLR, 2023d.
Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Theoretical linear convergence of unfolded ista and its practical weights and thresholds. Advances in Neural Information Processing Systems, 31, 2018.
Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for proximal algorithm for sampling. In Conference on Learning Theory, pp. 2984-3014 PMLR, 2022b.
Noam Chomsky. On certain formal properties of grammars. Information and control, 2(2):137-167, 1959.
Hugo Cui, Florent Krzakala, Eric Vanden-Eijnden, and Lenka 7deborová. Analysis of learning a flow-based generative model from limited sample complexity. arXiv preprint a7X55:2310.03575, 2023.
Michal Drozdzal, Eugene Vorontsov, Gabriel Chartrand, Samuel Kadoury, and Chris Pal. The im- portance of skip connections biomedical image segmentation. In International workshop on deep learning in medical image analysis, international workshop large-scale annotation of biomedical data and expert label synthesis, pp. 179-187. Springer, 2016.
Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from the sherrington-kirkpariick gibbs measure via algorithmic stochastic localization. 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS), pp. 323-334. IEEE, 2022.
Ronen Eldan. Thin shell implies spectral gap up polylog via stochastic localization scheme. Geometric and Functional Analysis, 23(2):532-569, 2013.
Fabian Falck, Christopher Williams, Dominic Danks, George Deligiannidis, Christopher Yau, Chris C Holmes, Arnaud Doucet, and Matthew Willetts. multi-rcsolution framework for u- nets with applications to hierarchical vaes. Advances in Neural Information Processing Systems, 35:15529-15544 2022.
Alessandro Favero, Francesco Cagnetta, and Matthieu Wyart. Locality defeats the curse of dimen- sionality in convolutional teacher-student scenarios. Advances in Neural Information Processing Systems, 34:9456-94488 2021.
Hengyu Fu, Zhuoran Yang, Mengdi Wang, and Minshuo Chen. Unveil conditional diffusion models with classifier-free guidance: A sharp statistical theory. arXiv preprint arXiv:2403.1 11968, 2024.
Jérôme Garnier-Brun, Mare Mézard, Emanuele Moscato, and Luca Saglietti. How transformers learn structured data: insights from hierarchical filtering. arXiv preprint arXiv:2408. 15138, 2024.
Davide Ghio, Yatin Dandi, Florent Krzakala, and Lenka Zdeborová. Sampling with flows, diffusion and autoregressive neural networks: A spin-glass perspective. arXiv preprint arXiv:2308.144 110330888.44.... 14085, 2023.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dim- itris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301. 13/96, 2023.
12

---

Published as a conference paper at ICLR 2025
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning. pp. 399-406, 2010.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 685810-6651, 2020.
Ya Jin and Stuart Geman. Context and hierarchy in a probabilistic image model. 2006 IEEE computer society conference on computer vision and patlern recognition (CVPR'00) (66), volume 2, pp. 2145-2152. IEEE, 2006.
Zahra Kadkhodaie, Florentin Guth, Stéphane Mallat, and Eero P Simoncelli. Learning multi-scale local conditional probability models images. arXiv preprint arXiv:2303.0295. 2023a.
Zahra Kadkhodaie, Florentin Guth, Eero P Simoncelli, and Stéphane Mallat. Generalization in diffusion models arises from geometry-adaptive harmonic representation. arXiv preprint arXiv:2310.02555 2023b.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence Jackel. Handwritten digit recognition with a back pree-proppgation network. Advances in neural information processing systems, 2, 1989.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 36(11):2278-3224 1998.
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 986-985.100 PMLR, 2023.
Lillian Lee. Learning of context-free languages: A survey of the literature. 1996.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejic Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.09251, 2023.
Jia Li, Robert M Gray, and Richard Olshen. Multiresolution image classification by hierarchical modeling with two-dimensional hidden markov models. IEEE transactions on information theory, 46(5):1826-1888 1841, 2000.
Zhiyuan Li, Yi Zhang, and Sanjeev Arora. Why are convolutional nets more sample-efficient than fully-connected nets? arXiv preprint arXiv:20/0.085/5. 2020.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. arXiv preprint arXiv:2310.08566 2023.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:22 10749, 2022a.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208. 0eti22288..4999 14699, 2022b.
Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang. and Neil Shah. A unified view graph neural networks as graph signal denoising. In Proceedings of the 30th ACM International Conference Information& & Knowledge Management, pp. 1202-1211, 2021.
13

---

Published as a conference paper at ICLR 2025
Eran Malach and Shai Shalev-Shwarte The implications of local correlation on learning some deep functions. Advances in Neural Information Processing Systems, 33:1322-133 1332, 2020.
Tanya Marwah, Zachary Lipton, and Andrej Risteski. Parametric complexity bounds for approx- imating pdes with neural networks. Advances Neural Information Processing Systems, 34: 15044-15055. 2021.
Tanya Marwah, Zachary Chase Lipton, Jianfeng Lu, and Andrej Risteski. Neural network approxi- mations of pdes beyond linearity: A representational perspective. In International Conference on Machine Learning, pp. 24139-24172 PMIR, 2023.
Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models high-dimensional graphical models. arXiv preprint arXiv:239.111200 2023.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pp. 3351-3418. PMLR, 2021.
Mare Mezard and Andrea Montanari. Information, physics, and computation. Oxford University Press, 2009.
Theodor Misiakiewicz and Song Mei. Learning with convolution and pooling operations in kernel methods. Advances in Neural Information Processing Systems, 35:29014-29023 2022.
Vishal Monga, Yuelong Li, and Yonina Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2):18-44, 2021.
Andrea Montanari. Sampling, diffusions, and stochastic localization.
arXiv preprint
arXiv:2305. 10690, 2023.
Andrea Montanari and Yuchen Wu. Posterior sampling from the spiked models via diffusion pro- cesses. arXiv preprint arXiv:2304.1144/ 2023.
Elchanan Mossel. Deep learning and hierarchal generative models.
arXiv preprint
arXiv: 1612.00077 2016.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribu- tion estimators. arXiv preprint arXiv2233.001611 2023.
Ozan Oktay, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori, Steven McDonagh, Nils Y Hammerla, Bernhard Kainz, et al. Attention u-net: Learning where to look for the pancreas. arXiv preprint 180.:v880440399 2018.
Vardan Papyan, Yaniv Romano, and Michael Elad. Convolutional neural networks analyzed via volutional sparse coding. The Journal of Machine Learning Research, 18(1):2887-295 2017.
Judea Pearl. Reverend bayes inference engines: A distributed hierarchical approach. In Proba- bilistic and Causal Inference: The Works of Judea Pearl, pp. 129-138. 1982.
Leonardo Petrini, Francesco Cagnetta, Umberto M Tomasini. Alessandro Favero, and Matthieu Wyart. How deep neural networks learn compositional data: The random hierarchy model. arXiv preprint arXiv:2307.022 199,. 2023.
Tomaso Poggio, Hrushikesh Mhaskar, Lorenzo Rosasco, Brando Miranda, and Qianli Liao. Why and when can deep-bur but not shallow-networks avoid the curse of dimensionality: a review. Inter- national Journal of Automation and Computing, 14(5):503-51,, 2017.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 234-241. Springer, 2015.
Meyer Scetbon and Zaid Harchaoui. Harmonic decompositions of convolutional networks. In Inter- national Conference Machine Learning, pp. 8522-8532. PMLR, 2020.
14

---

Published as a conference paper at ICLR 2025
Johannes Sccmidtt-Hibbed Nonparametric regression using deep neural networks with relu activation function. 2020.
Antonio Sclocchi, Alessandro Favero, and Matthieu Wyart. A phase transition in diffusion models reveals the hierarchical nature of data. arXiv preprint arXiv:2402.16999 2024.
Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the ddpm objec- tive. arXiv preprint arXiv:2307.01178, 2023.
Nahian Siddique, Sidike Paheding, Colin Elkin, and Vijay Devabhaktuni. U-net and its variants for medical image segmentation: review of theory and applications. leee Access, 9:82031-82057 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2265, PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1-9, 2015.
Umberto Tomasini and Matthicu Wyart. How deep networks learn sparse and hicrarchical data: the sparse random hierarchy model. arXiv preprint urXiv:2404.10727. 2024.
Martin J. Wainwright. High- High-DDiensional Statistics: A Non-Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019. doi: 10.1017/ 9781 108627771.
Martin Wainwright, Michael I Jordan, al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning, 1(1-2):1-305, 2008.
Haonan Wang, Peng Cao, Jiaqi Wang, and Osmar R Zaiane. Uctransnet: rethinking the skip con- nections in u-net from a channel-wise perspective with transformer. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pp. 2441-2449, 2022.
Zihao Wang and Lei Wu. Theoretical analysis of the inductive biases in deep convolutional networks. Advances in Neural Information Processing Systems, 36, 2024.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers. Advances Neural Information Processing Systems, 15:11201-2008 2022.
Christopher Williams, Fabian Falck, George Deligiannidis, Chris C Holmes, Arnaud Doucet, and Saifuddin Syed. A unified framework for u-net design and analysis. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Alan S Willsky. Multiresolution markov models for signal and image processing. Proceedings of the JEEE, 9068):1396-1458 1458, 2002.
Yuchen Wu, Minshuo Chen, Zihao Li. Mengdi Wang, and Yuting Wei. Theoretical insights for diffusion guidance: case study for gaussian mixture models. arXiv preprint 2403022403.01639, 2024.
Lechao Xiao. Eigenspace restructuring: a principle of space and frequency in neural networks. In Conference Learning Theory, pp. 4888 4944. PMLR, 2022.
15

---

Published as a conference paper at ICLR 2025
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D Haeffele, and Yi Ma. White-box transformers via sparse rate reduction. arXiv preprint arXiv:2306.01122 2023a.
Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma. Emergence of segmentation with minimalistic white-box transformers. arXiv preprint arXiv:2308.16211 2023b.
Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Minshuo Chen, and Mengdi Wang. Reward-directes conditional diffusion: Provable distribution estimation and reward improvement. arXiv preprint arXiv:2307.07055, 2023.
Jian Zhang and Bernard Ghanem. Ista-net: Interpretable optimization-inspired deep network for image compressive sensing. In Proceedings of the IEEE conference computer vision and paltern recognition, pp. 1828-1837, 2018.
Shuai Zheng, Sadeep Jayasumana, Bernardino Romera-Paredes Vibhav Vinect, Zhizhong Su, Da- long Du, Chang Huang, and Philip HS Torr. Conditional random fields recurrent neural net- works. In Proceedings of the IEEE international conference computer vision, pp. 1537, 2015.
Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang. Unet++: A nested u-net architecture for medical image segmentation. In Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, ML-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 20, 2018, Proceedings 4, pp. 3-11. Springer, 2018.
16

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2 The generative hierarchical model
3
The warm-up problem: Classification in GHMs
4 Denoising and diffusion in GHMs
5 Further related work
Conclusions and Discussions
10
A Proof strategy
18
A.1 Proof strategy: ConvNets approximate the belief propagation algorithm
18
A.2 Proof strategy: U-Nets approximate the belief propagation algorithm
19
Technical preliminaries
20
Proof of Proposition
D Proof of Theorem 4
21
D.I Auxillary lemmas
23
Proof of Theorem
28
E,1 Error decomposition
28
E.2 Results on generalization
28
E.3 Auxillary lemmas
29
Proof of Proposition 5
30
G Proof of Theorem 6
30
G.I Auxillary lemmas
31
Proof of Theorem
33
H.1 Error decomposition
H.2 Results on generalization
33
H.3 Auxillary lemmas
34
17

---

Published as a conference paper at ICLR 2025
A PROOF STRATEGY
A. PROOF STRATEGY: CONVNETS APPROXIMATE THE BELIEF PROPAGATION ALGORITHM
Lemma 17 introduces decomposition of (2 p.) into two components, approximation error and generalization error:
(M (4) inf WEW D3 R4) +2 sup
R((NN) R(UNN)
WEW
approximation ertor
generalization error
The bound of generalization error follows standard approach: employing chaining argument, the error is controlled by (gydn/)) where dp represents the number of parameters in the ConvNet class.
In the following, we describe our strategy to control the approximation enror: we first present the belief propagation and message passing algorithm for computing the Bayes classifier μlx, and then demonstrate that ConvNets are capable of effectively approximating this message passing algorithm.
The belief propagation and message passing algorithm.
The belief propagation algorithm operates input [s]d and iteratively calculates the beliefs {ut E (o...)) as follows:
(x( = ((r! = =rb,,
(x! x _()()))
IWEE(w)) (wlet))(ab)) ( x x ((tt)) (ett 8= 11..,
(wilo)
(BP(y)x) = (0 (y).
(BP-CLS) Classical results in graphical models verify that the belief propagation algorithm accurately com- putes the Bayes classifier in this tree graph.
Lemma (BP calculates the Bayes classifier exactly (Pearl, 1982; Wainwright et al., 2008; Mezard & Montanari, 2009)). When applying the belief propagation algorithm (BP-CLS) starting with æ e [S]6, it holds that μ.a (-1z) μBP [BP(-1a) (-1z).
The belief propagation algorithm can be streamlined into a message passing algorithm, starting with the initialization (L) Ty for each node in the highest layer (L) The operations are defined as follows:
( 1/00 (1(9) ERS,
i [L], v v(O,
h normalize E Evec()) qo E RS,
l [4], ve v(e-1)
(MP-CLS)
(Mr(y)æ) (1(t)ax(((()))
The functions f!()) [S] + R {t(( RR + RS JeE(Z-1 are defined as:
(x)) log (b( x),
E [S], E [S),
(22)
(h)) log Easist w/ (8)
(s, a),
h EERE R sE [T-
We note that the normalization operator in (MP-CLS) is non- non-essential and could be dropped; how- ever, we include it to ensure the formulation closely mirrors (ConvNet), offering a technical benefit.
The subsequent proposition affirms that message passing is essentially equivalent to belief propaga- tion:
Proposition 3 (BP reduces to MP). Consider the belief propagation algorithm and the message passing algorithm, both starting from x E [S]d, as in Eq. (BP-CLS) and (MP-CLS). Then we have Vo (.)) softmax(h!") (2) for all 0 < L 1 and v(9) In particular, we have (Bp()1z) (MP(-1x) = (-
The proof of Proposition 3 is presented in Section C.
18

---

Published as a conference paper at ICLR 2025
Approximating message passing with ConvNets. By comparing the message passing algorithm (MP-CLS) alongside the ConvNet (ConvNet), the primary distinction lies in the nonlinear functions used: p(() versus NN((), Given the expression f(
(h), (()( log Ee " a))"", it becomes ev- ident that approximating the logarithmic and exponential functions using one-hidden-layer networks enables f:() (h) be effectively approximated by a two-hidden layer neural network. This leads to the following theorem:
Theorem 4 (ConvNets approximation of Bayes classifier). Let Assumption I and hold. For any 8 > take
D= D=4 44
B = Poly(d, 3,3" ,/6).
Then there exists W E W.s.s.S....R in Eq. (9), such that defining PNN as in Eq. (ConvNet), we huve
n. (ylæz) - log (x)()) S 6.
max [E(
The proof of Theorem 4 is detailed in Section D.
A.2 PROOF STRATEGY: U-NETS APPROXIMATE THE BELIEF PROPAGATION ALGORITTIM
The proof strategy for the denoising task closely aligns with that of the classification task as detailed in Section A.1.
W Lemma 22 introduces a decomposition of (3m m.) into two components, approximation error and generalization error:
W (m m.) S inf WEW
W (m m.) +2
- R(m
sup WEW
approximation
generalization erree
The bound of generalization emror follows standard approach: employing chaining argument, the error controlled by Ö(Vdp/n), where dn represents the number of parameters in the U-Net class.
In the following, we describe our strategy to control the approximation error: we first present the belief propagation and message passing algorithm for computing the Bayes denoiser m., and then demonstrate that U-Nets are capable of effectively approximating this message passing algorithm.
The belief propagation and message passing algorithm. The belief propagation algorithm (1) operates input z E IRd and iteratively calculates the beliefs 1.0r 4,o ((..)))eeo.....)).eees as follows:
(L((c = (1(++11 20)); z0);
(1) (=(() Eneoy) IIeeC()) (u(ett) (0) (1+ (ett)) 1))
((0)) z.d
t= - I,..,,,
(o(o) xl,
(o E, 2 wi (1((22 (6-1) (2)
(m(t)) , E((pa(()) { (ompll (spatt) II 10 ( l p= L=
o (=(()) (1 (1)))))((+))) (o(f)) (u))
mBr(2), = E,ao 1o vo
(1) (1)
(I
(BP-DNS)
where s (((+1) (4) 2 == expl- - )223. Classical results in graphical models verify that the belief propagation algorithm accurately computes the Bayes denoiser in this tree graph.
Lemma 2 (BP calculates the Bayes denoiser exactly (Pearl, 1982; Wainwright et al., 2008; Mezard & Montanari, 2009)). When applying the belief propagation algorithm (BP-DNS) starting with z e Rdd it holds that μt, (x.lzz) = v/ (xx) for E V(I.), so that mBp (z) = m(x).
We remark that the downward wwwwardd-pwaard structure of belief propagation in generative hierarchical models has been pointed out in the literature (Sclocchi al., 2024).
19

---

Published as a conference paper at ICLR 2025
The belief propagation algorithm can be streamlined into message passing algorithm, starting with the initialization 1 (-(x =.)² /2) e IR' for each node in the highest layer V(L), The operations are defined as follows:
4. (L((o) (normalize( (2)) E R,,
(6)
E [4], VE vlo,
(6 Evec(E) (1)
4.00 E ER,
L [4]. We v(6-1),
te (ipu(e) E ERS, (with 6/ = 11 11
(6)
LE [4], We WE v(O,
(MP-DNS)
6 (0)
1e (((())) (normalizo(ug) (0)
(2)
h1. ERS
te [4] Ve y(o),
(2) = ExE[S] 8. softmax(b (b),))
vE v(L)
The functions ( p(t) :R$ RR "R." hee[n) are defined as
(1): = log Ene(s) v/ (s,a)e 1/ (h), = log Ene(s) v (a, s)cAa
h HERS, 8 E [S], [4] h hERS, 8 E [S] E [Z].
(23)
We note that the normalization operator in (MP-DNS) non-essential and could be dropped; how- we include it to ensure the formulation closely mirrors (UNet), offering technical benefit. The subsequent proposition affirms that message passing is essentially equivalent to belief propaga- tion:
Proposition 5 (BP reduces to MP). Consider the belief propagation algorithm and the message passing algorithm, both with input E Rdd as in Eq. (BP-DNS) and (MP-DNS). Then we have () = softmax(hi].) (f)
(,0) and v (.)()) = softmax(b), -if.)) and va (.) = softmax(b(.) (T)). In
particular, mmp(=)= = (z)(2)) m(z).
The proof Proposition 5 presented in Section
Approximating message passing with ConvNets. By comparing the message passing algorithm (MP-DNS) alongside the U-Net (UNet), the primary distinction lies in the nonlinear functions used: f0 versus NN(!) o.l. Notably, f entails a log-sum-exponential structure. This structure suggests that approximating the logarithmic and exponential functions with one- hidden-laye neural networks can allow f (h) to be effectively approximated by a two-hidden-layer neural network. This leads to the following theorem:
Theorem 6 (U-Nets approximation of Bayes denoiser). Let Assumption and 2 hold. For any 8 > take
D 4[S 18/, B = (oly S, K, 1,,,,/8) 1/8).
Then there exists W E L.S. in Eq. (18), such that defining m NNN in Eq. (UNet), have
|m(=) mNN(2))loc 6. =ERd
sup
The proof of Theorem 6 is detailed in Section G.
B TECHNICAL PRELIMINARIES
We here present a bound on the supremum of sub-Gaussian processes, whose proof was based on the chaining argument.
Lemma 3 (Proposition A.4 of (Bai et al., 2024)), Suppose that {Xw hwwe is zero-mean random process given by
X = Σ f(z; w) [E]]((2;
where "1.".. In are i.i.d samples from distribution P- such that the following assumption holds:
20

---

Published as a conference paper at ICLR 2025
(a) The index set θ is equipped with distance and diameter By- Further, assume that for some constant Ap. for any ball @' of radius r in 9 the covering number admits upper bound log (og,N p) dp Log(2Arr/A)) AA) for all 0 2
(b) For any fixed θ and sampled from Pz. the random variable f(z; w) [=[[(z;;)) is (f(":"""" a --sub-Gaussian random variable (B(e>lfts:))-Ba (f(s =w]l]] I 12 for any X E IR).
(c) For any w, w' E and & sampled from Pz, the random variable f(2; ()) - f(z:w) is a o' p(w, w') w')-sub-CCeee --uub-aasssaan random variable (EleA[((ww)-Ssaw))) S (" 1/2 for any IR).
Then with probability at least 1 17, it holds that
X S CCov dp 1og(2A,(1 + Bpo' (a)) log(1/n)
sup
WE
where C a universal constans.
We next present simple inequality used in the proof of Theorem 1.
Lemma 4 (From log ratio bound to square distance bound). Let and q be two probability measures on A(SS) such that
max (E(S)
log (gg((y) (()) 8.
Then we have
(p(s) - (s))) S (65 -1)2
Proof of Lemma 4. The lemma is by the fact that lp(y) ((v)) < - log (1))( 1)-p(y).
C PROOF OF PROPOSITION 3
Proof of Proposition 3. By Eq. (MP-CLS) and (22), defining 2. ( = softmax(h!") then for
L. < L 2, get
16) (x
II e'EC(u)
Σ
(1t11) (1! a) ()
[E[S]
Σ
II
(V(tt) (a!?) x ) (e-1) (a (
(e+1)
JC1 1) E'EC(v) Cce)
This coincides with the update rule in Eq. (BP-CLS). For = 1, we get
(((-1) (x) & II Σ
(x (1)
a)l(a -wrl x II
(LL) c(er) (a (1( zzr)..
d'E((1)
RE(S]
e'EC(")
This again coincides with the update rule in Eq. (BP-CLS). This finishes the proof of Proposition 3.
D PROOF OF THEOREM
Proof of Theorem 4. By Lemma 13, take M( 228kk34/01++ and M0 12 /87+1. Then there exists {(az, Wj, by)h JE(M(()) and {(a,, wg b,)} ) with
sup |ayl < 2SK, sup |uyl < 1, sup |b,l < SK, sup أا < 2, sup byy < sup [byl S ' Sdk? (8),
such that defining
m1 4100 log,. (x) = = M1 0j ReLUUWwe by), (xp = Ex āj R.R.U + bz).
21

---

Published as a conference paper at ICLR 2025
and defining 1. (h), E ER (1), - log Σラ v/ (1,j) exp(h;)) (h): - logs. Σj vsl expp) we have
E
sup (I) (h)]lo S 8
(24)
max, h1 0
In addition, by Lemma 15, take M(L) |3/ + and M = Then there exists (2)
((a))) _u(h ..1()) SE(M(() and {(a))) _0j (L)
_ sE(M with
sup (a))) 11) 1 2K, sup 1u (1) 11, sup 1 = SK, sup 10g |
s 4, sup m (D)
sup 10, 1 1,
$ 1,
$
j
such that defining
logs (1) = -E of
(L) ReLU(uf + +b(1), Ind(x) = Ene a. ReLU x B
and defining 1. (1), (b) E R by
1. (x), log Eÿ=1 v/ (1, [() = ;); (x): logs Ej=1 (t.))nddo Vi E [S),
we have
sup (x) - F.4 S 88/(d34). E||S]
(25)
By Eq. (24) and (25) and Lemma 5, taking 0
R$ to be as defined Eq. (MP-CLS) and ( E R to be as defined in Eq. (A-MP-CLS) with (Ft2) Jeel4) defined above, we have
[AAt)) |1 < [5/(d3*] x II (2ml (0) + +1)
ISESL
As consequence. we just need to show that the approximate version of message passing algorithm as in Eq. (A-MP-CLS) could be cast as a neural network.
Indeed, by Lemma 14, there exist two-hidden-layer neural networks (for E [L - 1] and & E [m (9)
NN WeO Whe. wilh wi? (h) w ReL ReLU(W!! (2) [h; 1])),
with w &SSSM w ( w!l E
and
||W/0 Imax S Poly(SKM M()),, | Imma log(4SK [| Imax 2SK,
such that
NN wit w w (h) (h), vh R such that max 0. Furthermore, by Lemma 16, there exist two-hidden layer neural networks (for E [m(
NN NN wit with (x) (1) = w(F (eI.U(W.!))) 2, ReLUUW!!! [m:1]))
with w(O RRSSSM whl ERSM ((MM) +1), whlf (SMM
|| I. Ilaax S 2K,
|| Imax & Poly(SKM( (1) M [W]]]] Imux S 1.
such that
NN WiL with wiyh (x) (x) = (x), Va E [S].
This proves that the approximate version of message passing in Eq. (A-MP-CLS) can be cast into the convolutional neural network as in Eq. (ConvNet) with proper choice of dimension
> (S (9) SM. + +1] S (12SK 1 + 1) + 1, EE(L)
and norm of the weights. This finishes the proof of Theorem 4.
22

---

Published as a conference paper at ICLR 2025
D.1 AUXILLARY LEMMAS
Lemma 5 (Error propagation of the approximate version of message passing in classification). As- sume have functions f(0 and f( such that
(x) - -1 I 6, Vo E [S},
(26)
!1! f.'0 S, wh VhE such that max h 0, A4 1
jE(S]
Furthermore, consider the following approximate version of message passing algorithm with initial- ization ((()) Zu for u V(L):
To (()) (1 ERS,
E [T. 10-11,, 1], NEDO
(A-MP-CLS)
h1 ((f-1) normalize (Evecto) E R$
EE 1----n 1], V e D((-1),
Taking hy R to be defined in Eq. (MP-CLS) CLS) and h CLS), we have
R to be defined in Eq. (A-MP-
| 1 S <8x x II
(2m (6) + +1).
ISPLL
Proof of Lemma 5. We prove this lemma by induction, aiming to show that for any e [L we have
[T n 1 2l,c S 2m(++1) I (km + 1)6, Ve V(O.
(27)
To prove the formula for l L 1, since 11 T(o by Eq. (26), we get
17 - 1 S o, Vo E v(t).
By Lemma 7, we get
| (L-1) 1 11:0 S $2 Ex -
- )
s 2m
Yo E ((t-))
This proves the formula (27) for El 1-
Assuming that (27) holds at the layer f, by the update formula, we have
1:!h) - 1 = [1 (h f(00(( (1o
$ 11 (1) (h!0) - f12 )1 + (1 (N! - {(() (lo)))
i + 2m (f+1) I1 -ik=et2 (2m + 1) < 12m 1)6,
where the middle inequality is by the assumption of f!( and by Lemma 6. By Lemma 7, we get 1)8, Vo VO
[ h |00 2m 11
This proves Lemma by the induction argument.
Lemma 6 (Non-expansiveness of log- log-sum- -e-eseetttop.iiaaa). For h E R and v e RSS define fh) S by
f(h), = log Σn= V, exp(hy), Vi [
Then for hi, h2 E S we have
- S 11/1 h2/120.
Proof of Lemma 6. Fix E [S]. We have
Vi exp(hy) Vns(h); = (I
EvE(s] Vin exp(hx)) SE[S]*
so that (4)ell = 1. By intermediate value theorem, we have (s(b)); f(h2), (V.((S)(( -ha)l (V.(()) (h I| halloo 1701 h2lboo- This proves Lemma 6.
23

---

Published as a conference paper at ICLR 2025
Lemma 7 (Lipschitzness of the normalization operator), For h E RR, define normalize(h) E R by
normalize(h); h4 max, hy, Vi [S].
Then for hi, h2 E R we have
|normalize(hi) normalize(ha)lle S 2 halloo-
Proof of Lemma 7, Note we have the following inequality
maxhiy max h2.51 S ll helloo.
3
so that
|normalize(h:): (orm) & llı h2 loc max h1.5 max h2 2 h22loo.
This completes the proof of Lemma 7.
Lemma (ReLU approximation of the exponential function). For any take M [1/0]+1] N. Then there exists {(a,, wj, 6ū SE(M with
sup lر,| 2, sup/wy) [pyl $1,
sup 6,و S log M,
(28)
such that defining expş R + R by
M exp, (a) -i=1 0j ReLUU(wy + by),
we have expş is non-decreasing on (-o0, 0). and
sup =E(-oc,0)
(exp(x) - (exp((x)) S exps(0) = 1.
Proof of Lemmaa 8. Define e, j/(M - 1), - -looee, for j [M 1], a, (ez (ey-e))/(b2
b1) and a₁ = (ej+1 - e,) (bj+1 by) - (e, - -e--1)/(b, bbj-) for < j I M 2. Furthermore, define
M-2
expa(x)= = Σ 1=1
a ReLU(x + b,) + ReLU(-u+ ReLU(-x)..
Then we have exp)(-b))) e, for j E [M 1]. and exp, is piece-wise linear and non-decreasing on (-o0, O). Note that also have exp( bb) ej for j [M 1]. and exp is increasing (-o0, 0j. This proves that sUPxe(-o<.0) exp(x) exps(x)) I 5. Furthermore, it is easy to see that (0) and exps is non-decreasing on (-oc, Oj. Finally, since exp is 1-Lipschitz, it is easy see that sup)e[M] (ayl < 2. It is also easy see the other parts of Eq. (28) are satisfied, and this proves Lemma 8.
Lemma 9 (ReLU approximation of the logarithm function). For any A > 0, O, take M [2A/0] + - N. Then there exists {(aj. Wy b,) be)]) with
sup |b, < A, sup | S 2A, sup |uyl 11,
(29)
,
such that defining loga R + R by
logs(x) Em aj ReLU(w,r + by),
have logg is non-decreasing [1/A, A), and
log(æ>) logs (x) < 8. (AA.AA
sup
Proof of Lemma 9, The proof of Lemma is similar to Lemma 8.
24

---

Published as a conference paper at ICLR 2025
Lemma 10 (ReLU approximation of indicator function). Define
Ind(x) 2ReLU(x 1/2) 2ReLU(x 1/2) (R)LU(x))
we have
1(x - - Ind(r - j), Vj 2.
Proof of Lemma 10. The lemma holds by direct calculation.
Lemma 11 (Log-sum-exponentaa) approximation). Assume logs. IR R and exps; :R + R are
such that,
sup =E[1/K,SK]
log(x) - logs, (x)) S 61,
sup =E(-00.0
|exp(>) - expo, (x) S 62.
exP6, (0) = exPb2 is non-decreasing on (-o0,0).
Assume that 1/K S minij Vij S maxij Vi S K. Define f(h), S4 (h) E R by
log f(h); = log Eÿ= 1 exp(hj): (h); = logg, Σi= Vi exps (hz): vi [S].
Then we have
sup 11f(M) - (,),, 51 + Sk max, 1,=0
Proof of Lemma 11. We have
10(h); (81.82(1); log V exp(h)) - logs, exps, (h))
$ llog(V..: exps, (h)) logs, V c. expo, (h)) + log v e. exp(h); log v.. expa, (h)))-
For the first term, since exp,, (h) S for all h S 0. we have v exps, (h)) < S max,j V1 S SK. Furthermore, since max; h, and exps, (0) 1, we have (Ve, exps, (1)) mini Vij 2 1/K. As consequence, by assumption, we have
Vo, exps, (h)) logs, (V.., expo, (h)) S 61.
For the second term, since both (Ve,, exp())) and (Ve, exps, (h)) are within [1/K, SK] on which log function has Lipschitz constant K, we have
log((V: i., exp(h)) log(U V.. exps, (b)) S SK max Vi |xexp(h) exps,(A)l S SK2
This finishes the proof of Lemma 11.
Lemma 12 (Log-Psi approximation). Assume log, IR. - R is such that,
sup [log(x) logg S 8. xE[I/KK]
Assume that 1/K < min;j Vi S max,j 1 S K. For E [S], define f(x), fe(x) e R by
f(z); log E5 U 1, 1xx = 5), fa(x), = logg Ex v Wizl(x 5). Vi [S].
Then we have
sup
[ fo( 8..
=E|S]
Proof of Lemma 12. For any fixed x E [S] and i E [S], we have
[f(x), fs(2)a log Vie log, Vix1 < 8,
where the last inequality is by assumption. This proves Lemma 12.
25

---

Published as a conference paper at ICLR 2025
Lemma 13 (ReLU approximation of log-sum-exponentia)) Assume that 1/K Σ mingy is 1
maxij Vij S K. For any 5 > 0, take M1 = [2SK/6] /6| + 1 and M2 [2SK2 /0] + 1. there exists {laj, Wy. bj) and {(āj, Wj. b5)aa with
Then
sup layl S 2SK, 3
sup layl S 1, sup suplbsl - SK, sup وما وراةة 2, sup أراا 1 sup 1byl s log(4SK2 16),
such that defining
Mo logs. (x) ΣØ= Mi Uj ReLU(W52 + bj), expo. (2) = =i== a, ReLUU@@e + b,),
and defining f(h), R$ by
s(h); = log Σك Vis exp(h;), fa(h); = logs. Σ expsx Vi E [S],
we have
sup max, h,=0
[1((h) fa(h)llo 5
Proof of Lemma 13. Lemma 9 implies that taking M1 [2SK/0] + 1, there exists
{aaj, wy [E[M) with
sup las1 s 2SK, supluyl 1, sup |byl SK, ;
such that defining
Mi 1j=1a;; 02 ReLU(wwar + by); logs, (x)
have
sup 1/($K)$=59K
log(z) logg (x) S (x) 6/2.
Lemma 8 implies that taking M2 = |2SK2/8 /8/ + 1, there exists {(ā,, ūū b,) JE(M))) with
sup ag < 2, sup 1 < 1, sup |ر, < log(4SKK) (8)),
such that defining
Mr expx (x) Ex aj ReLU bb),
we have exps, is non-decreasing on (-o0, 0J. and
|exp(r) - exps. (x)) S 6/(SK3), (0ps.(0)= =1 sup =E(-
As a consequence, the condition of Lemma 11 is satisfied with d1 8/2 and 82 8/(2SK3), so that we have
sup (1((h) - fs(h))l S + SK252 max; h;=0
This finishes the proof of Lemma 13.
Lemma 14 (Existence of ReLU network approximating aog-sum-expooentiin Let fa be the func- tion as defined in Lemma 13, Then there exists a two-hidden-lan layer neural network
NNw, W.WW, (h) W1 ReLU(W, ReLU [h; 1)),
with W1 E SX5M W2 E R SM, (SM=-11), Wa E ((5M++1)(($+)) and
[| mma < 2SK, ||22|max - Mo)) ||a||max S Log(4SK3 18).
such that
NN (h) = fo(h), Vh E R such that maxh, 0.
26

---

Published as a conference paper at ICLR 2025
Proof of Lemma 14. Define
a i
0
0
0
d M:M2
0
0
W₂
E IR
(5 (5M2))
0
112
0
Wi:Ma
b1:M,
W1:Ma
b1:Mz
W
W1:Ma
bi:Mo
then we have
(exp&+ (h); 1] W₂ ReLU(Ws [h; 1]).
Define
aiM
T
W1-
i:M1
R SxSM
di:M,
Wi:M
0
0
b1:
W2-
0
W1:M
0
b1:M5
E (M+x
W1:M,
b1:ae,
then we have
logs. (h) = Wi ReLU(W2 [h; 1)).
As consequence, we have
f. (h) = logs. Vp exps,(h) = W1 W.-ReLLU 2)- 2-ReeLU 1)) = Wr-ReLU(gg-RLLU(N 1])),
where we define W2 = W₂ diag( V, 1)- W2. 2. also direct to upper bound [[i]] ml W| mmx. and [[a Imax- This finishes the proof of Lemma 14.
Lemma 15 (ReLU approximation of log-Psi). Assume that 1/K S minsy Viy I max,, Vij S K. For any {lag, Wy b,) b5))) with
0, take M1 [K/8] + and M2 3. Then there exists {(aj. Wj, by) by) [E([M)] and
sup ر, S 1, sup |ayl & 2K, sup |ayl S 1, sup|byl SK, sup jāy1 s 4, sup 1
j
such that defining
M2 logs(x) = 1==1 aj ReLUU b,), Ind(x) = Ej āy ReLU(W,r + b,),
and defining [a)) E ERS by
logg Σ v LoyInd(o - j), Vi E [S]. f(z), log Σ 11 =() j), fo(x);
we have
sup
ll(()) fo(z) <
=Eis|
Proof of Lemma 15. The proof of the lemma is similar to the proof of Lemma 13, using a combina- tion of Lemma 10, 9, and 12.
Lemma 16 (Existence of ReLU network approximating log-Psi). Let fo be the function as defined in Lemma 15. Then there exists a awo-hidden- -ayeer neural network
NNw, WNw..wawws (x) = W1 ReLU(W, ReLU(W; 1])), with W1 E RSSMM, W2 E RSM. (SMM-11) Wa ((5M++1)x(5+1)) and
[| [max < 2K, ||a]|lax S Poly(SKM, [[a max S 1.
such that
NNW, (x) = fa(x), Yr E [S].
Proof of Lemma 16. The proof of Lemma 16 similar the proof of Lemma 14.
27

---

Published as a conference paper at ICLR 2025
E PROOF OF THEOREM 1
Proof of Theorem I. By Lemma 17, we have the error decomposition
D H.) I wew D w m.)) + suPwew R(uNN) R(UNN)
To control the first term (the approximation enror), by Theorem 4, there exists W E W..mm. I.S.S. D. as in Eq. (9) with norm bound B = Poly (d, K, K, D), such that defining IN as in Eq. (ConvNet), we have
W
log (, - log W NN (u/zz) SC. S2
max VE (S)+
D
Furthermore, by Lemma 4, when D 2 52k 34, we have
inf D2 U4) c(e 8 stnz -1) =c S"K wew
D²
To control the second term (the generalization error), by Proposition 7, with probability at least 1-17, 77, we have
< C.. LD(D +2 1) m
log (d m
DDB 34) log(1/1)
sup W.Wa... L.S.D,B
Combining the above two equations proves Theorem 1.
E.1 ERROR DECOMPOSITION
Lemma 17. Consider the setting of Theorem 1. We have decomposition
H.) I inf DH p.) + supwew R(uK)) R(mNK)
Proof of Lemma 17. We have that for any conditional distribution (I(-))). there is decomposition
s
D3 = Exenga
Σ
(1(() - (s(2))
s=1
S ((m()))z) 1(y- = 8})2] 1E
(m)()) - 1fv a}}²] = R(w1) -(1+).
Define
W, = arg min R(HW)) = arg min
WEW
WEW
Then we have
DE )4.) = R(MNN)) R(U))
= R(UNK) R( + R(UNK) R(AN)) + R(UN)) - R(HW)) R(uN)) R(n.)
2. sup R(UNK) R(ANK) + +D WEW
This proves Lemma 17.
E.2 RESULTS ON GENERALIZATION
Proposition 7 (Generalization error of the classification problem). Let L. be the set defined as in Eq. (9). Then, with probability at least 17, have
LD(D + 25 1) m log (d m
DSB 3*) (1gg))))
R - R(WN < C.
sup
W W.Waa.L.s.D.
28

---

Published as a conference paper at ICLR 2025
Proof of Proposition 7. In Lemma 3, can take (y.z), w W, θ Wd.m L.S.D.B p(w, w²) IW Wl and [(=;:) loss(y) y, NN (-(:)). Therefore, show Proposition 7, we just need to apply Lemma 3 by checking (a), (b), (c).
W.m.L.S.D.. equipped with p(u, w') IIW Will has diameter Bp == 2B. Further note that W...I.SS.D. has a dimension bounded by dp = D(D + 25 + m)]|m|]| According to Example 5.8 of (Wainwright, 2019), it holds that log WAm. II-I) S dp llo(( + 2r/ A) for any 0 < A < 2r. This verifies (a).
Check (a). We note that the index Q
Check (b). Since w(z;w) = loss(y, (MN((11=)) is 2-bounded. As a consequence, f(z, w) E. [= w)] is sub-Gaussian random variable with the sub-Gaussian parameter to be a universal constant.
Check (c). Lemma 20 implies that
[[(2;)]) [2; w2)l o'. |W, W21. d' mm (3B*) d. 58/2 (S + D).
As consequence, (2; w1) (2; w2) is o o p(w1, W2) o I W2l sub-Gaussian
Therefore, we apply Lemma to conclude the proof of Proposition 7.
E.3 AUXILLARY LEMMAS
Lemma 18 (Norm bound in the chain rule in classification settings). Consider the ConvNei in Eq. (ConvNet). Assume that [ S B. Then for any E, v, L, and [3]. we have
[ 2| S 53/2
||q 2| B² (11bg0 1|2 1),
[|hg 1/(4-1112 S 2m (0) max
119 112.
v'EC(e)
max ve [S]
V w opp S B3
max
V
w 2, ops
Vk 2 E 1+1,
iE|S
wingg olop S B2 ||b! 112.
max (E[S]
V with
( 1lop (0)
max max V w w v'El(r) SE|S]
4k 4k
sE[S]
max EE(S)
V
w:o w: sortmaxft ):llop S max | w:o 7, llop.
(E(S)
Proof of Lemma 18. The proof of the lemma uses the chain rule, the 1-Lipschitzness of ReLU, the 2- L-pisshhizeesss of normalize, and the 1-Lipschitzness of softmax.
Lemma 19. Consider the ConvNet as in Eq. (ConvNet). Assume that [| 1 < B. Then we have
max V wi! softmax(h?"))llll < (3B²)² b .d. 53/2
max @E[S| +E[3] DE
LE (E(M SE[S]
Proof of Lemma 19. This lemma is implied by Lemma 18 and an induction argument. Lemma 20. Consider the ConvNet as in Eq. (ComNer)). Assume that ||| I B. Then we have
max V.ue
UNN (y)l xx) MNN (v)z)) s 3 |m]l (imll((3/)))) (3B*)4 4 d S3/2 (S + D)
- |W-
WW].
Therefore, we have
loss(y, ( x)) loss(y, MNN (-1z)) < 12 m (3B3)L d. S3/2 (S + D) IW W..
Proof of Lemma 20. The first inequality is by the fact that
W (y/(z)
w (yl x))
max 3.
PNN
s Σ
Σ
Σ
Hin(nrow(W!!)) (Wol(W (oostmmaxhh)) (elop))))))lpppp)))))) (W - llop
A RE(L)
AE [m)0)
29

---

Published as a conference paper at ICLR 2025
where we have used the inequality that trace(A" B) (nrow(A), (col)) BBllpp To prove the second equation, we have
8 1(y= =) - (11)) 2
W (NN( x)) - loossssyy MNN (-12)) = Σ loss(y,
Σ [1 $s ( '
NNN My = 5} MNN (s|2)) + -1 (N (s|z) - -NNk
(o (s|zz) S 4- max
44-max (NN(( PNN (5|2z)
This completes the proof of Lemma 20.
F PROOF OF PROPOSITION 5
Proof of Proposition 5. By Eq. (MP-DNS) and (23), defining (1)()) softmax(h (h).)) then for
8 < L
get
(8) (x! x
II ('EC(v)
Σ
((tt)) ((o)) (at a)e (4
eE[s]
Σ
II
(ett)) ((0)
(o x y (6t))) (2(tt))
(2t1)
(f+1)
v'EC(o)
=c())
This coincides with the update rule of (6) 4. as in Eq. (BP-DNS).
Furthermore, defining (0)) (.0()) (.) softmax
(0) 1o
( then for l 0, 1,..,. I. we get
(x x
(i(() (b, 2 ju ((+1)
II
Σ
(b, ale a) (1
d(or)
BE(S)
vEN())
aE(S]
( w1 (1 ((pa(o)) )u (ttia(t)) (1(pet)))
II
(8) (x.9))
(2-13 pa(u) Mit)
VEN())
(0) This coincides with the update rule of 1.u as in Eq. (BP-DNS).
Finally, defining (1) (.) = softmax
(b
1.u we get
(x(()) (b)) x
((1) (b, x
(L)
3o (ap(( (6)
II
[ aE[s]
wi (a)ela)) x 110000
bE[s]
v'EN(o)
x (1)) f.v (z (o(()))i((+11) (2(1) (z) 2.).
This coincides with the formula of (L) in Eq. (BP-DNS).
This finishes the proof of Proposition
G PROOF OF THEOREM 6
Proof of Theorem 6. By Lemma 13, take Mi ==252Kd188 Kd18 /0|+1 and M₂ 1. Then there exists {(a,, wj, by)bsema and {(ā,, W,, by JE[M21] with
/81+
sup |a,1 < 2SK, i
sup juyl < 1, sup suplbyl < SK, sup ول < sup kuy S 1,
sup suplbyl S 1og(4-18" Sak2//))
j
j
such that defining
logs. (x) Σドー
(x)= aj ReLUU(we by),
(x) Ex aj ReLU( + bz).
30

---

Published as a conference paper at ICLR 2025
and defining 16. (o. & (1 by
(22))) log - j=1 v/ (i, exp(h)),
[ (hi) logó+ E (i, (kpp)(tg))
e [S].
(2) (j. 1) exp(h)), (42 log Σ 2
Jt (h); logs. Ei=1 - (j. expox
Vi [S].
have
sup max; R2
1 (h) for (h)ll S 88/(d8''
(30)
By Eq. (30) and Lemma 21, taking ( E RR to be as defined Eq. (MP-DNS) and ((4)) be as defined in Eq. (A-MP-DNS) with fo, hee[L].EE/mtn as defined above, we have 10 - 1. - (8/(d18))) x 18 = 8/S,
f," E R to
which gives
sup [m))) m S8. 11 b(4 0c S 8. EER
As consequence, we just need to show that the approximate version of message passing algorithm as in Eq. (A-MP-DNS) could be cast as a neural network.
Indeed, by Lemma 14, there exist two-hidden-layer neural networks (for E [h E E [L]. and (m)
NN w/12 1.0.. 2.... s....
w22 W... (h) = 1,0.c ReLU(W
2,0, ReLU(W,, 3,0, [4: 1]));
with (7(e)
Loo,s E RSSSM w(()
2,0,6 E R. (SM w 3,0, E ((($M)+))((++1), and
[ L,o, mmax 2SK,
| 2,, max I Poly(SKMMM)), || 3,,, log(45 d18 k2)8).
such that
NN wr
222 22. w0o (h) = fo)( 4h E R such that max h = 0.
3
This proves that the approximate version of message passing in Eq. (A-MP-DNS) coincides with the U-Net as in Eq. (UNet) with proper choice of dimension
2 max S. + 1 = x (125 /01 + +1) + rE[L]
and norm of the weights This finishes the proof of Theorem 6.
G.1 AUXILLARY LEMMAS
Lemma 21 (Error propagation of the approximate version of message passing in denoising). Assume we have functions UL9 ((0) and (8 f(()) such that
1 (h) 78. (h) 5, wh E R such that max hg = 0, (L
(31)
jE(S]
Furthermore, consider the following approximate version of message passing algorithm with initial- ization h (- =2)) E R for u E V(L), defined as below
(.10 Enlon (normalize (h R
E [4], Ve V(0)
h (--1) Evecco) L. ERS.
E [L], Ve v(e-1),
u T.u = (Tp(()) ERS. (with 6 = (
E [4], VE V(0)
7. to(l)) (normalize(af% (a (.21)) ) 1. E
(()) F(e)
(9)
le [L], VE V())
mMr (z), EvEls] 8 softmax ((L))
e V())
(A-MP-DNS)
Taking (L)) E S to be as defined in Eq. (MP-DNS) and (LL) 7,0
E R
to be as defined in Eq. (A-MP-
DNS), we have
-((2) S 56x 18 d. [b(L
max
"EV(L)
31

---

Published as a conference paper at ICLR 2025
Proof of Lemma 21.
Step 1. Downward induction. In the first step, we aim to show that for any & E [L - 1] we have
| N.olllo S m(++1) [Tk= + 1)ō, E V(0)
(32)
To prove the formula for l = L- 1, 1, since uormalizeeh( (4) (ormalize(h(..)) by Eq. (31), we get
1 -
(. Só, Vo E vin)
Hence we get
(L-1)
Vu E v(e-1), ( 1100 Eweeet) - - S mml
This proves the formula (32) for =
Assuming that (32) holds at the layer f, by the update formula, have
(2)
= 11 (normalizo(h)) (h. (u(()) (normalizo(h
$ ((0)
(.(t)) (normalizo(h1...)) f1. (normalizo (f..))) + Lo(o) (normalize(( (h)))) - -1(000 (hormalizo(h1(.)))
< + 2m (f+1)
2 + 1)0 - IIk=et1 (2m + 1)6,
where the middle inequality is by the assumption of /9 and by Lemma and Lemma Hence we get
[ - n1. ( llo S m(() (0m=++122m( + 1)6, Ve 210)
This proves Eq. (32) by the induction argument.
Step 2. Upward induction. The downward induction argument proves that, for II (2mlk)+
(k)
1), we have
L0 2. 1 10 IT, =e= = 0,1..... L, Vu V(A)
In this step, we aim show that for any 0,1, L, we have
[ T.e
(2 S 56. . Ve E V(0)
(33)
To prove this formula for 0, note that 6/0 h1p and b h [0)
we have
10 ( loo 1 ( 1lc .
This proves the formula (33) for 0.
Assuming that (33) holds at layer e 1, by the update formula, we have
16
((0
Te llo
1.0
(0)
(.(0)) (normalize (6-1)
(b((
((,)))
(a()) (normalized (e-1) (2)
x(()
(1. + [l
1.pa(r)
(((())
(normalizo (6-1)
(.01)
p(0)
ta(1)) (normalize (((1-)
1100
(a(())
T.pa(e)
T.pa(v)
(normalizo(B (1(-1)
f(()
(.(()) (normalizo T.pa(v)
((-1)
(.
t.pa(e)
< |normalize (b( (L-1)
pa(vv) ( - normalize (1-1)
c + 8
pp(())
< 6 5 - . 6º 8
This proves Eq. (33) by the induction argument. This proves the Lemma 21 by observing that T < 34 I (0) 34. d.
32

---

Published as a conference paper at ICLR 2025
H
PROOF OF THEOREM
Proof of Theorem 2. By Lemma 22, we have the enror decomposition
w Dmm m) S inf wew (m W m) + 2- supwew R(m R(m
To control the first term (the approximation enror), by Theorem 6, there exists W E L.S.L. as in Eq. (18) with norm bound B Poly(d, S, K, 18", D), such that defining m W NN as in Eq. (UNet),
NN
we have
sup [l((2) mNN(=)ll SC. S32dd 18 =ER4
Therefore, have
(m W m) supp|m((>) - mNN SC. 5 K' 182L
inf WEW
D2
To control the second term (the generalization error), by Proposition 8, with probability at least 1-17, have
- SC.S LD(D + 2S m log d ml
DSB 185) + log(1/1)
sup
W.Wa...L...D.
Combining the above two equations proves Theorem 2.
H.1 ERROR DECOMPOSITION
Lemma 22. Consider the setting of Theorem 2. We have decomposition
M m) S inf (3 W m) + 2 supwew R(mm))) - R(m
Proof of Lemma 22. We have that for any conditional expectation m1 there is decomposition
D2 m) BB d (mlmm m(2))l
R(m:) R(m). xx13] -x11 d [0'[mm1(2) (z) - x011 E, (d |1 (z)
Define
W. = arg min WEW
R(m W = arg min
M m.).
wew
Then we have
D2 km m. R(m W
w
R(m)
R(m R(m W.
+ R(m W
W; R(m R(m + R(m R(m)
2 sup
R(m wk) wew
+ D2 w.
NN: m)
This proves Lemma 22.
H.2 RESULTS ON GENTRALIZAAIIN
Proposition 8 (Generalization error the denoising problem). Let Wd..LL be the set defined as in Eq. (9). Then, with probability at least we have
R(W R(WN SC LD(D 25 + 1) m
log (d m DSB 18t) + log(1/n)
sup W W.Wam.L..DD.
33

---

Published as a conference paper at ICLR 2025
Proof of Proposition 8. In Lemma 3, can take (z, W, θ Wd.m L.S.D. D.S.D.B, p(w, w²) IW Il and [(= = | W ( Therefore, show Proposition 8, we just need to apply Lemma 3 by checking (a). (b), (c).
W.m.L.S.D.. equipped with p(u, w') I[W - W has diameter Bp == 2B. Further note that W.s.I..S.D. has a dimension bounded by dp 2D(D + 28 + 1)(mm|| 1. According Example 5.8 of (Wainwright, 2019), it holds that log Wam,Waamm........ II-1) S dp -og((1 + 2r/ A) for any 0 < A 2r. This verifies (a).
Check (a). We note that the index set Q
Check (b). Since f(z;; w) = | - (=) is bo-bbunded. As a consequence, f(z, w)
[= w)] is sub-Gaussian random variable with the sub-Gaussian parameter to be C-S*
Check (c). Lemma 25 implies that
[f(z wx) w2) S Lp |W1 Wall, 4,:- m 18 118BB" BoL d. d
Σ (S (5 (S D). DEVCL
Since A + y where (x, 9) μ. x (O Ia). and ||0||1 S Sd. |Igl| 1 is Cd-sub-Gauusian. Hence and hence f2; wi) fz; W2) is o' p(w1, W2) rub-Gaussian with
|| 11 is
o' Cl| 18 B6L de 5t (S D).
Therefore, we apply Lemma conclude the proof of Proposition 8.
H.3 AUXILLARY LEMMAS
Lemma 23 (Norm bound in the chain rule in denoising settings), Consider the U-Net as in Eq. (UNet) with modified input ((..) = (-x2/22 + x2 nev)ne(s] E IR$ (Since we will immediately nor- malize the imput, this input is effectively the same as the imput h (-(x =,))/2)zE[S) (RS), Assume that W S B. Thenfor for any L, v, i, and E [3], we have
||h 1,0 - 5 + 5
1q 14 < B (2- 10-llh 1.0 ll2 + 1):
1h ((-1)) ll2 S m (1)
(0) max 1a
e'C(()
[b/f0/12 |1 1l2
[|b[]] 12 < B3 (2 (()1)(-1)) ( ll2 + 2114 + 1) 11h..001..
pa(
max
V wi0
L. 2BB
max V wrlt 1.
AA. 2 E+1
#E
rE(S]
W L.o, s 22B ||h [1
max
*E(S)
max
1lop S m()
max max d'Ec(o) rE|S]
V| wii Lh wh.
4244
#E|5)
V W +.
2BB
max
V
with by b(t-1)
s.t. (.pa()))llopp
Ak Ak2E+1
:E(S'
vE[S]
V t T.m.alop
2BB
(1Ib (8-1)
||l + [1q ll2),
max *E(S)
max *E[S]
V Wilt Topp 5 <BB
max
V Wit b(t-1) illoo
iE(s)
max
V ( (ee.ellop)
(e)
+ max
V wis (0) vE(s] wis
Lealloop Wk [L],
iE[S]
|/ wilh) MNN ( S max
V
Wk [T]. (1.1) 1).
(E(s]
Proof of Lemma 23. The proof of the lemma uses the chain rule, the 1-Lipschitzness of RELU, the 2-1 2-Lipschitzness of normalize, and the 1-Lippchitzness of softmax.
34

---

Published as a conference paper at ICLR 2025
Lemma 24. Consider the U-Net as in Eq. (UNet). Assume that |W]| S B. Then for any V(t) we have
V wio (2)ellp S 184 BoL d. 5 (zol).
max max max max max oE(1.)) #ER* +E[3] (E((())
Proof of Lemma 24. This lemma is implied by Lemma and an induction argument.
Lemma 25. Consider the ConvNet in Eq. (UNer). Assume that || S B. Then for any
E V(L,, have
W (2) s <6|[m]]]188 18 B6L d (S + (ze)) (S + D) IW - WII. max (z)
Therefore, we have
max
|| ( | NN (2)1
S 12 18 BoL d. Sa Σ
(S (5 (S+ D) |W W
NEV(L)
Proof of Lemma 25. The first inequality is the fact that
(2)
W
W mNN (z)
< Σ ΣEWD Σ
nrow (W!?,,) ncol(W (Wo,), w!a..
OE (4.1) +E(3] RE(T. cE(m(0)
where we have used the inequality that trace (race(AFB) B) <
Bllop
To prove the second inequality, we have
x w (2) | (2)
2x W (2) (2)
W (=) (
max
< 2d max | (z) NNNN(22 (2)111.
This completes the proof of Lemma 25.
35


Reviews:
[{"review_id": "mhVKM8Knra", "reviewer": "Reviewer_iUQ6", "summary": "The paper examines the widely-used U-Net architecture under the scope of generative hierarchical models. The authors present a theoretical framework that treats the U-Net convolution and pooling layers as nodes of a hierarchical graphical model. Under their analysis, both classification and denoising with U-Nets can be seen as empirical risk minimization for a Bayes classifier/denoiser. The proposed interpretation can", "questions": "- Why are the established bounds useful for future neural network architecture development? Could they help us understand scaling laws for U-Nets? \n- What are the implications of interpreting each U-Net layer as a latent in the hierarchical model? The authors mention that \"[...] our theoretical findings generated a hypothesis of the functionality of each layer of the U-Nets\", but there are no actual hypotheses made. Can you conjecture something about the features learned in intermediate layers of the U-Net and how they have been used in other tasks [1]?\n\n[1] Baranchuk, Dmitry, et al. \"Label-Efficient Semantic Segmentation with Diffusion Models.\" International Conference on Learning Representations.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Interpreting U-Nets as hierarchical models is an interesting angle that could help explain some of the significant design choices behind training and inference in diffusion models.\n- The writing of the paper is clear and helps with understanding the intuition behind the hierarchical model interpretation of U-Nets.", "weaknesses": "- The goal of the analysis is slightly unclear. The main result of the paper establishes a bound for the difference between the U-Net denoiser and the optimal denoiser, based on a simplified U-Net architecture. However, there is no motivation on why such a bound could be useful. See questions.", "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": "Vgze9pd64t", "reviewer": "Reviewer_f34u", "summary": "U-Nets have shown promising performance across different computer vision tasks. In this paper, the authors aim to provide a theoretical explanation of the U-Net through the lens of belief propagation. Specifically, the authors have shown that the U-Net approximates the belief propagation denoising, which results in its specific model structure. In addition, the authors have shown that convolutional neural networks are well suited for classifications.", "questions": "- In Eq. 17 only ERM error is used while I think when these models are trained usually there's an extra regularization term, I was wondering does it influences the conclusion?\n\n- In neural network, the upper bound in the theorem has the tendency to be too loose, I was wondering how tight are the bounds?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Provide a rigorous theoretical explanation of U-Net and CNN through the lens of a graphical model", "weaknesses": "- Many hypotheses are made for the derivation. The conclusion would be strengthened if there were at least some toy-ish experiments to test these hypotheses\n- (minor) Given how dense the paper is, a lookup table for notations would be very helpful", "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": "besxu3wMYS", "reviewer": "Reviewer_UKfv", "summary": "This paper studies approximation and generalization of the classification function / score function in diffusion models. The data is assumed to satisfy a tree-structured Markov random field model with hierarchical latent variables. The authors show how respectively a (variant of) CNN and UNet architectures are tailored for this problem, due to their resemblance to the belief-propagation algorithm. They obtain polynomial rates for the approximation and generalization error in terms of the network width and number of samples, showing that these architectures together with the assumptions on the data are able to break the curse of dimensionality.", "questions": "- Why is the assumption that children are independent when conditioned on their parent necessary? Belief propagation would still be possible (with a computational complexity exponential in the branching factor, but this would remain manageable if it is small).\n- Do the authors believe that their model captures the relevant properties of image distributions? How would they go about testing it? And if no, what in their opinion are the missing features that should be studied in future work?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper studies an important problem, which is to understand what properties of the data distribution coupled with inductive biases of network architectures allow to break the curse of dimensionality when learning diffusion models. The observation that CNNs and UNet correspond to classical Bayesian inference algorithms in the case of tree-structured models is an interesting one. To the best of my knowledge, this is the first work that obtains such results for deep networks.", "weaknesses": "The main weakness of the paper is that the data assumptions are very restrictive. In particular, they exist simpler and more efficient classical algorithms to learn and sample from these distributions than diffusion models.", "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": "iLHaCMrTdw", "reviewer": "Reviewer_xjiq", "summary": "This paper presents a new perspective on the U-Net architecture, linking it to tree-structured models and the belief propagation denoising algorithm. The authors derive sample complexity bounds for U-Nets in learning the denoising function, employing a proof strategy similar to that used for ConvNets in classification tasks.", "questions": "See the Weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(1) Authors offer a novel interpretation of U-Nets by linking them to the belief propagation algorithm. It’s intriguing how this algorithm naturally gives rise to key U-Net components, such as skip connections and pooling layers.\n\n(2) Authors conduct a thorough complexity analysis, with extensive proofs provided in the appendix.\n\n(3) The paper is well-written, and the proof strategy in the ConvNet example aids in understanding the later sections on U-Nets.", "weaknesses": "(1) Although this is a theoretical paper, it lacks any form of empirical validation of its conclusions, even in a simplified format.\n\n(2) Why is the belief propagation algorithm applied exclusively to U-Nets? For instance, ResNet’s residual connections could also approximate “long skip connections” that link all nodes, disregarding the up/down sampling operations.\n\n(3) The title suggests that U-Nets are efficient for classification, yet the paper only analyzes classification in the context of ConvNets.", "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": "0qdq3vB55Y", "reviewer": "Reviewer_5xWe", "summary": "This paper provides a theoretical framework for understanding U-Nets. The authors present the connection between U-Nets and belief propagation algorithms within generative hierarchical models (GHMs). They show how the components of the U-Net model naturally implement belief propagation steps, which enable efficient denoising operations. In particular, they prove how they can approximate the message parsing algorithm and, therefore, also the belief propagation algorithm, using ConvNets and U-Nets for classification and denoising tasks, respectively. Finally, based on these proofs, this paper establishes sample complexity bounds for learning ConvNet classifiers and denoising functions with U-Nets.", "questions": "Can the authors clarify whether true convolution or cross-correlation is used, given that parameter sharing and local invariance properties are retained?\n\nWhat are the potential applications of the established bounds? How could the theoretical findings be directly applied to or inform improvements in state-of-the-art diffusion models?\n\nHow do the results theoretically align with empirical observations in commonly used U-Net-based models for diffusion, such as those in image generation tasks? Could the authors elaborate on any observed gaps between theory and practice?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is very well-written. It is well-organized, with a clear outline and structure that facilitates understanding. The presentation of the generative hierarchical model, followed by the classification problem with ConvNets and the Denoising Problem with U-Nets is impeccable. \n\nThe authors demonstrate strong theoretical rigor. They present detailed and thorough proofs in a way that encourages replication and verification. The paper leaves no gaps in its logical arguments and even addresses potential objections given some simplifying assumptions made by the authors.\n\nThe literature review is extensive and effectively situates the contributions of the papers within the current research, by highlighting its value of adding a theoretical foundation for U-Nets and presenting the first sample complexity results for learning a Bayes denoiser using U-Nets.", "weaknesses": "The paper makes simplifying assumptions regarding “convolutional layers” and the use of square loss instead of cross-entropy. While it is helpful that the authors clarify these deviations from practical implementations, referring to these layers as “convolutional” may be misleading, as they do not fully adhere to the standard convolution operations. However, it is noted that parameter sharing and local invariance properties are retained, though it remains unclear whether true convolution or cross-correlation is used.\n\nThe practical utility of the provided bounds is unclear, as their direct practical applicability for empirical settings is not fully established.", "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)