U-Nets as Belief Propagation: Efficient Classification, Denoising, and Diffusion in Generative Hierarchical Models

## Abstract

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 across language and image domains.

## 1 Introduction

U-Nets are one of the most prominent network architectures in computer vision, primarily employed for tasks such as image segmentation, denoising [, , , ], and diffusion modeling [, , , ]. These networks are structured as encoder-decoder convolutional neural networks equipped with long skip connections, and their input and output typically maintain the same dimensions. While U-Nets have demonstrated exceptional performance across a variety of applications, the theoretical foundations of their key components—including the encoder-decoder structure, long skip connections, and the pooling and up-sampling layers—remain inadequately understood. Notably, long skip connections have a significant impact on performance as shown in empirical studies [, ]. Existing explanations, often anecdotal, suggest their efficacy stems from improved information propagation and reduction of the vanishing gradient issue, but a thorough theoretical exploration is still lacking.

In this paper, we introduce a 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 a tree-structured graphical model, which has been widely employed in language and image generative modeling [, , , , , ]. We detail the precise definition of such generative hierarchical models in Section 2. A series of recent work [, , , ] have pioneered the use of generative hierarchical models in studying classification tasks and diffusion models. As noted by [], 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 this 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. Additionally, we discuss the broader implications of these findings for diffusion models in generative hierarchical models.

In addition to our main findings, in Section 3, we demonstrate that the standard architecture of convolutional neural networks (ConvNets) is well-suited for classification tasks within the same generative hierarchical model. We provide efficient sample complexity results to support this assertion. This offers a 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 ${\mathcal{T}}=({\mathcal{V}},{\mathcal{E}})$ with a height of $L$, where we conventionally designate the root of the tree ${\rm r}$ as layer 0 0. For each node $v\in{\mathcal{V}}$, we denote ${\rm pa}(v)$ as the parent of $v$, ${\mathcal{C}}(v)$ as the children of $v$, and ${\mathcal{N}}(v)$ as the siblings of $v$. We denote ${\mathcal{V}}^{(\ell)}$ as the set of nodes at layer $\ell$. We assume that for any $v\in{\mathcal{V}}^{(\ell-1)}$, the number of children is precisely ${m}^{(\ell)}$ for $\ell\in[L]$. The leaf nodes $v\in{\mathcal{V}}^{(L)}$ have no children. Additionally, for each $v\in{\mathcal{V}}^{(\ell)}$, we assume an ordering function (a bijection) $\iota:{\mathcal{C}}(v)\to[{m}^{(\ell)}]$, ensuring that any child $v^{\prime}\in{\mathcal{C}}(v)$ possesses a unique rank $\iota(v^{\prime})\in[{m}^{(\ell)}]$. We denote the number of nodes at layer $\ell$ as ${d}^{(\ell)}$, and the number of nodes at layer $L$ as ${d}={d}^{(L)}$. We further denote ${\underline{m}}=({m}^{(\ell)})_{\ell\in[L]}$, and $\|{\underline{m}}\|_{1}=\sum_{\ell=1}^{L}{m}^{(\ell)}$. By these definitions and assumptions, we have $d^{(\ell)}=\prod_{1\leq s\leq\ell}{m}^{(s)}$, and $1={d}^{(0)}\leq{d}^{(1)}\leq\cdots\leq{d}^{(L)}={d}$.

For each layer $\ell=0,\ldots,L$, every tree node $v\in{\mathcal{V}}^{(\ell)}$ is associated with a variable $x_{v}^{(\ell)}\in[{S}]$ for some ${S}\in{\mathbb{N}}_{\geq 2}$. (For simplicity, we use the same variable space $[{S}]$ across all layers, although our framework can accommodate variations across different layers $\ell$.) We denote ${\bm{x}}^{(\ell)}=(x_{v}^{(\ell)})_{v\in{\mathcal{V}}^{(\ell)}}\in[{S}]^{{d}_{
\ell}}$ as the variables at layer $\ell$. The variables at the leaves, ${\bm{x}}={\bm{x}}^{(L)}\in[{S}]^{{d}}$ are considered the observed covariates, exemplified by the pixel representation of an image. Conversely, the root node variable $y=x_{\rm r}^{(0)}\in[{S}]$ is treated as the associated label. Variables for the intermediate layers $\{{\bm{x}}^{(\ell)}\}_{1\leq\ell\leq L-1}$ remain unobserved.

#### The generative hierarchical model.

We consider a specific type of generative hierarchical model (GHM)(^1^11We define “generative hierarchical models” as general probabilistic models with a hierarchical structure. The specific model discussed in this paper is an instance of such generative hierarchical models. These models are also known by various other names, including “hierarchical generative models”, “latent hierarchical models”, “Bayesian hierarchical models”, “hierarchical Markov random fields”, among others. The reasons for choosing the term “generative hierarchical models” in this paper are detailed in Appendix H.) , which is a joint distribution ${\mu_{\star}}$ over variables

$$ $(y=x^{(0)}\in[{S}],~{}~{}~{}{\bm{x}}^{(1)}\in[{S}]^{{d}^{(1)}},~{}~{}~{}\ldots ,~{}~{}~{}{\bm{x}}^{(L-1)}\in[{S}]^{{d}^{(L-1)}},~{}~{}~{}{\bm{x}}^{(L)}={\bm{ x}}\in[{S}]^{{d}}),$ $$

associated with a set of functions $\{\psi^{(\ell)}:[{S}]\times[{S}]^{m^{(\ell)}}\to{\mathbb{R}}_{\geq 0}\}_{\ell
\in[L]}$, defined as

$$ $\displaystyle~{}{\mu_{\star}}(y,{\bm{x}}^{(1)},\ldots,{\bm{x}}^{(L-1)},{\bm{x}})$ (GHM) $\displaystyle\propto$ $\displaystyle~{}\textstyle\psi^{(1)}(y,{\bm{x}}^{(1)})\cdot\Big{(}\prod_{v\in{ \mathcal{V}}^{(1)}}\psi^{(2)}(x_{v}^{(1)},x_{{\mathcal{C}}(v)}^{(2)})\Big{)} \cdots\Big{(}\prod_{v\in{\mathcal{V}}^{(L-1)}}\psi^{(L)}(x_{v}^{(L-1)},x_{{ \mathcal{C}}(v)})\Big{)}.$ $$

The formula specifies that any two nodes $v_{1},v_{2}\in{\mathcal{V}}^{(\ell)}$ within the same level uses the same function $\psi^{(\ell)}$, thereby embedding specific invariance properties into ${\mu_{\star}}$. Consequently, this GHM ensures that $(x_{v})_{v\succeq v_{1}}\stackrel{{\scriptstyle d}}{{=}}(x_{v})_{v\succeq v_{2}}$ for any $v_{1},v_{2}\in{\mathcal{V}}^{(\ell)}$. The notation $v\succeq v_{1}$ denotes that $v$ is either identical to or a descendant of $v_{1}$. We will explore later how this invariance property interacts with the convolutional structure of the neural networks to be introduced.

Throughout the paper, we impose a specific factorization assumption on the $\psi$ functions of GHM.

###### Assumption 1 (Factorization of $\psi$ ) .

For each layer $\ell\in[L]$ and node $v\in{\mathcal{V}}^{(\ell-1)}$, we have

$$ $\textstyle\psi^{(\ell)}(x^{(\ell-1)}_{v},x^{(\ell)}_{{\mathcal{C}}(v)})=\prod_ {v^{\prime}\in{\mathcal{C}}(v)}\psi^{(\ell)}_{\iota(v^{\prime})}(x_{v}^{(\ell- 1)},x_{v^{\prime}}^{(\ell)}).$ (1) $$

We also state a technical assumption concerning the boundedness of these $\psi$ functions of GHM.

###### Assumption 2 (Boundedness of $\psi$ ) .

For any layer $\ell\in[L]$ and child node rank $\iota\in[{m}^{(\ell)}]$, the transition probabilities are bounded as follows:

$$ $1/K\leq\min_{x,x^{\prime}}\psi_{\iota}^{(\ell)}(x,x^{\prime})\leq\max_{x,x^{ \prime}}\psi_{\iota}^{(\ell)}(x,x^{\prime})\leq K.$ (2) $$

It is helpful to think about the joint distribution ${\mu_{\star}}$ as a tree-structured Markov process, which admits the factorization

$$ $\displaystyle~{}{\mu_{\star}}(y,{\bm{x}}^{(1)},\ldots,{\bm{x}}^{(L-1)},{\bm{x}})$ (3) $\displaystyle=$ $\displaystyle~{}\textstyle{\mu_{\star}}(y){\mu_{\star}}({\bm{x}}^{(1)}|y)\cdot \Big{(}\prod_{v\in{\mathcal{V}}^{(1)}}{\mu_{\star}}(x_{{\mathcal{C}}(v)}^{(2)} |x_{v}^{(1)})\Big{)}\cdots\Big{(}\prod_{v\in{\mathcal{V}}^{(L-1)}}{\mu_{\star} }(x_{{\mathcal{C}}(v)}|x_{v}^{(L-1)})\Big{)},$ $$

where we abuse the notation to denote ${\mu_{\star}}(x_{{\mathcal{C}}(v)}^{(\ell)}|x_{v}^{(\ell-1)})$ as the conditional probability of $x_{{\mathcal{C}}(v)}^{(\ell)}$ given $x_{v}^{(\ell-1)}$, and ${\mu_{\star}}(y)$ as the marginal probability of $y$. Indeed, any graphical model specified as in Eq. (1) can be cast into the form of (3). Furthermore, It can be checked that Eq. (3) coincides with (GHM) upon taking $\psi^{(\ell)}(x^{(\ell-1)}_{v},x^{(\ell)}_{{\mathcal{C}}(v)})={\mu_{\star}}(x_
{{\mathcal{C}}(v)}^{(\ell)}|x_{v}^{(\ell-1)})$ for $\ell\geq 1$, and $\psi^{(1)}(y,{\bm{x}}^{(1)})={\mu_{\star}}(y){\mu_{\star}}({\bm{x}}^{(1)}|y)$. In this scenario, the factorization assumption (1) implies that $(x_{v^{\prime}}^{\ell})_{v^{\prime}\in{\mathcal{C}}(v)}$ given $x_{v}^{(\ell-1)}$ are conditionally independent. We include a schematic plot of the generative hierarchical model with $3$ layers as in Figure 1(left).

Figure: Figure 1: Left: The generative hierarchical model with $3$ layers and ${m}^{(1)}=3$, ${m}^{(2)}=3$, and ${m}^{(3)}=2$ children in each layer. Right: A $3$-layer convolutional neural network.
Refer to caption: x1.png

###### Assumption 1 (Factorization of ψ 𝜓 \psi italic_ψ ) .

###### Assumption 2 (Boundedness of ψ 𝜓 \psi italic_ψ ) .

#### GHMs as natural models for languages and images.

In the field of linguistics, GHMs are very similar to context-free grammars (CFGs) [, , ]. The generative process of a context-free grammar involves creating valid strings or sentences based on a given set of production rules (the $\psi$ 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 to the terminal set, thus forming a complete sentence (the covariate ${\bm{x}}$).

In computer vision, GHMs are often utilized to model natural images [, , ], where they are sometimes referred to as multi-resolution Markov models. The hypothetical image generation process begins with a high-level concept of the image (represented by the label $y$), which is then iteratively refined using a production rule (the $\psi$ function). This refinement continues through successive resolution levels until the detail reaches the pixel level, resulting in the final image (the covariate ${\bm{x}}$).

We remark that a series of studies [, , ] 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 subsequent 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 a set of iid samples $\{({\bm{x}}_{i},y_{i})\}_{i\in[n]}$ drawn from ${\mu_{\star}}$ under the GHM. Our objective is to learn a probabilistic classifier $\hat{\mu}(y|{\bm{x}})$ from this dataset. With a suitable loss function, the optimal classifier is the Bayes classifier ${\mu_{\star}}(y|{\bm{x}})$, which represents the true conditional probability of $y$ given ${\bm{x}}$. We aim to examine the sample complexity of learning this classifier through empirical risk minimization over the class of convolutional neural networks.

#### The ConvNet architecture.

We here introduce the convolutional neural network (ConvNet) architecture used for classification, represented as $\mu_{{\rm NN}}(\cdot|{\bm{x}})\in\Delta([{S}])$ for input ${\bm{x}}\in[{S}]^{d}$. Initially, we set ${\mathsf{h}}^{(L)}_{v}=x_{v}\in[{S}]$ for each node $v\in{\mathcal{V}}^{(L)}$. The operational flow of the network unfolds as follows:

$$ $\displaystyle{\mathsf{q}}_{v}^{(\ell)}=$ $\displaystyle~{}{\rm NN}_{\iota(v)}^{(\ell)}({\mathsf{h}}_{v}^{(\ell)})\in{ \mathbb{R}}^{S},~{}~{}$ $\displaystyle\ell\in[L],~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (ConvNet) $\displaystyle{\mathsf{h}}_{v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle{\rm normalize}\Big{(}\sum_{v^{\prime}\in{\mathcal{C }}(v)}{\mathsf{q}}_{v^{\prime}}^{(\ell)}\Big{)}\in{\mathbb{R}}^{S},~{}~{}$ $\displaystyle\ell\in[L],~{}~{}v\in{\mathcal{V}}^{(\ell-1)},$ $\displaystyle\mu_{{\rm NN}}(\cdot|{\bm{x}})=$ $\displaystyle~{}{\rm softmax}({\mathsf{h}}_{{\rm r}}^{(0)})\in\Delta([{S}]).$ $$

The functions ${\rm NN}_{\iota(v)}^{(\ell)}:{\mathbb{R}}^{{S_{{\rm in}}^{(\ell)}}}\to{\mathbb
{R}}^{S}$ (which adjust their input dimensionality ${S_{{\rm in}}^{(\ell)}}$ based on layer depth, ${S_{{\rm in}}^{(\ell)}}=({S}+1)\cdot 1\{\ell\neq L\}+2\cdot 1\{\ell=L\}$) and ${\rm normalize}:{\mathbb{R}}^{S}\to{\mathbb{R}}^{S}$ are defined as follows:

$$ $\displaystyle{\rm NN}_{\iota}^{(\ell)}({\mathsf{h}})=$ $\displaystyle~{}W_{1,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{2,\iota}^{(\ell)}\cdot{ \rm ReLU}(W_{3,\iota}^{(\ell)}\cdot[{\mathsf{h}};1])),$ (4) $\displaystyle{\rm normalize}({h})_{s}=$ $\displaystyle~{}\textstyle{h}_{s}-\max_{s^{\prime}\in[{S}]}{h}_{s^{\prime}}.$ (5) $$

The dimensions of the weight matrices within the ConvNet are specified as follows:

$$ ${\bm{W}}=\Big{\{}\{W_{1,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}\times D}\}_{\iota \in[{m}^{(\ell)}]},\{W_{2,\iota}^{(\ell)}\in{\mathbb{R}}^{D\times D}\}_{\iota \in[{m}^{(\ell)}]},\{W_{3,\iota}^{(\ell)}\in{\mathbb{R}}^{D\times{S_{{\rm in}} ^{(\ell)}}}\}_{\iota\in[{m}^{(\ell)}]}\Big{\}}_{\ell\in[L]}.$ (6) $$

Furthermore, we denote $\mu_{{\rm NN}}^{{\bm{W}}}$ as the ConvNet classifier $\mu_{{\rm NN}}$ parameterized by ${\bm{W}}$. A schematic illustration of a ConvNet with $3$ layers is provided in Figure 1(right).

###### Remark 1 (An explanation of the “ConvNet” architecture) .

We remark that the neural network layers described in (ConvNet) are different from the “convolution operations” typically seen in practice. The convolution operations used in practice involve computing the inner products between convolutional filters and image patches, whereas in (ConvNet) and (4), a point-wise product is employed instead. Despite this, these layers are still referred to as convolutional layers because the mapping from $\{{\mathsf{h}}^{(\ell)}_{v}\}_{v\in{\mathcal{V}}^{(\ell)}}$ to $\{{\mathsf{q}}^{(\ell)}_{v}\}_{v\in{\mathcal{V}}^{(\ell)}}$, as per the first line of (ConvNet), preserves the translation-invariance property. Specifically, we use the same function ${\rm NN}_{\iota}^{(\ell)}$ across different inputs ${\mathsf{h}}^{(\ell)}_{v_{1}}$ and ${\mathsf{h}}^{(\ell)}_{v_{2}}$ as long as $\iota(v_{1})=\iota(v_{2})=\iota$.

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, (ConvNet) 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.

###### Remark 1 (An explanation of the “ConvNet” architecture) .

#### The ERM estimator.

In the classification task, we employ empirical risk minimization over ConvNets as outlined in the following equation:

$$ $\widehat{\bm{W}}=\arg\min_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S} ,D,B}}\Big{\{}{\widehat{\sf R}}(\mu_{{\rm NN}}^{{\bm{W}}})=\frac{1}{n}\sum_{i= 1}^{n}{\rm loss}(y_{i},\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}}_{i}))\Big{\}},$ (7) $$

where, for simplicity of analysis, we opt for the square loss rather than the more commonly used cross-entropy loss:

$$ $\textstyle{\rm loss}(y,\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}}))=\sum_{s=1}^{ {S}}\Big{(}1\{y=s\}-\mu_{{\rm NN}}^{{\bm{W}}}(s|{\bm{x}})\Big{)}^{2}.$ (8) $$

The parameter space for the ConvNets is defined as:

$$ $\displaystyle{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}:=\Big{\{}{\bm{W}} \text{ as defined in (\ref{eqn:W_set_classification})}:{|\!|\!|{{\bm{W}}}|\!| \!|}:=\max_{j\in[3]}\max_{\ell\in[L]}\max_{\iota\in[m^{(\ell)}]}\|W_{j,\iota}^ {(\ell)}\|_{{\rm op}}\leq B\Big{\}}.$ (9) $$

We anticipate that the empirical risk minimizer, $\mu_{{\rm NN}}^{\widehat{\bm{W}}}$, could learn the Bayes classifier ${\mu_{\star}}$, as the global minimizer of the population risk over all conditional distributions yields the Bayes classifier:

$$ ${\mu_{\star}}(\cdot|\cdot)=\arg\min_{\mu}\Big{\{}{\sf R}(\mu)={\mathbb{E}}[{ \rm loss}(y,\mu(\cdot|{\bm{x}}))]\Big{\}}.$ $$

In our theoretical analysis, we measure the discrepancy between $\mu_{\rm NN}^{\widehat{\bm{W}}}$ and ${\mu_{\star}}$ using the squared Euclidean distance:

$$ $\textstyle{\sf D_{2}^{2}}(\mu,{\mu_{\star}})={\mathbb{E}}_{{\bm{x}}\sim{\mu_{ \star}}}\Big{[}\sum_{s=1}^{{S}}\Big{(}\mu(s|{\bm{x}})-{\mu_{\star}}(s|{\bm{x}} )\Big{)}^{2}\Big{]}.$ (10) $$

#### Sample complexity bound.

The subsequent theorem establishes the bound of the ${\sf D_{2}^{2}}$-distance between the ConvNet estimator $\mu_{{\rm NN}}^{\widehat{\bm{W}}}$ and the true Bayes classifier ${\mu_{\star}}$.

###### Theorem 1 (Learning to classify using ConvNets) .

Let Assumption 1 and 2 hold. Let ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ be the set defined as in Eq. (9), where $D\geq{S}^{2}K^{2}{d}\cdot 3^{L}$ and $B={\rm Poly}({d},{S},K,3^{L},D)$. Let $\widehat{\bm{W}}$ be the empirical risk minimizer as in Eq. (7). Then with probability at least $1-\eta$, we have

$$ ${\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_{\star}})\leq C\cdot \Big{(}\frac{{S}^{4}K^{4}{d}^{2}\cdot 3^{2L}}{D^{2}}+\sqrt{\frac{LD(D+2S+1)\|{ \underline{m}}\|_{1}\log({d}\|{\underline{m}}\|_{1}D{S}\cdot 3^{L})+\log(1/ \eta)}{n}}\Big{)}.$ (11) $$

The proof of Theorem 1 is detailed in Section D.

###### Remark 2 .

To ensure the ${\sf D_{2}^{2}}$-distance is less than ${\varepsilon}^{2}$, Theorem 1 requires to take

$$ $D=\Theta({S}^{2}K^{2}{d}3^{L}/{\varepsilon}),~{}~{}~{}~{}~{}n=\tilde{\Theta}(L {S}^{4}K^{4}{d}^{2}3^{2L}\|{\underline{m}}\|_{1}/{\varepsilon}^{6}),$ (12) $$

where $\tilde{\Theta}$ hides a logarithmic factor $\log({d}\|{\underline{m}}\|_{1}{S}\cdot 3^{L}/(\eta{\varepsilon}))$. The dependency on any of the parameters $({S},{d},3^{L},K,\|{\underline{m}}\|_{1},{\varepsilon})$ could potentially be refined by imposing additional assumptions on the $\psi$ functions or through a 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 $m^{(\ell)}=m\geq 3$ for each $\ell\in[L]$, leading to ${d}=m^{L}$. In this setup, the sample complexity gives

$$ $n=\tilde{\Theta}(L^{2}m\cdot{d}^{2+2\log_{m}3}/{\varepsilon}^{6})\leq\tilde{ \Theta}(L^{2}m\cdot{d}^{4}/{\varepsilon}^{6}),$ (13) $$

exhibiting a polynomial dependence on ${d}$ and $1/{\varepsilon}$. Such polynomial scaling aligns with existing literature [, , , , ], which indicates that learning hierarchical models using multi-layer networks avoids the curse of dimensionality. Theorem 1 serves as a warm-up result in the classification context. In Section 4, we aim to extend similar methodologies to address denoising problems, employing analogous proof strategies.

###### Theorem 1 (Learning to classify using ConvNets) .

###### Remark 2 .

### 3.1 Proof strategy: ConvNets approximate the belief propagation algorithm

Lemma 17 introduces a decomposition of ${\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_{\star}})$ into two components, approximation error and generalization error:

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_ {\star}})\leq\underbrace{\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}(\mu_{{ \rm NN}}^{\bm{W}},{\mu_{\star}})}_{\text{approximation error}}+2\cdot \underbrace{\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{\sf R}}(\mu_{{\rm NN }}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big{|}}_{\text{generalization error}}.$ $$

The bound of generalization error follows a standard approach: employing a chaining argument, the error is controlled by $\tilde{O}(\sqrt{{d_{\rm p}}/n})$, where ${d_{\rm p}}$ represents the number of parameters in the ConvNet 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 classifier ${\mu_{\star}}$, 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 on input ${\bm{x}}\in[{S}]^{d}$ and iteratively calculates the beliefs $\{{\nu}_{v}^{(\ell)}\in\Delta([{S}])\}_{\ell\in\{0,\ldots,L\},v\in{\mathcal{V}
}^{(\ell)}}$ as follows:

$$ $\displaystyle{\nu}^{(L)}_{v}(x^{(L)}_{v})=$ $\displaystyle~{}1\{x^{(L)}_{v}=x_{v}\},$ (BP-CLS) $\displaystyle{\nu}^{(\ell)}_{v}(x^{(\ell)}_{v})\propto$ $\displaystyle~{}\textstyle\sum_{x^{(\ell+1)}_{{\mathcal{C}}(v)}}\prod_{v^{ \prime}\in{\mathcal{C}}(v)}\Big{(}\psi_{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{( \ell)},x^{(\ell+1)}_{v^{\prime}}){\nu}^{(\ell+1)}_{v^{\prime}}(x^{(\ell+1)}_{v ^{\prime}})\Big{)},$ $\displaystyle\ell=L-1,\ldots,0,$ $\displaystyle\mu_{{\rm BP}}(y|{\bm{x}})=$ $\displaystyle~{}\nu_{{\rm r}}^{(0)}(y).$ $$

Classical results in graphical models verify that the belief propagation algorithm accurately computes the Bayes classifier in this tree graph.

###### Lemma 1 (BP calculates the Bayes classifier exactly [ Pea82 , WJ + 08 , MM09 ] ) .

When applying the belief propagation algorithm (BP-CLS) starting with ${\bm{x}}\in[{S}]^{d}$, it holds that ${\mu_{\star}}(\cdot|{\bm{x}})=\mu_{{\rm BP}}(\cdot|{\bm{x}})$.

The belief propagation algorithm can be streamlined into a message passing algorithm, starting with the initialization ${h}_{v}^{(L)}=x_{v}$ for each node $v$ in the highest layer ${\mathcal{V}}^{(L)}$. The operations are defined as follows:

$$ $\displaystyle{q}_{v}^{(\ell)}=$ $\displaystyle~{}{f}_{\iota(v)}^{(\ell)}({h}_{v}^{(\ell)})\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (MP-CLS) $\displaystyle{h}_{v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle{\rm normalize}\Big{(}\sum_{v^{\prime}\in{\mathcal{C }}(v)}{q}_{v^{\prime}}^{(\ell)}\Big{)}\in{\mathbb{R}}^{{S}},~{}~{}~{}~{}~{}~{}$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell-1)},$ $\displaystyle\mu_{{\rm MP}}(y|{\bm{x}})=$ $\displaystyle~{}{\rm softmax}({h}^{(0)}_{{\rm r}}).$ $$

The functions ${f}_{\iota}^{(L)}:[{S}]\to{\mathbb{R}}^{{S}}$, $\{{f}_{\iota}^{(\ell)}:{\mathbb{R}}^{{S}}\to{\mathbb{R}}^{{S}}\}_{\ell\in[L-1]}$ are defined as:

$$ $\displaystyle{f}_{\iota}^{(L)}(x)_{s}=$ $\displaystyle~{}\log\psi_{\iota}^{(L)}(s,x),~{}~{}~{}$ $\displaystyle x\in[{S}],s\in[{S}],$ (14) $\displaystyle{f}_{\iota}^{(\ell)}({h})_{s}=$ $\displaystyle~{}\textstyle\log\sum_{a\in[{S}]}\psi_{\iota}^{(\ell)}(s,a)e^{{h} _{a}},~{}~{}~{}$ $\displaystyle{h}\in{\mathbb{R}}^{S},s\in[{S}],\ell\in[L-1].$ $$

We note that the normalization operator in (MP-CLS) is non-essential and could be dropped; however, 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 propagation:

###### Proposition 2 (BP reduces to MP) .

Consider the belief propagation algorithm and the message passing algorithm, both starting from ${\bm{x}}\in[{S}]^{d}$, as in Eq. (BP-CLS) and (MP-CLS). Then we have ${\nu}^{(\ell)}_{v}(\cdot)={\rm softmax}(h^{(\ell)}_{v})$ for all $0\leq\ell\leq L-1$ and $v\in{\mathcal{V}}^{(\ell)}$. In particular, we have $\mu_{{\rm BP}}(\cdot|{\bm{x}})=\mu_{{\rm MP}}(\cdot|{\bm{x}})={\mu_{\star}}(
\cdot|{\bm{x}})$.

The proof of Proposition 2 is presented in Section B.

###### Lemma 1 (BP calculates the Bayes classifier exactly [ Pea82 , WJ + 08 , MM09 ] ) .

###### Proposition 2 (BP reduces to MP) .

#### 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: ${f}_{\iota}^{(\ell)}$ versus ${\rm NN}_{\iota}^{(\ell)}$. Given the expression ${f}_{\iota}^{(\ell)}({h})_{s}=\log\sum_{a\in[{S}]}\psi_{\iota}^{(\ell)}(s,a)e^
{{h}_{a}}$, it becomes evident that approximating the logarithmic and exponential functions using one-hidden-layer networks enables ${f}_{\iota}^{(\ell)}({h})$ to be effectively approximated by a two-hidden-layer neural network. This leads to the following theorem:

###### Theorem 3 (ConvNets approximation of Bayes classifier) .

Let Assumption 1 and 2 hold. For any $\delta>0$, take

$$ $D=4\lceil{S}^{2}K^{2}{d}\cdot 3^{L}/\delta\rceil,~{}~{}~{}~{}~{}B={\rm Poly}({ d},{S},K,3^{L},1/\delta).$ $$

Then there exists ${\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ as in Eq. (9), such that defining $\mu_{{\rm NN}}^{\bm{W}}$ as in Eq. (ConvNet), we have

$$ $\max_{y\in[{S}],{\bm{x}}\in[{S}]^{d}}\Big{|}\log{\mu_{\star}}(y|{\bm{x}})-\log \mu_{{\rm NN}}^{\bm{W}}(y|{\bm{x}})\Big{|}\leq\delta.$ $$

The proof of Theorem 3 is detailed in Section C.

###### Theorem 3 (ConvNets approximation of Bayes classifier) .

## 4 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 $({\bm{z}},{\bm{x}})$, generated from the following: ${\bm{x}}\sim{\mu_{\star}}$ represents the clean covariates, and ${\bm{z}}={\bm{x}}+{\bm{g}}$ where ${\bm{g}}\sim{\mathcal{N}}({\bm{0}},{\mathbf{I}}_{d})$ 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 of $({\bm{z}},{\bm{x}})$ as ${\mu_{\star}}$.

We consider a scenario where a set of iid samples $\{({\bm{z}}_{i},{\bm{x}}_{i})\}_{i\in[n]}\sim_{iid}{\mu_{\star}}$ is drawn from the distribution. Our objective is to learn a denoiser ${\bm{m}}({\bm{z}})$ from this dataset. With a suitable loss function, the optimal denoiser is the Bayes denoiser ${{\bm{m}}_{\star}}({\bm{z}})={\mathbb{E}}_{({\bm{x}},{\bm{z}})\sim{\mu_{\star}
}}[{\bm{x}}|{\bm{z}}]$, which calculates the posterior expectation of ${\bm{x}}$ given ${\bm{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, represented as ${\bm{m}}_{\rm NN}({\bm{z}})\in{\mathbb{R}}^{{d}}$ for input ${\bm{z}}\in{\mathbb{R}}^{d}$. Initially, we set ${\mathsf{h}}_{\downarrow,v}^{(L)}=(-(x-z_{v})^{2}/2)_{x\in[{S}]}\in{\mathbb{R}
}^{{S}}$ for $v\in{\mathcal{V}}^{(L)}$ for each node $v\in{\mathcal{V}}^{(L)}$. The operational flow of the network unfolds as follows:

$$ $\displaystyle{\mathsf{q}}_{\downarrow,v}^{(\ell)}=$ $\displaystyle~{}{\rm NN}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize}({ \mathsf{h}}_{\downarrow,v}^{(\ell)}))\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (UNet) $\displaystyle{\mathsf{h}}_{\downarrow,v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle\sum_{v^{\prime}\in{\mathcal{C}}(v)}{\mathsf{q}}_{ \downarrow,v^{\prime}}^{(\ell)}\in{\mathbb{R}}^{{S}},~{}~{}~{}~{}~{}~{}$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell-1)},$ $\displaystyle{\mathsf{u}}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{\mathsf{b}}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}\in{\mathbb{R}}^ {{S}},~{}~{}~{}(\text{with }{\mathsf{b}}_{\uparrow,{\rm r}}^{(0)}={\mathsf{h}} _{\downarrow,{\rm r}}^{(0)})$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{\mathsf{b}}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{\rm NN}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({\mathsf {u}}_{\uparrow,v}^{(\ell)}-{\mathsf{q}}_{\downarrow,v}^{(\ell)}))+{\mathsf{h}} _{\downarrow,v}^{(\ell)}\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{\bm{m}}_{{\rm NN}}({\bm{z}})_{v}=$ $\displaystyle~{}\textstyle\sum_{s\in[{S}]}s\cdot{\rm softmax}({\mathsf{b}}_{ \uparrow,v}^{(L)})_{s},$ $\displaystyle v\in{\mathcal{V}}^{(L)}.$ $$

The functions $\{{\rm NN}_{\downarrow,\iota}^{(\ell)},{\rm NN}_{\uparrow,\iota}^{(\ell)}:{
\mathbb{R}}^{{S}}\to{\mathbb{R}}^{{S}}\}_{\ell\in[L]}$ and ${\rm normalize}:{\mathbb{R}}^{S}\to{\mathbb{R}}^{S}$ are defined as follows:

$$ $\displaystyle{\rm NN}_{\diamond,\iota}^{(\ell)}({\mathsf{h}})_{s}=$ $\displaystyle~{}\textstyle W_{1,\diamond,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{2, \diamond,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{3,\diamond,\iota}^{(\ell)}\cdot[{ \mathsf{h}};1])),~{}~{}s\in[{S}],\diamond\in\{\downarrow,\uparrow\},\ell\in[L] ,\iota\in[m^{(\ell)}],$ (15) $\displaystyle{\rm normalize}({h})_{s}=$ $\displaystyle~{}\textstyle{h}_{s}-\max_{s^{\prime}\in[{S}]}{h}_{s^{\prime}}.$ (16) $$

The dimensions of the weight matrices within the U-Net are specified as follows:

$$ ${\bm{W}}=\Big{\{}\{W_{1,\diamond,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}\times D} \}_{\iota\in[{m}^{(\ell)}]},\{W_{2,\diamond,\iota}^{(\ell)}\in{\mathbb{R}}^{D \times D}\}_{\iota\in[{m}^{(\ell)}]},\{W_{3,\diamond,\iota}^{(\ell)}\in{ \mathbb{R}}^{D\times({S}+1)}\}_{\iota\in[{m}^{(\ell)}]}\Big{\}}_{\ell\in[L], \diamond\in\{\downarrow,\uparrow\}}$ (17) $$

Furthermore, we denote ${\bm{m}}_{{\rm NN}}^{{\bm{W}}}$ as the U-Net denoiser ${\bm{m}}_{{\rm NN}}$ parameterized by ${\bm{W}}$. A schematic illustration of a U-Net with $3$ layers is provided in Figure 2.

###### Remark 3 (An explanation of the “U-Net” architecture) .

As noted in our discussion of the convolutional 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) is its encoder-decoder architecture and the inclusion of long skip connections, which closely mirror practical implementations. Specifically, the encoder sequence in (UNet) progresses as ${\mathsf{h}}_{\downarrow}^{(L)}\to{\mathsf{q}}_{\downarrow}^{(L)}\to{\mathsf{h
}}_{\downarrow}^{(L-1)}\to\cdots\to{\mathsf{q}}_{\downarrow}^{(1)}\to{\mathsf{
h}}_{\downarrow}^{(0)}$, consisting of a series of convolutional, average pooling, and normalization layers. This part of the architecture is the same as (ConvNet) for the classification task. The decoder sequence ascends as ${\mathsf{b}}_{\uparrow}^{(0)}\to{\mathsf{u}}_{\uparrow}^{(1)}\to{\mathsf{b}}_{
\uparrow}^{(1)}\to\cdots\to{\mathsf{b}}_{\uparrow}^{(L-1}\to{\mathsf{u}}_{
\uparrow}^{(L)}\to{\bm{m}}_{\rm NN}$, consisting of a series of convolutional, up-sampling, and normalization layers. Moreover, the computation of ${\mathsf{b}}_{\uparrow}^{(\ell)}$ utilizes ${\mathsf{u}}_{\uparrow}^{(\ell)}$ from the upward sequence, and $({\mathsf{h}}_{\downarrow}^{(\ell)},{\mathsf{q}}_{\downarrow}^{(\ell)})$ from the downward process, which is enabled by the long skip connections. This encoder-decoder architecture and the role of the long skip connections are effectively visualized in Figure 2.

Figure: Figure 2: A U-Net with $L=3$. “AP” stands for average-pooling. “US” stands for up-sampling. “SKIP” stands for long skip connections.
Refer to caption: x3.png

###### Remark 3 (An explanation of the “U-Net” architecture) .

#### The ERM estimator.

In the denoising task, we employ empirical risk minimization over U-Nets as outlined in the following equation:

$$ $\widehat{\bm{W}}=\arg\min_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S} ,D,B}}\Big{\{}{\widehat{\sf R}}({\bm{m}}_{{\rm NN}}^{\bm{W}})=\frac{1}{n{d}} \sum_{i=1}^{n}\|{\bm{x}}_{i}-{\bm{m}}_{{\rm NN}}^{\bm{W}}({\bm{z}}_{i})\|_{2}^ {2}\Big{\}}.$ (18) $$

The parameter space for the U-Nets is defined as:

$$ ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}:=\Big{\{}{\bm{W}}\text{ as defined in (\ref{eqn:W_set_denoising})}:{|\!|\!|{{\bm{W}}}|\!|\!|}:=\max_{ \diamond\in\{\downarrow,\uparrow\}}\max_{j\in[3]}\max_{\ell\in[L]}\max_{\iota \in[m^{(\ell)}]}\|W_{j,\iota}^{(\ell)}\|_{{\rm op}}\leq B\Big{\}}.$ (19) $$

We anticipate that the empirical risk minimizer, ${\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}}$, could learn the Bayes denoiser ${{\bm{m}}_{\star}}$, as the global minimizer of the population risk over all functions yields the Bayes denoiser:

$$ ${{\bm{m}}_{\star}}(\cdot)=\arg\min_{{\bm{m}}}\Big{\{}{\sf R}({\bm{m}})={ \mathbb{E}}_{({\bm{x}},{\bm{z}})\sim{\mu_{\star}}}[{d}^{-1}\|{\bm{x}}-{\bm{m}} ({\bm{z}})\|_{2}^{2}]\Big{\}}.$ $$

In our theoretical analysis, we measure the discrepancy between ${\bm{m}}_{\rm NN}^{\widehat{\bm{W}}}$ and ${{\bm{m}}_{\star}}$ using the squared Euclidean distance:

$$ $\textstyle{\sf D_{2}^{2}}({\bm{m}},{{\bm{m}}_{\star}})={\mathbb{E}}_{{\bm{z}} \sim{\mu_{\star}}}\Big{[}d^{-1}\|{\bm{m}}-{{\bm{m}}_{\star}}({\bm{z}})\|_{2}^{ 2}\Big{]}.$ (20) $$

#### Sample complexity bound.

The subsequent theorem establishes the bound of the ${\sf D_{2}^{2}}$-distance between the U-Net estimator ${\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}}$ and the true Bayes denoiser ${{\bm{m}}_{\star}}$.

###### Theorem 4 (Learning to denoise using U-Nets) .

Let Assumption 1 and 2 hold. Let ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ be the set defined as in Eq. (19), where $B={\rm Poly}({d},{S},K,18^{L},D)$. Let $\widehat{\bm{W}}$ be the empirical risk minimizer as in Eq. (18). Then with probability at least $1-\eta$, we have

$$ ${\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_{\star}})\leq C\cdot \Big{(}\frac{{S}^{6}K^{4}{d}^{2}\cdot 18^{2L}}{D^{2}}+{S}^{2}\cdot\sqrt{\frac{ LD(D+2S+1)\|{\underline{m}}\|_{1}\log({d}\|{\underline{m}}\|_{1}D{S}\cdot 18^{ L})+\log(1/\eta)}{n}}\Big{)}.$ $$

The proof of Theorem 4 is detailed in Section G.

###### Remark 4 .

To ensure the ${\sf D_{2}^{2}}$-distance is less than ${\varepsilon}^{2}$, Theorem 4 requires to take

$$ $D=\Theta({S}^{3}K^{2}{d}18^{L}/{\varepsilon}),~{}~{}~{}~{}~{}n=\tilde{\Theta}( L{S}^{10}K^{4}{d}^{2}18^{2L}\|{\underline{m}}\|_{1}/{\varepsilon}^{6}),$ (21) $$

where $\tilde{\Theta}$ hides a logarithmic factor $\log({d}\|{\underline{m}}\|_{1}{S}\cdot 18^{L}/(\eta{\varepsilon}))$. Similar to Theorem 1, the dependency on any of the parameters $({S},{d},18^{L},K,\|{\underline{m}}\|_{1},{\varepsilon})$ could potentially be refined by imposing additional assumptions on the $\psi$ functions or through a 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 $m^{(\ell)}=m\geq 3$ for each $\ell\in[L]$, leading to ${d}=m^{L}$. In this setup, the sample complexity gives

$$ $n=\tilde{\Theta}(L^{2}m\cdot{d}^{2+2\log_{m}18}/{\varepsilon}^{6})\leq\tilde{ \Theta}(L^{2}m\cdot{d}^{8}/{\varepsilon}^{6}),$ (22) $$

exhibiting a polynomial dependence on ${d}$ and $1/{\varepsilon}$. Although the degree of the polynomial is substantial and potentially improvable, this gives the first polynomial sample complexity result for learning the Bayes denoiser in a hierarchical model using the U-Nets.

###### Theorem 4 (Learning to denoise using U-Nets) .

###### Remark 4 .

#### Connection to the task of diffusion generative modeling.

The denoising task examined in this section is closely related to the diffusion model approach [, , , ] to generative modeling. Diffusion models involve learning a generative model from a dataset of $n$ independent and identically distributed samples $\{{\bm{x}}_{i}\}_{i=1}^{n}$, drawn from an unknown data distribution ${\mu_{\star}}\in{\mathcal{P}}([{S}]^{d})$. The goal is to generate new samples $\hat{{\bm{x}}}\sim\hat{\mu}$ that match the distribution ${\mu_{\star}}$. Most diffusion model formulations involve a series of steps that are closely related to the denoising task described earlier. Here, we illustrate how they work using a variant of diffusion models, the stochastic localization process [, , , , ]:

- •
Step 1. Fit approximate denoising functions $\hat{\bm{m}}_{t}:{\mathbb{R}}^{d}\to{\mathbb{R}}^{d}$ for $t\in[0,T]$. This is done by minimizing the empirical risk over a class of neural networks ${\mathcal{F}}$ (i.e., the denoising task discussed in this section):
$\hat{\bm{m}}_{t}\equiv\arg\min_{{\rm NN}_{t}\in{\mathcal{F}}}\frac{1}{n{d}}
\sum_{i=1}^{n}\|{\bm{x}}_{i}-{\rm NN}_{t}(t\cdot{\bm{x}}_{i}+\sqrt{t}\cdot{\bm
{g}}_{i})\|_{2}^{2},~{}~{}~{}~{}{\bm{g}}_{i}\sim_{iid}{\mathcal{N}}({\bm{0}},{
\mathbf{I}}_{d}).$
(ERM)
- •
Step 2. Simulate a discretized version of a stochastic differential equation (SDE) starting from zero, whose drift term gives the approximate denoising function:
${\rm d}{\bm{z}}_{t}=\hat{\bm{m}}_{t}({\bm{z}}_{t})\cdot{\rm d}t+{\rm d}{\bm{B}
}_{t},~{}~{}t\in[0,T],~{}~{}~{}~{}~{}{\bm{z}}_{0}={\bm{0}},$
(SDE)
and generate an approximate sample $\hat{\bm{x}}={\bm{z}}_{T}/T\in{\mathbb{R}}^{d}$ at the final time $T$.

Standard analysis shows that by replacing the fitted denoising functions $\hat{\bm{m}}_{t}({\bm{z}})$ with the true denoising functions ${\bm{m}}_{t}({\bm{z}})$ in Eq. (SDE) and allowing $T\rightarrow\infty$, we can effectively recover the original data distribution ${\mu_{\star}}$. Consequently, the quality of samples generated from diffusion models hinges on two critical factors:

- (1)
How well the fitted denoising functions $\hat{\bm{m}}_{t}$ (ERM) approximate the true denoising functions ${\bm{m}}_{t}$;
- (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 a reliable denoising function estimator is available [, , , , ]; and controlling the denoising function approximation error through neural networks [, , ]. 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 4, with standard SDE discretization error bounds, such as the result established in [], it is straightforward to derive an end-to-end error bound for the sampling process of diffusion models in GHMs, similar to the strategy of [, , ](^2^22We note that the stochastic localization formulation is equivalent to the DDPM diffusion model, differing only in parametrization []. 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, as per Tweedie’s formula. Consequently, Theorem 4 can be readily adapted to establish a sample complexity bound for learning this score function.).

### 4.1 Proof strategy: U-Nets approximate the belief propagation algorithm

The proof strategy for the denoising task closely aligns with that of the classification task as detailed in Section 3.1.

Lemma 22 introduces a decomposition of ${\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}},{{\bm{m}}_{\star}})$ into two components, approximation error and generalization error:

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}}, {{\bm{m}}_{\star}})\leq\underbrace{\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{ 2}}({\bm{m}}_{{\rm NN}}^{\bm{W}},{{\bm{m}}_{\star}})}_{\text{approximation error}}+2\cdot\underbrace{\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{\sf R }}({\bm{m}}_{{\rm NN}}^{\bm{W}})-{\sf R}({\bm{m}}_{{\rm NN}}^{\bm{W}})\Big{|}} _{\text{generalization error}}.$ $$

The bound of generalization error follows a standard approach: employing a chaining argument, the error is controlled by $\tilde{O}(\sqrt{{d_{\rm p}}/n})$, where ${d_{\rm p}}$ 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 ${{\bm{m}}_{\star}}$, 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 operates on input ${\bm{z}}\in{\mathbb{R}}^{d}$ and iteratively calculates the beliefs $\{{\nu}_{\downarrow,v}^{(\ell)},{\nu}_{\uparrow,v}^{(\ell)}\in\Delta([{S}])\}_
{\ell\in\{0,\ldots,L\},v\in{\mathcal{V}}^{(\ell)}}$ as follows:

$$ $\displaystyle{\nu}^{(L)}_{\downarrow,v}(x^{(L)}_{v})=$ $\displaystyle~{}\psi^{(L+1)}(x^{(L)}_{v},z_{v}),$ $\displaystyle{\nu}^{(\ell)}_{\downarrow,v}(x^{(\ell)}_{v})\propto$ $\displaystyle~{}\textstyle\sum_{x^{(\ell+1)}_{{\mathcal{C}}(v)}}\prod_{v^{ \prime}\in{\mathcal{C}}(v)}\Big{(}\psi_{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{( \ell)},x^{(\ell+1)}_{v^{\prime}}){\nu}^{(\ell+1)}_{\downarrow,v^{\prime}}(x^{( \ell+1)}_{v^{\prime}})\Big{)},$ $\displaystyle\ell=L-1,\ldots,0,$ $\displaystyle{\nu}_{\uparrow,{\rm r}}^{(0)}(x^{(0)}_{{\rm r}})\propto$ $\displaystyle~{}1,$ $\displaystyle{\nu}_{\uparrow,v}^{(\ell)}(x^{(\ell)}_{v})\propto$ $\displaystyle~{}\textstyle\sum_{x^{(\ell-1)}_{{\rm pa}(v)},x^{(\ell)}_{{ \mathcal{N}}(v)}}\psi^{(\ell)}(x^{(\ell-1)}_{{\rm pa}(v)},x^{(\ell)}_{{ \mathcal{C}}({\rm pa}(v))}){\nu}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}(x^{(\ell-1) }_{{\rm pa}(v)})\prod_{v^{\prime}\in{\mathcal{N}}(v)}{\nu}_{\downarrow,v^{ \prime}}^{(\ell)}(x^{(\ell)}_{v^{\prime}}),$ $\displaystyle\ell=1,\ldots,L,$ $\displaystyle{\nu}_{v}^{(L)}(x^{(L)}_{v})\propto$ $\displaystyle~{}{\nu}_{\uparrow,L}(x^{(L)}_{v})\psi^{(L+1)}(x^{(L)}_{v},z_{v}),$ $\displaystyle{\bm{m}}_{{\rm BP}}({\bm{z}})_{v}=$ $\displaystyle~{}\textstyle\sum_{x_{v}^{(L)}}x_{v}^{(L)}{\nu}_{v}^{(L)}(x^{(L)} _{v}),$ (BP-DNS) $$

where $\psi^{(L+1)}(x^{(L)}_{v},z_{v}):=\exp\{-(x^{(L)}_{v}-z_{v})^{2}/2\}$.
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 [ Pea82 , WJ + 08 , MM09 ] ) .

When applying the belief propagation algorithm (BP-DNS) starting with ${\bm{z}}\in{\mathbb{R}}^{d}$, it holds that ${\mu_{\star}}(x_{v}|{\bm{z}})={\nu}_{v}^{(L)}(x_{v})$ for $v\in{\mathcal{V}}^{(L)}$, so that ${\bm{m}}_{{\rm BP}}({\bm{z}})={\bm{m}}({\bm{z}})$.

We remark that the downward-upward structure of belief propagation in generative hierarchical models has been pointed out in the literature [].

The belief propagation algorithm can be streamlined into a message passing algorithm, starting with the initialization ${h}_{\downarrow,v}^{(L)}=(-(x-z_{v})^{2}/2)_{x\in[{S}]}\in{\mathbb{R}}^{{S}}$ for each node $v$ in the highest layer ${\mathcal{V}}^{(L)}$. The operations are defined as follows:

$$ $\displaystyle{q}_{\downarrow,v}^{(\ell)}=$ $\displaystyle~{}{f}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize}({h}_{ \downarrow,v}^{(\ell)}))\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (MP-DNS) $\displaystyle{h}_{\downarrow,v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle\sum_{v^{\prime}\in{\mathcal{C}}(v)}{q}_{\downarrow, v^{\prime}}^{(\ell)}\in{\mathbb{R}}^{{S}},~{}~{}~{}~{}~{}~{}$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell-1)},$ $\displaystyle{u}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{b}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}\in{\mathbb{R}}^{{S}},~{} ~{}~{}(\text{with }{b}_{\uparrow,{\rm r}}^{(0)}={h}_{\downarrow,{\rm r}}^{(0)})$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{b}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{f}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({u}_{\uparrow ,v}^{(\ell)}-{q}_{\downarrow,v}^{(\ell)}))+{h}_{\downarrow,v}^{(\ell)}\in{ \mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{\bm{m}}_{{\rm MP}}({\bm{z}})_{v}=$ $\displaystyle~{}\textstyle\sum_{s\in[{S}]}s\cdot{\rm softmax}({b}_{\uparrow,v} ^{(L)})_{s},$ $\displaystyle v\in{\mathcal{V}}^{(L)}.$ $$

The functions $\{{f}_{\downarrow,\iota}^{(\ell)},{f}_{\uparrow,\iota}^{(\ell)}:{\mathbb{R}}^{
{S}}\to{\mathbb{R}}^{{S}}\}_{\ell\in[L]}$ are defined as

$$ $\displaystyle{f}_{\downarrow,\iota}^{(\ell)}({h})_{s}=$ $\displaystyle~{}\textstyle\log\sum_{a\in[{S}]}\psi_{\iota}^{(\ell)}(s,a)e^{{h} _{a}},~{}~{}~{}$ $\displaystyle{h}\in{\mathbb{R}}^{S},s\in[{S}],\ell\in[L],$ (23) $\displaystyle{f}_{\uparrow,\iota}^{(\ell)}({h})_{s}=$ $\displaystyle~{}\textstyle\log\sum_{a\in[{S}]}\psi_{\iota}^{(\ell)}(a,s)e^{h_{ a}},~{}~{}~{}$ $\displaystyle{h}\in{\mathbb{R}}^{S},s\in[{S}],\ell\in[L].$ $$

We note that the normalization operator in (MP-DNS) is non-essential and could be dropped; however, we include it to ensure the formulation closely mirrors (UNet), offering a technical benefit.

The subsequent proposition affirms that message passing is essentially equivalent to belief propagation:

###### Proposition 5 (BP reduces to MP) .

Consider the belief propagation algorithm and the message passing algorithm, both with input ${\bm{z}}\in{\mathbb{R}}^{d}$, as in Eq. (BP-DNS) and (MP-DNS). Then we have ${\nu}^{(\ell)}_{\downarrow,v}(\cdot)={\rm softmax}({h}^{(\ell)}_{\downarrow,v})$ and ${\nu}^{(\ell)}_{\uparrow,v}(\cdot)={\rm softmax}({b}_{\uparrow,v}^{(\ell)}-{h}
^{(\ell)}_{\downarrow,v})$, and ${\nu}_{v}^{(L)}(\cdot)={\rm softmax}(b_{\uparrow,v}^{(L)})$. In particular, ${\bm{m}}_{{\rm MP}}({\bm{z}})={\bm{m}}_{{\rm BP}}({\bm{z}})={\bm{m}}({\bm{z}})$.

The proof of Proposition 5 is presented in Section E.

###### Lemma 2 (BP calculates the Bayes denoiser exactly [ Pea82 , WJ + 08 , MM09 ] ) .

###### Proposition 5 (BP reduces to MP) .

#### 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: ${f}_{\diamond,\iota}^{(\ell)}$ versus ${\rm NN}_{\diamond,\iota}^{(\ell)}$. Notably, ${f}_{\diamond,\iota}^{(\ell)}$ entails a log-sum-exponential structure. This structure suggests that approximating the logarithmic and exponential functions with one-hidden-layer neural networks can allow ${f}_{\diamond,\iota}^{(\ell)}({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 1 and 2 hold. For any $\delta>0$, take

$$ $\displaystyle D=4\lceil{S}^{3}K^{2}{d}\cdot 18^{L}/\delta\rceil,~{}~{}~{}~{}~{ }B={\rm Poly}({d},{S},K,18^{L},1/\delta).$ $$

Then there exists ${\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ as in Eq. (19), such that defining ${\bm{m}}_{{\rm NN}}^{\bm{W}}$ as in Eq. (UNet), we have

$$ $\sup_{{\bm{z}}\in{\mathbb{R}}^{d}}\|{\bm{m}}({\bm{z}})-{\bm{m}}_{{\rm NN}}({ \bm{z}})\|_{\infty}\leq\delta.$ $$

The proof of Theorem 6 is detailed in Section F.

###### Theorem 6 (U-Nets approximation of Bayes denoiser) .

## 5 Further related work

#### Generative hierarchical models.

Hierarchical modeling of data distributions has been proposed in a series of works [, , , , , , , ]. While the hierarchical models in [, , , ] remain deterministic, [, , , ] studied the generative version of hierarchical models. The theoretical and empirical evidence presented in [, , ] 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 [ PCT + 23 , SFW24 , TW24 ] .

The series of works on hierarchical generative models [, , ] inspired the current study. [] first pointed out that the belief propagation denoising algorithm of hierarchical models consists of downward and upward processes. Through mean-field analysis on a 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. [], on the other hand, first introduced these models in a classification context. [, ] demonstrated that learning hierarchical models using multi-layer networks circumvents the curse of dimensionality. Specifically, they theoretically and empirically characterized the sample complexity, showing that it remains polynomial in dimension when learning convolutional networks under random generative rules. On the other hand, in the absence of correlations, they showed that the sample complexity is 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 [, , , , ] have become the state-of-the-art architecture for image classification and have been the backbone for many computer vision tasks. U-Nets [, , , ] have been particularly well-suited for image segmentation and denoising tasks [], and have served as the backbone architecture for diffusion models [, , , ]. A series of theoretical works has explained the inductive bias of CNNs [, , , , , , , , , , , , , , , ]. 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 [, ], where they found that the U-Nets are conjugate to the ResNets. In contrast, we demonstrate that U-Nets can effectively approximate the belief propagation denoising algorithms of GHMs. We also note that [] has analyzed the learning dynamics for a simple U-Net in diffusion models.

#### Variational inference in graphical models.

Variational inference is commonly employed to approximate the marginal statistics of graphical models [, , , , ]. For high-dimensional statistical models, approximate message passing [, ] and equivalently TAP variational inference [, , , , , ] can achieve consistent estimation of the Bayes posterior. In this paper, the hierarchical model is represented as a tree graph, and classical results [] indicate that belief propagation performs exact inference in such structures.

#### Neural networks approximation of algorithms.

Classical neural network approximation theory employs a function approximation viewpoint [, , , ], which typically relies on assumptions that the target function is smooth or hierarchically smooth. A recent line of work has investigated the expressiveness of neural networks through an algorithm approximation viewpoint [, , , , , , , ]. In particular, [, , , , ] demonstrate that transformers can efficiently approximate several algorithm classes, such as gradient descent, reinforcement learning algorithms, and even Turing machines. In the context of diffusion models, [] shows that ResNets can efficiently approximate the score function of high-dimensional graphical models by approximating the variational inference algorithm. Our work is closely related to [], except that we study neural network approximation in a different statistical model and network architecture.

From a practical viewpoint, a line of work has focused on neural network denoising by unrolling iterative denoising algorithms into deep networks [, , , , , , , , , ]. These approaches include unrolling the Iterative Shrinkage-Thresholding Algorithm (ISTA) for LASSO into recurrent networks [, , , ], unrolling belief propagation for Markov random fields into recurrent networks [], and unrolling graph denoising algorithms into graph neural networks []. 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 [, , , ] have emerged as a leading approach for generative modeling. Neural network-based score function approximation has been recently studied from the function approximation viewpoint in [, , , , ], and from the algorithm approximation viewpoint in []. Theoretical studies of other aspects of diffusion models include [, , , , , , , , , , , , , , , , , ]. For a comprehensive introduction to the theory of diffusion models, see the recent review [].

## 6 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 ${\bm{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/\epsilon$ 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 belief 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, extending these results to include conditional denoising functions represents an exciting direction for future research. Finally, we hope that the insights provided in this paper could guide the design of innovative network architectures.

## Acknowledgement

This work is supported by NSF DMS-2210827, CCF-2315725, an NSF Career Award, and a Google Research Scholar Award. The author would like to thank Hui Xu, Yu Bai, and Yuchen Wu for valuable discussions.

## Appendix A 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 [ BCW + 24 ] ) .

Suppose that $\{X_{w}\}_{w\in\Theta}$ is a zero-mean random process given by

$$ $\displaystyle X_{w}\equiv\frac{1}{n}\sum_{i=1}^{n}f(z_{i};w)-{\mathbb{E}}_{z}[ f(z;w)],$ $$

where $z_{1},\cdots,z_{n}$ are i.i.d samples from a distribution ${\mathbb{P}}_{z}$ such that the following assumption holds:

- (a)
The index set $\Theta$ is equipped with a distance $\rho$ and diameter ${B_{\rm p}}$. Further, assume that for some constant ${A_{\rm p}}$, for any ball $\Theta^{\prime}$ of radius $r$ in $\Theta$, the covering number admits upper bound $\log N(\Delta;\Theta^{\prime},\rho)\leq{d_{\rm p}}\log(2{A_{\rm p}}r/\Delta)$ for all $0<\Delta\leq 2r$.
- (b)
For any fixed $w\in\Theta$ and $z$ sampled from ${\mathbb{P}}_{z}$, the random variable $f(z;w)-{\mathbb{E}}_{z}[f(z;w)]$ is a $\sigma$-sub-Gaussian random variable (${\mathbb{E}}[e^{\lambda[f(z;w)-{\mathbb{E}}_{z^{\prime}}[f(z^{\prime};w)]]}]
\leq e^{\lambda^{2}\sigma^{2}/2}$ for any $\lambda\in{\mathbb{R}}$).
- (c)
For any $w,w^{\prime}\in\Theta$ and $z$ sampled from ${\mathbb{P}}_{z}$, the random variable $f(z;w)-f(z;w^{\prime})$ is a $\sigma^{\prime}\rho(w,w^{\prime})$-sub-Gaussian random variable (${\mathbb{E}}[e^{\lambda[f(z;w)-f(z;w^{\prime})]}]\leq e^{\lambda^{2}(\sigma^{
\prime})^{2}\rho^{2}(w,w^{\prime})/2}$ for any $\lambda\in{\mathbb{R}}$).

Then with probability at least $1-\eta$, it holds that

$$ $\displaystyle\sup_{w\in\Theta}|X_{w}|\leq C\sigma\sqrt{\frac{{d_{\rm p}}\cdot \log(2{A_{\rm p}}(1+{B_{\rm p}}\sigma^{\prime}/\sigma))+\log(1/\eta)}{n}},$ $$

where $C$ is a universal constant.

We next present a simple inequality used in the proof of Theorem 1.

###### Lemma 4 (From log ratio bound to square distance bound) .

Let $p$ and $q$ be two probability measures on $\Delta([{S}])$ such that

$$ $\max_{y\in[{S}]}\Big{|}\log p(y)-\log q(y)\Big{|}\leq\delta.$ $$

Then we have

$$ $\textstyle\sum_{s=1}^{{S}}(p(s)-q(s))^{2}\leq(e^{\delta}-1)^{2}.$ $$

###### Proof of Lemma 4 .

The lemma is by the fact that $|p(y)-q(y)|\leq(\exp\{|\log p(y)-\log q(y)|\}-1)\cdot p(y)$.
∎

###### Lemma 3 (Proposition A.4 of [ BCW + 24 ] ) .

###### Lemma 4 (From log ratio bound to square distance bound) .

###### Proof of Lemma 4 .

## Appendix B Proof of Proposition 2

###### Proof of Proposition 2 .

By Eq. (MP-CLS) and (14), defining ${\nu}^{(\ell)}_{v}(\cdot)={\rm softmax}({h}^{(\ell)}_{v})$, then for $\ell\leq L-2$, we get

$$ $\displaystyle{\nu}^{(\ell)}_{v}(x_{v}^{(\ell)})\propto$ $\displaystyle~{}\prod_{v^{\prime}\in{\mathcal{C}}(v)}\Big{(}\sum_{a\in[S]}\psi _{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{(\ell)},a)e^{({h}_{v}^{(\ell+1)})_{a}} \Big{)}$ $\displaystyle\propto$ $\displaystyle~{}\sum_{x^{(\ell+1)}_{{\mathcal{C}}(v)}}\prod_{v^{\prime}\in{ \mathcal{C}}(v)}\Big{(}\psi_{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{(\ell)},x^{( \ell+1)}_{v^{\prime}}){\nu}^{(\ell+1)}_{v^{\prime}}(x^{(\ell+1)}_{v^{\prime}}) \Big{)}.$ $$

This coincides with the update rule in Eq. (BP-CLS). For $\ell=L-1$, we get

$$ $\displaystyle{\nu}^{(L-1)}_{v}(x_{v}^{(L-1)})\propto$ $\displaystyle~{}\prod_{v^{\prime}\in{\mathcal{C}}(v)}\Big{(}\sum_{a\in[S]}\psi _{\iota(v^{\prime})}^{(L)}(x_{v}^{(L-1)},a)1\{a=x_{v^{\prime}}\}\Big{)}\propto \prod_{v^{\prime}\in{\mathcal{C}}(v)}\psi_{\iota(v^{\prime})}^{(L)}(x_{v}^{(L- 1)},x_{v^{\prime}}).$ $$

This again coincides with the update rule in Eq. (BP-CLS). This finishes the proof of Proposition 2.
∎

###### Proof of Proposition 2 .

## Appendix C Proof of Theorem 3

###### Proof of Theorem 3 .

By Lemma 13, take $M_{1}^{(\ell)}=\lceil 2{S}K{d}3^{L}/\delta\rceil+1$ and $M_{2}^{(\ell)}=\lceil 2{S}K^{2}{d}3^{L}/\delta\rceil+1$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M_{1}^{(\ell)}]}$ and $\{(\bar{a}_{j},\bar{w}_{j},\bar{b}_{j})\}_{j\in[M_{2}^{(\ell)}]}$ with

$$ $\sup_{j}|a_{j}|\leq 2{S}K,~{}~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}~{}\sup_{ j}|b_{j}|\leq{S}K,~{}~{}~{}~{}\sup_{j}|\bar{a}_{j}|\leq 2,~{}~{}~{}~{}\sup_{j} |\bar{w}_{j}|\leq 1,~{}~{}~{}~{}\sup_{j}|\bar{b}_{j}|\leq\log(4\cdot 3^{L}{S}{ d}K^{2}/\delta),$ $$

such that defining

$$ $\textstyle\log_{\delta\star}(x)=\sum_{j=1}^{M_{1}^{(\ell)}}a_{j}\cdot{\rm ReLU }(w_{j}x+b_{j}),~{}~{}~{}\exp_{\delta\star}(x)=\sum_{j=1}^{M_{2}^{(\ell)}}\bar {a}_{j}\cdot{\rm ReLU}(\bar{w}_{j}x+\bar{b}_{j}),$ $$

and defining ${f}_{\iota}^{(\ell)}(h),{\bar{f}}_{\iota}^{(\ell)}(h)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}_{\iota}^{(\ell)}(h)_{i}=\log\sum_{j=1}^{S}\psi_{\iota}^{(\ell)}( i,j)\exp(h_{j}),~{}~{}~{}{\bar{f}}_{\iota}^{(\ell)}(h)_{i}=\log_{\delta\star} \sum_{j=1}^{S}\psi_{\iota}^{(\ell)}(i,j)\exp_{\delta\star}(h_{j}),~{}~{}~{}~{} ~{}\forall i\in[{S}],$ $$

we have

$$ $\sup_{\max_{i}h_{i}=0}\|{f}_{\iota}^{(\ell)}(h)-{\bar{f}}_{\iota}^{(\ell)}(h) \|_{\infty}\leq\delta/({d}3^{L}).$ (24) $$

In addition, by Lemma 15, take $M_{1}^{(L)}=\lceil 3^{L}{d}K/\delta\rceil+1$ and $M_{2}^{(L)}=3$. Then there exists $\{(a_{j}^{(L)},w_{j}^{(L)},b_{j}^{(L)})\}_{j\in[M_{1}^{(L)}]}$ and $\{(\bar{a}_{j}^{(L)},\bar{w}_{j}^{(L)},\bar{b}_{j}^{(L)})\}_{j\in[M_{2}^{(L)}]}$ with

$$ $\sup_{j}|a_{j}^{(L)}|\leq 2K,~{}~{}~{}~{}\sup_{j}|w_{j}^{(L)}|\leq 1,~{}~{}~{} ~{}\sup_{j}|b_{j}^{(L)}|\leq{S}K,~{}~{}~{}~{}\sup_{j}|\bar{a}_{j}^{(L)}|\leq 4 ,~{}~{}~{}~{}\sup_{j}|\bar{w}_{j}^{(L)}|\leq 1,~{}~{}~{}~{}\sup_{j}|\bar{b}_{j }^{(L)}|\leq 1,$ $$

such that defining

$$ $\textstyle\log_{\delta}(x)=\sum_{j=1}^{M_{1}^{(L)}}a_{j}^{(L)}\cdot{\rm ReLU}( w_{j}^{(L)}x+b_{j}^{(L)}),~{}~{}~{}{\rm Ind}(x)=\sum_{j=1}^{M_{2}^{(L)}}\bar{a }_{j}^{(L)}\cdot{\rm ReLU}(\bar{w}_{j}^{(L)}x+\bar{b}_{j}^{(L)}),$ $$

and defining ${f}_{\iota}^{(L)}(h),{\bar{f}}_{\iota}^{(L)}(h)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}_{\iota}^{(L)}(x)_{i}=\log\sum_{j=1}^{S}\psi_{\iota}^{(L)}(i,j)1( x=j),~{}~{}~{}{\bar{f}}_{\iota}^{(L)}(x)_{i}=\log_{\delta}\sum_{j=1}^{S}\psi_{ \iota}^{(L)}(i,j){\rm Ind}(x-j),~{}~{}~{}~{}~{}\forall i\in[{S}],$ $$

we have

$$ $\sup_{x\in[{S}]}\|{f}_{\iota}^{(L)}(x)-{\bar{f}}_{\iota}^{(L)}(x)\|_{\infty} \leq\delta/({d}3^{L}).$ (25) $$

By Eq. (24) and (25) and Lemma 5, taking ${h}_{{\rm r}}^{(0)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (MP-CLS) and ${\bar{h}}_{{\rm r}}^{(0)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (A-MP-CLS) with $\{{\bar{f}}_{\iota}^{(\ell)}\}_{\ell\in[L],\iota\in[{m}^{(\ell)}]}$ as defined above, we have

$$ $\|{h}_{{\rm r}}^{(0)}-{\bar{h}}_{{\rm r}}^{(0)}\|_{\infty}\leq[\delta/({d}3^{L })]\times\prod_{1\leq\ell\leq L}(2{m}^{(\ell)}+1)\leq\delta.$ $$

As a 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 $\ell\in[L-1]$ and $\iota\in[{m}^{(\ell)}]$)

$$ ${\rm NN}_{W_{1,\iota}^{(\ell)},W_{2,\iota}^{(\ell)},W_{3,\iota}^{(\ell)}}(h)=W _{1,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{2,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{3, \iota}^{(\ell)}\cdot[h;1])),$ $$

with $W_{1,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}\times{S}M_{1}^{(\ell)}}$, $W_{2,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}M_{1}^{(\ell)}\times({S}M_{2}^{(\ell)}
+1)}$, $W_{3,\iota}^{(\ell)}\in{\mathbb{R}}^{({S}M_{2}^{(\ell)}+1)\times({S}+1)}$, and

$$ $\|W_{1,\iota}^{(\ell)}\|_{\max}\leq 2{S}K,~{}~{}~{}\|W_{2,\iota}^{(\ell)}\|_{ \max}\leq{\rm Poly}({S}KM_{1}^{(\ell)}M_{2}^{(\ell)}),~{}~{}~{}\|W_{3,\iota}^{ (\ell)}\|_{\max}\leq\log(4{S}K^{2}/\delta).$ $$

such that

$$ ${\rm NN}_{W_{1,\iota}^{(\ell)},W_{2,\iota}^{(\ell)},W_{3,\iota}^{(\ell)}}(h)={ \bar{f}}_{\iota}^{(\ell)}(h),~{}~{}~{}\forall h\in{\mathbb{R}}^{S}\text{ such that }\max_{j}h_{j}=0.$ $$

Furthermore, by Lemma 16, there exist two-hidden-layer neural networks (for $\iota\in[{m}^{(L)}]$)

$$ ${\rm NN}_{W_{1,\iota}^{(L)},W_{2,\iota}^{(L)},W_{3,\iota}^{(L)}}(x)=W_{1,\iota }^{(L)}\cdot{\rm ReLU}(W_{2,\iota}^{(L)}\cdot{\rm ReLU}(W_{3,\iota}^{(L)}\cdot [x;1])),$ $$

with $W_{1,\iota}^{(L)}\in{\mathbb{R}}^{{S}\times{S}M_{1}^{(L)}}$, $W_{2,\iota}^{(L)}\in{\mathbb{R}}^{{S}M_{1}^{(L)}\times({S}M_{2}^{(L)}+1)}$, $W_{3,\iota}^{(L)}\in{\mathbb{R}}^{({S}M_{2}^{(L)}+1)\times 2}$,

$$ $\|W_{1,\iota}^{(L)}\|_{\max}\leq 2K,~{}~{}~{}\|W_{2,\iota}^{(L)}\|_{\max}\leq{ \rm Poly}({S}KM_{1}^{(L)}M_{2}^{(L)}),~{}~{}~{}\|W_{3,\iota}^{(L)}\|_{\max} \leq 1.$ $$

such that

$$ ${\rm NN}_{W_{1,\iota}^{(L)},W_{2,\iota}^{(L)},W_{3,\iota}^{(L)}}(x)={\bar{f}}_ {\iota}^{(L)}(x),~{}~{}~{}\forall x\in[{S}].$ $$

This proves that the approximate version of message passing as in Eq. (A-MP-CLS) can be cast into the convolutional neural network as in Eq. (ConvNet) with proper choice of dimension

$$ $D\geq\max_{\ell\in[L]}\{SM_{1}^{(\ell)},SM_{2}^{(\ell)}+1\}={S}\times\Big{(} \lceil 2{S}K^{2}{d}3^{L}/\delta\rceil+1\Big{)}+1,$ $$

and norm of the weights. This finishes the proof of Theorem 3.
∎

###### Proof of Theorem 3 .

### C.1 Auxillary lemmas

###### Lemma 5 (Error propagation of the approximate version of message passing in classification) .

Assume we have functions ${f}_{\iota}^{(\ell)}$ and ${\bar{f}}_{\iota}^{(\ell)}$ such that

$$ $\displaystyle\|{f}_{\iota}^{(L)}(x)-{\bar{f}}_{\iota}^{(L)}(x)\|_{\infty}\leq$ $\displaystyle~{}\delta,~{}~{}~{}\forall x\in[{S}],$ (26) $\displaystyle\|{f}_{\iota}^{(\ell)}(h)-{\bar{f}}_{\iota}^{(\ell)}(h)\|_{\infty}\leq$ $\displaystyle~{}\delta,~{}~{}~{}\forall h\in{\mathbb{R}}^{S}\text{ such that } \max_{j\in[{S}]}h_{j}=0,~{}~{}\forall\ell\leq L-1.$ $$

Furthermore, consider the following approximate version of message passing algorithm with initialization ${\bar{h}}_{v}^{(L)}=x_{v}$ for $v\in{\mathcal{V}}^{(L)}$:

$$ $\displaystyle{\bar{q}}_{v}^{(\ell)}=$ $\displaystyle~{}{\bar{f}}_{\iota(v)}^{(\ell)}({\bar{h}}_{v}^{(\ell)})\in{ \mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L-1],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (A-MP-CLS) $\displaystyle{\bar{h}}_{v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle{\rm normalize}\Big{(}\sum_{v^{\prime}\in{\mathcal{C }}(v)}{\bar{q}}_{v^{\prime}}^{(\ell)}\Big{)}\in{\mathbb{R}}^{{S}},~{}~{}~{}~{} ~{}~{}$ $\displaystyle\ell\in[L-1],~{}~{}~{}v\in{\mathcal{V}}^{(\ell-1)}.$ $$

Taking ${h}_{{\rm r}}^{(0)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (MP-CLS) and ${\bar{h}}_{{\rm r}}^{(0)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (A-MP-CLS), we have

$$ $\|{h}_{{\rm r}}^{(0)}-{\bar{h}}_{{\rm r}}^{(0)}\|_{\infty}\leq\delta\times \prod_{1\leq\ell\leq L}(2{m}^{(\ell)}+1).$ $$

###### Proof of Lemma 5 .

We prove this lemma by induction, aiming to show that for any $\ell\in[L-1]$ we have

$$ $\textstyle\|{\bar{h}}_{v}^{(\ell)}-{h}_{v}^{(\ell)}\|_{\infty}\leq 2{m}^{(\ell +1)}\prod_{k=\ell+2}^{L}(2{m}^{(k)}+1)\delta,~{}~{}~{}~{}\forall v\in{\mathcal {V}}^{(\ell)}.$ (27) $$

To prove the formula for $\ell=L-1$, since ${h}_{v}^{(L)}={\bar{h}}_{v}^{(L)}$, by Eq. (26), we get

$$ $\|{\bar{q}}_{v}^{(L)}-{q}_{v}^{(L)}\|_{\infty}\leq\delta,~{}~{}~{}~{}\forall v \in{\mathcal{V}}^{(L)}.$ $$

By Lemma 7, we get

$$ $\textstyle\|{h}_{v}^{(L-1)}-{\bar{h}}_{v}^{(L-1)}\|_{\infty}\leq 2\Big{\|}\sum _{v^{\prime}\in{\mathcal{C}}(v)}({\bar{q}}_{v^{\prime}}^{(L)}-{q}_{v^{\prime}} ^{(L)})\Big{\|}_{\infty}\leq 2{m}^{(L)}\delta,~{}~{}~{}~{}\forall v\in{ \mathcal{V}}^{(L-1)}.$ $$

This proves the formula (27) for $\ell=L-1$.

Assuming that (27) holds at the layer $\ell$, by the update formula, we have

$$ $\displaystyle~{}\|{\bar{q}}_{v}^{(\ell)}-{q}_{v}^{(\ell)}\|_{\infty}=\|{\bar{f }}_{\iota(v)}^{(\ell)}({\bar{h}}_{v}^{(\ell)})-{f}_{\iota(v)}^{(\ell)}({h}_{v} ^{(\ell)})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|{\bar{f}}_{\iota(v)}^{(\ell)}({\bar{h}}_{v}^{(\ell)})-{f}_{ \iota(v)}^{(\ell)}({\bar{h}}_{v}^{(\ell)})\|_{\infty}+\|{f}_{\iota(v)}^{(\ell) }({h}_{v}^{(\ell)})-{f}_{\iota(v)}^{(\ell)}({\bar{h}}_{v}^{(\ell)})\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\textstyle\delta+2{m}^{(\ell+1)}\prod_{k=\ell+2}^{L}(2{m}^{(k) }+1)\delta\leq\prod_{k=\ell+1}^{L}(2{m}^{(k)}+1)\delta,$ $$

where the middle inequality is by the assumption of ${f}_{\iota}^{(\ell)}$ and by Lemma 6. By Lemma 7, we get

$$ $\textstyle\|{\bar{h}}_{v}^{(\ell-1)}-{h}_{v}^{(\ell-1)}\|_{\infty}\leq 2{m}^{( \ell)}\prod_{k=\ell+1}^{L}(2{m}^{(k)}+1)\delta,~{}~{}~{}~{}\forall v\in{ \mathcal{V}}^{(\ell)}.$ $$

This proves Lemma 5 by the induction argument.
∎

###### Lemma 6 (Non-expansiveness of log-sum-exponential) .

For $h\in{\mathbb{R}}^{S}$ and $\Psi\in{\mathbb{R}}^{{S}\times{S}}$, define ${f}(h)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}(h)_{i}=\log\sum_{j=1}^{S}\Psi_{ij}\exp(h_{j}),~{}~{}~{}\forall i \in[{S}].$ $$

Then for $h_{1},h_{2}\in{\mathbb{R}}^{S}$, we have

$$ $\|{f}(h_{1})-{f}(h_{2})\|_{\infty}\leq\|h_{1}-h_{2}\|_{\infty}.$ $$

###### Proof of Lemma 6 .

Fix $i\in[{S}]$. We have

$$ $\nabla_{h}{f}(h)_{i}=\Big{(}\frac{\Psi_{ij}\exp(h_{j})}{\sum_{k\in[{S}]}\Psi_{ ik}\exp(h_{k})}\Big{)}_{j\in[{S}]},$ $$

so that $\|\nabla_{h}{f}(h)_{i}\|_{1}=1$. By intermediate value theorem, we have

$$ $|{f}(h_{1})_{i}-{f}(h_{2})_{i}|=|\nabla_{h}{f}(\xi)_{i}^{\mathsf{T}}(h_{1}-h_{ 2})|\leq\|\nabla_{h}{f}(\xi)_{i}^{\mathsf{T}}\|_{1}\|h_{1}-h_{2}\|_{\infty}=\| h_{1}-h_{2}\|_{\infty}.$ $$

This proves Lemma 6.
∎

###### Lemma 7 (Lipschitzness of the normalization operator) .

For $h\in{\mathbb{R}}^{S}$, define ${\rm normalize}(h)\in{\mathbb{R}}^{S}$ by

$$ $\displaystyle\textstyle{\rm normalize}(h)_{i}=$ $\displaystyle~{}\textstyle h_{i}-\max_{j}h_{j},~{}~{}~{}\forall i\in[{S}].$ $$

Then for $h_{1},h_{2}\in{\mathbb{R}}^{S}$, we have

$$ $\|{\rm normalize}(h_{1})-{\rm normalize}(h_{2})\|_{\infty}\leq 2\|h_{1}-h_{2} \|_{\infty}.$ $$

###### Proof of Lemma 7 .

Note we have the following inequality

$$ $|\max_{j}h_{1,j}-\max_{j}h_{2,j}|\leq\|h_{1}-h_{2}\|_{\infty},$ $$

so that

$$ $|{\rm normalize}(h_{1})_{i}-{\rm normalize}(h_{2})_{i}|\leq\|h_{1}-h_{2}\|_{ \infty}+|\max_{j}h_{1,j}-\max_{j}h_{2,j}|\leq 2\|h_{1}-h_{2}\|_{\infty}.$ $$

This completes the proof of Lemma 7.
∎

###### Lemma 8 (ReLU approximation of the exponential function) .

For any $\delta>0$, take $M=\lceil 1/\delta\rceil+1\in{\mathbb{N}}$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M]}$ with

$$ $\sup_{j}|a_{j}|\leq 2,~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}\sup_{j}|b_{j}| \leq\log M,$ (28) $$

such that defining $\exp_{\delta}:{\mathbb{R}}\to{\mathbb{R}}$ by

$$ $\textstyle\exp_{\delta}(x)=\sum_{j=1}^{M}a_{j}\cdot{\rm ReLU}(w_{j}x+b_{j}),$ $$

we have $\exp_{\delta}$ is non-decreasing on $(-\infty,0]$, and

$$ $\sup_{x\in(-\infty,0]}|\exp(x)-\exp_{\delta}(x)|\leq\delta,~{}~{}~{}~{}\exp_{ \delta}(0)=1.$ $$

###### Proof of Lemma 8 .

Define $e_{j}=j/(M-1)$, $b_{j}=-\log(e_{j})$ for $j\in[M-1]$, $a_{1}=(e_{2}-e_{1})/(b_{2}-b_{1})$ and $a_{j}=(e_{j+1}-e_{j})/(b_{j+1}-b_{j})-(e_{j}-e_{j-1})/(b_{j}-b_{j-1})$ for $2\leq j\leq M-2$. Furthermore, define

$$ $\exp_{\delta}(x)=\sum_{j=1}^{M-2}a_{j}{\rm ReLU}(x+b_{j})+{\rm ReLU}(-x+e_{1}) -{\rm ReLU}(-x).$ $$

Then we have $\exp_{\delta}(-b_{j})=e_{j}$ for $j\in[M-1]$, and $\exp_{\delta}$ is piece-wise linear and non-decreasing on $(-\infty,0]$. Note that we also have $\exp(-b_{j})=e_{j}$ for $j\in[M-1]$, and $\exp$ is increasing on $(-\infty,0]$. This proves that $\sup_{x\in(-\infty,0]}|\exp(x)-\exp_{\delta}(x)|\leq\delta$. Furthermore, it is easy to see that $\exp_{\delta}(0)=1$ and $\exp_{\delta}$ is non-decreasing on $(-\infty,0]$. Finally, since $\exp$ is $1$-Lipschitz, it is easy to see that $\sup_{j\in[M]}|a_{j}|\leq 2$. It is also easy to 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$, $\delta>0$, take $M=\lceil 2A/\delta\rceil+1\in{\mathbb{N}}$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M]}$ with

$$ $\sup_{j}|a_{j}|\leq 2A,~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}\sup_{j}|b_{j}| \leq A,$ (29) $$

such that defining $\log_{\delta}:{\mathbb{R}}\to{\mathbb{R}}$ by

$$ $\textstyle\log_{\delta}(x)=\sum_{j=1}^{M}a_{j}\cdot{\rm ReLU}(w_{j}x+b_{j}),$ $$

we have $\log_{\delta}$ is non-decreasing on $[1/A,A]$, and

$$ $\sup_{x\in[1/A,A]}|\log(x)-\log_{\delta}(x)|\leq\delta.$ $$

###### Proof of Lemma 9 .

The proof of Lemma 9 is similar to Lemma 8.
∎

###### Lemma 10 (ReLU approximation of indicator function) .

Define

$$ $\textstyle{\rm Ind}(x)=2{\rm ReLU}(x-1/2)+2{\rm ReLU}(x+1/2)-4{\rm ReLU}(x),$ $$

we have

$$ $1(x=j)={\rm Ind}(x-j),~{}~{}~{}~{}~{}\forall j,x\in{\mathbb{Z}}.$ $$

###### Proof of Lemma 10 .

The lemma holds by direct calculation.
∎

###### Lemma 11 (Log-sum-exponential approximation) .

Assume $\log_{\delta_{1}}:{\mathbb{R}}\to{\mathbb{R}}$ and $\exp_{\delta_{2}}:{\mathbb{R}}\to{\mathbb{R}}$ are such that,

$$ $\displaystyle\sup_{x\in[1/K,{S}K]}|\log(x)-\log_{\delta_{1}}(x)|\leq$ $\displaystyle~{}\delta_{1},~{}~{}~{}\sup_{x\in(-\infty,0]}|\exp(x)-\exp_{ \delta_{2}}(x)|\leq\delta_{2},$ $\displaystyle\exp_{\delta_{2}}(0)=$ $\displaystyle~{}1,~{}~{}~{}~{}~{}~{}~{}\exp_{\delta_{2}}\text{ is non- decreasing on }(-\infty,0].$ $$

Assume that $1/K\leq\min_{ij}\Psi_{ij}\leq\max_{ij}\Psi_{ij}\leq K$. Define ${f}(h),{f}_{\delta_{1},\delta_{2}}(h)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}(h)_{i}=\log\sum_{j=1}^{S}\Psi_{ij}\exp(h_{j}),~{}~{}~{}{f}_{ \delta_{1},\delta_{2}}(h)_{i}=\log_{\delta_{1}}\sum_{j=1}^{S}\Psi_{ij}\exp_{ \delta_{2}}(h_{j}),~{}~{}~{}~{}~{}\forall i\in[{S}].$ $$

Then we have

$$ $\sup_{\max_{i}h_{i}=0}\|{f}(h)-{f}_{\delta_{1},\delta_{2}}(h)\|_{\infty}\leq \delta_{1}+{S}K^{2}\delta_{2}.$ $$

###### Proof of Lemma 11 .

We have

$$ $\displaystyle~{}|{f}(h)_{i}-{f}_{\delta_{1},\delta_{2}}(h)_{i}|=|\log\langle \Psi_{i:},\exp(h)\rangle-\log_{\delta_{1}}\langle\Psi_{i:},\exp_{\delta_{2}}(h )\rangle|$ $\displaystyle\leq$ $\displaystyle~{}|\log\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle-\log_{\delta _{1}}\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle|+|\log\langle\Psi_{i:},\exp( h)\rangle-\log\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle|.$ $$

For the first term, since $\exp_{\delta_{2}}(h)\leq 1$ for all $h\leq 0$, we have $\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle\leq{S}\max_{ij}\Psi_{ij}\leq{S}K$. Furthermore, since $\max_{i}h_{i}=0$ and $\exp_{\delta_{2}}(0)=1$, we have $\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle\geq\min_{ij}\Psi_{ij}\geq 1/K$. As a consequence, by assumption, we have

$$ $|\log\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle-\log_{\delta_{1}}\langle\Psi _{i:},\exp_{\delta_{2}}(h)\rangle|\leq\delta_{1}.$ $$

For the second term, since both $\langle\Psi_{i:},\exp(h)\rangle$ and $\langle\Psi_{i:},\exp_{\delta_{2}}(h)\rangle$ are within $[1/K,{S}K]$ on which $\log$ function has Lipschitz constant $K$, we have

$$ $|\log\langle\Psi_{i:},\exp(h)\rangle-\log\langle\Psi_{i:},\exp_{\delta_{2}}(h) \rangle|\leq{S}K\max_{ij}\Psi_{ij}\cdot\|\exp(h)-\exp_{\delta_{2}}(h)\|_{ \infty}\leq{S}K^{2}\delta_{2}.$ $$

This finishes the proof of Lemma 11.
∎

###### Lemma 12 (Log-Psi approximation) .

Assume $\log_{\delta}:{\mathbb{R}}\to{\mathbb{R}}$ is such that,

$$ $\sup_{x\in[1/K,K]}|\log(x)-\log_{\delta}(x)|\leq\delta.$ $$

Assume that $1/K\leq\min_{ij}\Psi_{ij}\leq\max_{ij}\Psi_{ij}\leq K$. For $x\in[{S}]$, define ${f}(x),{f}_{\delta}(x)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}(x)_{i}=\log\sum_{j=1}^{S}\Psi_{ij}1(x=j),~{}~{}~{}{f}_{\delta}(x )_{i}=\log_{\delta}\sum_{j=1}^{S}\Psi_{ij}1(x=j),~{}~{}~{}~{}~{}\forall i\in[{ S}].$ $$

Then we have

$$ $\sup_{x\in[{S}]}\|{f}(x)-{f}_{\delta}(x)\|_{\infty}\leq\delta.$ $$

###### Proof of Lemma 12 .

For any fixed $x\in[{S}]$ and $i\in[{S}]$, we have

$$ $|{f}(x)_{i}-{f}_{\delta}(x)_{i}|=|\log\Psi_{ix}-\log_{\delta}\Psi_{ix}|\leq\delta,$ $$

where the last inequality is by assumption. This proves Lemma 12.
∎

###### Lemma 13 (ReLU approximation of log-sum-exponential) .

Assume that $1/K\leq\min_{ij}\Psi_{ij}\leq\max_{ij}\Psi_{ij}\leq K$. For any $\delta>0$, take $M_{1}=\lceil 2{S}K/\delta\rceil+1$ and $M_{2}=\lceil 2{S}K^{2}/\delta\rceil+1$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M_{1}]}$ and $\{(\bar{a}_{j},\bar{w}_{j},\bar{b}_{j})\}_{j\in[M_{2}]}$ with

$$ $\sup_{j}|a_{j}|\leq 2{S}K,~{}~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}~{}\sup_{ j}|b_{j}|\leq{S}K,~{}~{}~{}~{}\sup_{j}|\bar{a}_{j}|\leq 2,~{}~{}~{}~{}\sup_{j} |\bar{w}_{j}|\leq 1,~{}~{}~{}~{}\sup_{j}|\bar{b}_{j}|\leq\log(4{S}K^{2}/\delta),$ $$

such that defining

$$ $\textstyle\log_{\delta\star}(x)=\sum_{j=1}^{M_{1}}a_{j}\cdot{\rm ReLU}(w_{j}x+ b_{j}),~{}~{}~{}\exp_{\delta\star}(x)=\sum_{j=1}^{M_{2}}\bar{a}_{j}\cdot{\rm ReLU }(\bar{w}_{j}x+\bar{b}_{j}),$ $$

and defining ${f}(h),{f}_{\delta}(h)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}(h)_{i}=\log\sum_{j=1}^{S}\Psi_{ij}\exp(h_{j}),~{}~{}~{}{f}_{ \delta}(h)_{i}=\log_{\delta\star}\sum_{j=1}^{S}\Psi_{ij}\exp_{\delta\star}(h_{ j}),~{}~{}~{}~{}~{}\forall i\in[{S}],$ $$

we have

$$ $\sup_{\max_{i}h_{i}=0}\|{f}(h)-{f}_{\delta}(h)\|_{\infty}\leq\delta.$ $$

###### Proof of Lemma 13 .

Lemma 9 implies that taking $M_{1}=\lceil 2{S}K/\delta\rceil+1$, there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M_{1}]}$ with

$$ $\sup_{j}|a_{j}|\leq 2{S}K,~{}~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}~{}\sup_{ j}|b_{j}|\leq{S}K,$ $$

such that defining

$$ $\textstyle\log_{\delta\star}(x)=\sum_{j=1}^{M_{1}}a_{j}\cdot{\rm ReLU}(w_{j}x+ b_{j}),$ $$

we have

$$ $\sup_{1/({S}K)\leq x\leq{S}K}|\log(x)-\log_{\delta\star}(x)|\leq\delta/2.$ $$

Lemma 8 implies that taking $M_{2}=\lceil 2{S}K^{2}/\delta\rceil+1$, there exists $\{(\bar{a}_{j},\bar{w}_{j},\bar{b}_{j})\}_{j\in[M_{2}]}$ with

$$ $\sup_{j}|\bar{a}_{j}|\leq 2,~{}~{}~{}~{}\sup_{j}|\bar{w}_{j}|\leq 1,~{}~{}~{}~ {}\sup_{j}|\bar{b}_{j}|\leq\log(4{S}K^{2}/\delta),$ $$

such that defining

$$ $\textstyle\exp_{\delta\star}(x)=\sum_{j=1}^{M_{2}}\bar{a}_{j}\cdot{\rm ReLU}( \bar{w}_{j}x+\bar{b}_{j}),$ $$

we have $\exp_{\delta\star}$ is non-decreasing on $(-\infty,0]$, and

$$ $\sup_{x\in(-\infty,0]}|\exp(x)-\exp_{\delta\star}(x)|\leq\delta/(2{S}K^{2}),~{ }~{}~{}\exp_{\delta\star}(0)=1.$ $$

As a consequence, the condition of Lemma 11 is satisfied with $\delta_{1}=\delta/2$ and $\delta_{2}=\delta/(2{S}K^{2})$, so that we have

$$ $\sup_{\max_{i}h_{i}=0}\|{f}(h)-{f}_{\delta}(h)\|_{\infty}\leq\delta_{1}+{S}K^{ 2}\delta_{2}=\delta.$ $$

This finishes the proof of Lemma 13.
∎

###### Lemma 14 (Existence of ReLU network approximating log-sum-exponential) .

Let ${f}_{\delta}$ be the function as defined in Lemma 13. Then there exists a two-hidden-layer neural network

$$ ${\rm NN}_{W_{1},W_{2},W_{3}}(h)=W_{1}\cdot{\rm ReLU}(W_{2}\cdot{\rm ReLU}(W_{3 }\cdot[h;1])),$ $$

with $W_{1}\in{\mathbb{R}}^{{S}\times{S}M_{1}}$, $W_{2}\in{\mathbb{R}}^{{S}M_{1}\times({S}M_{2}+1)}$, $W_{3}\in{\mathbb{R}}^{({S}M_{2}+1)\times({S}+1)}$, and

$$ $\|W_{1}\|_{\max}\leq 2{S}K,~{}~{}~{}\|W_{2}\|_{\max}\leq{\rm Poly}({S}KM_{1}M_ {2}),~{}~{}~{}\|W_{3}\|_{\max}\leq\log(4{S}K^{2}/\delta).$ $$

such that

$$ ${\rm NN}_{W_{1},W_{2},W_{3}}(h)={f}_{\delta}(h),~{}~{}~{}\forall h\in{\mathbb{ R}}^{S}\text{ such that }\max_{j}h_{j}=0.$ $$

###### Proof of Lemma 14 .

Define

$$ $\displaystyle\overline{W}_{2}=$ $\displaystyle~{}\begin{bmatrix}\bar{a}_{1:M_{2}}^{\mathsf{T}}&0&\cdots&0&0\\ 0&\bar{a}_{1:M_{2}}^{\mathsf{T}}&\cdots&0&0\\ \cdots&\cdots&\cdots&\cdots&\cdots\\ 0&0&\cdots&\bar{a}_{1:M_{2}}^{\mathsf{T}}&0\\ 0&0&\cdots&0&1\\ \end{bmatrix}\in{\mathbb{R}}^{({S}+1)\times({S}M_{2}+1)},$ $\displaystyle W_{3}=$ $\displaystyle~{}\begin{bmatrix}\bar{w}_{1:M_{2}}&0&\cdots&0&\bar{b}_{1:M_{2}} \\ 0&\bar{w}_{1:M_{2}}&\cdots&0&\bar{b}_{1:M_{2}}\\ \cdots&\cdots&\cdots&\cdots&\cdots\\ 0&0&\cdots&\bar{w}_{1:M_{2}}&\bar{b}_{1:M_{2}}\\ 0&0&\cdots&0&1\\ \end{bmatrix}\in{\mathbb{R}}^{({S}M_{2}+1)\times({S}+1)},$ $$

then we have

$$ $[\exp_{\delta\star}(h);1]=\overline{W}_{2}\cdot{\rm ReLU}(W_{3}\cdot[h;1]).$ $$

Define

$$ $\displaystyle W_{1}=$ $\displaystyle~{}\begin{bmatrix}a_{1:M_{1}}^{\mathsf{T}}&0&\cdots&0\\ 0&a_{1:M_{1}}^{\mathsf{T}}&\cdots&0\\ \cdots&\cdots&\cdots&\cdots\\ 0&0&\cdots&a_{1:M_{1}}^{\mathsf{T}}\\ \end{bmatrix}\in{\mathbb{R}}^{{S}\times{S}M_{1}},$ $\displaystyle\tilde{W}_{2}=$ $\displaystyle~{}\begin{bmatrix}w_{1:M_{1}}&0&\cdots&0&b_{1:M_{1}}\\ 0&w_{1:M_{1}}&\cdots&0&b_{1:M_{1}}\\ \cdots&\cdots&\cdots&\cdots&\cdots\\ 0&0&\cdots&w_{1:M_{1}}&\bar{b}_{1:M_{1}}\\ \end{bmatrix}\in{\mathbb{R}}^{{S}M_{1}\times({S}+1)},$ $$

then we have

$$ $\log_{\delta\star}(h)=W_{1}\cdot{\rm ReLU}(\tilde{W}_{2}\cdot[h;1]).$ $$

As a consequence, we have

$$ ${f}_{\star}(h)=\log_{\delta\star}\Psi\exp_{\delta\star}(h)=W_{1}\cdot{\rm ReLU }(\tilde{W}_{2}\cdot{\rm diag}(\Psi,1)\cdot\overline{W}_{2}\cdot{\rm ReLU}(W_{ 3}\cdot[h;1]))=W_{1}\cdot{\rm ReLU}(W_{2}\cdot{\rm ReLU}(W_{3}\cdot[h;1])),$ $$

where we define $W_{2}=\tilde{W}_{2}\cdot{\rm diag}(\Psi,1)\cdot\overline{W}_{2}$. It is also direct to upper bound $\|W_{1}\|_{\max}$, $\|W_{2}\|_{\max}$, and $\|W_{3}\|_{\max}$. This finishes the proof of Lemma 14.
∎

###### Lemma 15 (ReLU approximation of log-Psi) .

Assume that $1/K\leq\min_{ij}\Psi_{ij}\leq\max_{ij}\Psi_{ij}\leq K$. For any $\delta>0$, take $M_{1}=\lceil K/\delta\rceil+1$ and $M_{2}=3$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M_{1}]}$ and $\{(\bar{a}_{j},\bar{w}_{j},\bar{b}_{j})\}_{j\in[M_{2}]}$ with

$$ $\sup_{j}|a_{j}|\leq 2K,~{}~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}~{}\sup_{j}| b_{j}|\leq{S}K,~{}~{}~{}~{}\sup_{j}|\bar{a}_{j}|\leq 4,~{}~{}~{}~{}\sup_{j}| \bar{w}_{j}|\leq 1,~{}~{}~{}~{}\sup_{j}|\bar{b}_{j}|\leq 1,$ $$

such that defining

$$ $\textstyle\log_{\delta}(x)=\sum_{j=1}^{M_{1}}a_{j}\cdot{\rm ReLU}(w_{j}x+b_{j} ),~{}~{}~{}{\rm Ind}(x)=\sum_{j=1}^{M_{2}}\bar{a}_{j}\cdot{\rm ReLU}(w_{j}x+b_ {j}),$ $$

and defining ${f}(x),{f}_{\delta}(x)\in{\mathbb{R}}^{S}$ by

$$ $\textstyle{f}(x)_{i}=\log\sum_{j=1}^{S}\Psi_{ij}1(x=j),~{}~{}~{}{f}_{\delta}(x )_{i}=\log_{\delta}\sum_{j=1}^{S}\Psi_{ij}{\rm Ind}(x-j),~{}~{}~{}~{}~{} \forall i\in[{S}],$ $$

we have

$$ $\sup_{x\in[{S}]}\|{f}(x)-{f}_{\delta}(x)\|_{\infty}\leq\delta.$ $$

###### Proof of Lemma 15 .

The proof of the lemma is similar to the proof of Lemma 13, using a combination of Lemma 10, 9, and 12.
∎

###### Lemma 16 (Existence of ReLU network approximating log-Psi) .

Let ${f}_{\delta}$ be the function as defined in Lemma 15. Then there exists a two-hidden-layer neural network

$$ ${\rm NN}_{W_{1},W_{2},W_{3}}(x)=W_{1}\cdot{\rm ReLU}(W_{2}\cdot{\rm ReLU}(W_{3 }\cdot[x;1])),$ $$

with $W_{1}\in{\mathbb{R}}^{{S}\times{S}M_{1}}$, $W_{2}\in{\mathbb{R}}^{{S}M_{1}\times({S}M_{2}+1)}$, $W_{3}\in{\mathbb{R}}^{({S}M_{2}+1)\times({S}+1)}$, and

$$ $\|W_{1}\|_{\max}\leq 2K,~{}~{}~{}\|W_{2}\|_{\max}\leq{\rm Poly}({S}KM_{1}M_{2} ),~{}~{}~{}\|W_{3}\|_{\max}\leq 1.$ $$

such that

$$ ${\rm NN}_{W_{1},W_{2},W_{3}}(x)={f}_{\delta}(x),~{}~{}~{}\forall x\in[{S}].$ $$

###### Proof of Lemma 16 .

The proof of Lemma 16 is similar to the proof of Lemma 14.
∎

###### Lemma 5 (Error propagation of the approximate version of message passing in classification) .

###### Proof of Lemma 5 .

###### Lemma 6 (Non-expansiveness of log-sum-exponential) .

###### Proof of Lemma 6 .

###### Lemma 7 (Lipschitzness of the normalization operator) .

###### Proof of Lemma 7 .

###### Lemma 8 (ReLU approximation of the exponential function) .

###### Proof of Lemma 8 .

###### Lemma 9 (ReLU approximation of the logarithm function) .

###### Proof of Lemma 9 .

###### Lemma 10 (ReLU approximation of indicator function) .

###### Proof of Lemma 10 .

###### Lemma 11 (Log-sum-exponential approximation) .

###### Proof of Lemma 11 .

###### Lemma 12 (Log-Psi approximation) .

###### Proof of Lemma 12 .

###### Lemma 13 (ReLU approximation of log-sum-exponential) .

###### Proof of Lemma 13 .

###### Lemma 14 (Existence of ReLU network approximating log-sum-exponential) .

###### Proof of Lemma 14 .

###### Lemma 15 (ReLU approximation of log-Psi) .

###### Proof of Lemma 15 .

###### Lemma 16 (Existence of ReLU network approximating log-Psi) .

###### Proof of Lemma 16 .

## Appendix D Proof of Theorem 1

###### Proof of Theorem 1 .

By Lemma 17, we have the error decomposition

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_ {\star}})\leq\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{ \bm{W}},{\mu_{\star}})+2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{ \sf R}}(\mu_{{\rm NN}}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big{|}.$ $$

To control the first term (the approximation error), by Theorem 3, there exists ${\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ as in Eq. (9) with norm bound $B={\rm Poly}({d},{S},K,3^{L},D)$, such that defining $\mu_{{\rm NN}}^{\bm{W}}$ as in Eq. (ConvNet), we have

$$ $\max_{y\in[{S}],{\bm{x}}\in[{S}]^{d}}\Big{|}\log{\mu_{\star}}(y|{\bm{x}})-\log \mu_{{\rm NN}}^{\bm{W}}(y|{\bm{x}})\Big{|}\leq C\cdot\frac{S^{2}K^{2}{d}\cdot 3 ^{L}}{D}.$ $$

Furthermore, by Lemma 4, when $D\geq S^{2}K^{2}{d}\cdot 3^{L}$, we have

$$ $\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\bm{W}},{\mu_{ \star}})\leq C\Big{(}e^{\frac{S^{2}K^{2}{d}\cdot 3^{L}}{D}}-1\Big{)}^{2}\leq C \frac{S^{4}K^{4}{d}^{2}\cdot 3^{2L}}{D^{2}}.$ $$

To control the second term (the generalization error), by Proposition 7, with probability at least $1-\eta$, we have

$$ $\sup_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}}\Big{|}{ \widehat{\sf R}}(\mu_{{\rm NN}}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big {|}\leq C\cdot\sqrt{\frac{LD(D+2S+1)\|{\underline{m}}\|_{1}\log({d}\|{ \underline{m}}\|_{1}D{S}B\cdot 3^{L})+\log(1/\eta)}{n}}.$ $$

Combining the above two equations proves Theorem 1.
∎

###### Proof of Theorem 1 .

### D.1 Error decomposition

###### Lemma 17 .

Consider the setting of Theorem 1. We have decomposition

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_ {\star}})\leq\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{ \bm{W}},{\mu_{\star}})+2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{ \sf R}}(\mu_{{\rm NN}}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big{|}.$ $$

###### Proof of Lemma 17 .

We have that for any conditional distribution $\mu_{1}(\cdot|\cdot)$, there is decomposition

$$ $\displaystyle~{}{\sf D_{2}^{2}}(\mu_{1},{\mu_{\star}})={\mathbb{E}}_{{\bm{x}} \sim{\mu_{\star}}}\Big{[}\sum_{s=1}^{{S}}\Big{(}\mu_{1}(s|{\bm{x}})-{\mu_{ \star}}(s|{\bm{x}})\Big{)}^{2}\Big{]}$ $\displaystyle=$ $\displaystyle~{}{\mathbb{E}}_{({\bm{x}},y)\sim{\mu_{\star}}}\Big{[}\sum_{s=1}^ {S}(\mu_{1}(s|{\bm{x}})-1\{y=s\})^{2}\Big{]}-{\mathbb{E}}_{({\bm{x}},y)\sim{ \mu_{\star}}}\Big{[}\sum_{s=1}^{S}({\mu_{\star}}(s|{\bm{x}})-1\{y=s\})^{2}\Big {]}={\sf R}(\mu_{1})-{\sf R}({\mu_{\star}}).$ $$

Define

$$ ${\bm{W}}_{\star}=\arg\min_{{\bm{W}}\in{\mathcal{W}}}{\sf R}(\mu_{{\rm NN}}^{{ \bm{W}}})=\arg\min_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{{ \bm{W}}},{\mu_{\star}}).$ $$

Then we have

$$ $\displaystyle~{}{\sf D_{2}^{2}}(\mu_{{\rm NN}}^{\widehat{\bm{W}}},{\mu_{\star} })={\sf R}(\mu_{{\rm NN}}^{\widehat{\bm{W}}})-{\sf R}({\mu_{\star}})$ $\displaystyle=$ $\displaystyle~{}{\sf R}(\mu_{{\rm NN}}^{\widehat{\bm{W}}})-{\widehat{\sf R}}( \mu_{{\rm NN}}^{\widehat{\bm{W}}})+{\widehat{\sf R}}(\mu_{{\rm NN}}^{\widehat{ \bm{W}}})-{\widehat{\sf R}}(\mu_{{\rm NN}}^{{\bm{W}}_{\star}})+{\widehat{\sf R }}(\mu_{{\rm NN}}^{{\bm{W}}_{\star}})-{\sf R}(\mu_{{\rm NN}}^{{\bm{W}}_{\star} })+{\sf R}(\mu_{{\rm NN}}^{{\bm{W}}_{\star}})-{\sf R}({\mu_{\star}})$ $\displaystyle\leq$ $\displaystyle~{}2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\sf R}(\mu_{{\rm NN }}^{{\bm{W}}})-{\widehat{\sf R}}(\mu_{{\rm NN}}^{{\bm{W}}})\Big{|}+{\sf D_{2}^ {2}}(\mu_{{\rm NN}}^{{\bm{W}}_{\star}},{\mu_{\star}})$ $$

This proves Lemma 17.
∎

###### Lemma 17 .

###### Proof of Lemma 17 .

### D.2 Results on generalization

###### Proposition 7 (Generalization error of the classification problem) .

Let ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ be the set defined as in Eq. (9). Then, with probability at least $1-\eta$, we have

$$ $\sup_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}}\Big{|}{ \widehat{\sf R}}(\mu_{{\rm NN}}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big {|}\leq C\cdot\sqrt{\frac{LD(D+2S+1)\|{\underline{m}}\|_{1}\log({d}\|{ \underline{m}}\|_{1}D{S}B\cdot 3^{L})+\log(1/\eta)}{n}}.$ $$

###### Proof of Proposition 7 .

In Lemma 3, we can take $z=(y,{\bm{x}})$, $w={\bm{W}}$, $\Theta={\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$, $\rho(w,w^{\prime})={|\!|\!|{{\bm{W}}-{\bm{W}}^{\prime}}|\!|\!|}$, and $f(z_{i};w)={\rm loss}(y,\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}}))$. Therefore, to show Proposition 7, we just need to apply Lemma 3 by checking (a), (b), (c).

Check (a). We note that the index set $\Theta:={\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ equipped with $\rho(w,w^{\prime}):={|\!|\!|{{\bm{W}}-{\bm{W}}^{\prime}}|\!|\!|}$ has diameter ${B_{\rm p}}:=2B$. Further note that ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ has a dimension bounded by ${d_{\rm p}}:=D(D+2S+1)\|{\underline{m}}\|_{1}$. According to Example 5.8 of [], it holds that $\log N(\Delta;{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B},{|\!|\!|{\cdot}|\!
|\!|})\leq{d_{\rm p}}\cdot\log(1+2r/\Delta)$ for any $0<\Delta\leq 2r$. This verifies (a).

Check (b). Since $f(z_{i};w)={\rm loss}(y,\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}}))$ is $2$-bounded. As a consequence, $f(z,w)-{\mathbb{E}}_{z}[f(z,w)]$ is a sub-Gaussian random variable with the sub-Gaussian parameter to be a universal constant.

Check (c). Lemma 20 implies that

$$ $\displaystyle|f(z;w_{1})-f(z;w_{2})|\leq\sigma^{\prime}\cdot{|\!|\!|{{\bm{W}}_ {1}-{\bm{W}}_{2}}|\!|\!|},~{}~{}~{}\sigma^{\prime}:=12\|{\underline{m}}\|_{1}( 3B^{3})^{L}\cdot{d}\cdot S^{3/2}\cdot(S+D).$ $$

As a consequence, $f(z;w_{1})-f(z;w_{2})$ is $\sigma^{\prime}\rho(w_{1},w_{2})=\sigma^{\prime}{|\!|\!|{{\bm{W}}_{1}-{\bm{W}}
_{2}}|\!|\!|}$ sub-Gaussian.

Therefore, we apply Lemma 3 to conclude the proof of Proposition 7.
∎

###### Proposition 7 (Generalization error of the classification problem) .

###### Proof of Proposition 7 .

### D.3 Auxillary lemmas

###### Lemma 18 (Norm bound in the chain rule in classification settings) .

Consider the ConvNet as in Eq. (ConvNet). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then for any $\ell$, $v$, $\iota$, and $\star\in[3]$, we have

$$ $\displaystyle\|{\mathsf{h}}_{v}^{(L)}\|_{2}\leq$ $\displaystyle~{}S^{3/2},$ $\displaystyle\|{\mathsf{q}}_{v}^{(\ell)}\|_{2}\leq$ $\displaystyle~{}B^{3}\cdot(\|{\mathsf{h}}_{v}^{(\ell)}\|_{2}+1),$ $\displaystyle\|{\mathsf{h}}_{v}^{(\ell-1)}\|_{2}\leq$ $\displaystyle~{}2m^{(\ell)}\cdot\max_{v^{\prime}\in{\mathcal{C}}(v)}\|{\mathsf {q}}_{v^{\prime}}^{(\ell)}\|_{2},$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\iota}}{\mathsf{q}}_{v,i }^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}B^{3}\cdot\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\iota}}{ \mathsf{h}}_{v,i}^{(\ell)}\|_{{\rm op}},$ $\displaystyle\forall k\geq\ell+1,$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\iota}}{\mathsf{q}}_{ v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}B^{2}\cdot\|{\mathsf{h}}_{v}^{(\ell)}\|_{2},$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\iota}}{\mathsf{h}}_{v,i }^{(\ell-1)}\|_{{\rm op}}\leq$ $\displaystyle~{}2m^{(\ell)}\cdot\max_{v^{\prime}\in{\mathcal{C}}(v)}\max_{i\in [{S}]}\|\nabla_{W^{(k)}_{\star,\iota}}{\mathsf{q}}_{v^{\prime},i}^{(\ell)}\|_{ {\rm op}},~{}~{}~{}~{}$ $\displaystyle\forall k\geq\ell,$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\iota}}{\rm softmax}( {\mathsf{h}}^{(0)}_{{\rm r}})_{i}\|_{{\rm op}}\leq$ $\displaystyle~{}\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\iota}}{\mathsf{h} }^{(0)}_{{\rm r},i}\|_{{\rm op}}.$ $$

###### Proof of Lemma 18 .

The proof of the lemma uses the chain rule, the $1$-Lipschitzness of ${\rm ReLU}$, the $2$-Lipschitzness of ${\rm normalize}$, and the $1$-Lipschitzness of ${\rm softmax}$.
∎

###### Lemma 19 .

Consider the ConvNet as in Eq. (ConvNet). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then we have

$$ $\max_{{\bm{x}}\in[{S}]^{d}}\max_{\star\in[3]}\max_{\ell\in[L]}\max_{\iota\in[m ^{(\ell)}]}\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\iota}}{\rm softmax}({ \mathsf{h}}^{(0)}_{{\rm r}})_{i}\|_{{\rm op}}\leq(3B^{3})^{L}\cdot{d}\cdot{S}^ {3/2}.$ $$

###### Proof of Lemma 19 .

This lemma is implied by Lemma 18 and an induction argument.
∎

###### Lemma 20 .

Consider the ConvNet as in Eq. (ConvNet). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then we have

$$ $\displaystyle\max_{y,{\bm{x}}}\Big{|}\mu_{{\rm NN}}^{{\bm{W}}}(y|{\bm{x}})-\mu _{{\rm NN}}^{\overline{\bm{W}}}(y|{\bm{x}})\Big{|}\leq$ $\displaystyle~{}3\|{\underline{m}}\|_{1}(3B^{3})^{L}\cdot{d}\cdot{S}^{3/2} \cdot(S+D)\cdot{|\!|\!|{{\bm{W}}-\overline{\bm{W}}}|\!|\!|}.$ $$

Therefore, we have

$$ $\Big{|}{\rm loss}(y,\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}}))-{\rm loss}(y, \mu_{{\rm NN}}^{\overline{\bm{W}}}(\cdot|{\bm{x}}))\Big{|}\leq 12\|{\underline {m}}\|_{1}(3B^{3})^{L}\cdot{d}\cdot{S}^{3/2}\cdot(S+D)\cdot{|\!|\!|{{\bm{W}}- \overline{\bm{W}}}|\!|\!|}.$ $$

###### Proof of Lemma 20 .

The first inequality is by the fact that

$$ $\displaystyle~{}\max_{y,{\bm{x}}}\Big{|}\mu_{{\rm NN}}^{{\bm{W}}}(y|{\bm{x}})- \mu_{{\rm NN}}^{\overline{\bm{W}}}(y|{\bm{x}})\Big{|}$ $\displaystyle\leq$ $\displaystyle~{}\sum_{\star\in[3]}\sum_{\ell\in[L]}\sum_{\iota\in[m^{(\ell)}]} \min\{{\rm nrow}(W^{(\ell)}_{\star,\iota}),{\rm ncol}(W^{(\ell)}_{\star,\iota} )\}\|\nabla_{W^{(\ell)}_{\star,\iota}}{\rm softmax}({\mathsf{h}}^{(0)}_{{\rm r }})_{i}\|_{{\rm op}}\|W^{(\ell)}_{\star,\iota}-\overline{W}^{(\ell)}_{\star, \iota}\|_{{\rm op}},$ $$

where we have used the inequality that ${\rm trace}(A^{\mathsf{T}}B)\leq\{{\rm nrow}(A),{\rm ncol}(A)\}\|A\|_{{\rm op}
}\|B\|_{{\rm op}}$.

To prove the second equation, we have

$$ $\displaystyle~{}\Big{|}{\rm loss}(y,\mu_{{\rm NN}}^{{\bm{W}}}(\cdot|{\bm{x}})) -{\rm loss}(y,\mu_{{\rm NN}}^{\overline{\bm{W}}}(\cdot|{\bm{x}}))\Big{|}=\Big{ |}\sum_{s=1}^{{S}}\Big{(}1\{y=s\}-\mu_{{\rm NN}}^{{\bm{W}}}(s|{\bm{x}})\Big{)} ^{2}-\sum_{s=1}^{{S}}\Big{(}1\{y=s\}-\mu_{{\rm NN}}^{\overline{\bm{W}}}(s|{\bm {x}})\Big{)}^{2}\Big{|}$ $\displaystyle\leq$ $\displaystyle~{}\sum_{s=1}^{{S}}\Big{|}1\{y=s\}-\mu_{{\rm NN}}^{{\bm{W}}}(s|{ \bm{x}})+1\{y=s\}-\mu_{{\rm NN}}^{\overline{\bm{W}}}(s|{\bm{x}})\Big{|}\cdot \Big{|}\mu_{{\rm NN}}^{{\bm{W}}}(s|{\bm{x}})-\mu_{{\rm NN}}^{\overline{\bm{W}} }(s|{\bm{x}})\Big{|}\leq 4\cdot\max_{s}\Big{|}\mu_{{\rm NN}}^{{\bm{W}}}(s|{\bm {x}})-\mu_{{\rm NN}}^{\overline{\bm{W}}}(s|{\bm{x}})\Big{|}.$ $$

This completes the proof of Lemma 20.
∎

###### Lemma 18 (Norm bound in the chain rule in classification settings) .

###### Proof of Lemma 18 .

###### Lemma 19 .

###### Proof of Lemma 19 .

###### Lemma 20 .

###### Proof of Lemma 20 .

## Appendix E Proof of Proposition 5

###### Proof of Proposition 5 .

By Eq. (MP-DNS) and (23), defining ${\nu}^{(\ell)}_{\downarrow,v}(\cdot)={\rm softmax}({h}^{(\ell)}_{\downarrow,v})$, then for $\ell\leq L-1$, we get

$$ $\displaystyle{\nu}^{(\ell)}_{\downarrow,v}(x_{v}^{(\ell)})\propto$ $\displaystyle~{}\prod_{v^{\prime}\in{\mathcal{C}}(v)}\Big{(}\sum_{a\in[S]}\psi _{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{(\ell)},a)e^{({h}_{\downarrow,v}^{(\ell +1)})_{a}}\Big{)}$ $\displaystyle\propto$ $\displaystyle~{}\sum_{x^{(\ell+1)}_{{\mathcal{C}}(v)}}\prod_{v^{\prime}\in{ \mathcal{C}}(v)}\Big{(}\psi_{\iota(v^{\prime})}^{(\ell+1)}(x_{v}^{(\ell)},x^{( \ell+1)}_{v^{\prime}}){\nu}^{(\ell+1)}_{\downarrow,v^{\prime}}(x^{(\ell+1)}_{v ^{\prime}})\Big{)}.$ $$

This coincides with the update rule of ${\nu}^{(\ell)}_{\downarrow,v}$ as in Eq. (BP-DNS).

Furthermore, defining ${\nu}^{(\ell)}_{\uparrow,v}(\cdot)={\rm softmax}({b}_{\uparrow,v}^{(\ell)}-{h}
^{(\ell)}_{\downarrow,v})$, then for $\ell=0,1,\ldots,L$, we get

$$ $\displaystyle{\nu}^{(\ell)}_{\uparrow,v}(x_{v}^{(\ell)})\propto$ $\displaystyle~{}\sum_{b\in[{S}]}\psi_{\iota(v)}^{(\ell+1)}(b,x_{v}^{(\ell)}){ \nu}^{(\ell+1)}_{\uparrow,{\rm pa}(v)}(b)\prod_{v^{\prime}\in{\mathcal{N}}(v)} \Big{(}\sum_{a\in[S]}\psi_{\iota(v^{\prime})}^{(\ell+1)}(b,a)e^{({h}_{ \downarrow,v}^{(\ell+1)})_{a}}\Big{)}$ $\displaystyle\propto$ $\displaystyle~{}\sum_{x^{(\ell-1)}_{{\rm pa}(v)},x^{(\ell)}_{{\mathcal{N}}(v)} }\psi^{(\ell)}(x^{(\ell-1)}_{{\rm pa}(v)},x^{(\ell)}_{{\mathcal{C}}({\rm pa}(v ))}){\nu}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}(x^{(\ell-1)}_{{\rm pa}(v)})\prod_{ v^{\prime}\in{\mathcal{N}}(v)}{\nu}_{\downarrow,v^{\prime}}^{(\ell)}(x^{(\ell) }_{v^{\prime}}).$ $$

This coincides with the update rule of ${\nu}^{(\ell)}_{\uparrow,v}$ as in Eq. (BP-DNS).

Finally, defining ${\nu}_{v}^{(L)}(\cdot)={\rm softmax}(b_{\uparrow,v}^{(L)})$, we get

$$ $\displaystyle{\nu}^{(L)}_{v}(x_{v}^{(L)})\propto$ $\displaystyle~{}\sum_{b\in[{S}]}\psi_{\iota(v)}^{(L)}(b,x_{v}^{(L)}){\nu}^{(L) }_{\uparrow,{\rm pa}(v)}(b)\prod_{v^{\prime}\in{\mathcal{N}}(v)}\Big{(}\sum_{a \in[S]}\psi_{\iota(v^{\prime})}^{(L)}(b,a)e^{({h}_{\downarrow,v}^{(L)})_{a}} \Big{)}\times{\nu}^{(L)}_{\downarrow,v}(x_{v}^{(L)})$ $\displaystyle\propto$ $\displaystyle~{}{\nu}^{(L)}_{\uparrow,v}(x_{v}^{(L)})\psi^{(L+1)}(x^{(L)}_{v}, z_{v}).$ $$

This coincides with the formula of ${\nu}^{(L)}_{v}$ as in Eq. (BP-DNS).

This finishes the proof of Proposition 5.
∎

###### Proof of Proposition 5 .

## Appendix F Proof of Theorem 6

###### Proof of Theorem 6 .

By Lemma 13, take $M_{1}=\lceil 2{S}^{2}K{d}18^{L}/\delta\rceil+1$ and $M_{2}=\lceil 2{S}^{2}K^{2}{d}18^{L}/\delta\rceil+1$. Then there exists $\{(a_{j},w_{j},b_{j})\}_{j\in[M_{1}]}$ and $\{(\bar{a}_{j},\bar{w}_{j},\bar{b}_{j})\}_{j\in[M_{2}]}$ with

$$ $\sup_{j}|a_{j}|\leq 2{S}K,~{}~{}~{}~{}\sup_{j}|w_{j}|\leq 1,~{}~{}~{}~{}\sup_{ j}|b_{j}|\leq{S}K,~{}~{}~{}~{}\sup_{j}|\bar{a}_{j}|\leq 2,~{}~{}~{}~{}\sup_{j} |\bar{w}_{j}|\leq 1,~{}~{}~{}~{}\sup_{j}|\bar{b}_{j}|\leq\log(4\cdot 18^{L}{S} ^{2}{d}K^{2}/\delta),$ $$

such that defining

$$ $\textstyle\log_{\delta\star}(x)=\sum_{j=1}^{M_{1}}a_{j}\cdot{\rm ReLU}(w_{j}x+ b_{j}),~{}~{}~{}\exp_{\delta\star}(x)=\sum_{j=1}^{M_{2}}\bar{a}_{j}\cdot{\rm ReLU }(\bar{w}_{j}x+\bar{b}_{j}),$ $$

and defining ${f}_{\diamond,\iota}^{(\ell)}(h),{\bar{f}}_{\diamond,\iota}^{(\ell)}(h)\in{
\mathbb{R}}^{S}$ for $\diamond\in\{\downarrow,\uparrow\}$ by

$$ $\displaystyle\textstyle{f}_{\downarrow,\iota}^{(\ell)}(h)_{i}=\log\sum_{j=1}^{ S}\psi_{\iota}^{(\ell)}(i,j)\exp(h_{j}),~{}~{}~{}{\bar{f}}_{\downarrow,\iota}^ {(\ell)}(h)_{i}=\log_{\delta\star}\sum_{j=1}^{S}\psi_{\iota}^{(\ell)}(i,j)\exp _{\delta\star}(h_{j}),~{}~{}~{}~{}~{}\forall i\in[{S}],$ $\displaystyle\textstyle{f}_{\uparrow,\iota}^{(\ell)}(h)_{i}=\log\sum_{j=1}^{S} \psi_{\iota}^{(\ell)}(j,i)\exp(h_{j}),~{}~{}~{}{\bar{f}}_{\uparrow,\iota}^{( \ell)}(h)_{i}=\log_{\delta\star}\sum_{j=1}^{S}\psi_{\iota}^{(\ell)}(j,i)\exp_{ \delta\star}(h_{j}),~{}~{}~{}~{}~{}\forall i\in[{S}],$ $$

we have

$$ $\sup_{\max_{i}h_{i}=0}\|{f}_{\diamond,\iota}^{(\ell)}(h)-{\bar{f}}_{\diamond, \iota}^{(\ell)}(h)\|_{\infty}\leq\delta/({d}18^{L}{S}).$ (30) $$

By Eq. (30) and Lemma 21, taking ${b}_{\uparrow,v}^{(L)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (MP-DNS) and ${\bar{b}}_{\uparrow,v}^{(L)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (A-MP-DNS) with $\{{\bar{f}}_{\diamond,\iota}^{(\ell)}\}_{\ell\in[L],\iota\in[{m}^{(\ell)}]}$ as defined above, we have

$$ $\|{b}_{\uparrow,v}^{(L)}-{\bar{b}}_{\uparrow,v}^{(L)}\|_{\infty}\leq[\delta/({ d}18^{L}{S})]\times 18^{L}{d}=\delta/{S},$ $$

which gives

$$ $\sup_{{\bm{z}}\in{\mathbb{R}}^{d}}\|{\bm{m}}({\bm{z}})-{\bm{m}}_{{\rm NN}}({ \bm{z}})\|_{\infty}\leq{S}\cdot\|{b}_{\uparrow,v}^{(L)}-{\bar{b}}_{\uparrow,v} ^{(L)}\|_{\infty}\leq\delta.$ $$

As a 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 $\diamond\in\{\downarrow,\uparrow\}$, $\ell\in[L]$, and $\iota\in[{m}^{(\ell)}]$)

$$ ${\rm NN}_{W_{1,\diamond,\iota}^{(\ell)},W_{2,\diamond,\iota}^{(\ell)},W_{3, \diamond,\iota}^{(\ell)}}(h)=W_{1,\diamond,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{2 ,\diamond,\iota}^{(\ell)}\cdot{\rm ReLU}(W_{3,\diamond,\iota}^{(\ell)}\cdot[h; 1])),$ $$

with $W_{1,\diamond,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}\times{S}M_{1}}$, $W_{2,\diamond,\iota}^{(\ell)}\in{\mathbb{R}}^{{S}M_{1}\times({S}M_{2}+1)}$, $W_{3,\diamond,\iota}^{(\ell)}\in{\mathbb{R}}^{({S}M_{2}+1)\times({S}+1)}$, and

$$ $\|W_{1,\diamond,\iota}^{(\ell)}\|_{\max}\leq 2{S}K,~{}~{}~{}\|W_{2,\diamond, \iota}^{(\ell)}\|_{\max}\leq{\rm Poly}({S}KM_{1}M_{2}),~{}~{}~{}\|W_{3, \diamond,\iota}^{(\ell)}\|_{\max}\leq\log(4{S}^{2}{d}18^{L}K^{2}/\delta).$ $$

such that

$$ ${\rm NN}_{W_{1,\diamond,\iota}^{(\ell)},W_{2,\diamond,\iota}^{(\ell)},W_{3, \diamond,\iota}^{(\ell)}}(h)={\bar{f}}_{\diamond,\iota}^{(\ell)}(h),~{}~{}~{} \forall h\in{\mathbb{R}}^{S}\text{ such that }\max_{j}h_{j}=0.$ $$

This proves that the approximate version of message passing as in Eq. (A-MP-DNS) coincides with the U-Net as in Eq. (UNet) with proper choice of dimension

$$ $D\geq\max_{\ell\in[L]}\{SM_{1}^{(\ell)},SM_{2}^{(\ell)}+1\}={S}\times\Big{(} \lceil 2{S}^{2}K^{2}{d}18^{L}/\delta\rceil+1\Big{)}+1$ $$

and norm of the weights. This finishes the proof of Theorem 6.
∎

###### Proof of Theorem 6 .

### F.1 Auxillary lemmas

###### Lemma 21 (Error propagation of the approximate version of message passing in denoising) .

Assume we have functions $\{{f}_{\downarrow,\iota}^{(\ell)},{f}_{\uparrow,\iota}^{(\ell)}\}$ and $\{{\bar{f}}_{\downarrow,\iota}^{(\ell)},{\bar{f}}_{\uparrow,\iota}^{(\ell)}\}$ such that

$$ $\displaystyle\|{f}_{\diamond,\iota}^{(\ell)}(h)-{\bar{f}}_{\diamond,\iota}^{( \ell)}(h)\|_{\infty}\leq$ $\displaystyle~{}\delta,~{}~{}~{}\forall h\in{\mathbb{R}}^{S}\text{ such that } \max_{j\in[{S}]}h_{j}=0,~{}~{}\diamond\in\{\downarrow,\uparrow\}.$ (31) $$

Furthermore, consider the following approximate version of message passing algorithm with initialization ${h}_{\downarrow,v}^{(L)}=(-(x-z_{v})^{2}/2)_{x\in[{S}]}\in{\mathbb{R}}^{{S}}$ for $v\in{\mathcal{V}}^{(L)}$, defined as below

$$ $\displaystyle{\bar{q}}_{\downarrow,v}^{(\ell)}=$ $\displaystyle~{}{\bar{f}}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar {h}}_{\downarrow,v}^{(\ell)}))\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ (A-MP-DNS) $\displaystyle{\bar{h}}_{\downarrow,v}^{(\ell-1)}=$ $\displaystyle~{}\textstyle\sum_{v^{\prime}\in{\mathcal{C}}(v)}{\bar{q}}_{ \downarrow,v^{\prime}}^{(\ell)}\in{\mathbb{R}}^{{S}},~{}~{}~{}~{}~{}~{}$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell-1)},$ $\displaystyle{\bar{u}}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{\bar{b}}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}\in{\mathbb{R}}^{{S }},~{}~{}~{}(\text{with }{\bar{b}}_{\uparrow,{\rm r}}^{(0)}={\bar{h}}_{ \downarrow,{\rm r}}^{(0)})$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{\bar{b}}_{\uparrow,v}^{(\ell)}=$ $\displaystyle~{}{\bar{f}}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar{u }}_{\uparrow,v}^{(\ell)}-{\bar{q}}_{\downarrow,v}^{(\ell)}))+{\bar{h}}_{ \downarrow,v}^{(\ell)}\in{\mathbb{R}}^{{S}},$ $\displaystyle\ell\in[L],~{}~{}~{}v\in{\mathcal{V}}^{(\ell)},$ $\displaystyle{\bar{\bm{m}}}_{{\rm MP}}({\bm{z}})_{v}=$ $\displaystyle~{}\textstyle\sum_{s\in[{S}]}s\cdot{\rm softmax}({\bar{b}}_{ \uparrow,v}^{(L)})_{s},$ $\displaystyle v\in{\mathcal{V}}^{(L)}.$ $$

Taking ${b}_{\uparrow,v}^{(L)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (MP-DNS) and ${\bar{b}}_{\uparrow,v}^{(L)}\in{\mathbb{R}}^{{S}}$ to be as defined in Eq. (A-MP-DNS), we have

$$ $\max_{v\in{\mathcal{V}}^{(L)}}\|{b}_{\uparrow,v}^{(L)}-{\bar{b}}_{\uparrow,v}^ {(L)}\|_{\infty}\leq\delta\times 18^{L}\cdot{d}.$ $$

###### Proof of Lemma 21 .

Step 1. Downward induction. In the first step, we aim to show that for any $\ell\in[L-1]$ we have

$$ $\textstyle\|{\bar{h}}_{\downarrow,v}^{(\ell)}-{h}_{\downarrow,v}^{(\ell)}\|_{ \infty}\leq{m}^{(\ell+1)}\prod_{k=\ell+2}^{L}(2{m}^{(k)}+1)\delta,~{}~{}~{}~{} \forall v\in{\mathcal{V}}^{(\ell)}.$ (32) $$

To prove the formula for $\ell=L-1$, since ${\rm normalize}({h}_{\downarrow,v}^{(L)})={\rm normalize}({\bar{h}}_{
\downarrow,v}^{(L)})$, by Eq. (31), we get

$$ $\|{\bar{q}}_{\downarrow,v}^{(L)}-{q}_{\downarrow,v}^{(L)}\|_{\infty}\leq\delta ,~{}~{}~{}~{}\forall v\in{\mathcal{V}}^{(L)}.$ $$

Hence we get

$$ $\textstyle\|{h}_{\downarrow,v}^{(L-1)}-{\bar{h}}_{\downarrow,v}^{(L-1)}\|_{ \infty}=\Big{\|}\sum_{v^{\prime}\in{\mathcal{C}}(v)}({\bar{q}}_{\downarrow,v^{ \prime}}^{(L)}-{q}_{\downarrow,v^{\prime}}^{(L)})\Big{\|}_{\infty}\leq{m}^{(L) }\delta,~{}~{}~{}~{}\forall v\in{\mathcal{V}}^{(L-1)}.$ $$

This proves the formula (32) for $\ell=L-1$.

Assuming that (32) holds at the layer $\ell$, by the update formula, we have

$$ $\displaystyle~{}\|{\bar{q}}_{\downarrow,v}^{(\ell)}-{q}_{\downarrow,v}^{(\ell) }\|_{\infty}=\|{\bar{f}}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar{ h}}_{\downarrow,v}^{(\ell)}))-{f}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize }({h}_{\downarrow,v}^{(\ell)}))\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|{\bar{f}}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize}({ \bar{h}}_{\downarrow,v}^{(\ell)}))-{f}_{\downarrow,\iota(v)}^{(\ell)}({\rm normalize }({\bar{h}}_{\downarrow,v}^{(\ell)}))\|_{\infty}+\|{f}_{\downarrow,\iota(v)}^{ (\ell)}({\rm normalize}({h}_{\downarrow,v}^{(\ell)}))-{f}_{\downarrow,\iota(v) }^{(\ell)}({\rm normalize}({\bar{h}}_{\downarrow,v}^{(\ell)}))\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\textstyle\delta+2{m}^{(\ell+1)}\prod_{k=\ell+2}^{L}(2{m}^{(k) }+1)\delta\leq\prod_{k=\ell+1}^{L}(2{m}^{(k)}+1)\delta,$ $$

where the middle inequality is by the assumption of ${f}_{\downarrow,\iota}^{(\ell)}$ and by Lemma 6 and Lemma 7. Hence we get

$$ $\textstyle\|{\bar{h}}_{\downarrow,v}^{(\ell-1)}-{h}_{\downarrow,v}^{(\ell-1)} \|_{\infty}\leq{m}^{(\ell)}\prod_{k=\ell+1}^{L}(2{m}^{(k)}+1)\delta,~{}~{}~{}~ {}\forall v\in{\mathcal{V}}^{(\ell)}.$ $$

This proves Eq. (32) by the induction argument.

Step 2. Upward induction. The downward induction argument proves that, for $\Gamma=\prod_{k=1}^{L}(2{m}^{(k)}+1)$, we have

$$ $\textstyle\|{\bar{q}}_{\downarrow,v}^{(\ell)}-{q}_{\downarrow,v}^{(\ell)}\|_{ \infty},\|{\bar{h}}_{\downarrow,v}^{(\ell)}-{h}_{\downarrow,v}^{(\ell)}\|_{ \infty}\leq\Gamma\delta,~{}~{}~{}~{}\forall\ell=0,1,\ldots,L,~{}~{}\forall v \in{\mathcal{V}}^{(\ell)}.$ $$

In this step, we aim to show that for any $\ell=0,1,\ldots,L$, we have

$$ $\textstyle\|{\bar{b}}_{\uparrow,v}^{(\ell)}-{b}_{\uparrow,v}^{(\ell)}\|_{ \infty}\leq 6^{\ell}\cdot\Gamma\cdot\delta,~{}~{}~{}~{}\forall v\in{\mathcal{V }}^{(\ell)}.$ (33) $$

To prove this formula for $\ell=0$, note that ${b}_{\uparrow,{\rm r}}^{(0)}={h}_{\downarrow,{\rm r}}^{(0)}$ and ${\bar{b}}_{\uparrow,{\rm r}}^{(0)}={\bar{h}}_{\downarrow,{\rm r}}^{(0)}$, we have

$$ $\|{\bar{b}}_{\uparrow,{\rm r}}^{(0)}-{b}_{\uparrow,{\rm r}}^{(0)}\|_{\infty}= \|{\bar{h}}_{\downarrow,{\rm r}}^{(0)}-{h}_{\downarrow,{\rm r}}^{(0)}\|_{ \infty}\leq\Gamma\cdot\delta.$ $$

This proves the formula (33) for $\ell=0$.

Assuming that (33) holds at layer $\ell-1$, by the update formula, we have

$$ $\displaystyle~{}\|{b}_{\uparrow,v}^{(\ell)}-{\bar{b}}_{\uparrow,v}^{(\ell)}\|_ {\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|{f}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({b}_{ \uparrow,{\rm pa}(v)}^{(\ell-1)}-{q}_{\downarrow,v}^{(\ell)}))-{\bar{f}}_{ \uparrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar{b}}_{\uparrow,{\rm pa}(v)}^{ (\ell-1)}-{\bar{q}}_{\downarrow,v}^{(\ell)}))\|_{\infty}+\|{h}_{\downarrow,v}^ {(\ell)}-{\bar{h}}_{\downarrow,v}^{(\ell)}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|{f}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({b}_{ \uparrow,{\rm pa}(v)}^{(\ell-1)}-{q}_{\downarrow,v}^{(\ell)}))-{f}_{\uparrow, \iota(v)}^{(\ell)}({\rm normalize}({\bar{b}}_{\uparrow,{\rm pa}(v)}^{(\ell-1)} -{\bar{q}}_{\downarrow,v}^{(\ell)}))\|_{\infty}$ $\displaystyle~{}+\|{f}_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar{b}}_ {\uparrow,{\rm pa}(v)}^{(\ell-1)}-{\bar{q}}_{\downarrow,v}^{(\ell)}))-{\bar{f} }_{\uparrow,\iota(v)}^{(\ell)}({\rm normalize}({\bar{b}}_{\uparrow,{\rm pa}(v) }^{(\ell-1)}-{\bar{q}}_{\downarrow,v}^{(\ell)}))\|_{\infty}+\|{h}_{\downarrow, v}^{(\ell)}-{\bar{h}}_{\downarrow,v}^{(\ell)}\|_{\infty}$ $\displaystyle\leq$ $\displaystyle~{}\|{\rm normalize}({b}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}-{q}_{ \downarrow,v}^{(\ell)})-{\rm normalize}({\bar{b}}_{\uparrow,{\rm pa}(v)}^{( \ell-1)}-{\bar{q}}_{\downarrow,v}^{(\ell)})\|_{\infty}+\delta+\Gamma\cdot\delta$ $\displaystyle\leq$ $\displaystyle~{}4\cdot 6^{\ell-1}\cdot\Gamma\cdot\delta+\delta+\Gamma\cdot \delta\leq 6^{\ell}\cdot\Gamma\cdot\delta.$ $$

This proves Eq. (33) by the induction argument. This proves the Lemma 21 by observing that $\Gamma\leq 3^{L}\prod_{\ell=1}^{L}m^{(\ell)}=3^{L}\cdot{d}$.
∎

###### Lemma 21 (Error propagation of the approximate version of message passing in denoising) .

###### Proof of Lemma 21 .

## Appendix G Proof of Theorem 4

###### Proof of Theorem 4 .

By Lemma 22, we have the error decomposition

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}}, {\bm{m}})\leq\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN} }^{\bm{W}},{\bm{m}})+2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{\sf R }}({\bm{m}}_{{\rm NN}}^{\bm{W}})-{\sf R}({\bm{m}}_{{\rm NN}}^{\bm{W}})\Big{|}.$ $$

To control the first term (the approximation error), by Theorem 6, there exists ${\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ as in Eq. (19) with norm bound $B={\rm Poly}({d},{S},K,18^{L},D)$, such that defining ${\bm{m}}_{{\rm NN}}^{\bm{W}}$ as in Eq. (UNet), we have

$$ $\sup_{{\bm{z}}\in{\mathbb{R}}^{d}}\|{\bm{m}}({\bm{z}})-{\bm{m}}_{{\rm NN}}({ \bm{z}})\|_{\infty}\leq C\cdot\frac{S^{3}K^{2}{d}\cdot 18^{L}}{D}.$ $$

Therefore, we have

$$ $\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\bm{W}},{ \bm{m}})\leq\sup_{{\bm{z}}}\|{\bm{m}}({\bm{z}})-{\bm{m}}_{{\rm NN}}({\bm{z}}) \|_{\infty}^{2}\leq C\cdot\frac{S^{6}K^{4}{d}^{2}\cdot 18^{2L}}{D^{2}}.$ $$

To control the second term (the generalization error), by Proposition 8, with probability at least $1-\eta$, we have

$$ $\sup_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}}\Big{|}{ \widehat{\sf R}}({\bm{m}}_{{\rm NN}}^{\bm{W}})-{\sf R}({\bm{m}}_{{\rm NN}}^{ \bm{W}})\Big{|}\leq C\cdot{S}^{2}\cdot\sqrt{\frac{LD(D+2S+1)\|{\underline{m}} \|_{1}\log({d}\|{\underline{m}}\|_{1}D{S}B\cdot 18^{L})+\log(1/\eta)}{n}}.$ $$

Combining the above two equations proves Theorem 4.
∎

###### Proof of Theorem 4 .

### G.1 Error decomposition

###### Lemma 22 .

Consider the setting of Theorem 4. We have decomposition

$$ $\displaystyle\textstyle{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}}, {\bm{m}})\leq\inf_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN} }^{\bm{W}},{\bm{m}})+2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\widehat{\sf R }}({\bm{m}}_{{\rm NN}}^{\bm{W}})-{\sf R}({\bm{m}}_{{\rm NN}}^{\bm{W}})\Big{|}.$ $$

###### Proof of Lemma 22 .

We have that for any conditional expectation ${\bm{m}}_{1}({\bm{z}})$, there is decomposition

$$ $\displaystyle~{}{\sf D_{2}^{2}}({\bm{m}}_{1},{\bm{m}})={\mathbb{E}}_{({\bm{x}} ,{\bm{z}})\sim{\mu_{\star}}}\Big{[}{d}^{-1}\|{\bm{m}}_{1}({\bm{z}})-{\bm{m}}({ \bm{z}})\|_{2}^{2}\Big{]}$ $\displaystyle=$ $\displaystyle~{}{\mathbb{E}}_{({\bm{x}},{\bm{z}})\sim{\mu_{\star}}}\Big{[}{d}^ {-1}\|{\bm{m}}_{1}({\bm{z}})-{\bm{x}}\|_{2}^{2}\Big{]}-{\mathbb{E}}_{({\bm{x}} ,{\bm{z}})\sim{\mu_{\star}}}\Big{[}{d}^{-1}\|{\bm{m}}_{1}({\bm{z}})-{\bm{x}}\| _{2}^{2}\Big{]}={\sf R}({\bm{m}}_{1})-{\sf R}({\bm{m}}).$ $$

Define

$$ ${\bm{W}}_{\star}=\arg\min_{{\bm{W}}\in{\mathcal{W}}}{\sf R}({\bm{m}}_{{\rm NN} }^{{\bm{W}}})=\arg\min_{{\bm{W}}\in{\mathcal{W}}}{\sf D_{2}^{2}}({\bm{m}}_{{ \rm NN}}^{{\bm{W}}},{\bm{m}}).$ $$

Then we have

$$ $\displaystyle~{}{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}},{\bm{m} })={\sf R}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}})-{\sf R}({\bm{m}})$ $\displaystyle=$ $\displaystyle~{}{\sf R}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}})-{\widehat{\sf R }}({\bm{m}}_{{\rm NN}}^{\widehat{\bm{W}}})+{\widehat{\sf R}}({\bm{m}}_{{\rm NN }}^{\widehat{\bm{W}}})-{\widehat{\sf R}}({\bm{m}}_{{\rm NN}}^{{\bm{W}}_{\star} })+{\widehat{\sf R}}({\bm{m}}_{{\rm NN}}^{{\bm{W}}_{\star}})-{\sf R}({\bm{m}}_ {{\rm NN}}^{{\bm{W}}_{\star}})+{\sf R}({\bm{m}}_{{\rm NN}}^{{\bm{W}}_{\star}}) -{\sf R}({\bm{m}})$ $\displaystyle\leq$ $\displaystyle~{}2\cdot\sup_{{\bm{W}}\in{\mathcal{W}}}\Big{|}{\sf R}({\bm{m}}_{ {\rm NN}}^{{\bm{W}}})-{\widehat{\sf R}}({\bm{m}}_{{\rm NN}}^{{\bm{W}}})\Big{|} +{\sf D_{2}^{2}}({\bm{m}}_{{\rm NN}}^{{\bm{W}}_{\star}},{\bm{m}})$ $$

This proves Lemma 22.
∎

###### Lemma 22 .

###### Proof of Lemma 22 .

### G.2 Results on generalization

###### Proposition 8 (Generalization error of the denoising problem) .

Let ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ be the set defined as in Eq. (9). Then, with probability at least $1-\eta$, we have

$$ $\sup_{{\bm{W}}\in{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}}\Big{|}{ \widehat{\sf R}}(\mu_{{\rm NN}}^{\bm{W}})-{\sf R}(\mu_{{\rm NN}}^{\bm{W}})\Big {|}\leq C\cdot{S}^{2}\cdot\sqrt{\frac{LD(D+2S+1)\|{\underline{m}}\|_{1}\log({d }\|{\underline{m}}\|_{1}D{S}B\cdot 18^{L})+\log(1/\eta)}{n}}.$ $$

###### Proof of Proposition 8 .

In Lemma 3, we can take $z=({\bm{z}},{\bm{x}})$, $w={\bm{W}}$, $\Theta={\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$, $\rho(w,w^{\prime})={|\!|\!|{{\bm{W}}-{\bm{W}}^{\prime}}|\!|\!|}$, and $f(z_{i};w)=\|{\bm{x}}-{\bm{m}}_{\rm NN}^{\bm{W}}({\bm{z}})\|_{2}^{2}$. Therefore, to show Proposition 8, we just need to apply Lemma 3 by checking (a), (b), (c).

Check (a). We note that the index set $\Theta:={\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ equipped with $\rho(w,w^{\prime}):={|\!|\!|{{\bm{W}}-{\bm{W}}^{\prime}}|\!|\!|}$ has diameter ${B_{\rm p}}:=2B$. Further note that ${\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B}$ has a dimension bounded by ${d_{\rm p}}:=2D(D+2S+1)\|{\underline{m}}\|_{1}$. According to Example 5.8 of [], it holds that $\log N(\Delta;{\mathcal{W}}_{{d},{\underline{m}},L,{S},D,B},{|\!|\!|{\cdot}|\!
|\!|})\leq{d_{\rm p}}\cdot\log(1+2r/\Delta)$ for any $0<\Delta\leq 2r$. This verifies (a).

Check (b). Since $f(z_{i};w)={d}^{-1}\|{\bm{x}}-{\bm{m}}_{\rm NN}^{\bm{W}}({\bm{z}})\|_{2}^{2}$ is ${S}^{2}$-bounded. As a consequence, $f(z,w)-{\mathbb{E}}_{z}[f(z,w)]$ is a sub-Gaussian random variable with the sub-Gaussian parameter to be $C\cdot{S}^{2}$.

Check (c). Lemma 25 implies that

$$ $\displaystyle|f(z;w_{1})-f(z;w_{2})|\leq{L_{\rm p}}\cdot{|\!|\!|{{\bm{W}}_{1}- {\bm{W}}_{2}}|\!|\!|},~{}~{}~{}{L_{\rm p}}:=12\|{\underline{m}}\|_{1}18^{L}B^{ 6L}\cdot{d}\cdot{S}^{3}\Big{(}\sum_{v\in{\mathcal{V}}^{(L)}}({S}+|z_{v}|)\Big{ )}\cdot(S+D).$ $$

Since ${\bm{z}}\stackrel{{\scriptstyle d}}{{=}}{\bm{x}}+{\bm{g}}$ where $({\bm{x}},{\bm{g}})\sim{\mu_{\star}}\times{\mathcal{N}}({\bm{0}},{\mathbf{I}}_
{d})$, and $\|{\bm{x}}\|_{1}\leq{S}{d}$, $\|{\bm{g}}\|_{1}$ is $C{d}$-sub-Gaussian. Hence $\|{\bm{z}}\|_{1}$ is $C{S}{d}$-sub-Gausssian, and hence $f(z;w_{1})-f(z;w_{2})$ is $\sigma^{\prime}\rho(w_{1},w_{2})$ sub-Gaussian with

$$ $\sigma^{\prime}=C\|{\underline{m}}\|_{1}18^{L}B^{6L}\cdot{d}^{2}\cdot{S}^{4} \cdot(S+D).$ $$

Therefore, we apply Lemma 3 to conclude the proof of Proposition 8.
∎

###### Proposition 8 (Generalization error of the denoising problem) .

###### Proof of Proposition 8 .

### G.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 ${\mathsf{h}}_{\downarrow,v}^{(L)}=(-x^{2}/2+xz_{v})_{x\in[{S}]}\in{\mathbb{R}}
^{{S}}$ (Since we will immediately normalize the input, this input is effectively the same as the input ${\mathsf{h}}_{\downarrow,v}^{(L)}=(-(x-z_{v})^{2}/2)_{x\in[{S}]}\in{\mathbb{R}
}^{{S}}$). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then for any $\ell$, $v$, $\iota$, and $\star\in[3]$, we have

$$ $\displaystyle\|{\mathsf{h}}_{\downarrow,v}^{(L)}\|_{2}\leq$ $\displaystyle~{}{S}^{3}+{S}^{2}|z_{v}|,$ $\displaystyle\|{\mathsf{q}}_{\downarrow,v}^{(\ell)}\|_{2}\leq$ $\displaystyle~{}B^{3}\cdot(2\cdot\|{\mathsf{h}}_{\downarrow,v}^{(\ell)}\|_{2}+ 1),$ $\displaystyle\|{\mathsf{h}}_{\downarrow,v}^{(\ell-1)}\|_{2}\leq$ $\displaystyle~{}m^{(\ell)}\cdot\max_{v^{\prime}\in{\mathcal{C}}(v)}\|{\mathsf{ q}}_{\downarrow,v^{\prime}}^{(\ell)}\|_{2},$ $\displaystyle\|{\mathsf{b}}_{\uparrow,{\rm r}}^{(0)}\|_{2}=$ $\displaystyle~{}\|{\mathsf{h}}_{\downarrow,{\rm r}}^{(0)}\|_{2},$ $\displaystyle\|{\mathsf{b}}_{\uparrow,v}^{(\ell)}\|_{2}\leq$ $\displaystyle~{}B^{3}\cdot(2\|{\mathsf{b}}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}\| _{2}+2\|{\mathsf{q}}_{\downarrow,v}^{(\ell)}\|_{2}+1)+\|{\mathsf{h}}_{ \downarrow,v}^{(\ell)}\|_{2},$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\downarrow,\iota}}{ \mathsf{q}}_{\downarrow,v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}2B^{3}\cdot\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\downarrow ,\iota}}{\mathsf{h}}_{\downarrow,v,i}^{(\ell)}\|_{{\rm op}},$ $\displaystyle\forall k\geq\ell+1,$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\downarrow,\iota}}{ \mathsf{q}}_{\downarrow,v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}2B^{2}\cdot\|{\mathsf{h}}_{\downarrow,v}^{(\ell)}\|_{2},$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\downarrow,\iota}}{ \mathsf{h}}_{\downarrow,v,i}^{(\ell-1)}\|_{{\rm op}}\leq$ $\displaystyle~{}m^{(\ell)}\cdot\max_{v^{\prime}\in{\mathcal{C}}(v)}\max_{i\in[ {S}]}\|\nabla_{W^{(k)}_{\star,\downarrow,\iota}}{\mathsf{q}}_{\downarrow,v^{ \prime},i}^{(\ell)}\|_{{\rm op}},~{}~{}~{}~{}$ $\displaystyle\forall k\geq\ell,$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\uparrow,\iota}}{\mathsf {b}}_{\uparrow,v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}2B^{3}\cdot\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\uparrow, \iota}}{\mathsf{b}}_{\uparrow,{\rm pa}(v),i}^{(\ell-1)}\|_{{\rm op}},$ $\displaystyle\forall k\geq\ell+1,$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(\ell)}_{\star,\uparrow,\iota}}{ \mathsf{b}}_{\uparrow,v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}2B^{2}\cdot(\|{\mathsf{b}}_{\uparrow,{\rm pa}(v)}^{(\ell-1)}\| _{2}+\|{\mathsf{q}}_{\downarrow,v}^{(\ell)}\|_{2}),$ $\displaystyle\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\downarrow,\iota}}{ \mathsf{b}}_{\uparrow,v,i}^{(\ell)}\|_{{\rm op}}\leq$ $\displaystyle~{}2B^{3}\cdot\Big{(}\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star, \downarrow,\iota}}{\mathsf{b}}_{\uparrow,{\rm pa}(v),i}^{(\ell-1)}\|_{{\rm op}}$ $\displaystyle~{}+\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\downarrow,\iota}}{ \mathsf{q}}_{\downarrow,v,i}^{(\ell)}\|_{{\rm op}}\Big{)}+\max_{i\in[{S}]}\| \nabla_{W^{(k)}_{\star,\downarrow,\iota}}{\mathsf{h}}_{\downarrow,v,i}^{(\ell) }\|_{{\rm op}},$ $\displaystyle\forall k\in[L],$ $\displaystyle\|\nabla_{W^{(k)}_{\star,\diamond,\iota}}{\bm{m}}_{{\rm NN}}^{\bm {W}}({\bm{z}})_{v}\|_{{\rm op}}\leq$ $\displaystyle~{}{S}\cdot\max_{i\in[{S}]}\|\nabla_{W^{(k)}_{\star,\diamond, \iota}}{\mathsf{b}}_{\uparrow,v,i}^{(L)}\|_{{\rm op}},$ $\displaystyle\forall k\in[L],\diamond\in\{\downarrow,\uparrow\}.$ $$

###### Proof of Lemma 23 .

The proof of the lemma uses the chain rule, the $1$-Lipschitzness of ${\rm ReLU}$, the $2$-Lipschitzness of ${\rm normalize}$, and the $1$-Lipschitzness of ${\rm softmax}$.
∎

###### Lemma 24 .

Consider the U-Net as in Eq. (UNet). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then for any $v\in{\mathcal{V}}^{(L)}$, we have

$$ $\max_{\diamond\in\{\downarrow,\uparrow\}}\max_{{\bm{z}}\in{\mathbb{R}}^{d}} \max_{\star\in[3]}\max_{\ell\in[L]}\max_{\iota\in[m^{(\ell)}]}\|\nabla_{W^{( \ell)}_{\star,\diamond,\iota}}{\bm{m}}_{{\rm NN}}^{\bm{W}}({\bm{z}})_{v}\|_{{ \rm op}}\leq 18^{L}B^{6L}\cdot{d}\cdot{S}^{3}({S}+|z_{v}|).$ $$

###### Proof of Lemma 24 .

This lemma is implied by Lemma 23 and an induction argument.
∎

###### Lemma 25 .

Consider the ConvNet as in Eq. (UNet). Assume that ${|\!|\!|{{\bm{W}}}|\!|\!|}\leq B$. Then for any $v\in{\mathcal{V}}^{(L)}$, we have

$$ $\displaystyle\max_{{\bm{z}}}\Big{|}{\bm{m}}_{{\rm NN}}^{{\bm{W}}}({\bm{z}})_{v }-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{z}})_{v}\Big{|}\leq$ $\displaystyle~{}6\|{\underline{m}}\|_{1}18^{L}B^{6L}\cdot{d}\cdot{S}^{3}({S}+| z_{v}|)\cdot(S+D)\cdot{|\!|\!|{{\bm{W}}-\overline{\bm{W}}}|\!|\!|}.$ $$

Therefore, we have

$$ $\displaystyle\max_{{\bm{z}}}{d}^{-1}\Big{|}\|{\bm{x}}-{\bm{m}}_{{\rm NN}}^{{ \bm{W}}}({\bm{z}})\|_{2}^{2}-\|{\bm{x}}-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}} }({\bm{z}})\|_{2}^{2}\Big{|}\leq$ $\displaystyle~{}12\|{\underline{m}}\|_{1}18^{L}B^{6L}\cdot{d}\cdot{S}^{3}\Big{ (}\sum_{v\in{\mathcal{V}}^{(L)}}({S}+|z_{v}|)\Big{)}\cdot(S+D)\cdot{|\!|\!|{{ \bm{W}}-\overline{\bm{W}}}|\!|\!|}.$ $$

###### Proof of Lemma 25 .

The first inequality is by the fact that

$$ $\displaystyle~{}\max_{{\bm{z}}}\Big{|}{\bm{m}}_{{\rm NN}}^{{\bm{W}}}({\bm{z}}) _{v}-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{z}})_{v}\Big{|}$ $\displaystyle\leq$ $\displaystyle~{}\sum_{\diamond\in\{\downarrow,\uparrow\}}\sum_{\star\in[3]} \sum_{\ell\in[L]}\sum_{\iota\in[m^{(\ell)}]}\min\{{\rm nrow}(W^{(\ell)}_{\star ,\diamond,\iota}),{\rm ncol}(W^{(\ell)}_{\star,\diamond,\iota})\}\|\nabla_{W^{ (\ell)}_{\star,\diamond,\iota}}{\bm{m}}_{{\rm NN}}^{\tilde{\bm{W}}}({\bm{z}})_ {v}\|_{{\rm op}}\|W^{(\ell)}_{\star,\diamond,\iota}-\overline{W}^{(\ell)}_{ \star,\diamond,\iota}\|_{{\rm op}},$ $$

where we have used the inequality that ${\rm trace}(A^{\mathsf{T}}B)\leq\{{\rm nrow}(A),{\rm ncol}(A)\}\|A\|_{{\rm op}
}\|B\|_{{\rm op}}$.

To prove the second inequality, we have

$$ $\displaystyle~{}\max_{{\bm{z}}}\Big{|}\|{\bm{x}}-{\bm{m}}_{{\rm NN}}^{{\bm{W}} }({\bm{z}})\|_{2}^{2}-\|{\bm{x}}-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{ z}})\|_{2}^{2}\Big{|}$ $\displaystyle\leq$ $\displaystyle~{}\max_{{\bm{z}}}\|2{\bm{x}}-{\bm{m}}_{{\rm NN}}^{{\bm{W}}}({\bm {z}})-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{z}})\|_{\infty}\|{\bm{m}}_{ {\rm NN}}^{{\bm{W}}}({\bm{z}})-{\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{z} })\|_{1}$ $\displaystyle\leq$ $\displaystyle~{}2{d}\max_{{\bm{z}}}\|{\bm{m}}_{{\rm NN}}^{{\bm{W}}}({\bm{z}})- {\bm{m}}_{{\rm NN}}^{\overline{\bm{W}}}({\bm{z}})\|_{1}.$ $$

This completes the proof of Lemma 25.
∎

###### Lemma 23 (Norm bound in the chain rule in denoising settings) .

###### Proof of Lemma 23 .

###### Lemma 24 .

###### Proof of Lemma 24 .

###### Lemma 25 .

###### Proof of Lemma 25 .

## Appendix H The name “Generative Hierarchical Models”

There are several alternative names for “generative hierarchical models” in the literature, each of which we have opted not to use for specific reasons:

- •
Hierarchical generative models. Although the term “Hierarchical Generative Models” also appears in [, , , ], we avoid using “generative models” as it is commonly associated with a particular task or statistical method, which could lead to confusion.
- •
Bayesian hierarchical models/Hierarchical Bayesian models. The term “Bayesian Hierarchical Models” is prevalent in Bayesian statistics; however, it primarily refers to models that incorporate prior distributions on hyperparameters, which differs from the focus of our paper. Additionally, there is no “Bayesian” component in our model, making this terminology inappropriate.
- •
Latent hierarchical models/Hierarchical latent variable models. The use of “latent” in the name might mislead readers about the model’s structure and properties.
- •
Hierarchical Markov random fields. We find this term overly lengthy and not particularly informative regarding the specific characteristics of our model.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: - Why are the established bounds useful for future neural network architecture development? Could they help us understand scaling laws for U-Nets? 
- 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]?

[1] Baranchuk, Dmitry, et al. "Label-Efficient Semantic Segmentation with Diffusion Models." International Conference on Learning Representations.
soundness: 3
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.
- The writing of the paper is clear and helps with understanding the intuition behind the hierarchical model interpretation of U-Nets.
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
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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
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?

- In neural network, the upper bound in the theorem has the tendency to be too loose, I was wondering how tight are the bounds?
soundness: 3
strengths: - Provide a rigorous theoretical explanation of U-Net and CNN through the lens of a graphical model
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.
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
- (minor) Given how dense the paper is, a lookup table for notations would be very helpful

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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).
- 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?
soundness: 4
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.
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.
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.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: See the Weakness.
soundness: 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.

(2) Authors conduct a thorough complexity analysis, with extensive proofs provided in the appendix.

(3) The paper is well-written, and the proof strategy in the ConvNet example aids in understanding the later sections on U-Nets.
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.
weaknesses: (1) Although this is a theoretical paper, it lacks any form of empirical validation of its conclusions, even in a simplified format.

(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.

(3) The title suggests that U-Nets are efficient for classification, yet the paper only analyzes classification in the context of ConvNets.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
questions: Can the authors clarify whether true convolution or cross-correlation is used, given that parameter sharing and local invariance properties are retained?

What 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?

How 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?
soundness: 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. 

The 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.

The 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.
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.
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.

The practical utility of the provided bounds is unclear, as their direct practical applicability for empirical settings is not fully established.

Meta-review: This work examines graphical models and the U-Net architecture in simplified form in order to theoretically relate the two. In particular, or a restricted setting, the U-Net architecture and computation can be connected with belief propagation and approximation bounds are derived. This is a purely theoretical paper, and relates to prior work in this direction, while providing novel insight into the design and working of diffusion models with U-Net architectures.

The main strengths and weaknesses are:

- Pro: Important problem and novel theoretical connection (UKfv: "studies an important problem [...] to break the curse of dimensionality when learning diffusion models")
- Pro: The theoretical content is insightful, new, and able to be extended further (f34u: "rigorous theoretical explanation of U-Net and CNN through the lens of a graphical model", UKfv: "The restrictions in the theoretical analysis help simplify the exposition and could be relaxed in future work")
- Pro: While there are no experiments in this work, this weakness is tempered by evidence from existing work (UKfv: "I do not think that this paper requires experimental evaluations to be complete, as its message is (partly) supported by experimental evidence from the literature."). Furthermore experiments are not required.
- Con: The theoretical content may not be valid w.r.t. practice because the operations studied do not reflect convolution as it is computed in deep networks (5xWe: "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")
- Con: No experiments to support the theoretical hypotheses and conclusions (f34u: "conclusion would be strengthened if there were at least some toy-ish experiments to test these hypotheses", xjiq: "it lacks any form of empirical validation of its conclusions").
- Con: Questions about the purpose and value of the analysis and what precisely the interpretation of features as beliefs of latent nodes implies (iUQ6: "[...] there is no motivation on why such a bound could be useful.").

The strongest arguments for acceptance are the theoretical connection between deep nets and graphical models (in the given restricted scope) and the importance of the topic given the popularity of diffusion. While the connection is not complete, in that it does not extend to deep nets as they are used in practice, it has been judged as theoretically sound and can inform future work. Furthermore the work is well-connected to the recent literature, and carries on from the previous steps in this direction.

The strongest arguments for rejection are 1. the disconnect from networks studied and deep networks (specifically: U-Nets and ConvNets) as they are usually defined and used 2. the lack of experiments to justify theoretical hypotheses and check conclusions given the distance between the theory and practice as already noted.

Five expert reviewers agree on acceptance, although 4/5 of the reviewers vote only for marginal accept (rating: 6). However the reviewer UKfv votes for clear acceptance (rating: 8). UKfv is the most confident, has the most related background, and defends the paper during the discussion. The meta-reviewer sides with acceptance, though cautions the authors to include experiments and a discussion of the gaps between theory and practice in the final revision in order to achieve greater impact. This suggestion is in line with the authors' own statement that "we plan to include experimental results in the final version". This is merely a suggestion to emphasize the comments of multiple reviewers, and not a condition on acceptance, as the submission is sufficiently informative theoretically for publication per the evaluation of the reviewers.

Decision: Accept (Poster)